visualization of sobol indices

Author: matthiaskoenig

src/sbmlsim/sensitivity/analysis.py

@@ -249,16 +249,34 @@ def sensitivity_df(self, key="normalized") -> pd.DataFrame:
             index=self.sensitivity[key].coords["parameter"]
         )
 
+    def plot_sensitivity(
+        self,
+        key: str, cutoff=0.1,
+        cluster_rows: bool = True,
+        title: Optional[str] = None,
+        cmap: str = "seismic",
+        **kwargs
+    ) -> None:
+        df = self.sensitivity_df(key=key)
+        heatmap(
+            df=df,
+            parameter_labels={p.uid: p.name for p in self.parameters},
+            output_labels={q.uid: q.name for q in self.outputs},
+            cutoff=cutoff,
+            cluster_rows=cluster_rows,
+            title=title,
+            cmap=cmap,
+            **kwargs
+        )
+
 import os
 
 def run_simulation(
     params_tuple
 ):
     """Pass all required arguments as parameter tuple."""
     sensitivity_simulation, r, chunked_changes = params_tuple
-
     outputs = []
-
     for kc in track(range(len(chunked_changes)), description=f"Simulate samples PID={os.getpid()}"):
         changes = chunked_changes[kc]
         # console.print(f"PID={os.getpid()} | k={kc}")
