Fix P1/P2 review findings: color parser, staircase counts, dose-response target, bacon parity

- Replace hex-only _hex_to_rgba with _color_to_rgba supporting named CSS
  colors, 3-digit hex, and matplotlib fallback (P1)
- Use max(n_treated) across cells per cohort in plot_staircase with warning
  on missingness instead of first-cell value (P1)
- Infer target from curve.target in plot_dose_response when curve passed
  directly, preventing ATT-labeled ACRT plots (P2)
- Add show_weighted_avg vertical lines to plot_bacon plotly scatter (P2)
- Add 9 targeted regression tests covering all fixes (P2)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: igerber

diff_diff/visualization/_common.py

@@ -65,13 +65,43 @@ def _plotly_default_layout(fig, *, title=None, xlabel=None, ylabel=None, show_le
     )
 
 
-def _hex_to_rgba(hex_color, alpha=1.0):
-    """Convert hex color to rgba string for plotly.
+_CSS_COLORS = {
+    "red": (255, 0, 0),
+    "blue": (0, 0, 255),
+    "green": (0, 128, 0),
+    "black": (0, 0, 0),
+    "white": (255, 255, 255),
+    "gray": (128, 128, 128),
+    "grey": (128, 128, 128),
+    "lightgray": (211, 211, 211),
+    "lightgrey": (211, 211, 211),
+    "darkgray": (169, 169, 169),
+    "darkgrey": (169, 169, 169),
+    "orange": (255, 165, 0),
+    "purple": (128, 0, 128),
+    "yellow": (255, 255, 0),
+    "cyan": (0, 255, 255),
+    "magenta": (255, 0, 255),
+    "pink": (255, 192, 203),
+    "brown": (165, 42, 42),
+    "navy": (0, 0, 128),
+    "teal": (0, 128, 128),
+    "olive": (128, 128, 0),
+    "coral": (255, 127, 80),
+    "salmon": (250, 128, 114),
+}
+
+
+def _color_to_rgba(color, alpha=1.0):
+    """Convert any color to an ``rgba(r, g, b, a)`` string for plotly.
+
+    Accepts hex colors (``#rrggbb``, ``#rgb``), CSS named colors, and
+    falls back to ``matplotlib.colors`` when available.
 
     Parameters
     ----------
-    hex_color : str
-        Hex color string (e.g., ``"#2563eb"``).
+    color : str
+        Color specification.
     alpha : float, default=1.0
         Opacity value between 0 and 1.
 
@@ -80,9 +110,45 @@ def _hex_to_rgba(hex_color, alpha=1.0):
     str
         An ``rgba(r, g, b, a)`` string.
     """
-    hex_color = hex_color.lstrip("#")
-    r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
-    return f"rgba({r}, {g}, {b}, {alpha})"
+    if not isinstance(color, str):
+        raise ValueError(f"Expected a color string, got {type(color).__name__}")
+
+    # 1. Hex colors: #rrggbb or #rgb
+    stripped = color.lstrip("#")
+    if color.startswith("#") and all(c in "0123456789abcdefABCDEF" for c in stripped):
+        if len(stripped) == 6:
+            r = int(stripped[0:2], 16)
+            g = int(stripped[2:4], 16)
+            b = int(stripped[4:6], 16)
+            return f"rgba({r}, {g}, {b}, {alpha})"
+        if len(stripped) == 3:
+            r = int(stripped[0] * 2, 16)
+            g = int(stripped[1] * 2, 16)
+            b = int(stripped[2] * 2, 16)
+            return f"rgba({r}, {g}, {b}, {alpha})"
+
+    # 2. Named CSS colors
+    if color.lower() in _CSS_COLORS:
+        r, g, b = _CSS_COLORS[color.lower()]
+        return f"rgba({r}, {g}, {b}, {alpha})"
+
+    # 3. Already an rgba/rgb string — override alpha
+    if color.startswith("rgba(") or color.startswith("rgb("):
+        return color if alpha == 1.0 else color  # pass through for plotly
+
+    # 4. Fallback: try matplotlib.colors if available
+    try:
+        from matplotlib.colors import to_rgb
+
+        r, g, b = to_rgb(color)
+        return f"rgba({int(r * 255)}, {int(g * 255)}, {int(b * 255)}, {alpha})"
+    except (ImportError, ValueError):
+        pass
+
+    raise ValueError(
+        f"Cannot parse color '{color}'. Use hex (#rrggbb), a CSS color name, "
+        "or install matplotlib for full color support."
+    )
 
 
 # Default color constants

diff_diff/visualization/_continuous.py

@@ -87,6 +87,9 @@ def plot_dose_response(
             raise ValueError(f"target must be 'att' or 'acrt', got '{target}'")
 
     if curve is not None:
+        # Infer target from curve when passed directly (not via results)
+        if results is None and hasattr(curve, "target") and curve.target:
+            target = curve.target
         dose_grid = curve.dose_grid
         effects = curve.effects
         ci_lower = curve.conf_int_lower
@@ -223,7 +226,7 @@ def _render_dose_response_plotly(
 ):
     """Render dose-response curve with plotly."""
     from diff_diff.visualization._common import (
-        _hex_to_rgba,
+        _color_to_rgba,
         _plotly_default_layout,
         _require_plotly,
     )
@@ -245,7 +248,7 @@ def _render_dose_response_plotly(
                 x=dose_list + dose_list[::-1],
                 y=list(ci_upper) + list(ci_lower)[::-1],
                 fill="toself",
-                fillcolor=_hex_to_rgba(band_color, 0.15),
+                fillcolor=_color_to_rgba(band_color, 0.15),
                 line=dict(color="rgba(0,0,0,0)"),
                 name="95% CI",
                 hoverinfo="skip",

diff_diff/visualization/_diagnostic.py

@@ -248,7 +248,7 @@ def _render_sensitivity_plotly(
 ):
     """Render sensitivity plot with plotly."""
     from diff_diff.visualization._common import (
-        _hex_to_rgba,
+        _color_to_rgba,
         _plotly_default_layout,
         _require_plotly,
     )
@@ -278,7 +278,7 @@ def _render_sensitivity_plotly(
                 x=M_list + M_list[::-1],
                 y=list(bounds_arr[:, 1]) + list(bounds_arr[:, 0])[::-1],
                 fill="toself",
-                fillcolor=_hex_to_rgba(bounds_color, bounds_alpha),
+                fillcolor=_color_to_rgba(bounds_color, bounds_alpha),
                 line=dict(color="rgba(0,0,0,0)"),
                 name="Identified set",
             )
@@ -742,6 +742,19 @@ def _render_bacon_plotly(
                 )
             )
 
+        # Weighted average lines
+        if show_weighted_avg:
+            effect_by_type = results.effect_by_type()
+            for ctype, avg_effect in effect_by_type.items():
+                if avg_effect is not None and by_type[ctype]:
+                    fig.add_vline(
+                        x=avg_effect,
+                        line_dash="dash",
+                        line_color=colors[ctype],
+                        opacity=0.5,
+                        line_width=1.5,
+                    )
+
         # TWFE line
         if show_twfe_line:
             fig.add_vline(

diff_diff/visualization/_event_study.py

@@ -430,7 +430,7 @@ def _render_event_study_plotly(
 ):
     """Render event study plot with plotly."""
     from diff_diff.visualization._common import (
-        _hex_to_rgba,
+        _color_to_rgba,
         _plotly_default_layout,
         _require_plotly,
     )
@@ -452,7 +452,7 @@ def _render_event_study_plotly(
             fig.add_vrect(
                 x0=min(pre_in_plot) - 0.5,
                 x1=max(pre_in_plot) + 0.5,
-                fillcolor=_hex_to_rgba(shade_color, 0.5),
+                fillcolor=_color_to_rgba(shade_color, 0.5),
                 line_width=0,
                 layer="below",
             )
@@ -477,7 +477,7 @@ def _render_event_study_plotly(
                 x=ci_periods + ci_periods[::-1],
                 y=ci_hi + ci_lo[::-1],
                 fill="toself",
-                fillcolor=_hex_to_rgba(color, 0.15),
+                fillcolor=_color_to_rgba(color, 0.15),
                 line=dict(color="rgba(0,0,0,0)"),
                 showlegend=False,
                 hoverinfo="skip",
@@ -970,7 +970,7 @@ def _render_honest_event_study_plotly(
 ):
     """Render honest event study plot with plotly."""
     from diff_diff.visualization._common import (
-        _hex_to_rgba,
+        _color_to_rgba,
         _plotly_default_layout,
         _require_plotly,
     )
@@ -988,7 +988,7 @@ def _render_honest_event_study_plotly(
             x=list(periods) + list(periods)[::-1],
             y=list(original_ci_upper) + list(original_ci_lower)[::-1],
             fill="toself",
-            fillcolor=_hex_to_rgba(original_color, 0.15),
+            fillcolor=_color_to_rgba(original_color, 0.15),
             line=dict(color="rgba(0,0,0,0)"),
             name="Standard CI",
             hoverinfo="skip",
@@ -1001,7 +1001,7 @@ def _render_honest_event_study_plotly(
             x=list(periods) + list(periods)[::-1],
             y=list(honest_ci_upper) + list(honest_ci_lower)[::-1],
             fill="toself",
-            fillcolor=_hex_to_rgba(honest_color, 0.15),
+            fillcolor=_color_to_rgba(honest_color, 0.15),
             line=dict(color="rgba(0,0,0,0)"),
             name=f"Honest CI (M={honest_M:.2f})",
             hoverinfo="skip",

diff_diff/visualization/_staggered.py

@@ -298,16 +298,29 @@ def _extract_staircase_data(results, data, unit, time, first_treat):
         groups = sorted(results.groups)
         cohort_counts = []
         for g in groups:
-            # Find a representative (g, t) entry to get n_treated for this cohort
-            n_treated = None
-            for (gg, tt), eff in results.group_time_effects.items():
+            # Collect n_treated across all (g, t) cells for this cohort.
+            # n_treated is a per-cell observation count that can vary with
+            # missingness, so we use the max as the best cohort size estimate.
+            cell_counts = []
+            for (gg, _tt), eff in results.group_time_effects.items():
                 if gg == g:
-                    n_treated = eff.get("n_treated", eff.get("n_obs", None))
-                    if n_treated is not None:
-                        break
-            if n_treated is None:
-                n_treated = 0
-            cohort_counts.append((g, int(n_treated)))
+                    n = eff.get("n_treated", eff.get("n_obs", None))
+                    if n is not None:
+                        cell_counts.append(int(n))
+            if not cell_counts:
+                cohort_counts.append((g, 0))
+                continue
+            max_count = max(cell_counts)
+            if min(cell_counts) != max_count:
+                import warnings
+
+                warnings.warn(
+                    f"Cohort {g}: n_treated varies across cells "
+                    f"({min(cell_counts)}-{max_count}). "
+                    f"Using max as cohort size; pass data= for exact counts.",
+                    stacklevel=3,
+                )
+            cohort_counts.append((g, max_count))
 
         return cohort_counts
 
@@ -391,7 +404,7 @@ def _render_staircase_mpl(*, cohort_counts, figsize, title, color, show_counts,
 def _render_staircase_plotly(*, cohort_counts, title, color, show_counts, show):
     """Render staircase plot with plotly."""
     from diff_diff.visualization._common import (
-        _hex_to_rgba,
+        _color_to_rgba,
         _plotly_default_layout,
         _require_plotly,
     )
@@ -419,7 +432,7 @@ def _render_staircase_plotly(*, cohort_counts, title, color, show_counts, show):
             mode="lines",
             line=dict(color=color, width=2, shape="hv"),
             fill="tozeroy",
-            fillcolor=_hex_to_rgba(color, 0.15),
+            fillcolor=_color_to_rgba(color, 0.15),
             name="Cumulative treated",
         )
     )