⚡️ Refactor SobolSanalysis to enforce calc_second_order consistency and improve error handling in sensitivity analysis tests

Author: BaptisteDE

corrai/sensitivity.py

@@ -176,7 +176,7 @@ def analyze(
                 prefix=method,
             )
         else:
-            if method or agg_method_kwarg:
+            if agg_method_kwarg is not None or (method is not None and method != "mean"):
                 warnings.warn(
                     "'method' or 'agg_method_kwarg' was provided but Model is static."
                     " Arguments will be ignored"
@@ -406,14 +406,37 @@ class SobolSanalysis(Sanalysis):
     sensitivity analysis following the Sobol method. Sampling of the parameter
     space is performed using the Saltelli scheme, which ensures efficient
     estimation of first-order, second-order, and total-order Sobol indices.
+
+    Parameters
+    ----------
+    parameters : list of Parameter
+        Parameters that define the sampling space.
+    model : Model
+        Model instance to be simulated.
+    simulation_options : dict, optional
+        Options passed to the model simulation.
+    calc_second_order : bool, default=True
+        Whether to compute second-order (interaction) Sobol indices.
+        This value is fixed at construction and used consistently for
+        both sampling and analysis. Setting it to False reduces the
+        required sample size from N*(2D+2) to N*(D+2).
+
+    Notes
+    -----
+    Sobol analysis requires a single call to :meth:`add_sample`. Calling it
+    more than once would concatenate independent Saltelli matrices, which
+    invalidates the variance decomposition. A ``ValueError`` is raised on
+    any subsequent call.
     """
 
     def __init__(
         self,
         parameters: list[Parameter],
         model: Model,
         simulation_options: dict = None,
+        calc_second_order: bool = True,
     ):
+        self._calc_second_order = calc_second_order
         super().__init__(parameters, model, simulation_options)
 
     def _set_sampler(
@@ -430,15 +453,19 @@ def add_sample(
         N: int,
         simulate: bool = True,
         n_cpu: int = 1,
-        *,
-        calc_second_order: bool = True,
         **sample_kwargs,
     ):
+        if len(self.sampler.sample) > 0:
+            raise ValueError(
+                "SobolSanalysis does not support incremental sampling. "
+                f"The sample already contains {len(self.sampler.sample)} rows. "
+                "Create a new SobolSanalysis instance to generate a new sample."
+            )
         super().add_sample(
             N=N,
             simulate=simulate,
             n_cpu=n_cpu,
-            calc_second_order=calc_second_order,
+            calc_second_order=self._calc_second_order,
             **sample_kwargs,
         )
 
@@ -449,7 +476,6 @@ def analyze(
         agg_method_kwarg: dict = None,
         reference_time_series: pd.Series = None,
         freq: str | pd.Timedelta | dt.timedelta = None,
-        calc_second_order: bool = True,
         **analyse_kwargs,
     ):
         return super().analyze(
@@ -458,7 +484,7 @@ def analyze(
             agg_method_kwarg=agg_method_kwarg,
             reference_time_series=reference_time_series,
             freq=freq,
-            calc_second_order=calc_second_order,
+            calc_second_order=self._calc_second_order,
             **analyse_kwargs,
         )
 
@@ -468,7 +494,6 @@ def plot_bar(
         sensitivity_metric: str = "ST",
         method: str = "mean",
         reference_time_series: pd.Series = None,
-        calc_second_order: bool = True,
         unit: str = "",
         agg_method_kwarg: dict = None,
         title: str = None,
@@ -485,7 +510,6 @@ def plot_bar(
             agg_method_kwarg=agg_method_kwarg,
             title=title,
             plot_kwargs=plot_kwargs,
-            calc_second_order=calc_second_order,
             **analyse_kwargs,
         )
 
@@ -498,7 +522,6 @@ def plot_dynamic_metric(
         reference_time_series: pd.Series = None,
         unit: str = "",
         agg_method_kwarg: dict = None,
-        calc_second_order: bool = True,
         title: str = None,
         plot_kwargs: dict = None,
     ):
@@ -511,7 +534,6 @@ def plot_dynamic_metric(
             unit=unit,
             agg_method_kwarg=agg_method_kwarg,
             reference_time_series=reference_time_series,
-            calc_second_order=calc_second_order,
             stacked=True,
             title=title,
             plot_kwargs=plot_kwargs,
@@ -528,6 +550,11 @@ def plot_s2_matrix(
         plot_kwargs: dict = None,
         **analyse_kwargs,
     ):
+        if not self._calc_second_order:
+            raise ValueError(
+                "plot_s2_matrix() requires second-order indices. "
+                "Set calc_second_order=True when creating SobolSanalysis."
+            )
         return super().salib_plot_matrix(
             indicator=indicator,
             sensitivity_method_name="Sobol",

tests/test_sensitivity.py

@@ -1,5 +1,8 @@
+import warnings
+
 import numpy as np
 import pandas as pd
+import pytest
 
 from corrai.base.parameter import Parameter
 from corrai.base.model import IshigamiDynamic, Ishigami
@@ -29,55 +32,82 @@
 
 
 class TestSensitivity:
-    def test_sanalysis_sobol_with_sobol_sampler(self):
-        # sobol_analysis = SobolSanalysis(
-        #     parameters=PARAMETER_LIST,
-        #     model=IshigamiDynamic(),
-        #     simulation_options=SIMULATION_OPTIONS,
-        # )
-        #
-        # sobol_analysis.add_sample(N=1000, n_cpu=1, calc_second_order=True, seed=42)
-        # res = sobol_analysis.analyze("res", calc_second_order=True, seed=42)
-        #
-        # np.testing.assert_almost_equal(
-        #     res["mean_res"]["S1"],
-        #     np.array([0.33080399, 0.44206835, 0.00946747]),
-        # )
-        #
-        # res = sobol_analysis.analyze("res", freq="h", calc_second_order=True, seed=42)
-        # assert res.index.tolist() == [
-        #     pd.Timestamp("2009-01-01 00:00:00"),
-        #     pd.Timestamp("2009-01-01 01:00:00"),
-        #     pd.Timestamp("2009-01-01 02:00:00"),
-        #     pd.Timestamp("2009-01-01 03:00:00"),
-        #     pd.Timestamp("2009-01-01 04:00:00"),
-        #     pd.Timestamp("2009-01-01 05:00:00"),
-        # ]
-        #
-        # sobol_analysis.plot_sample_hist(
-        #     "res", bins=10, reference_value=10, reference_label="ref"
-        # )
-        #
-        # np.testing.assert_almost_equal(
-        #     res["2009-01-01 00:00:00"]["S1"],
-        #     np.array([0.33080399, 0.44206835, 0.00946747]),
-        #     decimal=3,
-        # )
-        #
+    def test_sanalysis_sobol_static(self):
         sobol_analysis = SobolSanalysis(
             parameters=PARAMETER_LIST,
             model=Ishigami(),
+            calc_second_order=True,
+        )
+
+        sobol_analysis.add_sample(N=1000, n_cpu=1, seed=42)
+        res = sobol_analysis.analyze("res", seed=42)
+
+        np.testing.assert_almost_equal(
+            res["mean_res"]["S1"],
+            np.array([0.33080399, 0.44206835, 0.00946747]),
+        )
+
+    def test_sanalysis_sobol_dynamic(self):
+        sobol_analysis = SobolSanalysis(
+            parameters=PARAMETER_LIST,
+            model=IshigamiDynamic(),
+            simulation_options=SIMULATION_OPTIONS,
+            calc_second_order=True,
         )
 
-        sobol_analysis.add_sample(N=1000, n_cpu=1, calc_second_order=True, seed=42)
-        res = sobol_analysis.analyze("res", calc_second_order=True, seed=42)
+        sobol_analysis.add_sample(N=1000, n_cpu=1, seed=42)
+        res = sobol_analysis.analyze("res", seed=42)
 
         np.testing.assert_almost_equal(
             res["mean_res"]["S1"],
             np.array([0.33080399, 0.44206835, 0.00946747]),
         )
 
-        assert True
+        res_freq = sobol_analysis.analyze("res", freq="h", seed=42)
+        assert res_freq.index.tolist() == [
+            pd.Timestamp("2009-01-01 00:00:00"),
+            pd.Timestamp("2009-01-01 01:00:00"),
+            pd.Timestamp("2009-01-01 02:00:00"),
+            pd.Timestamp("2009-01-01 03:00:00"),
+            pd.Timestamp("2009-01-01 04:00:00"),
+            pd.Timestamp("2009-01-01 05:00:00"),
+        ]
+        np.testing.assert_almost_equal(
+            res_freq["2009-01-01 00:00:00"]["S1"],
+            np.array([0.33080399, 0.44206835, 0.00946747]),
+            decimal=3,
+        )
+
+    def test_sobol_incremental_sampling_raises(self):
+        sobol_analysis = SobolSanalysis(
+            parameters=PARAMETER_LIST,
+            model=Ishigami(),
+        )
+        sobol_analysis.add_sample(N=64, n_cpu=1, seed=42)
+        with pytest.raises(ValueError, match="does not support incremental sampling"):
+            sobol_analysis.add_sample(N=64, n_cpu=1, seed=0)
+
+    def test_sobol_plot_s2_matrix_raises_when_no_second_order(self):
+        sobol_analysis = SobolSanalysis(
+            parameters=PARAMETER_LIST,
+            model=Ishigami(),
+            calc_second_order=False,
+        )
+        sobol_analysis.add_sample(N=64, n_cpu=1, seed=42)
+        with pytest.raises(ValueError, match="requires second-order indices"):
+            sobol_analysis.plot_s2_matrix()
+
+    def test_sobol_static_model_no_spurious_warning(self):
+        sobol_analysis = SobolSanalysis(
+            parameters=PARAMETER_LIST,
+            model=Ishigami(),
+        )
+        sobol_analysis.add_sample(N=64, n_cpu=1, seed=42)
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            sobol_analysis.analyze("res", seed=42)
+        our_warnings = [w for w in caught if "sensitivity" in str(w.filename)]
+        assert len(our_warnings) == 0, f"Unexpected warnings from sensitivity.py: {our_warnings}"
 
     def test_sanalysis_morris(self):
         morris_analysis = MorrisSanalysis(
@@ -202,8 +232,9 @@ def test_sobol_s2_matrix(self):
             parameters=PARAMETER_LIST,
             model=IshigamiDynamic(),
             simulation_options=SIMULATION_OPTIONS,
+            calc_second_order=True,
         )
-        sobol_analysis.add_sample(N=2**2, n_cpu=1, calc_second_order=True)
+        sobol_analysis.add_sample(N=2**2, n_cpu=1)
         fig_matrix = sobol_analysis.plot_s2_matrix()
         assert fig_matrix["layout"]["title"]["text"] == (
             "Sobol mean res " "- 2nd order interactions"
@@ -321,7 +352,7 @@ def test_sobol_plot_bar_plot_kwargs(self):
             parameters=PARAMETER_LIST,
             model=Ishigami(),
         )
-        sobol_analysis.add_sample(N=2**4, n_cpu=1, calc_second_order=True, seed=42)
+        sobol_analysis.add_sample(N=2**4, n_cpu=1, seed=42)
         fig = sobol_analysis.plot_bar(
             plot_kwargs={
                 "title": "My Custom Title",