@@ -271,16 +289,6 @@ def run_simulation(
     return outputs
 
 
-
-
-
-
-
-
-
-
-
-
 class LocalSensitivityAnalysis(SensitivityAnalysis):
     """Local sensitivity analysis based on local differences.
 
@@ -371,35 +379,7 @@ def calculate_sensitivity(self):
                 sensitivity_normalized[kp, ko] = sensitivity_raw[kp, ko] * p_ref/q_ref
 
 
-    def plot_sensitivity(self, cutoff=0.1, cluster_rows: bool = True, title: Optional[str] = None):
-        df = self.sensitivity_df(key="normalized")
-        self.plot_sensitivity_df(
-            df=df,
-            parameter_labels={p.uid: p.name for p in self.parameters},
-            output_labels={q.uid: q.name for q in self.outputs},
-            cutoff=cutoff,
-            cluster_rows=cluster_rows,
-            title=title
-        )
-
-    @staticmethod
-    def plot_sensitivity_df(
-        df: pd.DataFrame,
-        parameter_labels: dict[str, str],
-        output_labels: dict[str, str],
-        cutoff=0.1, cluster_rows: bool = True,
-        title: Optional[str] = None,
-    ):
-        console.print(df)
 
-        heatmap(
-            df,
-            parameter_labels=parameter_labels,
-            output_labels=output_labels,
-            cutoff=cutoff,
-            cluster_rows=False,
-            title=title,
-        )
 
 
 @dataclass
@@ -457,50 +437,41 @@ def create_samples(self, N: int=1024):
 
 
     def calculate_sensitivity(self):
-        # transfer results in libsa results format
+        """Calculate the sensitivity matrices."""
 
         Y = self.results.values
         self.ssa_problem.set_results(Y)
 
+        # num_parameters x num_outputs
+        sensitivity_keys = ["S1", "ST", "S1_conf", "ST_conf"]
+        for key in sensitivity_keys:
+            self.sensitivity[key] = xr.DataArray(
+                np.full((self.num_parameters, self.num_outputs), np.nan),
+                dims=["parameter", "output"],
+                coords={"parameter": self.parameter_ids,
+                        "output": self.output_ids},
+                name=key
+            )
+
         # Perform Analysis
         # Si is a Python dict-like with the keys "S1", "S2", "ST",
         # "S1_conf", "S2_conf", and "ST_conf".
         # The _conf keys store the corresponding confidence intervals,
         # typically with a confidence level of 95%.
 
         # Calculate Sobol indices for every output
-        Si_all = []
         for ko in range(self.num_outputs):
             Yo = Y[:, ko]
             Si = SALib.analyze.sobol.analyze(
                 self.ssa_problem, Yo,
                 calc_second_order=True,
                 print_to_console=True,
             )
-            Si_all.append(Si)
-
-            Si.plot()
-            from matplotlib import pyplot as plt
-            plt.show()
-
-        # Si = SALib.analyze.sobol.analyze(
-        #     self.ssa_problem, Y,
-        #     calc_second_order=True,
-        #     print_to_console=True,
-        # )
-
-        # Store the sensitivity matrices
+            console.print("S1")
+            console.print(Si["S1"])
+            for key in sensitivity_keys:
+                self.sensitivity[key][:, ko] = Si[key]
 
-        sensitivity_total = Si['ST']
-        sensitivity_first = Si['S1']
-        print(Si['S1'])
-        print(Si['ST'])
-
-
-
-        Si.plot()
-        from matplotlib import pyplot as plt
-        plt.show()
 
     def plot(self):
         Si.plot()

src/sbmlsim/sensitivity/plots.py

@@ -17,12 +17,15 @@ def heatmap(
     df: pd.DataFrame,
     parameter_labels: Optional[dict[str, str]] = None,
     output_labels: Optional[dict[str, str]] = None,
-    cutoff: float=0.01,
+    cutoff: float=0.1,
     annotate_values=True,
     cluster_rows: bool = True, # cluster parameters
     cluster_cols: bool = False, # cluster outputs
-    transpose: bool=False,
     title: Optional[str] = None,
+    cmap: str = "seismic",
+    vcenter: float = 0.0,
+    vmin: float = -2.0,
+    vmax: float = 2.0,
 ):
     """Creates heatmap of model sensitivity"""
 
@@ -47,7 +50,9 @@ def calculate_subset(df, cutoff=0.01) -> pd.DataFrame:
 
     if cutoff > 0:
         df_subset = calculate_subset(df, cutoff=cutoff)
-        df_subset_mask = calculate_mask(df_subset, cutoff)
+    else:
+        df_subset = df
+    df_subset_mask = calculate_mask(df_subset, cutoff)
 
     # outputs
     xticklabels = [qid for qid in df_subset.columns]
@@ -62,19 +67,18 @@ def calculate_subset(df, cutoff=0.01) -> pd.DataFrame:
 
     n_outputs = df_subset.shape[1]
     n_parameters = df_subset.shape[0]
-    figsize = (int(n_outputs/n_parameters*30), 15)
+    figsize = (int(n_outputs/n_parameters*15), 15)
 
-    colorbar_range = 2.0
 
     # plot heatmap
     cg = sns.clustermap(
         df_subset,
-        center=0,
-        vmin=-colorbar_range,
-        vmax=colorbar_range,
+        center=vcenter,
+        vmin=vmin,
+        vmax=vmax,
         xticklabels=xticklabels,
         yticklabels=yticklabels,
-        cmap="seismic",
+        cmap=cmap,
         # cbar_pos=(0.0, 0.0, 0.6, 0.05), #  (left, bottom, width, height),
         cbar_pos=(0.0, 0.4, 0.03, 0.2),  # (left, bottom, width, height),
         cbar_kws={
@@ -85,8 +89,8 @@ def calculate_subset(df, cutoff=0.01) -> pd.DataFrame:
         fmt="1.2f",
         annot_kws={"size": 11},
         mask=df_subset_mask,
-        col_cluster=False,
-        row_cluster=True,
+        col_cluster=cluster_cols,
+        row_cluster=cluster_rows,
         method="single",
         figsize=figsize,
     )
@@ -96,7 +100,7 @@ def calculate_subset(df, cutoff=0.01) -> pd.DataFrame:
         horizontalalignment="right",
         size=20,
     )
-    label_fontsize=10
+    label_fontsize=13
     plt.setp(cg.ax_heatmap.get_yticklabels(), size=label_fontsize)
     plt.setp(cg.ax_heatmap.get_xticklabels(), size=label_fontsize)
     cg.ax_cbar.tick_params(labelsize=label_fontsize)
@@ -106,50 +110,3 @@ def calculate_subset(df, cutoff=0.01) -> pd.DataFrame:
     if title:
         plt.suptitle(title)
 
-    # for label in cg.ax_heatmap.get_xticklabels():
-    #     label.set_bbox(dict(facecolor='tab:blue', edgecolor='black', alpha=0.8))
-    #
-    # for label in cg.ax_heatmap.get_yticklabels():
-    #     label.set_bbox(dict(facecolor='tab:orange', edgecolor='black', alpha=0.8))
-
-    # create custom legend containing yticklabels and their description
-    # handles = [t.get_text() for t in ax.ax_heatmap.get_yticklabels()]
-    # labels = [pnames[pid]["label"] for pid in handles]
-    #
-    # # FIXME: update after defining labels
-    # idx = [pnames[pid]["idx"] for pid in handles]
-    # # idx = [k for k, pid in enumerate(handles)]
-    #
-    # labels = [label for _, label in sorted(zip(idx, labels))]
-    # handles = [f"{handle}:" for _, handle in sorted(zip(idx, handles))]
-    # handles = [handle.replace("_", "\_") for handle in handles]
-
-    # mid = int(np.ceil(len(handles) / 2))
-    # legend1 = plt.legend(
-    #     handles[:mid],
-    #     labels[:mid],
-    #     handler_map={str: LegendTitle({"fontsize": 16})},
-    #     fontsize=16,
-    #     frameon=False,
-    #     bbox_to_anchor=(1.2, -0.6),
-    #     loc="upper left",
-    #     handlelength=14,
-    # )
-    # legend2 = plt.legend(
-    #     handles[mid:],
-    #     labels[mid:],
-    #     handler_map={str: LegendTitle({"fontsize": 16})},
-    #     fontsize=16,
-    #     frameon=False,
-    #     bbox_to_anchor=(13, -0.6),
-    #     loc="upper left",
-    #     handlelength=19,
-    # )
-    # plt.gca().add_artist(legend1)
-
-    # plt.savefig(
-    #     results_dir / "parameter.sensitivity_cluster.png", dpi=300, bbox_inches="tight"
-    # )
-    # plt.savefig(results_dir / "parameter.sensitivity_cluster.svg", bbox_inches="tight")
-
-    # plt.show()