From 6c69c8e1a663d17a405c6e7471355838deb809a8 Mon Sep 17 00:00:00 2001
From: Runpeng Luo <rl6004@princeton.edu>
Date: Thu, 12 Nov 2026 00:48:34 -0500
Subject: [PATCH] 0.1.1b3: cn_palette option with a cnamaste (a, b) palette

- New "cnamaste" palette for (a, b) states, from cnamaste.palette.COPY_STATES:
  hue by minor-allele fraction, darker for more copies. The existing palette
  is named "hatchet" and stays the default.
- cn_palette= on get_cn_cmap, get_mixcn_cmap and plot_cnv_profile (passed on
  to its key); the figure helpers forward it through profile_kwargs.
- CN_PALETTES lists the names, for an argparse choices; an unknown name
  raises ValueError.
- The key's overflow box is labelled from the palette: >7 under hatchet,
  >6 under cnamaste.
- examples/plot_gallery.py takes --cn_palette {hatchet,cnamaste} and --outdir.
- Version 0.1.1b3.
---
 CHANGELOG.md                  |  8 ++++
 README.md                     |  5 +++
 VERSION                       |  2 +-
 examples/plot_gallery.py      | 30 ++++++++++++---
 src/cnplot/__init__.py        |  2 +
 src/cnplot/cnplot_colormap.py | 70 ++++++++++++++++++++++++++++++-----
 src/cnplot/cnplot_intcnp.py   | 20 +++++++---
 tests/test_intcnp.py          | 19 +++++++++-
 tests/test_utils_colormap.py  | 36 ++++++++++++++++++
 9 files changed, 170 insertions(+), 22 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ee6c9fc..5957c77 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,14 @@
 ## Unreleased
 
 ### cnplot package
+- A second (a, b) state palette, `"cnamaste"` (`cnamaste.palette.COPY_STATES`: hue by
+  minor-allele fraction, darker for more copies), next to the default, now named
+  `"hatchet"`. Picked with `cn_palette=` on `get_cn_cmap`, `get_mixcn_cmap` and
+  `plot_cnv_profile` (which passes it to its key); `CN_PALETTES` lists the names, for an
+  argparse `choices`. States a palette does not list take the default color, and the
+  key's trailing box is labelled from the palette: `>7` under `"hatchet"`, `>6` under
+  `"cnamaste"`. An unknown name raises `ValueError`. `examples/plot_gallery.py` takes
+  `--cn_palette {hatchet,cnamaste}` and `--outdir`.
 - Clone names are set in the regular weight: the clone labels beside the integer-CN profile
   and the `<label> (<pct>%)` row-block labels on the heatmap. The CN key's `CNA` title is
   regular as well.
diff --git a/README.md b/README.md
index e7b59bf..2285a1e 100644
--- a/README.md
+++ b/README.md
@@ -55,6 +55,11 @@ fig.savefig("profile.png", bbox_inches="tight")   # returns None; you own the fi
 
 ![profile](examples/profile.png)
 
