fix(visual): improve contrast of interaction heatmaps (#236)

python/interpret-core/interpret/glassbox/_ebm.py

@@ -2243,20 +2243,40 @@ def explain_global(self, name: str | None = None) -> EBMExplanation:
         # Obtain min/max for model scores
         lower_bound = np.inf
         upper_bound = -np.inf
+        main_lower_bound = np.inf
+        main_upper_bound = -np.inf
+        inter_lower_bound = np.inf
+        inter_upper_bound = -np.inf
         std_devs = (
             self.standard_deviations_
             if self.standard_deviations_ is not None
             else [None] * len(self.term_scores_)
         )
-        for scores, errors in zip(self.term_scores_, std_devs):
+        for scores, errors, feature_idxs in zip(
+            self.term_scores_, std_devs, self.term_features_
+        ):
             if errors is None:
-                lower_bound = min(lower_bound, float(np.min(scores)))
-                upper_bound = max(upper_bound, float(np.max(scores)))
+                term_lower_bound = float(np.min(scores))
+                term_upper_bound = float(np.max(scores))
             else:
-                lower_bound = min(lower_bound, float(np.min(scores - errors)))
-                upper_bound = max(upper_bound, float(np.max(scores + errors)))
+                term_lower_bound = float(np.min(scores - errors))
+                term_upper_bound = float(np.max(scores + errors))
+            lower_bound = min(lower_bound, term_lower_bound)
+            upper_bound = max(upper_bound, term_upper_bound)
+            if len(feature_idxs) == 1:
+                main_lower_bound = min(main_lower_bound, term_lower_bound)
+                main_upper_bound = max(main_upper_bound, term_upper_bound)
+            elif len(feature_idxs) == 2:
+                inter_lower_bound = min(inter_lower_bound, term_lower_bound)
+                inter_upper_bound = max(inter_upper_bound, term_upper_bound)
 
         bounds = (lower_bound, upper_bound)
+        main_bounds = (main_lower_bound, main_upper_bound)
+        if np.isinf(inter_lower_bound) or np.isinf(inter_upper_bound):
+            # No interaction terms, so fall back to the main effects bounds
+            inter_bounds = main_bounds
+        else:
+            inter_bounds = (inter_lower_bound, inter_upper_bound)
 
         mod_weights = remove_last(self.bin_weights_, self.bin_weights_)
         mod_term_scores = remove_last(self.term_scores_, self.bin_weights_)
@@ -2359,7 +2379,7 @@ def explain_global(self, name: str | None = None) -> EBMExplanation:
                     "type": "univariate",
                     "names": bin_labels,
                     "scores": scores,
-                    "scores_range": bounds,
+                    "scores_range": main_bounds,
                     "upper_bounds": upper_bounds,
                     "lower_bounds": lower_bounds,
                 }
@@ -2370,7 +2390,7 @@ def explain_global(self, name: str | None = None) -> EBMExplanation:
                     "type": "univariate",
                     "names": bin_labels,
                     "scores": model_graph,
-                    "scores_range": bounds,
+                    "scores_range": main_bounds,
                     "upper_bounds": None if errors is None else model_graph + errors,
                     "lower_bounds": None if errors is None else model_graph - errors,
                     "density": {
@@ -2454,7 +2474,8 @@ def explain_global(self, name: str | None = None) -> EBMExplanation:
                         "left_names": bin_labels_left,
                         "right_names": bin_labels_right,
                         "scores": model_graph,
-                        "scores_range": bounds,
+                        "scores_range": inter_bounds,
+                        "global_scores_range": bounds,
                     }
                     feature_list.append(feature_dict)
                     density_list.append({})
@@ -2464,7 +2485,8 @@ def explain_global(self, name: str | None = None) -> EBMExplanation:
                         "left_names": bin_labels_left,
                         "right_names": bin_labels_right,
                         "scores": model_graph,
-                        "scores_range": bounds,
+                        "scores_range": inter_bounds,
+                        "global_scores_range": bounds,
                     }
                     data_dicts.append(data_dict)
             else:  # pragma: no cover

python/interpret-core/interpret/visual/plot.py

