nan handling for continuous and categorical coloring (#427)

Co-authored-by: Tim Treis <tim.treis@stud.uni-heidelberg.de>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: 3 people

src/spatialdata_plot/pl/basic.py

@@ -213,7 +213,7 @@ def render_shapes(
             `fill_alpha` will overwrite the value present in the cmap.
         groups : list[str] | str | None
             When using `color` and the key represents discrete labels, `groups` can be used to show only a subset of
-            them. Other values are set to NA. If elment is None, broadcasting behaviour is attempted (use the same
+            them. Other values are set to NA. If element is None, broadcasting behaviour is attempted (use the same
             values for all elements).
         palette :  list[str] | str | None
             Palette for discrete annotations. List of valid color names that should be used for the categories. Must

src/spatialdata_plot/pl/render.py

@@ -983,6 +983,40 @@ def _infer_color_data_kind(
     return "numeric", pd.to_numeric(series, errors="coerce")
 
 
+def _build_alignment_dtype_hint(
+    sdata: sd.SpatialData | None,
+    element: object,
+    color_series: pd.Series,
+    table_name: str | None,
+) -> str:
+    """Build a diagnostic hint string for dtype mismatches between element and table indices."""
+    hints: list[str] = []
+    color_index_dtype = getattr(color_series.index, "dtype", None)
+    element_index_dtype = getattr(getattr(element, "index", None), "dtype", None) if element is not None else None
+
+    table_instance_dtype = None
+    instance_key = None
+    if table_name is not None and sdata is not None and table_name in sdata.tables:
+        table = sdata.tables[table_name]
+        try:
+            _, _, instance_key = get_table_keys(table)
+        except (KeyError, ValueError, TypeError, AttributeError):
+            instance_key = None
+        if instance_key is not None and hasattr(table, "obs") and instance_key in table.obs:
+            table_instance_dtype = table.obs[instance_key].dtype
+
+    if (
+        element_index_dtype is not None
+        and table_instance_dtype is not None
+        and element_index_dtype != table_instance_dtype
+    ):
+        hints.append(f"element index dtype is {element_index_dtype}, '{instance_key}' dtype is {table_instance_dtype}")
+    if color_index_dtype is not None and element_index_dtype is not None and color_index_dtype != element_index_dtype:
+        hints.append(f"color index dtype is {color_index_dtype}, element index dtype is {element_index_dtype}")
+
+    return f" (hint: {'; '.join(hints)})" if hints else ""
+
+
 def _set_color_source_vec(
     sdata: sd.SpatialData,
     element: SpatialElement | None,
@@ -1012,7 +1046,8 @@ def _set_color_source_vec(
 
     if len(origins) > 1:
         raise ValueError(
-            f"Color key '{value_to_plot}' for element '{element_name}' been found in multiple locations: {origins}."
+            f"Color key '{value_to_plot}' for element '{element_name}' was found in multiple locations: {origins}. "
+            "Please keep it in exactly one place (preferably on the points parquet for speed) to avoid ambiguity."
         )
 
     if len(origins) == 1 and value_to_plot is not None:
@@ -1035,6 +1070,17 @@ def _set_color_source_vec(
             color_source_vector if isinstance(color_source_vector, pd.Series) else pd.Series(color_source_vector)
         )
 
+        if color_series.isna().all():
+            element_label = _format_element_name(element_name)
+            location = f"table '{table_name}'" if table_name is not None else "the element"
+            dtype_hint = _build_alignment_dtype_hint(sdata, element, color_series, table_name)
+            raise ValueError(
+                f"Column '{value_to_plot}' for element '{element_label}' contains only missing values after aligning "
+                f"with {location}. This usually means the instance ids/indices could not be aligned or converted, so "
+                "colors cannot be determined. Please ensure the table annotates the element with matching instance ids."
+                f"{dtype_hint}"
+            )
+
         kind, processed = _infer_color_data_kind(
             series=color_series,
             value_to_plot=value_to_plot,
@@ -1059,6 +1105,9 @@ def _set_color_source_vec(
             return None, numeric_vector, False
 
         assert isinstance(processed, pd.Categorical)
+        if not processed.ordered:
+            # ensure deterministic category order when the source is unordered (e.g., from a Python set)
+            processed = processed.reorder_categories(sorted(processed.categories))
         color_source_vector = processed  # convert, e.g., `pd.Series`
 
         # Use the provided table_name parameter, fall back to only one present
@@ -1138,6 +1187,12 @@ def _set_color_source_vec(
         # (e.g. two categories share a color). Wrapping back in pd.Categorical ensures
         # downstream consumers always receive a Categorical for categorical data.
         color_vector = pd.Categorical(color_source_vector.map(color_mapping, na_action="ignore"))
+        # nan handling: only add the NA category if needed, and store it as a hex string
+        na_color_hex = na_color.get_hex_with_alpha() if isinstance(na_color, Color) else str(na_color)
+        if color_vector.isna().any():
+            if na_color_hex not in color_vector.categories:
+                color_vector = color_vector.add_categories(na_color_hex)
+            color_vector[pd.isna(color_vector)] = na_color_hex
 
         return color_source_vector, color_vector, True
 
@@ -1165,15 +1220,18 @@ def _map_color_seg(
 
     if isinstance(color_vector.dtype, pd.CategoricalDtype):
         # Case A: users wants to plot a categorical column
-        if np.any(color_source_vector.isna()):
-            cell_id[color_source_vector.isna()] = 0
         val_im: ArrayLike = map_array(seg.copy(), cell_id, color_vector.codes + 1)
         cols = colors.to_rgba_array(color_vector.categories)
     elif pd.api.types.is_numeric_dtype(color_vector.dtype):
         # Case B: user wants to plot a continous column
         if isinstance(color_vector, pd.Series):
             color_vector = color_vector.to_numpy()
-        cols = cmap_params.cmap(cmap_params.norm(color_vector))
+        # normalize only the not nan values, else the whole array would contain only nan values
+        normed_color_vector = color_vector.copy().astype(float)
+        normed_color_vector[~np.isnan(normed_color_vector)] = cmap_params.norm(
+            normed_color_vector[~np.isnan(normed_color_vector)]
+        )
+        cols = cmap_params.cmap(normed_color_vector)
         val_im = map_array(seg.copy(), cell_id, cell_id)
     else:
         # Case C: User didn't specify any colors
@@ -2688,6 +2746,7 @@ def _validate_col_for_column_table(
     elif table_name is not None:
         tables = get_element_annotators(sdata, element_name)
         if table_name not in tables:
+            logger.warning(f"Table '{table_name}' does not annotate element '{element_name}'.")
             raise KeyError(f"Table '{table_name}' does not annotate element '{element_name}'.")
         if col_for_color not in sdata[table_name].obs.columns and col_for_color not in sdata[table_name].var_names:
             raise KeyError(
@@ -3084,7 +3143,7 @@ def _prepare_transformation(
 def _datashader_map_aggregate_to_color(
     agg: DataArray,
     cmap: str | list[str] | ListedColormap,
-    color_key: None | list[str] = None,
+    color_key: list[str] | dict[str, str] | None = None,
     min_alpha: float = 40,
     span: None | list[float] = None,
     clip: bool = True,