+The (a, b) state colors come from one of `CN_PALETTES`: `"hatchet"` (default) or
+`"cnamaste"`. Pass the same `cn_palette=` to `plot_cnv_profile` (or through
+`profile_kwargs` of the figure helpers) and to `get_mixcn_cmap`, so points, profile
+and key agree. The gallery script takes it as `--cn_palette {hatchet,cnamaste}`.
+
 ### 2. Genome-wide RDR and BAF (multi-sample)
 We plots the per-bin read-depth ratio (RDR) and B-allele frequency (BAF) for both samples over the shared clonal profile.
 Reference: [`obs_df`](docs/reference.md#obs_df-bin-level-observations), [`expected_df` (1D)](docs/reference.md#expected_df-1d-overlay).
diff --git a/VERSION b/VERSION
index 696eb95..67fbf65 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-0.1.1b2
+0.1.1b3
diff --git a/examples/plot_gallery.py b/examples/plot_gallery.py
index 8bf98c6..8b724b1 100644
--- a/examples/plot_gallery.py
+++ b/examples/plot_gallery.py
@@ -7,11 +7,14 @@
 #   cnplot, and the tests/ simulator (tests/simulate.py + tests/data/)
 #
 # Usage:
-#   python examples/plot_gallery.py
+#   python examples/plot_gallery.py [--cn_palette {hatchet,cnamaste}] [--outdir DIR]
+#           --cn_palette    # (a, b) state palette; default hatchet
+#           --outdir        # directory for the PNGs; default examples/
 #
 # Outputs:
-#   examples/{profile,scatter_1d,scatter_2d,fcn_ab,heatmap_rdr,heatmap_baf}.png
+#   <outdir>/{profile,scatter_1d,scatter_2d,fcn_ab,heatmap_rdr,heatmap_baf}.png
 
+import argparse
 import os
 import sys
 
@@ -28,6 +31,7 @@
 
 from cnplot import (  # noqa: E402
     BAF_LIM,
+    CN_PALETTES,
     FS_TITLE,
     GenomeAxis,
     get_baf_cmap,
@@ -42,9 +46,16 @@
 )
 from simulate import CELL_POSTERIORS, load_dataset, reference  # noqa: E402
 
+parser = argparse.ArgumentParser(description="Render the README gallery figures.")
+parser.add_argument("--cn_palette", choices=CN_PALETTES, default="hatchet")
+parser.add_argument("--outdir", default=HERE)
+args = parser.parse_args()
+OUTDIR = os.path.abspath(args.outdir)
+os.makedirs(OUTDIR, exist_ok=True)
+
 
 def save(fig, name):
-    path = os.path.join(HERE, f"{name}.png")
+    path = os.path.join(OUTDIR, f"{name}.png")
     fig.savefig(path, dpi=100, bbox_inches="tight")
     plt.close(fig)
     print(path)
@@ -62,12 +73,19 @@ def suptitle(fig, text, y=1.0):
 palette = get_mixcn_cmap(
     list(sim.obs["cnp"].unique()),
     [float(sim.seg_ucn[f"u_{c}"].iloc[0]) for c in sim.clones],
+    cn_palette=args.cn_palette,
 )
 
 # 1. integer copy-number profile
 fig, (ax, lg) = plt.subplots(2, 1, figsize=(20, 3.5), height_ratios=[3, 1])
 plot_cnv_profile(
-    ax, sim.seg_ucn, genome_axis, sample_id="S1", ax_leg=lg, show_pi_viol=False
+    ax,
+    sim.seg_ucn,
+    genome_axis,
+    sample_id="S1",
+    ax_leg=lg,
+    show_pi_viol=False,
+    cn_palette=args.cn_palette,
 )
 suptitle(fig, "Integer copy-number profile", y=1.18)
 save(fig, "profile")
@@ -87,6 +105,7 @@ def suptitle(fig, text, y=1.0):
     seg_df=sim.seg_ucn,
     mb_ticks=True,
     show_pi_viol=False,
+    cn_palette=args.cn_palette,
 )
 suptitle(fig, "Genome-wide RDR and BAF (multi-sample)", y=1.0)
 save(fig, "scatter_1d")
@@ -124,6 +143,7 @@ def suptitle(fig, text, y=1.0):
     seg_df=sim.seg_ucn,
     mb_ticks=True,
     show_pi_viol=False,
+    cn_palette=args.cn_palette,
 )
 suptitle(fig, "Allele-specific fractional copy number (FCN-A / FCN-B)", y=1.0)
 save(fig, "fcn_ab")
@@ -169,6 +189,6 @@ def suptitle(fig, text, y=1.0):
         strips=strips,
         figsize=(20, 9),
         mb_ticks=True,
-        profile_kwargs={"show_pi_viol": False},
+        profile_kwargs={"show_pi_viol": False, "cn_palette": args.cn_palette},
     )
     save(fig, out)