@@ -703,8 +703,30 @@ def mli_plot_horizontal_bar(
 
 
 def plot_pairwise_heatmap(
-    data_dict, title="", xtitle="", ytitle="", transform_vals=True
+    data_dict,
+    title="",
+    xtitle="",
+    ytitle="",
+    transform_vals=True,
+    scale: str = "interaction",
 ):
+    """Plot a pairwise interaction heatmap.
+
+    Args:
+        data_dict: Dictionary with 'left_names', 'right_names' and 'scores'.
+        title: Plot title.
+        xtitle: X-axis title.
+        ytitle: Y-axis title.
+        transform_vals: Transpose the scores before plotting.
+        scale: Color scale range. 'interaction' uses the shared range of all
+            interaction terms ('scores_range'), 'term' uses the minimum and
+            maximum of this term's scores, and 'global' uses the shared range
+            of all terms ('global_scores_range', falling back to
+            'scores_range' when unavailable).
+
+    Returns:
+        A Plotly figure of the heatmap, or None if there are no scores.
+    """
     if data_dict.get("scores", None) is None:  # pragma: no cover
         return None
 
@@ -716,7 +738,17 @@ def plot_pairwise_heatmap(
         bin_vals = np.ascontiguousarray(np.transpose(bin_vals, (1, 0)))
 
     heatmap = go.Heatmap(z=bin_vals, x=bin_labels_left, y=bin_labels_right)
-    if data_dict.get("scores_range", None) is not None:
+    if scale == "term":
+        heatmap["zmin"] = float(np.min(bin_vals))
+        heatmap["zmax"] = float(np.max(bin_vals))
+    elif scale == "global":
+        global_scores_range = data_dict.get("global_scores_range", None)
+        if global_scores_range is None:
+            global_scores_range = data_dict.get("scores_range", None)
+        if global_scores_range is not None:
+            heatmap["zmin"] = global_scores_range[0]
+            heatmap["zmax"] = global_scores_range[1]
+    elif data_dict.get("scores_range", None) is not None:
         heatmap["zmin"] = data_dict["scores_range"][0]
         heatmap["zmax"] = data_dict["scores_range"][1]
 

python/interpret-core/tests/visual/test_plot.py

@@ -3,6 +3,7 @@
 
 from interpret.visual.plot import plot_line
 from interpret.visual.plot import plot_density
+from interpret.visual.plot import plot_pairwise_heatmap
 
 
 def test_plot_line_bounds_smoke():
@@ -30,3 +31,50 @@ def test_plot_density_large_numbers():
     # The x-axis tick text should show ranges using our new formatting
     assert "90T - 100T" in figure.layout.xaxis.ticktext[0]
     assert "100T - 1.00e+15" in figure.layout.xaxis.ticktext[1]
+
+
+def _heatmap_data_dict():
+    return {
+        "left_names": ["a", "b", "c"],
+        "right_names": ["x", "y"],
+        "scores": [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]],
+        "scores_range": (-2.0, 2.0),
+        "global_scores_range": (-10.0, 10.0),
+    }
+
+
+def test_plot_pairwise_heatmap_interaction_scale():
+    data_dict = _heatmap_data_dict()
+
+    figure = plot_pairwise_heatmap(data_dict)
+
+    assert figure.data[0].zmin == -2.0
+    assert figure.data[0].zmax == 2.0
+
+
+def test_plot_pairwise_heatmap_term_scale():
+    data_dict = _heatmap_data_dict()
+
+    figure = plot_pairwise_heatmap(data_dict, scale="term")
+
+    assert figure.data[0].zmin == 1.0
+    assert figure.data[0].zmax == 6.0
+
+
+def test_plot_pairwise_heatmap_global_scale():
+    data_dict = _heatmap_data_dict()
+
+    figure = plot_pairwise_heatmap(data_dict, scale="global")
+
+    assert figure.data[0].zmin == -10.0
+    assert figure.data[0].zmax == 10.0
+
+
+def test_plot_pairwise_heatmap_global_scale_falls_back_to_scores_range():
+    data_dict = _heatmap_data_dict()
+    del data_dict["global_scores_range"]
+
+    figure = plot_pairwise_heatmap(data_dict, scale="global")
+
+    assert figure.data[0].zmin == -2.0
+    assert figure.data[0].zmax == 2.0