diff --git a/src/pycea/pl/_utils.py b/src/pycea/pl/_utils.py index 04f7fe9..5c98a1d 100755 --- a/src/pycea/pl/_utils.py +++ b/src/pycea/pl/_utils.py @@ -365,7 +365,10 @@ def _get_colors( if data.dtype.kind in ["i", "f"]: # Numeric norm = _get_norm(vmin=vmin, vmax=vmax, data=data) color_map = plt.get_cmap(cmap) - colors = [color_map(norm(data[i])) if i in data.index else na_color for i in indicies] + # Vectorized: reindex to align with indicies (NaN for missing), then apply colormap in bulk + values = data.reindex(indicies) + color_map.set_bad(na_color) + colors = color_map(norm(np.ma.masked_invalid(values.values.astype(float)))) legend = _cbar_legend(key, color_map, norm) n_categories = 0 else: # Categorical diff --git a/tests/test_plot_utils.py b/tests/test_plot_utils.py index bec54ed..0edea3e 100755 --- a/tests/test_plot_utils.py +++ b/tests/test_plot_utils.py @@ -240,9 +240,9 @@ def test_get_colors_numeric(): data = pd.Series([0, 1, 2], index=["a", "b", "c"]) indices = ["a", "b", "c", "d"] colors, legend, ncat = _get_colors(tdata, "num", data, indices, cmap="viridis") - assert isinstance(colors, list) + assert isinstance(colors, np.ndarray) assert len(colors) == 4 - assert colors[-1] == "lightgrey" + np.testing.assert_allclose(colors[-1], mcolors.to_rgba("lightgrey")) assert ncat == 0 assert isinstance(legend, dict)