diff --git a/src/cnplot/__init__.py b/src/cnplot/__init__.py
index 2e3dc9f..4ff5e0d 100644
--- a/src/cnplot/__init__.py
+++ b/src/cnplot/__init__.py
@@ -18,6 +18,7 @@
 from .cnplot_colormap import (
     BAF_COLORS,
     CELLTYPE_CMAP,
+    CN_PALETTES,
     DATASET_CMAP,
     NA_COLOR,
     NA_LABELS,
@@ -116,6 +117,7 @@ def _read_version() -> str:
     "BAF_COLORS",
     "BAF_LIM",
     "CELLTYPE_CMAP",
+    "CN_PALETTES",
     "CN_PREFIX",
     "DATASET_CMAP",
     "DEFAULT_EXCLUDED_CHROMS",
diff --git a/src/cnplot/cnplot_colormap.py b/src/cnplot/cnplot_colormap.py
index dbc545c..34cd619 100644
--- a/src/cnplot/cnplot_colormap.py
+++ b/src/cnplot/cnplot_colormap.py
@@ -3,7 +3,8 @@
 Three independent palettes:
 
 - joint CN states (a, b) -> :func:`get_cn_cmap`, shared by profiles, legends,
-  and scatter points so all three agree.
+  and scatter points so all three agree. ``cn_palette`` picks one of
+  :data:`CN_PALETTES`, "hatchet" by default.
 - per-allele CN -> :func:`get_ascn_cmap`, a sequential ramp for A/B panels.
 - categorical labels -> :func:`get_multiclass_cmap`, which colors several
   annotation strips from one shared palette and keeps "normal" and missing
@@ -23,6 +24,7 @@
 __all__ = [
     "BAF_COLORS",
     "CELLTYPE_CMAP",
+    "CN_PALETTES",
     "DATASET_CMAP",
     "NA_COLOR",
     "NA_LABELS",
@@ -72,10 +74,10 @@
 _TUMOR_COLORS = [mcolors.to_hex(c) for c in plt.get_cmap("tab10").colors]
 _TUMOR_COLORS_20 = [mcolors.to_hex(c) for c in plt.get_cmap("tab20").colors]
 
-# Integer CN palette, keyed by (major, minor). Hue tracks total CN and shade
+# "hatchet" integer CN palette, keyed by (major, minor). Hue tracks total CN and shade
 # tracks allelic imbalance, so LOH reads darker than a balanced state of the
 # same total.
-_CN_PALETTE = {
+_HATCHET_PALETTE = {
     (0, 0): "darkblue",
     (1, 0): "lightblue",
     (1, 1): "lightgray",
@@ -98,7 +100,31 @@
     (7, 0): "#4b0082",
 }
 
-# teal-green, distinct from every CN <= 7 color above
+# NB: cnamaste.palette.COPY_STATES; hue by minor-allele fraction, darker for more copies
+_CNAMASTE_PALETTE = {
+    (1, 1): "#EEEEEE",
+    (2, 2): "#FC9272",
+    (3, 3): "#CB181D",
+    (1, 0): "#66D9E8",
+    (2, 0): "#22B5D2",
+    (3, 0): "#1A86C2",
+    (4, 0): "#0D56A0",
+    (5, 0): "#08306B",
+    (6, 0): "#041A3A",
+    (2, 1): "#A1D99B",
+    (4, 2): "#31A354",
+    (3, 1): "#FEC44F",
+    (4, 1): "#6A51A3",
+    (5, 1): "#AE017E",
+    (3, 2): "#F16913",
+    (0, 0): "#000000",
+}
+
+_CN_PALETTE_MAPS = {"hatchet": _HATCHET_PALETTE, "cnamaste": _CNAMASTE_PALETTE}
+
+CN_PALETTES = tuple(_CN_PALETTE_MAPS)
+
+# NB: teal-green, distinct from every color of either palette above
 _CN_DEFAULT_COLOR = "#00cc99"
 
 _CN_COPY_STATES = [
@@ -252,22 +278,40 @@ def _clone_order_key(label: str, na_labels) -> tuple:
 # =============================================================================
 
 
-def get_cn_cmap() -> tuple:
+def get_cn_cmap(cn_palette: str = "hatchet") -> tuple:
     """Return the integer allele-specific copy-number color map and legend order.
 
     Keyed by (a, b) under both orderings, so (2, 1) and (1, 2) render alike.
-    States above total CN 7 fall through to the default color.
+    States the palette does not list fall through to the default color: total
+    CN above 7 under "hatchet", above 6 under "cnamaste".
+
+    Args:
+        cn_palette: One of :data:`CN_PALETTES`. "hatchet" gives hue by total CN
+            and shade by allelic imbalance. "cnamaste" is
+            ``cnamaste.palette.COPY_STATES``: hue by minor-allele fraction
+            (balanced reds from a light-gray (1, 1), LOH cyan to navy, 1/3
+            greens, 1/4 yellow, 1/5 purple, 1/6 magenta, 2/5 orange), darker
+            for more copies.
 
     Returns:
         (state_style, tcn_states): the {(a, b): color} map with a "default"
-        key, and {total_cn: [(a, b), ...]} in the order legends draw them.
+        key, and {total_cn: [(a, b), ...]} in the order legends draw them,
+        holding only the states the palette colors.
+
+    Raises:
+        ValueError: If ``cn_palette`` is not in :data:`CN_PALETTES`.
     """
+    if cn_palette not in _CN_PALETTE_MAPS:
+        raise ValueError(f"cn_palette must be one of {CN_PALETTES}, got {cn_palette!r}")
+    palette = _CN_PALETTE_MAPS[cn_palette]
+
     tcn_states = {}
     for a, b in _CN_COPY_STATES:
-        tcn_states.setdefault(int(a + b), []).append((a, b))
+        if (max(a, b), min(a, b)) in palette:
+            tcn_states.setdefault(int(a + b), []).append((a, b))
 
     state_style = {}
-    for (major, minor), color in _CN_PALETTE.items():
+    for (major, minor), color in palette.items():
         state_style[(major, minor)] = color
         state_style[(minor, major)] = color
     state_style["default"] = _CN_DEFAULT_COLOR
@@ -303,6 +347,7 @@ def get_mixcn_cmap(
     clone_states,
     clone_props=None,
     display_min_clone_prop: float | None = None,
+    cn_palette: str = "hatchet",
 ) -> dict:
     """Map each joint-clone mixture CNP string to a scatter color.
 
@@ -317,11 +362,16 @@ def get_mixcn_cmap(
             normal. None treats every clone as visible.
         display_min_clone_prop: Minimum proportion for a tumor clone to count as
             visible. None treats every clone as visible.
+        cn_palette: Integer-CN palette, one of :data:`CN_PALETTES`; pass the
+            same one to the profile so points and profile agree.
 
     Returns:
         {cnp_string: color} covering every entry of ``clone_states``.
+
+    Raises:
+        ValueError: If ``cn_palette`` is not in :data:`CN_PALETTES`.
     """
-    state_style, _ = get_cn_cmap()
+    state_style, _ = get_cn_cmap(cn_palette)
 
     def visible(j):
         """Test whether clone j is above the display threshold.
diff --git a/src/cnplot/cnplot_intcnp.py b/src/cnplot/cnplot_intcnp.py
index b2e4bab..61da0ee 100644
--- a/src/cnplot/cnplot_intcnp.py
+++ b/src/cnplot/cnplot_intcnp.py
@@ -370,6 +370,7 @@ def plot_cnv_profile(
     contain: bool = True,
     legend_kwargs: dict | None = None,
     rasterized: bool = True,
+    cn_palette: str = "hatchet",
 ) -> None:
     """Draw a stacked integer copy-number profile, one row per clone.
 
@@ -415,9 +416,13 @@ def plot_cnv_profile(
         legend_kwargs: Extra styling forwarded to the legend, ignored without
             ``ax_leg``.
         rasterized: Rasterize the rectangles, keeping vector output small.
+        cn_palette: (a, b) palette for the bins and the legend, one of
+            :data:`~cnplot.cnplot_colormap.CN_PALETTES`; see
+            :func:`~cnplot.cnplot_colormap.get_cn_cmap`.
 
     Raises:
-        ValueError: If ``seg_df`` is empty or has no clone left to draw.
+        ValueError: If ``seg_df`` is empty, has no clone left to draw, or
+            ``cn_palette`` is unknown.
     """
     if seg_df.empty:
         raise ValueError("seg_df is empty")
@@ -432,7 +437,7 @@ def plot_cnv_profile(
     props = get_clone_proportions(seg_df, clones) if show_prop else None
     pi_viol = get_pi_viol(seg_df) if show_pi_viol else None
 
-    state_style, _ = get_cn_cmap()
+    state_style, _ = get_cn_cmap(cn_palette)
     mirrored = _mirrored_bins(states) if show_mirror else None
     num_clones = len(clones)
     h = height / num_clones
@@ -508,6 +513,7 @@ def plot_cnv_profile(
         _plot_cnv_legend(
             ax_leg,
             has_mirror=mirrored is not None and bool(mirrored.any()),
+            cn_palette=cn_palette,
             **(legend_kwargs or {}),
         )
 
@@ -559,11 +565,13 @@ def _plot_cnv_legend(
     pair_w: float = 2.0,
     pair_h: float = 0.6,
     fontsize: int = FS_LEGEND,
+    cn_palette: str = "hatchet",
 ) -> None:
     """Draw the integer copy-number legend, grouped by total copy number.
 
     Mirrored pairs share a swatch since (1, 2) and (2, 1) share a color; the "/"
-    hatch distinguishes them. A trailing box covers total CN above 7.
+    hatch distinguishes them. A trailing box covers the totals above the
+    palette's largest, ">7" under "hatchet" and ">6" under "cnamaste".
 
     Each state is labelled ``(a,b)`` under its swatch. The key takes as many
     rows as it needs for the widest label to fit under one swatch at
@@ -578,9 +586,11 @@ def _plot_cnv_legend(
         pair_w: Width of each swatch.
         pair_h: Height of each swatch.
         fontsize: Swatch label font size.
+        cn_palette: (a, b) palette, one of
+            :data:`~cnplot.cnplot_colormap.CN_PALETTES`.
 
     """
-    state_style, tcn_states = get_cn_cmap()
+    state_style, tcn_states = get_cn_cmap(cn_palette)
     ax.axis("off")
 
     gap_groups = 0.75 * pair_w
@@ -594,7 +604,7 @@ def _plot_cnv_legend(
         groups.append(
             [(f"({a},{b})", state_style[(a, b)], None) for a, b in uniq_pairs]
         )
-    groups.append([(">7", state_style["default"], None)])
+    groups.append([(f">{max(tcn_states)}", state_style["default"], None)])
     label_pt = max(_text_width_pt(ax, lab, fontsize) for g in groups for lab, _, _ in g)
     lead = [0.0] * len(groups)
     if has_mirror:
diff --git a/tests/test_intcnp.py b/tests/test_intcnp.py
index 2604e2b..423cd08 100644
--- a/tests/test_intcnp.py
+++ b/tests/test_intcnp.py
@@ -1,9 +1,10 @@
 """Integer copy-number profile and its self-drawn legend."""
 
+import matplotlib.colors as mcolors
 import matplotlib.pyplot as plt
 import pandas as pd
 
-from cnplot import plot_cnv_profile
+from cnplot import get_cn_cmap, plot_cnv_profile
 from cnplot.cnplot_intcnp import get_clone_names, get_clone_states, has_mirror
 
 
@@ -94,3 +95,19 @@ def test_unknown_sample_raises(axis, sim):
         fig, ax = plt.subplots()
         plot_cnv_profile(ax, sim.seg_ucn, axis, sample_id="NOPE")
     plt.close("all")
+
+
+def test_cnamaste_palette_colors_profile_and_key(axis, sim):
+    style, _ = get_cn_cmap("cnamaste")
+    allowed = {mcolors.to_hex(c) for c in style.values()}
+    fig, (ax, lg) = plt.subplots(2, 1)
+    plot_cnv_profile(
+        ax, sim.seg_ucn, axis, sample_id="S1", ax_leg=lg, cn_palette="cnamaste"
+    )
+    faces = [p.get_facecolor() for p in ax.patches]
+    bins = {mcolors.to_hex(c) for c in faces if c[3] > 0}
+    assert bins and bins <= allowed
+    labels = [t.get_text() for t in lg.texts]
+    assert "(2,0)" in labels and ">6" in labels
+    assert ">7" not in labels and "(4,3)" not in labels
+    plt.close(fig)
diff --git a/tests/test_utils_colormap.py b/tests/test_utils_colormap.py
index 0eb6d47..d2ea0a3 100644
--- a/tests/test_utils_colormap.py
+++ b/tests/test_utils_colormap.py
@@ -3,15 +3,18 @@
 import matplotlib.colors as mcolors
 import matplotlib.pyplot as plt
 import numpy as np
+import pytest
 
 from cnplot import (
     BAF_LIM,
+    CN_PALETTES,
     MARKER_SIZE_LARGE,
     MARKER_SIZE_SMALL,
     get_baf_cmap,
     get_categorical_cmap,
     get_cn_cmap,
     get_log2rdr_cmap,
+    get_mixcn_cmap,
     get_multiclass_cmap,
     resolve_colors,
     resolve_marker_size,
@@ -26,6 +29,39 @@ def test_cn_palette_is_allele_symmetric():
     assert "default" in style
 
 
+def test_cn_palette_defaults_to_hatchet():
+    assert CN_PALETTES == ("hatchet", "cnamaste")
+    assert get_cn_cmap() == get_cn_cmap("hatchet")
+    style, tcn = get_cn_cmap()
+    assert style[(3, 1)] == "orange"
+    assert max(tcn) == 7
+
+
+def test_cnamaste_palette_covers_total_cn_to_6():
+    style, tcn = get_cn_cmap("cnamaste")
+    assert style[(1, 1)] == "#EEEEEE"
+    assert style[(3, 1)] == style[(1, 3)] == "#FEC44F"
+    assert style[(0, 0)] == "#000000"
+    assert (4, 3) not in style
+    assert max(tcn) == 6
+    assert all(len(tcn[t]) == t + 1 for t in range(1, 7))
+    assert style["default"] == get_cn_cmap()[0]["default"]
+
+
+def test_unknown_cn_palette_raises():
+    with pytest.raises(ValueError, match="cn_palette"):
+        get_cn_cmap("viridis")
+
+
+def test_mixcn_clonal_state_takes_the_chosen_palette():
+    cnps = ["1|1;2|0;2|0", "1|1;2|0;1|1"]
+    hatchet = get_mixcn_cmap(cnps)
+    cnamaste = get_mixcn_cmap(cnps, cn_palette="cnamaste")
+    assert hatchet[cnps[0]] == "dimgray"
+    assert cnamaste[cnps[0]] == "#22B5D2"
+    assert cnamaste[cnps[1]] not in get_cn_cmap("cnamaste")[0].values()
+
+
 def test_marker_size_switch():
     assert resolve_marker_size(100) == MARKER_SIZE_LARGE
     assert resolve_marker_size(4999) == MARKER_SIZE_LARGE