local sensitivity analysis

Author: matthiaskoenig

src/sbmlsim/sensitivity/analysis.py

@@ -48,6 +48,16 @@
 from sbmlsim.sensitivity.outputs import SensitivityOutput
 import pandas as pd
 
+@dataclass
+class SensitivityOutput:
+    """Output for SensitivityAnalysis"""
+    uid: str
+    name: str
+    unit: Optional[str] = None
+
+    def __hash__(self):
+        return hash(self.uid)
+
 @dataclass
 class SensitivitySimulation:
     """Base class for sensitivity calculation.
@@ -60,10 +70,10 @@ class SensitivitySimulation:
     model_path: Path
     selections: list[str]
     rr: roadrunner.RoadRunner = None
-    outputs: list[str] = None
+    outputs: list[SensitivityOutput] = None
     changes_simulation: dict[str, float] = None
 
-    def __init__(self, model_path: Path, selections: list[str], changes_simulation: dict[str, float]):
+    def __init__(self, model_path: Path, selections: list[str], changes_simulation: dict[str, float], outputs: list[SensitivityOutput]):
         self.model_path = model_path
         self.selections = selections
         self.rr: roadrunner.RoadRunner = roadrunner.RoadRunner(str(model_path))
@@ -74,10 +84,15 @@ def __init__(self, model_path: Path, selections: list[str], changes_simulation:
 
         # store the simulation changes
         self.changes_simulation = changes_simulation
+        self.outputs: list[SensitivityOutput] = outputs
 
-        # get the outputs from the simulation
+        # validate the outputs from the simulation
         y = self.simulate(changes={})
-        self.outputs = list(y.keys())
+        outputs_dict = {q.uid for q in self.outputs}
+        for key in y:
+            if key not in outputs_dict:
+                raise ValueError(f"Key '{key}' missing in outputs dictionary: '{outputs_dict}")
+
 
 
     # def output_definitions(self) -> list[SensitivityOutput]:
@@ -118,17 +133,13 @@ def apply_changes(self, changes: dict[str, float], reset_all: bool=True) -> None
 
 
 
-@dataclass
+
 class SensitivityAnalysis:
     """Parent class for all sensitivity analysis.
 
     TODO: additional metadata for the outputs and the parameters; i.e. name, units, bounds, ....
     """
 
-    sensitivity_simulation: SensitivitySimulation
-    parameters: list[SensitivityParameter]
-    outputs: list[str]
-
     def __init__(self, sensitivity_simulation: SensitivitySimulation,
                  parameters: SensitivityParameter) -> None:
         """Create a sensitivity analysis for given parameter ids.
@@ -139,8 +150,11 @@ def __init__(self, sensitivity_simulation: SensitivitySimulation,
 
         # parameters to vary; shape: (num_parameters,)
         self.parameters: list[SensitivityParameter] = parameters
+        self.parameter_ids: list[str] = [p.uid for p in self.parameters]
+
         # outputs to calculate sensitivity on; shape: (num_outputs,)
-        self.outputs: list[output] = sensitivity_simulation.outputs
+        self.outputs: list[SensitivityOutput] = sensitivity_simulation.outputs
+        self.output_ids: list[str] = [q.uid for q in self.outputs]
 
         # parameter samples for sensitivity; shape: (num_samples x num_parameters)
         self.samples: Optional[xr.DataArray] = None
@@ -235,7 +249,7 @@ def create_samples(self) -> None:
         samples = xr.DataArray(
             np.full((num_samples, self.num_parameters), np.nan),
             dims=["sample", "parameter"],
-            coords={"sample": range(num_samples), "parameter": self.parameters},
+            coords={"sample": range(num_samples), "parameter": [p.uid for p in self.parameters]},
             name="samples"
         )
 
@@ -262,15 +276,15 @@ def calculate_sensitivity(self):
         self.sensitivity = xr.DataArray(
             np.full((self.num_parameters, self.num_outputs), np.nan),
             dims=["parameter", "output"],
-            coords={"parameter": [p.uid for p in self.parameters],
-                    "output": self.outputs},
+            coords={"parameter": self.parameter_ids,
+                    "output": self.output_ids},
             name="sensitivity"
         )
         self.sensitivity_normalized = xr.DataArray(
             np.full((self.num_parameters, self.num_outputs), np.nan),
             dims=["parameter", "output"],
-            coords={"parameter": [p.uid for p in self.parameters],
-                    "output": self.outputs},
+            coords={"parameter": self.parameter_ids,
+                    "output": self.output_ids},
             name="sensitivity"
         )
 
@@ -302,21 +316,33 @@ def sensitivity_df(self) -> pd.DataFrame:
 
     def plot_sensitivity(self):
         df = self.sensitivity_df
-        self.plot_sensitivity_df(df)
+        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},
+        )
 
     @staticmethod
