Update docs, change FOE/SOE to magic_binning, refactor print_indices (#61)

Author: justushelo

docs/tutorials.rst

@@ -5,3 +5,6 @@ Executable tutorials
    :maxdepth: 1
 
    notebooks/structural_reliability
+   Carbon Footprint <notebooks/3_Carbon_footprint>
+   Investment Model <notebooks/3_Investment_model>
+   Steel Structures <notebooks/3_Steel_structures>

src/simdec/sensitivity_indices.py

@@ -1,12 +1,53 @@
 from dataclasses import dataclass
+import warnings
 
 import numpy as np
 import pandas as pd
 from scipy import stats
 
-
 __all__ = ["sensitivity_indices"]
 
+try:
+    from IPython.display import display
+
+    HAS_IPYTHON = True
+except ImportError:
+    HAS_IPYTHON = False
+
+
+def _quantile_edges(x: np.ndarray, n_bins: int) -> np.ndarray:
+    """Bin edges holding approximately the same number of points.
+
+    Bins are defined by value, so identical values always fall in the same
+    bin. Duplicated edges are dropped: a factor with many ties gets fewer
+    than ``n_bins`` bins, and a constant factor gets a single one.
+
+    Discrete variables (few unique values relative to n_bins) get one bin
+    per unique value instead of quantile-based edges, to avoid collapsing
+    minority categories into a majority bin.
+    """
+    unique_vals = np.unique(x[~np.isnan(x)])
+    if len(unique_vals) <= n_bins:
+        # midpoints between consecutive unique values, so each value gets
+        # its own bin
+        if unique_vals.size == 1:
+            return np.array([unique_vals[0], np.nextafter(unique_vals[0], np.inf)])
+        midpoints = (unique_vals[:-1] + unique_vals[1:]) / 2
+        return np.concatenate([[unique_vals[0]], midpoints, [unique_vals[-1]]])
+
+    edges = np.unique(np.nanquantile(x, np.linspace(0, 1, n_bins + 1)))
+    if edges.size < 2:
+        edges = np.array([edges[0], np.nextafter(edges[0], np.inf)])
+    return edges
+
+
+def _conditional_var(sample: np.ndarray, y: np.ndarray, edges: list) -> float:
+    """Var(E[Y | bins]), each bin weighted by its number of points."""
+    mean, *_ = stats.binned_statistic_dd(sample, y, statistic="mean", bins=edges)
+    count, *_ = stats.binned_statistic_dd(sample, y, statistic="count", bins=edges)
+    valid = count > 0
+    return _weighted_var(mean[valid], weights=count[valid])
+
 
 def number_of_bins(n_runs: int, n_factors: int) -> tuple[int, int]:
     """Optimal number of bins for first & second-order sensitivity_indices indices.
@@ -131,60 +172,33 @@ def sensitivity_indices(
     foe = np.empty(n_factors)
     soe = np.zeros((n_factors, n_factors))
 
+    edges_foe = [
+        _quantile_edges(inputs[:, k], int(n_bins_foe)) for k in range(n_factors)
+    ]
+    edges_soe = [
+        _quantile_edges(inputs[:, k], int(n_bins_soe)) for k in range(n_factors)
+    ]
+
+    # Marginal Var(E[Y|Xk]) on the SOE binning, identical for every pair
+    var_marginal = np.array(
+        [
+            _conditional_var(inputs[:, [k]], output, [edges_soe[k]])
+            for k in range(n_factors)
+        ]
+    )
+
     for i in range(n_factors):
         # 1. First-order effects (FOE)
-        xi = inputs[:, i]
-
-        bin_avg, _, binnumber = stats.binned_statistic(
-            x=xi, values=output, bins=n_bins_foe, statistic="mean"
-        )
-
-        # Filter empty bins and get weights (counts)
-        mask_foe = ~np.isnan(bin_avg)
-        mean_i_foe = bin_avg[mask_foe]
-        # binnumber starts at 1; 0 is for values outside range
-        bin_counts_foe = np.unique(binnumber[binnumber > 0], return_counts=True)[1]
+        foe[i] = _conditional_var(inputs[:, [i]], output, [edges_foe[i]]) / var_y
 
-        foe[i] = _weighted_var(mean_i_foe, weights=bin_counts_foe) / var_y
-
-        # 2. Second-order effects (SOE)
         for j in range(n_factors):
             if j <= i:
                 continue
 
-            xj = inputs[:, j]
-
-            # 2D Binned Statistic for Var(E[Y|Xi, Xj])
-            bin_avg_ij, x_edges, y_edges, binnumber_ij = stats.binned_statistic_2d(
-                x=xi, y=xj, values=output, bins=n_bins_soe, expand_binnumbers=False
+            var_ij = _conditional_var(
+                inputs[:, [i, j]], output, [edges_soe[i], edges_soe[j]]
             )
-
-            mask_ij = ~np.isnan(bin_avg_ij)
-            mean_ij = bin_avg_ij[mask_ij]
-            counts_ij = np.unique(binnumber_ij[binnumber_ij > 0], return_counts=True)[1]
-            var_ij = _weighted_var(mean_ij, weights=counts_ij)
-
-            # Marginal Var(E[Y|Xi]) using n_bins_soe to match MATLAB logic
-            bin_avg_i_soe, _, binnumber_i_soe = stats.binned_statistic(
-                x=xi, values=output, bins=n_bins_soe, statistic="mean"
-            )
-            mask_i = ~np.isnan(bin_avg_i_soe)
-            counts_i = np.unique(
-                binnumber_i_soe[binnumber_i_soe > 0], return_counts=True
-            )[1]
-            var_i_soe = _weighted_var(bin_avg_i_soe[mask_i], weights=counts_i)
-
-            # Marginal Var(E[Y|Xj]) using n_bins_soe to match MATLAB logic
-            bin_avg_j_soe, _, binnumber_j_soe = stats.binned_statistic(
-                x=xj, values=output, bins=n_bins_soe, statistic="mean"
-            )
-            mask_j = ~np.isnan(bin_avg_j_soe)
-            counts_j = np.unique(
-                binnumber_j_soe[binnumber_j_soe > 0], return_counts=True
-            )[1]
-            var_j_soe = _weighted_var(bin_avg_j_soe[mask_j], weights=counts_j)
-
-            soe[i, j] = (var_ij - var_i_soe - var_j_soe) / var_y
+            soe[i, j] = (var_ij - var_marginal[i] - var_marginal[j]) / var_y
 
     # Mirror SOE and calculate Combined Effect (SI)
     # SI is FOE + half of all interactions associated with that variable
@@ -193,11 +207,18 @@ def sensitivity_indices(
         si[k] = foe[k] + (soe[:, k].sum() / 2)
 
     if print_indices:
-        df_foe = pd.DataFrame(foe, index=var_names, columns=["First-order effect"])
-        df_soe = pd.DataFrame(soe, index=var_names, columns=var_names)
-        df_si = pd.DataFrame(si, index=var_names, columns=["Combined effect"])
+        if not HAS_IPYTHON:
+            warnings.warn(
+                "print_indices=True requires ipython to be installed. "
+                "Install it with: pip install simdec[display]. Table skipped.",
+                stacklevel=2,
+            )
+        else:
+            df_foe = pd.DataFrame(foe, index=var_names, columns=["First-order effect"])
+            df_soe = pd.DataFrame(soe, index=var_names, columns=var_names)
+            df_si = pd.DataFrame(si, index=var_names, columns=["Combined effect"])
 
-        df_indices = pd.concat([df_foe, df_soe, df_si], axis=1)
-        print(f"\n{df_indices}\n")
+            df_indices = pd.concat([df_foe, df_soe, df_si], axis=1)
+            display(df_indices)
 
     return SensitivityAnalysisResult(si, foe, soe)

tests/test_sensitivity_indices.py

@@ -75,10 +75,10 @@ def test_sensitivity_indices(ishigami_ref_indices):
 @pytest.mark.parametrize(
     "fname, foe_ref, si_ref",
     [
-        (path_data / "stress.csv", [0.04, 0.50, 0.11, 0.28], [0.04, 0.51, 0.10, 0.35]),
+        (path_data / "stress.csv", [0.037, 0.49, 0.11, 0.28], [0.04, 0.52, 0.10, 0.35]),
         (
             path_data / "crying.csv",
-            [0.25, 0.22, 0.0, 0.0, 0.01, 0.38],
+            [0.24, 0.22, 0.0, 0.0, 0.01, 0.38],
             [0.28, 0.25, 0.01, 0.01, 0.01, 0.44],
         ),
     ],