Fix plotly heatmap masking, cmap passthrough, and __all__ exports

- Plotly mask_insignificant now uses grey overlay trace instead of NaN
  replacement, preserving distinction between insignificant and missing
- Remove cmap name swapping — pass through unchanged to plotly
- Add 4 new plot functions to diff_diff.__all__
- Add regression tests for heatmap masking, colorscale, and exports

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

Author: igerber

diff_diff/__init__.py

@@ -259,6 +259,10 @@
     "plot_group_effects",
     "plot_sensitivity",
     "plot_honest_event_study",
+    "plot_synth_weights",
+    "plot_staircase",
+    "plot_dose_response",
+    "plot_group_time_heatmap",
     # Parallel trends testing
     "check_parallel_trends",
     "check_parallel_trends_robust",

diff_diff/visualization/_staggered.py

@@ -757,11 +757,9 @@ def _render_group_time_heatmap_plotly(
 
     go = _require_plotly()
 
-    # Map matplotlib cmap names to plotly
+    # Pass cmap name through to plotly unchanged — plotly supports the same
+    # diverging colorscale names as matplotlib (RdBu, RdBu_r, etc.)
     plotly_cmap = cmap
-    cmap_mapping = {"RdBu_r": "RdBu", "RdBu": "RdBu_r", "coolwarm": "RdBu"}
-    if cmap in cmap_mapping:
-        plotly_cmap = cmap_mapping[cmap]
 
     # Build text annotations
     text = None
@@ -777,10 +775,11 @@ def _render_group_time_heatmap_plotly(
                     row.append(f"{val:{fmt}}")
             text.append(row)
 
-    display = effect_matrix.copy()
+    # Build significance mask for overlay (do NOT replace with NaN — that
+    # conflates "insignificant" with "missing cell")
+    sig_mask = None
     if mask_insignificant and p_matrix is not None:
         sig_mask = p_matrix > alpha
-        display = np.where(sig_mask, np.nan, display)
 
     # Center the colorscale
     finite_vals = effect_matrix[np.isfinite(effect_matrix)]
@@ -791,9 +790,10 @@ def _render_group_time_heatmap_plotly(
     else:
         zmin, zmax = -1, 1
 
+    # Main heatmap — always shows all values (insignificant cells greyed via opacity)
     fig = go.Figure(
         data=go.Heatmap(
-            z=display,
+            z=effect_matrix,
             x=[str(t) for t in time_labels],
             y=[str(g) for g in group_labels],
             colorscale=plotly_cmap,
@@ -805,6 +805,20 @@ def _render_group_time_heatmap_plotly(
         )
     )
 
+    # Grey overlay for insignificant cells (preserves underlying value)
+    if sig_mask is not None and np.any(sig_mask):
+        grey_z = np.where(sig_mask, 1.0, np.nan)
+        fig.add_trace(
+            go.Heatmap(
+                z=grey_z,
+                x=[str(t) for t in time_labels],
+                y=[str(g) for g in group_labels],
+                colorscale=[[0, "rgba(255,255,255,0.6)"], [1, "rgba(255,255,255,0.6)"]],
+                showscale=False,
+                hoverinfo="skip",
+            )
+        )
+
     _plotly_default_layout(
         fig,
         title=title,

tests/test_visualization_plotly.py

@@ -238,3 +238,58 @@ def test_css_color_outside_original_subset(self):
             show=False,
         )
         assert isinstance(fig, go.Figure)
+
+
+class TestPlotlyHeatmapMasking:
+    """Regression: mask_insignificant must grey out, not NaN-ify cells."""
+
+    def test_mask_preserves_values(self, cs_results):
+        from diff_diff.visualization import plot_group_time_heatmap
+
+        fig = plot_group_time_heatmap(
+            cs_results, mask_insignificant=True, backend="plotly", show=False
+        )
+        assert isinstance(fig, go.Figure)
+        # Should have 2 traces: main heatmap + grey overlay
+        assert len(fig.data) == 2
+        # Main heatmap should NOT have NaN where cells were insignificant
+        import numpy as np
+
+        main_z = fig.data[0].z
+        assert np.any(np.isfinite(main_z))
+
+    def test_no_mask_single_trace(self, cs_results):
+        from diff_diff.visualization import plot_group_time_heatmap
+
+        fig = plot_group_time_heatmap(
+            cs_results, mask_insignificant=False, backend="plotly", show=False
+        )
+        assert isinstance(fig, go.Figure)
+        assert len(fig.data) == 1  # Only main heatmap, no overlay
+
+    def test_cmap_not_swapped(self, cs_results):
+        """RdBu_r should not be swapped — last color should be warm (red)."""
+        from diff_diff.visualization import plot_group_time_heatmap
+
+        fig = plot_group_time_heatmap(cs_results, cmap="RdBu_r", backend="plotly", show=False)
+        assert isinstance(fig, go.Figure)
+        # Plotly resolves named colorscales to tuples. RdBu_r ends with red.
+        cs = fig.data[0].colorscale
+        # Last entry should be a reddish color (high R value)
+        last_color = cs[-1][1]  # e.g. "rgb(103,0,31)"
+        assert "103" in last_color or "178" in last_color  # dark red end of RdBu_r
+
+
+class TestTopLevelExports:
+    """Regression: new plot functions must be in diff_diff.__all__."""
+
+    def test_new_plots_in_all(self):
+        import diff_diff
+
+        for name in [
+            "plot_synth_weights",
+            "plot_staircase",
+            "plot_dose_response",
+            "plot_group_time_heatmap",
+        ]:
+            assert name in diff_diff.__all__, f"{name} missing from diff_diff.__all__"