From 58a177663b300e829f16d0ffb55b0424d9a19601 Mon Sep 17 00:00:00 2001
From: Jeremy Muhlich <jmuhlich@bitflood.org>
Date: Wed, 6 Sep 2023 15:05:22 -0400
Subject: [PATCH] Track LayerAligner discard reasons and show in report

---
 ashlar/reg.py    |  9 ++++---
 ashlar/report.py | 63 ++++++++++++++++++++++++++++++------------------
 2 files changed, 45 insertions(+), 27 deletions(-)

diff --git a/ashlar/reg.py b/ashlar/reg.py
index 7e5a5b9f..7ff73ad4 100644
--- a/ashlar/reg.py
+++ b/ashlar/reg.py
@@ -901,9 +901,10 @@ def constrain_positions(self):
         # positions to reference aligner positions, due to strong
         # self-correlation of the sensor dark current pattern which dominates in
         # low-signal images.
-        discard = (position_diffs == 0).all(axis=1)
+        self.discard_camera_bg = (position_diffs == 0).all(axis=1)
         # Discard any tile registration that error is infinite
-        discard |= np.isinf(self.errors)
+        self.discard_error = np.isinf(self.errors)
+        discard = self.discard_camera_bg | self.discard_error
         # Take the median of registered shifts to determine the offset
         # (translation) from the reference image to this one.
         if discard.all():
@@ -917,9 +918,9 @@ def constrain_positions(self):
         # replacing it with the relevant model prediction.
         distance = np.linalg.norm(self.positions - predictions - offset, axis=1)
         max_dist = self.max_shift_pixels
-        extremes = distance > max_dist
+        self.discard_distance = distance > max_dist
         # Recalculate the mean shift, also ignoring the extreme values.
-        discard |= extremes
+        discard |= self.discard_distance
         self.discard = discard
         if discard.all():
             self.offset = 0
diff --git a/ashlar/report.py b/ashlar/report.py
index 08423ddc..fa8b8ddd 100644
--- a/ashlar/report.py
+++ b/ashlar/report.py
@@ -288,7 +288,8 @@ def plot_layer_map(
     node_cmap = mcm.Greens.with_extremes(over="#333333")
     g = aligner.neighbors_graph
     pos = np.fliplr(centers)
-    qlen = np.min(aligner.metadata.size) * 0.45
+    qlen = np.repeat(np.min(aligner.metadata.size) * 0.45, len(pos))
+    qlen[aligner.discard_camera_bg | aligner.discard_error] = 0
     q_angles = np.rad2deg(np.arctan2(*(aligner.shifts * [-1, 1]).T))
     reference_offset = (
         aligner.cycle_offset
@@ -315,7 +316,7 @@ def plot_layer_map(
     ax.quiver(
         pos[:, 0],
         pos[:, 1],
-        [qlen] * len(pos),
+        qlen,
         [0] * len(pos),
         shifts,
         cmap=node_cmap,
@@ -327,28 +328,42 @@ def plot_layer_map(
         headlength=1,
         headaxislength=1,
     )
-    nx.draw_networkx_nodes(
-        g,
-        pos,
-        ax=ax,
-        cmap=node_cmap,
-        vmin=smin,
-        vmax=smax,
-        node_color=node_values,
-        node_size=node_size,
-        edgecolors=None,
-        **nx_kwargs,
-    )
-    draw_labels = functools.partial(
-        nx.draw_networkx_labels,
-        pos=pos,
-        ax=ax,
-        font_size=font_size,
-        **nx_kwargs,
-    )
-    draw_labels(g.subgraph(np.nonzero(aligner.discard)[0]), font_color="gray")
-    draw_labels(g.subgraph(np.nonzero(~aligner.discard)[0]), font_color="k")
+
+    def subgraph(cond):
+        return g.subgraph(np.nonzero(cond)[0])
+
+    def draw_nodes(cond, **kwargs):
+        nx.draw_networkx_nodes(
+            subgraph(cond),
+            pos=pos,
+            ax=ax,
+            cmap=node_cmap,
+            vmin=smin,
+            vmax=smax,
+            node_color=node_values[cond[np.array(aligner.neighbors_graph)]],
+            node_size=node_size,
+            edgecolors=None,
+            **nx_kwargs,
+            **kwargs,
+        )
+
+    def draw_labels(cond, **kwargs):
+        nx.draw_networkx_labels(
+            subgraph(cond),
+            pos=pos,
+            ax=ax,
+            font_size=font_size,
+            **nx_kwargs,
+            **kwargs,
+        )
+
+    draw_nodes(aligner.discard_camera_bg, node_shape="s")
+    draw_nodes(aligner.discard_error, node_shape="D")
+    draw_nodes(~aligner.discard | aligner.discard_distance, node_shape="o")
+    draw_labels(aligner.discard, font_color="gray")
+    draw_labels(~aligner.discard, font_color="k")
     draw_borders(ax, aligner, pos)
+
     cbar = fig.colorbar(
         mcm.ScalarMappable(mcolors.Normalize(smin, smax), node_cmap),
         extend="max",
@@ -357,9 +372,11 @@ def plot_layer_map(
         shrink=0.5,
         ax=ax,
     )
+
     ax.set_frame_on(False)
     ax.margins(0)
     fig.tight_layout()
+
     return fig