fix(labels): align instance_id with colour vector for partial tables (#775) (#776)

Author: timtreis

src/spatialdata_plot/pl/render.py

@@ -20,7 +20,7 @@
 from matplotlib import patheffects
 from matplotlib.cm import ScalarMappable
 from matplotlib.colors import BoundaryNorm, Colormap, ListedColormap, Normalize, to_rgba_array
-from spatialdata import get_extent, get_values
+from spatialdata import get_element_instances, get_extent, get_values
 from spatialdata.models import PointsModel, ShapesModel, get_table_keys
 from spatialdata.transformations import set_transformation
 from spatialdata.transformations.transformations import Identity
@@ -2340,6 +2340,12 @@ def _render_labels(
         # limits already clip to the box, so return before the instance-overlap/colour machinery.
         return
 
+    # instance_id must line up with the per-instance colour/outline vectors only when one of these
+    # drives the render; the same gate guards the rasterize reconciliation below.
+    needs_aligned_instances = (
+        col_for_color is not None or render_params.col_for_outline_color is not None or render_params.as_points
+    )
+
     if table_name is None:
         instance_id = unique_labels
         table = None
@@ -2355,8 +2361,19 @@ def _render_labels(
                 "instance_id=0 before plotting."
             )
 
-        # get instance id based on subsetted table
-        instance_id = np.unique(table.obs[instance_key].values)
+        if needs_aligned_instances:
+            # Restrict to instances that actually exist in the element (canonical scale-0),
+            # matching the colour vector's basis from get_values; table rows for absent
+            # instances are dropped instead of misaligning the mask. Rasterize/multiscale
+            # display drops are reconciled against both vectors by the block below.
+            # ponytail: this rescans the canonical raster get_values also scans; thread
+            # get_values' instance index out of resolve_color if the double pass ever bites.
+            instance_id = np.intersect1d(
+                np.unique(table.obs[instance_key].values),
+                np.asarray(get_element_instances(sdata_filt[element])),
+            )
+        else:
+            instance_id = np.unique(table.obs[instance_key].values)
 
     trans, trans_data = _prepare_transformation(label, coordinate_system, ax)
 
@@ -2408,7 +2425,7 @@ def _render_labels(
 
     # rasterize/downsampling can drop labels from the raster; remove their (now-absent) instance ids
     # so per-instance colors stay aligned and as_points does not emit dots for dropped cells.
-    if rasterize and (col_for_color is not None or col_for_outline_color is not None or render_params.as_points):
+    if rasterize and needs_aligned_instances:
         mask = np.isin(instance_id, unique_labels)
         instance_id = instance_id[mask]
         if col_for_color is not None:

tests/pl/test_render_labels.py

@@ -823,6 +823,60 @@ def test_render_labels_disjoint_instance_ids_clear_error():
         plt.close(fig)
 
 
+def _labels_with_partial_table(instance_id, *, scale_factors=None):
+    # labels contain instances 1, 2; the table's instance_id set is caller-controlled so tests can
+    # add phantom rows (instances absent from the raster) or omit a present instance.
+    arr = np.zeros((20, 20), dtype=np.int32)
+    arr[3:8, 3:8] = 1
+    arr[12:17, 12:17] = 2
+    obs = pd.DataFrame(
+        {
+            "instance_id": instance_id,
+            "region": pd.Categorical(["lbl"] * len(instance_id)),
+            "cat": pd.Categorical([c for c, _ in zip("ABCDEFG", instance_id, strict=False)]),
+        }
+    )
+    obs.index = obs.index.astype(str)
+    table = TableModel.parse(
+        AnnData(X=np.zeros((len(instance_id), 1)), obs=obs),
+        region=["lbl"],
+        region_key="region",
+        instance_key="instance_id",
+    )
+    table.obs["value"] = np.arange(len(instance_id), dtype=float)
+    labels = Labels2DModel.parse(arr, dims=["y", "x"], scale_factors=scale_factors)
+    return SpatialData(labels={"lbl": labels}, tables={"t": table})
+
+
+@pytest.mark.parametrize("color", ["value", "cat"])
+@pytest.mark.parametrize(
+    "instance_id",
+    [
+        [1, 2, 3],  # phantom row: instance 3 is absent from the raster (#775)
+        [1],  # present instance 2 has no table row
+    ],
+)
+def test_render_labels_partial_table_does_not_raise(color, instance_id):
+    # Regression test for #775: a table annotating instances that are not present in the labels
+    # (or missing a present instance) must render as missing, not raise IndexError.
+    sdata = _labels_with_partial_table(instance_id)
+    fig, ax = plt.subplots()
+    try:
+        sdata.pl.render_labels("lbl", color=color, table_name="t").pl.show(ax=ax)
+    finally:
+        plt.close(fig)
+
+
+def test_render_labels_phantom_row_survives_rasterization():
+    # #775: the alignment must also hold once rasterization/multiscale drops labels from the raster.
+    sdata = _labels_with_partial_table([1, 2, 3], scale_factors=[2])
+    fig, ax = plt.subplots()
+    try:
+        sdata.pl.render_labels("lbl", color="value", table_name="t").pl.show(ax=ax)
+    finally:
+        plt.close(fig)
+
+
 @pytest.mark.parametrize("scale_factors", [None, [2]])
 def test_render_labels_raises_on_3d(scale_factors):
     # Regression test for #608: 3D labels must raise a clear ValueError, not crash