-    def plot_sensitivity_df(df: pd.DataFrame, cutoff=0.1, cluster_rows: bool = True):
+    def plot_sensitivity_df(
+        df: pd.DataFrame,
+        parameter_labels: dict[str, str],
+        output_labels: dict[str, str],
+        cutoff=0.1, cluster_rows: bool = True
+    ):
         from sbmlsim.sensitivity.plots import heatmap
         console.print(df)
 
         # TODO: labels of parameters
         # TODO: labels of outputs
         # TODO: better position of colorbar
 
-        heatmap(df, cutoff=cutoff, cluster_rows=False)
-
-
-
+        heatmap(
+            df,
+            parameter_labels=parameter_labels,
+            output_labels=output_labels,
+            cutoff=cutoff,
+            cluster_rows=False
+        )
 
 @dataclass
 class SamplingSensitivityAnalysis(SensitivityAnalysis):

src/sbmlsim/sensitivity/parameters.py

@@ -24,6 +24,8 @@ def __hash__(self):
         return hash(self.uid)
 
 
+
+
 def parameters_for_sensitivity_analysis(
     sbml_path: Path,
     exclude_ids: Optional[set[str]] = None,
@@ -113,7 +115,7 @@ def parameter_from_sbase(sbase: libsbml.SBase) -> SensitivityParameter:
         else:
             parameters_filtered.append(sp)
 
-    console.print(f"Excluded parameters: {parameters_excluded}")
+    console.print(f"Excluded parameters: {[sp.uid for sp in parameters_excluded]}")
 
     return parameters_filtered
 
@@ -128,5 +130,4 @@ def parameter_from_sbase(sbase: libsbml.SBase) -> SensitivityParameter:
         }
     )
 
-    console.print("finished")
     console.print(parameters)

src/sbmlsim/sensitivity/plots.py

@@ -4,15 +4,19 @@
 https://stackoverflow.com/questions/59381273/heatmap-with-circles-indicating-size-of-population
 """
 import xarray as xr
-
+from typing import Optional
 
 from matplotlib import pyplot as plt
 import seaborn as sns
 import numpy as np
 import pandas as pd
+from sbmlutils.console import console
+
 
 def heatmap(
     df: pd.DataFrame,
+    parameter_labels: Optional[dict[str, str]] = None,
+    output_labels: Optional[dict[str, str]] = None,
     cutoff: float=0.01,
     annotate_values=True,
     cluster_rows: bool = True, # cluster parameters
@@ -44,12 +48,20 @@ def calculate_subset(df, cutoff=0.01) -> pd.DataFrame:
         df_subset = calculate_subset(df, cutoff=cutoff)
         df_subset_mask = calculate_mask(df_subset, cutoff)
 
+    # outputs
+    xticklabels = [qid for qid in df_subset.columns]
+    if output_labels:
+        console.print(output_labels)
+        xticklabels = [output_labels[qid] for qid in xticklabels]
+
+    # parameters
     yticklabels = [pid for pid in df_subset.index]
-    xticklabels = [pid for pid in df_subset.columns]
+    if parameter_labels:
+        yticklabels = [f"{pid}: {parameter_labels[pid]}" for pid in yticklabels]
 
     n_outputs = df_subset.shape[1]
     n_parameters = df_subset.shape[0]
-    figsize = (7, int(n_parameters / n_outputs * 7)/2)
+    figsize = (10, 15)
 
     colorbar_range = 2.0
 
@@ -66,7 +78,7 @@ def calculate_subset(df, cutoff=0.01) -> pd.DataFrame:
         cbar_pos=(0.0, 0.4, 0.03, 0.2),  # (left, bottom, width, height),
         cbar_kws={
             "orientation": "vertical",
-            "label": "sensitivity"
+            # "label": "sensitivity"
         },
         annot=annotate_values,
         fmt="1.2f",
@@ -83,8 +95,10 @@ def calculate_subset(df, cutoff=0.01) -> pd.DataFrame:
         horizontalalignment="right",
         size=20,
     )
-    plt.setp(ax.ax_heatmap.get_yticklabels(), size=20)
-    ax.ax_cbar.tick_params(labelsize=20)
+    label_fontsize=10
+    plt.setp(ax.ax_heatmap.get_yticklabels(), size=label_fontsize)
+    plt.setp(ax.ax_heatmap.get_xticklabels(), size=label_fontsize)
+    ax.ax_cbar.tick_params(labelsize=label_fontsize)
     ax.ax_row_dendrogram.set_visible(False)
     ax.ax_col_dendrogram.set_visible(False)