From a3dd9aed5ab800d0b0ae445c95be27b6ed9590ca Mon Sep 17 00:00:00 2001 From: gnopik <37065157+gnopik@users.noreply.github.com> Date: Thu, 17 Sep 2026 13:09:03 +0300 Subject: [PATCH 1/4] Update heterogeneity_indices.py --- src/simdec/heterogeneity_indices.py | 753 +++++++++++++++++++++------- 1 file changed, 574 insertions(+), 179 deletions(-) diff --git a/src/simdec/heterogeneity_indices.py b/src/simdec/heterogeneity_indices.py index d6bd457..3345562 100644 --- a/src/simdec/heterogeneity_indices.py +++ b/src/simdec/heterogeneity_indices.py @@ -1,247 +1,642 @@ from dataclasses import dataclass -import logging +import math +import warnings import matplotlib.pyplot as plt import numpy as np import pandas as pd -import simdec as sd +from .sensitivity_indices import sensitivity_indices -logger = logging.getLogger(__name__) +__all__ = [ + "HeterogeneityDetail", + "HeterogeneityResult", + "heterogeneity_indices", +] -__all__ = ["heterogeneity_indices", "plot_heterogeneity"] +_DEFAULT_N_REGIONS = 5 +# Computational safeguard, not a guarantee of statistical stability. +_MIN_REGION_SIZE = 100 @dataclass -class HeterogeneityResult: - summary: pd.DataFrame - regional_profiles: pd.DataFrame - split_name: str +class HeterogeneityDetail: + """Detailed results for one partitioning variable. + Attributes + ---------- + raw_profiles : pandas.DataFrame + Raw regional combined sensitivity profiles. Rows are regions and + columns are model inputs. + normalized_profiles : pandas.DataFrame + Regional profiles normalized to sum to one within each region. These + profiles are used to calculate H. + regional_sums : pandas.Series + Sum of the raw combined sensitivity indices in each region. + region_counts : pandas.Series + Number of observations in each region. + individual_contributions : pandas.Series + Input-level contributions C_{i,Z}. These sum to H_Z. For a categorical + input partition, the partitioning input itself is omitted because it is + constant within each of its own categories. + """ -def heterogeneity_indices( - output: pd.Series, - inputs: pd.DataFrame, - split_variable: str | pd.Series, - n_subdivisions: int | None = None, - plot: bool = False, -) -> HeterogeneityResult: - """Heterogeneity indices. + raw_profiles: pd.DataFrame + normalized_profiles: pd.DataFrame + regional_sums: pd.Series + region_counts: pd.Series + individual_contributions: pd.Series + + def __repr__(self) -> str: + n_regions, n_inputs = self.raw_profiles.shape + return ( + "HeterogeneityDetail(" + f"regions={n_regions}, inputs={n_inputs}; " + "raw_profiles, normalized_profiles, regional_sums, " + "region_counts, individual_contributions)" + ) - Compute sensitivity-based heterogeneity across subdivisions - of a variable. - Parameters - ---------- - output : pd.Series - Model output vector. - inputs : pd.DataFrame - Input/feature matrix. - split_variable : str or pd.Series - Variable to split on. If string, must be a column in 'inputs'. - n_subdivisions : int, optional - Number of regions for continuous variables. Defaults to 4. - plot : bool, default False - If True, displays a stacked bar chart of regional sensitivity profiles - by calling :func:`plot_heterogeneity`. The chart shows variance - contributions of each input across subdivisions of ``split_variable``, - ranked by global sensitivity indices. To capture the returned - ``matplotlib.axes.Axes`` object, call :func:`plot_heterogeneity` - directly on the result instead. +@dataclass +class HeterogeneityResult: + """Results returned by :func:`heterogeneity_indices`. - Returns - ------- - res : HeterogeneityResult - An object with attributes: + The main heterogeneity estimates are available through ``indices``. + Complete regional diagnostics are stored in ``details``. + """ + + indices: pd.Series + details: dict + _input_order: tuple + + def __getitem__(self, partition): + """Return H for one partition, e.g. ``H[\"X1\"]``.""" + return self.indices.loc[partition] + + def __repr__(self) -> str: + if self.indices.empty: + return "HeterogeneityResult\n\nNo heterogeneity indices were computed." + table = self.indices.rename("H").to_frame() + return "HeterogeneityResult\n\n" + table.to_string() + + def plot(self, partition=None, ax=None): + """Plot raw regional sensitivity profiles. + + Parameters + ---------- + partition : str or sequence of str, optional + Partition(s) to plot. If omitted, all available partitions are + plotted. Examples: ``H.plot(\"X1\")`` or + ``H.plot([\"Y\", \"X1\", \"X2\"])``. + ax : matplotlib.axes.Axes, optional + Existing axes. Only valid when plotting one partition. + + Returns + ------- + matplotlib.axes.Axes or numpy.ndarray + Axes containing the plot(s). + """ + return _plot_heterogeneity(self, partition=partition, ax=ax) + + +def _as_dataframe(inputs) -> pd.DataFrame: + X = pd.DataFrame(inputs).reset_index(drop=True).copy() + if X.shape[1] == 0: + raise ValueError("'inputs' must contain at least one input variable.") + if X.columns.has_duplicates: + raise ValueError("'inputs' must have unique column names.") + return X + + +def _as_series(values, name=None) -> pd.Series: + if isinstance(values, pd.Series): + s = values.reset_index(drop=True).copy() + else: + s = pd.Series(values).reset_index(drop=True) + if name is not None and s.name is None: + s.name = name + return s - summary : DataFrame - A summary of calculated heterogeneity indices. - regional_profiles : DataFrame - Regional sensitivity indices for each input across subdivisions. - split_name : str - The name of the variable used to split the data. +def _is_categorical(series: pd.Series) -> bool: + """Identify natural categorical partitions. + + Object, string, boolean, pandas categorical, and binary variables are + treated as categorical. Numeric variables with more than two levels are + treated as continuous unless explicitly stored with pandas ``category`` + dtype. This avoids silently interpreting low-cardinality integer-valued + continuous variables as categories. """ - y = pd.Series(output).reset_index(drop=True) - X = pd.DataFrame(inputs).reset_index(drop=True) - - if isinstance(split_variable, str): - if split_variable not in X.columns: - raise ValueError(f"'{split_variable}' not found in inputs.") - z = X[split_variable].reset_index(drop=True) - split_name = split_variable - else: - z = pd.Series(split_variable).reset_index(drop=True) - split_name = getattr(split_variable, "name", "split_variable") - - unique_vals = z.dropna().unique() - n_unique = len(unique_vals) - - # Determine if variable is categorical/binary - is_categorical = ( - isinstance(z.dtype, pd.CategoricalDtype) - or pd.api.types.is_object_dtype(z) - or pd.api.types.is_string_dtype(z) - or pd.api.types.is_bool_dtype(z) + non_missing = series.dropna() + n_unique = non_missing.nunique() + return ( + isinstance(series.dtype, pd.CategoricalDtype) + or pd.api.types.is_object_dtype(series) + or pd.api.types.is_string_dtype(series) + or pd.api.types.is_bool_dtype(series) or n_unique <= 2 ) - if is_categorical: - regions = z.astype("category") + +def _make_regions( + z: pd.Series, + n_regions: int, + name: str, +) -> tuple[pd.Series, bool]: + categorical = _is_categorical(z) + + if categorical: + regions = z.astype("category").cat.remove_unused_categories() else: - q = n_subdivisions if n_subdivisions is not None else 4 try: - regions = pd.qcut(z, q=q, duplicates="drop") - except ValueError as e: + regions = pd.qcut(z, q=n_regions, duplicates="drop") + except ValueError as exc: + raise ValueError( + f"Failed to divide '{name}' into {n_regions} equal-frequency " + f"regions: {exc}" + ) from exc + + if not hasattr(regions, "cat") or len(regions.cat.categories) < 2: + raise ValueError( + f"At least two regions are required to compute heterogeneity for " + f"'{name}'." + ) + + return regions, categorical + + +def _safe_output_variance(y: pd.Series) -> float | None: + try: + values = pd.to_numeric(y, errors="raise").to_numpy(dtype=float) + except (TypeError, ValueError): + return None + if len(values) < 2: + return 0.0 + return float(np.var(values, ddof=1)) + + +def _mean_pairwise_tv( + normalized_profiles: pd.DataFrame, +) -> tuple[float, pd.Series]: + """Mean pairwise total-variation distance and input contributions.""" + profiles = normalized_profiles.to_numpy(dtype=float) + n_regions, n_inputs = profiles.shape + + if n_regions < 2: + raise ValueError("At least two normalized regional profiles are required.") + + contribution_sum = np.zeros(n_inputs, dtype=float) + n_pairs = 0 + + for r in range(n_regions - 1): + for s in range(r + 1, n_regions): + contribution_sum += 0.5 * np.abs(profiles[r] - profiles[s]) + n_pairs += 1 + + contributions = contribution_sum / n_pairs + H = float(contributions.sum()) + + # For nonnegative compositional profiles, H is theoretically in [0, 1]. + if np.nanmin(profiles) >= -1e-12: + H = float(np.clip(H, 0.0, 1.0)) + + return H, pd.Series( + contributions, + index=normalized_profiles.columns, + dtype=float, + name="C_i", + ) + + +def _compute_partition( + *, + y: pd.Series, + X: pd.DataFrame, + z: pd.Series, + partition_name, + n_regions: int, + remove_partition_input: bool, +) -> tuple[float, HeterogeneityDetail]: + regions, partition_is_categorical = _make_regions( + z, + n_regions=n_regions, + name=str(partition_name), + ) + + X_profile = X + if remove_partition_input and partition_is_categorical: + X_profile = X.drop(columns=[partition_name]) + if X_profile.shape[1] == 0: raise ValueError( - f"Failed to bin '{split_name}' into {q} quantiles: {e}" - ) from e + f"Cannot compute H for categorical input '{partition_name}' " + "because no other inputs remain in the regional sensitivity " + "profile." + ) - regional_profiles = [] - skipped = [] + profile_rows = [] + counts = {} + # Every intended region must be admissible. H is not computed on a silently + # reduced subset of the requested partition. for region in regions.cat.categories: mask = regions == region - n_in_region = mask.sum() + n_in_region = int(mask.sum()) + counts[region] = n_in_region - if n_in_region < 10: - # Need enough samples for meaningful sensitivity indices - skipped.append((region, n_in_region, "too few samples (< 10)")) - continue + if n_in_region < _MIN_REGION_SIZE: + raise ValueError( + f"Region {region!r} of '{partition_name}' contains only " + f"{n_in_region} observations; at least {_MIN_REGION_SIZE} are " + "required for regional sensitivity analysis." + ) - X_sub = X.loc[mask] + X_sub = X_profile.loc[mask] y_sub = y.loc[mask] - # Skip if output has zero or near-zero variance in this region - if y_sub.var() < 1e-12: - skipped.append((region, n_in_region, "output variance ≈ 0")) - continue + variance = _safe_output_variance(y_sub) + if variance is not None and variance < 1e-12: + raise ValueError( + f"The output is constant, or approximately constant, in region " + f"{region!r} of '{partition_name}'. Regional variance-based " + "sensitivity indices and H are therefore undefined for this " + "partition." + ) try: - res = sd.sensitivity_indices(inputs=X_sub, output=y_sub) - si_vals = np.asarray(res.si).ravel() + result = sensitivity_indices(inputs=X_sub, output=y_sub) + si_values = np.asarray(result.si, dtype=float).ravel() + except Exception as exc: + raise ValueError( + f"Regional sensitivity calculation failed in region {region!r} " + f"of '{partition_name}': {exc}" + ) from exc + + if len(si_values) != X_profile.shape[1]: + raise ValueError( + f"Sensitivity calculation for region {region!r} of " + f"'{partition_name}' returned {len(si_values)} indices for " + f"{X_profile.shape[1]} inputs." + ) - # Guard against NaN/Inf from degenerate sensitivity computation - if not np.all(np.isfinite(si_vals)): - skipped.append((region, n_in_region, "non-finite SI values")) - continue + if not np.all(np.isfinite(si_values)): + raise ValueError( + f"Non-finite sensitivity indices were returned for region " + f"{region!r} of '{partition_name}'." + ) + + # Remove numerical dust only. Material negative estimates are preserved + # and flagged below rather than silently altered. + si_values[np.abs(si_values) < 1e-12] = 0.0 - si_region = pd.Series(si_vals, index=X.columns, name=region) - regional_profiles.append(si_region) + profile_rows.append( + pd.Series(si_values, index=X_profile.columns, name=region, dtype=float) + ) - except Exception as e: - skipped.append((region, n_in_region, f"exception: {e}")) - continue + raw = pd.DataFrame(profile_rows) + raw.index.name = "region" - if skipped: - logger.info("Skipped %d region(s) of '%s':", len(skipped), split_name) - for reg, n, reason in skipped: - logger.info(" - region=%r, n=%d, reason=%s", reg, n, reason) + regional_sums = raw.sum(axis=1) + regional_sums.name = "sum_S_i" - if len(regional_profiles) < 2: - total_regions = len(regions.cat.categories) - valid = len(regional_profiles) + bad_sums = (~np.isfinite(regional_sums)) | (np.abs(regional_sums) < 1e-12) + if bad_sums.any(): + bad_regions = list(regional_sums.index[bad_sums]) raise ValueError( - f"Not enough valid subdivisions to compute heterogeneity: " - f"{valid}/{total_regions} regions passed all checks for '{split_name}'.\n" - f"Skipped regions:\n" - "\n".join(f" {r!r}: n={n}, {reason} " for r, n, reason in skipped), - "\n\nTry: (1) reducing n_subdivisions, " - "(2) using a different split_variable, or " - "(3) ensuring more samples per region.", - ) - - regional_si = pd.concat(regional_profiles, axis=1) - - res_global = sd.sensitivity_indices(inputs=X, output=y) - overall_si = pd.Series( - np.asarray(res_global.si).ravel(), - index=X.columns, - name="Overall_SI", - ) + f"Regional sensitivity profiles for '{partition_name}' cannot be " + "normalized because their sensitivity-index sum is zero or " + f"non-finite in region(s) {bad_regions}." + ) - # Heterogeneity = 2 × population std dev across regions - hetero_scores = 2 * regional_si.std(axis=1, ddof=0) - total_hetero = hetero_scores.mean() + normalized = raw.div(regional_sums, axis=0) - hetero_col_name = f"Heterogeneity (across {split_name})" - summary = pd.DataFrame( - {"Overall_SI": overall_si, hetero_col_name: hetero_scores} - ).sort_values(by=hetero_col_name, ascending=False) - summary.loc["SUM / TOTAL"] = [overall_si.sum(), total_hetero] + if (normalized.to_numpy(dtype=float) < -1e-12).any(): + warnings.warn( + f"Negative combined sensitivity indices were found in regional " + f"profiles for '{partition_name}'. H is still calculated from the " + "normalized profiles, but the strict [0, 1] total-variation " + "interpretation assumes nonnegative sensitivity compositions.", + RuntimeWarning, + stacklevel=3, + ) - result = HeterogeneityResult(summary, regional_si, split_name) + H, contributions = _mean_pairwise_tv(normalized) - if plot: - plot_heterogeneity(result) + region_counts = pd.Series(counts, dtype=int, name="n") + region_counts.index.name = "region" - return result + detail = HeterogeneityDetail( + raw_profiles=raw, + normalized_profiles=normalized, + regional_sums=regional_sums, + region_counts=region_counts, + individual_contributions=contributions, + ) + return H, detail -def plot_heterogeneity(result: HeterogeneityResult, ax: plt.Axes = None) -> plt.Axes: - """Plot regional sensitivity profiles. +def heterogeneity_indices( + output, + inputs, + n_regions: int = _DEFAULT_N_REGIONS, + custom_partition=None, +) -> HeterogeneityResult: + """Calculate heterogeneity in regional sensitivity profiles. + + The heterogeneity index H measures how much the relative sensitivity + profile of a model changes across regions. H is calculated as the mean + pairwise total-variation distance between normalized regional combined + sensitivity profiles. + + With no ``custom_partition``, the function performs a full standard scan: + H_Y is calculated across regions of the output and H_Xi across regions of + every input. Continuous partitioning variables are divided into + ``n_regions`` equal-frequency regions; categorical variables use their + observed categories. + + If ``custom_partition`` is supplied, only H_Z for that partition is + calculated. A continuous custom partition is divided into ``n_regions`` + equal-frequency regions. A categorical custom partition uses its natural + categories, in which case ``n_regions`` is ignored. Parameters ---------- - result : HeterogeneityResult - The result object from heterogeneity_indices. - ax : matplotlib.axes.Axes, optional - Existing axes to plot on. + output : pandas.Series or array-like + Model output vector. + inputs : pandas.DataFrame or array-like + Model input matrix. DataFrame column names are used as input names. + n_regions : int, default 5 + Number of equal-frequency regions used for continuous partitioning + variables. + custom_partition : pandas.Series or array-like, optional + User-defined partitioning variable Z. When provided, only H_Z is + calculated. The partition is used only to assign observations to + regions and is not added to the model inputs. Returns ------- - ax : matplotlib.axes.Axes - The axes with the plot. - + HeterogeneityResult + ``H.indices`` contains the calculated heterogeneity indices. + ``H.details`` contains raw and normalized regional profiles, regional + sensitivity-index sums, observation counts, and input-level + contributions for every calculated partition. + + Notes + ----- + For a categorical input X_i, X_i is removed from its own regional + sensitivity profiles because it is constant within each category. + + H_Y is not calculated for a categorical output. Partitioning a categorical + output by its categories makes the output constant within every region, so + regional variance-based sensitivity indices are undefined. H_Xi is still + calculated for all inputs for which the regional analysis is admissible. + + Multi-level numeric categorical variables should use pandas ``category`` + dtype. Binary numeric variables are recognized as categorical + automatically. """ - summary = result.summary - regional_si = result.regional_profiles - split_name = result.split_name + if not isinstance(n_regions, (int, np.integer)) or n_regions < 2: + raise ValueError("'n_regions' must be an integer greater than or equal to 2.") - hetero_col_name = [c for c in summary.columns if "Heterogeneity" in c][0] - total_hetero = summary.loc["SUM / TOTAL", hetero_col_name] + y = _as_series(output, name="Y") + X = _as_dataframe(inputs) - plot_order = summary.index[summary.index != "SUM / TOTAL"] - plot_order = ( - summary.loc[plot_order].sort_values(by="Overall_SI", ascending=False).index - ) + if len(y) != len(X): + raise ValueError( + "'output' and 'inputs' must contain the same number of observations " + f"({len(y)} != {len(X)})." + ) - cmap = plt.colormaps["terrain"] - colors = [cmap(i) for i in np.linspace(0.05, 0.95, len(regional_si.index))] + input_order = tuple(X.columns) + indices = {} + details = {} - data_to_plot = regional_si.loc[plot_order].T + # ------------------------------------------------------------ + # Custom partition: calculate only H_Z. + # ------------------------------------------------------------ + if custom_partition is not None: + z = _as_series(custom_partition) + if len(z) != len(y): + raise ValueError( + "'custom_partition' must contain the same number of observations " + "as output and inputs." + ) + + partition_name = z.name if z.name is not None else "Z" + custom_is_categorical = _is_categorical(z) + + if custom_is_categorical and n_regions != _DEFAULT_N_REGIONS: + warnings.warn( + "'n_regions' is ignored because 'custom_partition' is " + "categorical. Its observed categories are used as regions.", + UserWarning, + stacklevel=2, + ) + + H_value, detail = _compute_partition( + y=y, + X=X, + z=z, + partition_name=partition_name, + n_regions=n_regions, + remove_partition_input=False, + ) + indices[partition_name] = H_value + details[partition_name] = detail - if ax is None: - _, ax = plt.subplots(figsize=(10, 6)) + return HeterogeneityResult( + indices=pd.Series(indices, dtype=float, name="H").rename_axis("partition"), + details=details, + _input_order=input_order, + ) - data_to_plot.plot( - kind="bar", - stacked=True, - ax=ax, - color=colors, - edgecolor="white", - width=0.8, + # ------------------------------------------------------------ + # Standard scan: H_Y and H_Xi for all inputs. + # ------------------------------------------------------------ + if _is_categorical(y): + warnings.warn( + "H_Y was not calculated because the output is categorical. " + "Partitioning a categorical output by its own categories makes " + "the output constant within each region, so regional variance-based " + "sensitivity indices are undefined. H_Xi will still be calculated " + "for all inputs where the output varies within the resulting " + "regions.", + UserWarning, + stacklevel=2, + ) + else: + try: + H_y, detail_y = _compute_partition( + y=y, + X=X, + z=y, + partition_name="Y", + n_regions=n_regions, + remove_partition_input=False, + ) + indices["Y"] = H_y + details["Y"] = detail_y + except ValueError as exc: + warnings.warn( + f"H_Y could not be calculated: {exc}", + UserWarning, + stacklevel=2, + ) + indices["Y"] = np.nan + + for input_name in X.columns: + try: + H_x, detail_x = _compute_partition( + y=y, + X=X, + z=X[input_name], + partition_name=input_name, + n_regions=n_regions, + remove_partition_input=True, + ) + indices[input_name] = H_x + details[input_name] = detail_x + except ValueError as exc: + warnings.warn( + f"H for input '{input_name}' could not be calculated: {exc}", + UserWarning, + stacklevel=2, + ) + indices[input_name] = np.nan + + return HeterogeneityResult( + indices=pd.Series(indices, dtype=float, name="H").rename_axis("partition"), + details=details, + _input_order=input_order, ) - ax.set_title( - f"Sensitivity Profiles across {split_name}\n" - f"(Total Heterogeneity: {total_hetero:.3f})", - fontsize=10, - ) - ax.set_ylabel("Variance Contribution", fontsize=8) - ax.set_xlabel(f"Regions of {split_name}", fontsize=8) +def _format_region_label(region) -> str: + if isinstance(region, pd.Interval): + left = f"{region.left:.3g}" + right = f"{region.right:.3g}" + return f"{left}-{right}" + if isinstance(region, (float, np.floating)): + return f"{region:.3g}" + return str(region) + + +def _plot_heterogeneity( + result: HeterogeneityResult, + partition=None, + ax=None, +): + available = list(result.details) + if not available: + raise ValueError("The result contains no regional profiles to plot.") + + if partition is None: + names = available + elif isinstance(partition, str) or not hasattr(partition, "__iter__"): + names = [partition] + else: + names = list(partition) - ax.legend( - title="Inputs (Ranked by Global SI)", - bbox_to_anchor=(1.05, 1), - loc="upper left", - ) + missing = [name for name in names if name not in result.details] + if missing: + raise KeyError( + f"Partition(s) {missing} are not available. Choose from {available}." + ) - ax.tick_params(axis="x", labelrotation=45) - ax.grid(axis="y", linestyle="--", alpha=0.7) + if ax is not None and len(names) != 1: + raise ValueError("'ax' can only be supplied when plotting one partition.") - if plt.get_backend().lower() != "agg": - plt.tight_layout() + input_order = list(result._input_order) + cmap = plt.colormaps["terrain"] + colors = { + name: cmap(pos) + for name, pos in zip( + input_order, + np.linspace(0.05, 0.95, max(len(input_order), 1)), + ) + } + + if len(names) == 1: + if ax is None: + _, ax = plt.subplots(figsize=(7.0, 4.8)) + axes = np.array([ax], dtype=object) + else: + ncols = min(3, len(names)) + nrows = math.ceil(len(names) / ncols) + _, grid = plt.subplots( + nrows, + ncols, + figsize=(5.2 * ncols, 4.2 * nrows), + squeeze=False, + ) + axes = grid.ravel() + + used_inputs = [] + + for axis, name in zip(axes, names): + detail = result.details[name] + raw = detail.raw_profiles + order = [item for item in input_order if item in raw.columns] + order += [item for item in raw.columns if item not in order] + data = raw.loc[:, order] + used_inputs.extend(item for item in order if item not in used_inputs) + + data.plot( + kind="bar", + stacked=True, + ax=axis, + color=[colors.get(item, cmap(0.5)) for item in order], + edgecolor="white", + linewidth=0.4, + width=0.82, + legend=False, + ) - return ax + H_value = result.indices.loc[name] + axis.set_title(f"{name}: H = {H_value:.3f}") + axis.set_ylabel(r"Combined sensitivity index, $S_i$") + axis.set_xlabel("Region") + axis.set_xticklabels( + [_format_region_label(region) for region in raw.index], + rotation=45, + ha="right", + ) + axis.grid(axis="y", linestyle=":", alpha=0.30) + axis.set_axisbelow(True) + + for axis in axes[len(names):]: + axis.remove() + + handles = [ + plt.Rectangle( + (0, 0), + 1, + 1, + facecolor=colors.get(name, cmap(0.5)), + edgecolor="white", + label=str(name), + ) + for name in used_inputs + ] + + if len(names) == 1: + if handles: + axes[0].legend( + handles=handles, + title="Inputs", + bbox_to_anchor=(1.02, 1), + loc="upper left", + ) + axes[0].figure.tight_layout() + return axes[0] + + if handles: + axes[0].figure.legend( + handles=handles, + title="Inputs", + loc="lower center", + bbox_to_anchor=(0.5, 0.01), + ncol=min(5, len(handles)), + frameon=False, + ) + axes[0].figure.subplots_adjust(bottom=0.16, hspace=0.48, wspace=0.28) + return axes From 6b1ba97f6159ebcaa288a709141dc25022550b28 Mon Sep 17 00:00:00 2001 From: gnopik <37065157+gnopik@users.noreply.github.com> Date: Thu, 17 Sep 2026 13:10:08 +0300 Subject: [PATCH 2/4] Update test_heterogeneity_indices.py --- tests/test_heterogeneity_indices.py | 281 +++++++++++++++++++--------- 1 file changed, 190 insertions(+), 91 deletions(-) diff --git a/tests/test_heterogeneity_indices.py b/tests/test_heterogeneity_indices.py index 5da9cb1..4eca395 100644 --- a/tests/test_heterogeneity_indices.py +++ b/tests/test_heterogeneity_indices.py @@ -1,140 +1,239 @@ -import pathlib -import pytest +import importlib +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt import numpy as np import pandas as pd -import matplotlib.pyplot as plt +import pytest -import simdec as sd +hi = importlib.import_module("simdec.heterogeneity_indices") -path_data = pathlib.Path(__file__).parent / "data" +class _SensitivityResult: + def __init__(self, si): + self.si = np.asarray(si, dtype=float) -@pytest.fixture(autouse=True) -def close_plots(): - yield - plt.close("all") +def _fake_sensitivity_indices(*, inputs, output): + """Deterministic stand-in for API and result-structure tests.""" + X = pd.DataFrame(inputs).reset_index(drop=True) + y = pd.Series(output).reset_index(drop=True) + yy = pd.to_numeric(y, errors="raise").to_numpy(dtype=float) + scores = [] + for column in X.columns: + x = X[column] + if not pd.api.types.is_numeric_dtype(x): + xx = pd.factorize(x)[0].astype(float) + else: + xx = pd.to_numeric(x, errors="raise").to_numpy(dtype=float) -@pytest.fixture -def dummy_data(): - rng = np.random.default_rng(42) - n = 200 + if np.std(xx) < 1e-12 or np.std(yy) < 1e-12: + score = 0.0 + else: + score = float(np.corrcoef(xx, yy)[0, 1] ** 2) + if not np.isfinite(score): + score = 0.0 + + # Positive floor keeps profiles normalizable in deliberately simple tests. + scores.append(score + 0.02) - inputs = pd.DataFrame( + return _SensitivityResult(scores) + + +@pytest.fixture +def example_data(): + rng = np.random.default_rng(123) + n = 1000 + category = pd.Series( + np.where(np.arange(n) % 2 == 0, "A", "B"), + dtype="category", + name="category", + ) + X = pd.DataFrame( { - "x1": rng.random(n), - "x2": rng.random(n), - "x3": rng.random(n), - "cat_var": rng.choice(["A", "B", "C"], size=n), + "x1": rng.normal(size=n), + "category": category, + "x3": rng.normal(size=n), } ) + y = pd.Series( + X["x1"] + + (category == "B").astype(float) * 2.0 * X["x3"] + + rng.normal(scale=0.2, size=n), + name="Y", + ) + return y, X - # Create a dummy output dependent on x1 and x2 - y = 2.0 * inputs["x1"] + 0.5 * inputs["x2"] + rng.normal(0, 0.1, n) - return inputs, y +def test_public_package_export(): + import simdec -def test_heterogeneity_categorical_str(dummy_data): - """Test splitting by a string column name (categorical).""" - inputs, y = dummy_data + assert callable(simdec.heterogeneity_indices) - res = sd.heterogeneity_indices(output=y, inputs=inputs, split_variable="cat_var") - # Check object structure - assert hasattr(res, "summary") - assert hasattr(res, "regional_profiles") - assert res.split_name == "cat_var" +def test_standard_call_computes_y_and_all_inputs(monkeypatch, example_data): + monkeypatch.setattr(hi, "sensitivity_indices", _fake_sensitivity_indices) + y, X = example_data - # Check DataFrame structures - assert "Overall_SI" in res.summary.columns - assert "Heterogeneity (across cat_var)" in res.summary.columns - assert "SUM / TOTAL" in res.summary.index + H = hi.heterogeneity_indices(output=y, inputs=X) - # 3 categories - assert res.regional_profiles.shape[1] == 3 - assert list(res.regional_profiles.index) == ["x1", "x2", "x3", "cat_var"] + assert list(H.indices.index) == ["Y", "x1", "category", "x3"] + assert H.indices.notna().all() + assert set(H.details) == set(H.indices.index) + assert H["x1"] == pytest.approx(H.indices["x1"]) -def test_heterogeneity_continuous_series(dummy_data): - """Test splitting by passing a pandas Series (continuous).""" - inputs, y = dummy_data - split_series = inputs["x1"] +def test_profiles_are_normalized_and_contributions_sum_to_h( + monkeypatch, example_data +): + monkeypatch.setattr(hi, "sensitivity_indices", _fake_sensitivity_indices) + y, X = example_data + + H = hi.heterogeneity_indices(output=y, inputs=X) + + for name, detail in H.details.items(): + assert np.allclose(detail.normalized_profiles.sum(axis=1), 1.0) + assert np.isclose(detail.individual_contributions.sum(), H.indices[name]) + assert detail.regional_sums.index.equals(detail.region_counts.index) - res = sd.heterogeneity_indices( - output=y, inputs=inputs, split_variable=split_series, n_subdivisions=4 - ) - assert res.split_name == "x1" - assert res.regional_profiles.shape[1] == 4 # 4 quantiles +def test_categorical_input_is_removed_from_its_own_profiles( + monkeypatch, example_data +): + monkeypatch.setattr(hi, "sensitivity_indices", _fake_sensitivity_indices) + y, X = example_data + H = hi.heterogeneity_indices(output=y, inputs=X) -def test_heterogeneity_missing_column(dummy_data): - """Test that a ValueError is raised when split_variable is not in inputs.""" - inputs, y = dummy_data + detail = H.details["category"] + assert "category" not in detail.raw_profiles.columns + assert "category" not in detail.normalized_profiles.columns + assert "category" not in detail.individual_contributions.index - with pytest.raises(ValueError, match="'missing_col' not found in inputs"): - sd.heterogeneity_indices(output=y, inputs=inputs, split_variable="missing_col") + # The same categorical input remains available in profiles for other partitions. + assert "category" in H.details["Y"].raw_profiles.columns + assert "category" in H.details["x1"].raw_profiles.columns -def test_heterogeneity_too_few_regions(): - """Test that a ValueError is raised when there are not enough valid subdivisions.""" - inputs = pd.DataFrame({"x1": [1, 2, 3, 4, 5], "cat": ["A", "B", "C", "D", "E"]}) - y = pd.Series([1.0, 2.0, 3.0, 4.0, 5.0]) +def test_categorical_output_skips_h_y_but_continues(monkeypatch, example_data): + monkeypatch.setattr(hi, "sensitivity_indices", _fake_sensitivity_indices) + y, X = example_data + y_binary = (y > y.median()).astype(int) - with pytest.raises(ValueError, match="Not enough valid subdivisions"): - sd.heterogeneity_indices(output=y, inputs=inputs, split_variable="cat") + with pytest.warns(UserWarning, match="output is categorical"): + H = hi.heterogeneity_indices(output=y_binary, inputs=X) + assert "Y" not in H.indices.index + assert "Y" not in H.details + assert list(H.indices.index) == ["x1", "category", "x3"] + assert H.indices.notna().all() -def test_heterogeneity_plot_argument(dummy_data): - """Test that setting plot=True works without throwing an error.""" - inputs, y = dummy_data - res = sd.heterogeneity_indices( - output=y, inputs=inputs, split_variable="cat_var", plot=True +def test_continuous_custom_partition_uses_n_regions(monkeypatch, example_data): + monkeypatch.setattr(hi, "sensitivity_indices", _fake_sensitivity_indices) + y, X = example_data + z = pd.Series(np.linspace(0.0, 1.0, len(y)), name="temperature") + + H = hi.heterogeneity_indices( + output=y, + inputs=X, + n_regions=5, + custom_partition=z, ) - assert res is not None - # Figure exists in the active pyplot state - assert len(plt.get_fignums()) > 0 + assert list(H.indices.index) == ["temperature"] + detail = H.details["temperature"] + assert len(detail.region_counts) == 5 + assert detail.region_counts.sum() == len(y) + assert (detail.region_counts >= 100).all() + + +def test_categorical_custom_partition_uses_categories_and_warns_if_n_regions_changed( + monkeypatch, example_data +): + monkeypatch.setattr(hi, "sensitivity_indices", _fake_sensitivity_indices) + y, X = example_data + z = pd.Series( + np.where(np.arange(len(y)) % 2 == 0, "OK", "Flood"), + dtype="category", + name="Flood regime", + ) + with pytest.warns(UserWarning, match="n_regions.*ignored"): + H = hi.heterogeneity_indices( + output=y, + inputs=X, + n_regions=10, + custom_partition=z, + ) + + assert list(H.indices.index) == ["Flood regime"] + detail = H.details["Flood regime"] + assert set(detail.region_counts.index.astype(str)) == {"OK", "Flood"} + assert sorted(detail.region_counts.tolist()) == [500, 500] + + +def test_minimum_region_size_is_100(monkeypatch, example_data): + monkeypatch.setattr(hi, "sensitivity_indices", _fake_sensitivity_indices) + y, X = example_data + z = pd.Series( + ["rare"] * 99 + ["common"] * (len(y) - 99), + dtype="category", + name="regime", + ) -def test_plot_heterogeneity(dummy_data): - """Test the independent plot_heterogeneity function.""" - inputs, y = dummy_data + with pytest.raises(ValueError, match="at least 100"): + hi.heterogeneity_indices( + output=y, + inputs=X, + custom_partition=z, + ) - res = sd.heterogeneity_indices(output=y, inputs=inputs, split_variable="cat_var") - ax = sd.plot_heterogeneity(res) +def test_region_counts_are_exposed(monkeypatch, example_data): + monkeypatch.setattr(hi, "sensitivity_indices", _fake_sensitivity_indices) + y, X = example_data - assert isinstance(ax, plt.Axes) + H = hi.heterogeneity_indices(output=y, inputs=X, n_regions=5) - # Calculate the expected title format - hetero_col_name = [c for c in res.summary.columns if "Heterogeneity" in c][0] - total_hetero = res.summary.loc["SUM / TOTAL", hetero_col_name] - expected_title = ( - f"Sensitivity Profiles across cat_var\n" - f"(Total Heterogeneity: {total_hetero:.3f})" + for name, detail in H.details.items(): + assert detail.region_counts.sum() == len(y) + assert len(detail.region_counts) >= 2 + + +def test_binary_analytical_tv_reference(): + # Controlled-model a=4 reference used in the manuscript: H_K = 6/13. + profiles = pd.DataFrame( + [ + [0.5, 0.5], + [1.0 / 26.0, 25.0 / 26.0], + ], + index=["K=0", "K=1"], + columns=["A", "B"], ) - assert ax.get_title() == expected_title - assert ax.get_ylabel() == "Variance Contribution" - assert ax.get_xlabel() == "Regions of cat_var" + H, contributions = hi._mean_pairwise_tv(profiles) + assert H == pytest.approx(6.0 / 13.0) + assert contributions.sum() == pytest.approx(H) -def test_heterogeneity_real_data(): - """Integration test using the real stress.csv dataset from the project.""" - fname = path_data / "stress.csv" - data = pd.read_csv(fname) - output_name, *v_names = list(data.columns) - inputs, output = data[v_names], data[output_name] +def test_plot_uses_stored_results_without_recalculation(monkeypatch, example_data): + monkeypatch.setattr(hi, "sensitivity_indices", _fake_sensitivity_indices) + y, X = example_data + H = hi.heterogeneity_indices(output=y, inputs=X) - res = sd.heterogeneity_indices( - output=output, inputs=inputs, split_variable="R", n_subdivisions=2 - ) + def _should_not_run(*args, **kwargs): + raise AssertionError("sensitivity_indices should not be called by plot()") + + monkeypatch.setattr(hi, "sensitivity_indices", _should_not_run) + + ax = H.plot("x1") + assert ax is not None + plt.close(ax.figure) - assert res.split_name == "R" - assert not res.summary.empty - assert res.regional_profiles.shape[1] == 2 + axes = H.plot(["Y", "x1"]) + assert len(np.ravel(axes)) >= 2 + plt.close(np.ravel(axes)[0].figure) From 6f68aa1dfcdb3279cebd106b3f68ebf82c33a427 Mon Sep 17 00:00:00 2001 From: gnopik <37065157+gnopik@users.noreply.github.com> Date: Thu, 17 Sep 2026 13:39:35 +0300 Subject: [PATCH 3/4] Update test_heterogeneity_indices.py --- tests/test_heterogeneity_indices.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_heterogeneity_indices.py b/tests/test_heterogeneity_indices.py index 4eca395..5b53bcb 100644 --- a/tests/test_heterogeneity_indices.py +++ b/tests/test_heterogeneity_indices.py @@ -1,12 +1,12 @@ import importlib -import matplotlib -matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np import pandas as pd import pytest +plt.switch_backend("Agg") + hi = importlib.import_module("simdec.heterogeneity_indices") From 796970b5e14d44d32838998634b6e83e1ce90df9 Mon Sep 17 00:00:00 2001 From: gnopik <37065157+gnopik@users.noreply.github.com> Date: Thu, 17 Sep 2026 10:52:08 +0000 Subject: [PATCH 4/4] Fix formatting and linting --- src/simdec/heterogeneity_indices.py | 2 +- tests/test_heterogeneity_indices.py | 8 ++------ 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/src/simdec/heterogeneity_indices.py b/src/simdec/heterogeneity_indices.py index 3345562..0ba493e 100644 --- a/src/simdec/heterogeneity_indices.py +++ b/src/simdec/heterogeneity_indices.py @@ -603,7 +603,7 @@ def _plot_heterogeneity( axis.grid(axis="y", linestyle=":", alpha=0.30) axis.set_axisbelow(True) - for axis in axes[len(names):]: + for axis in axes[len(names) :]: axis.remove() handles = [ diff --git a/tests/test_heterogeneity_indices.py b/tests/test_heterogeneity_indices.py index 5b53bcb..2886afd 100644 --- a/tests/test_heterogeneity_indices.py +++ b/tests/test_heterogeneity_indices.py @@ -85,9 +85,7 @@ def test_standard_call_computes_y_and_all_inputs(monkeypatch, example_data): assert H["x1"] == pytest.approx(H.indices["x1"]) -def test_profiles_are_normalized_and_contributions_sum_to_h( - monkeypatch, example_data -): +def test_profiles_are_normalized_and_contributions_sum_to_h(monkeypatch, example_data): monkeypatch.setattr(hi, "sensitivity_indices", _fake_sensitivity_indices) y, X = example_data @@ -99,9 +97,7 @@ def test_profiles_are_normalized_and_contributions_sum_to_h( assert detail.regional_sums.index.equals(detail.region_counts.index) -def test_categorical_input_is_removed_from_its_own_profiles( - monkeypatch, example_data -): +def test_categorical_input_is_removed_from_its_own_profiles(monkeypatch, example_data): monkeypatch.setattr(hi, "sensitivity_indices", _fake_sensitivity_indices) y, X = example_data