Skip to content

Vectorize continuous color computation in pl.nodes() for large speedup #48

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
colganwi merged 2 commits into from
Feb 25, 2026
Merged

Vectorize continuous color computation in pl.nodes() for large speedup #48

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b13c4e1
Vectorize continuous color computation in pl.nodes()
colganwi Feb 25, 2026
ec2bc80
Update test_get_colors_numeric for vectorized ndarray return type
colganwi Feb 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion src/pycea/pl/_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))))
Copy link

@chatgpt-codex-connector chatgpt-codex-connector bot Feb 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Avoid treating ±inf as missing in numeric color mapping

Using np.ma.masked_invalid(...) here masks both NaN and ±inf, so any infinite numeric values are now rendered with na_color instead of being mapped through the continuous colormap. This changes result semantics for workflows where overflow/div-by-zero values appear in node/edge metrics (they become visually indistinguishable from truly missing values), whereas the previous per-element path only treated absent indices as missing.

Useful? React with 👍 / 👎.

Copy link

@chatgpt-codex-connector chatgpt-codex-connector bot Feb 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Preserve infinite numeric values during color vectorization

Using np.ma.masked_invalid(...) here treats +/-inf as missing values, not just absent/NaN entries. In the previous path, infinite numeric values were still passed through Normalize and mapped via the colormap (typically to over/under extremes), but this change forces them to na_color, which can silently miscolor real overflow/outlier data in pl.nodes()/pl.branches() when annotations contain infinities.

Useful? React with 👍 / 👎.

legend = _cbar_legend(key, color_map, norm)
n_categories = 0
else: # Categorical
Expand Down
4 changes: 2 additions & 2 deletions tests/test_plot_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down