diff --git a/python/interpret-core/interpret/glassbox/_ebm.py b/python/interpret-core/interpret/glassbox/_ebm.py index bb8284324..d058e9e27 100644 --- a/python/interpret-core/interpret/glassbox/_ebm.py +++ b/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 diff --git a/python/interpret-core/interpret/visual/plot.py b/python/interpret-core/interpret/visual/plot.py index f96d8bb94..33dd1e63a 100644 --- a/python/interpret-core/interpret/visual/plot.py +++ b/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] diff --git a/python/interpret-core/tests/visual/test_plot.py b/python/interpret-core/tests/visual/test_plot.py index dd30902ef..752cb2fbe 100644 --- a/python/interpret-core/tests/visual/test_plot.py +++ b/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