adding heatmap plot for sensitivity analysis #632

Author: crangelsmith

autoemulate/experimental/sensitivity_analysis.py

@@ -341,6 +341,30 @@ def top_n_sobol_params(
             .nlargest(top_n)
             .index.tolist()
         )
+    def plot_sa_heatmap(
+        self,
+        results: pd.DataFrame,
+        index: str = "ST",
+        top_n: int | None = None,
+        cmap: str = "coolwarm",
+        normalize: bool = True,
+        figsize: tuple | None = None,
+    ):
+        """
+        Plot a normalized Sobol sensitivity analysis heatmap.
+
+        Parameters:
+            si_df (pd.DataFrame): Sensitivity index dataframe with columns ['index', 'parameter', 'output', 'value'].
+            index (str): The type of sensitivity index to plot (e.g., 'ST').
+            top_n (int, optional): Number of top parameters to include. Defaults to all.
+            cmap (str, optional): Matplotlib colormap. Defaults to 'coolwarm'.
+            normalize (bool, optional): Whether to normalize values to [0, 1]. Defaults to True.
+            figsize (tuple, optional): Figure size as (width, height) in inches. If None,
+        """
+        # Determine which parameters to include
+        parameter_list = self.top_n_sobol_params(results, top_n=len(results['parameter'].unique()) if top_n is None else top_n)
+
+        return _plot_sa_heatmap(results, index, parameter_list, cmap, normalize, fig_size=figsize)
 
 
 def _sobol_results_to_df(results: dict[str, ResultDict]) -> pd.DataFrame:
@@ -754,3 +778,72 @@ def _create_morris_plot(
     ax.set_ylabel("μ* (Modified Mean)")
     ax.set_title(f"Output: {output_name}")
     ax.grid(True, alpha=0.3)
+
+
+
+def _plot_sa_heatmap(si_df, index, parameters, cmap='coolwarm', normalize=True, fig_size=None):
+    """
+    Plot a sensitivity analysis heatmap for a given index.
+
+    Parameters:
+        si_df (pd.DataFrame): Sensitivity index dataframe with columns ['index', 'parameter', 'output', 'value'].
+        index (str): The type of sensitivity index to plot (e.g., 'ST').
+        parameters (list): List of  parameters to include in the grid.
+        cmap (str, optional): Matplotlib colormap. Defaults to 'coolwarm'.
+        normalize (bool, optional): Whether to normalize values to [0, 1]. Defaults to True.
+        fig_size (tuple, optional): Figure size as (width, height) in inches. If None,
+    """
+
+    # Filter the dataframe for the specified index
+    df = si_df[si_df['index'] == index]
+
+    # Pivot the dataframe to get a matrix: rows = outputs, cols = parameters
+    heatmap_df = (
+        df[df['parameter'].isin(parameters)]
+        .pivot_table(index='output', columns='parameter', values='value', fill_value=np.nan)
+        .reindex(columns=parameters)  # Ensure column order
+    )
+
+    # Normalize if requested
+    if normalize:
+        min_value = heatmap_df.min().min()
+        max_value = heatmap_df.max().max()
+        value_range = max_value - min_value if max_value != min_value else 1
+        heatmap_df = (heatmap_df - min_value) / value_range
+
+    # Convert to NumPy array
+    data_np = heatmap_df.to_numpy()
+
+    # layout - add space for legend
+    n_rows, n_cols = _calculate_layout(data_np.shape[1], data_np.shape[0])
+    fig_size = fig_size or (4.5 * n_cols, 4.5 * n_rows + 2)  # Extra width for legend
+
+    # Plotting
+    fig, ax = plt.subplots(figsize=fig_size)
+    cax = ax.imshow(data_np, cmap=cmap, aspect='auto')
+
+    # Colorbar
+    cbar = fig.colorbar(cax, ax=ax)
+    cbar_label = 'Normalized Sensitivity' if normalize else 'Sensitivity'
+    cbar.set_label(cbar_label, rotation=270, labelpad=15)
+
+    # Labels and ticks
+    ax.set_title(f"{index} Sensitivity Analysis Heatmap", fontsize=14, pad=12)
+    ax.set_xlabel("Parameters", fontsize=12)
+    ax.set_ylabel("Outputs", fontsize=12)
+
+    ax.set_xticks(np.arange(len(parameters)))
+    ax.set_xticklabels(parameters, rotation=45, ha='right')
+    ax.set_yticks(np.arange(len(heatmap_df.index)))
+    ax.set_yticklabels(heatmap_df.index)
+
+    # Gridlines
+    ax.set_xticks(np.arange(-0.5, len(parameters), 1), minor=True)
+    ax.set_yticks(np.arange(-0.5, len(heatmap_df.index), 1), minor=True)
+    ax.grid(which='minor', color='w', linestyle='-', linewidth=2)
+    ax.tick_params(which='minor', bottom=False, left=False)
+
+    plt.tight_layout()
+
+    return display_figure(fig)
+