causarray 0.0.8__tar.gz → 0.0.9__tar.gz

{causarray-0.0.8 → causarray-0.0.9}/.gitignore

@@ -184,6 +184,8 @@ docs/source/tutorial/replogle/replogle_subset.h5ad
 docs/source/tutorial/replogle/replogle_subset_norm*.h5ad
 docs/source/tutorial/replogle/replogle_normed.h5ad
 docs/source/tutorial/replogle/replogle_results.h5
+docs/source/tutorial/replogle/replogle_supt5h_go_*.csv
+docs/source/tutorial/replogle/replogle_propensity_batch12*
 
 # Unrelated tutorial directory (not part of causarray package docs)
 docs/source/tutorial/scp/

{causarray-0.0.8 → causarray-0.0.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: causarray
-Version: 0.0.8
+Version: 0.0.9
 Summary: causarray is a Python module for simultaneous causal inference with an array of outcomes.
 Author-email: Jin-Hong Du <jinhongd@hku.hk>, Maya Shen <myshen@andrew.cmu.edu>, Hansruedi Mathys <mathysh@pitt.edu>, Kathryn Roeder <jinhongd@hku.hk>
 Maintainer-email: Jin-Hong Du <jinhongd@hku.hk>
@@ -98,6 +98,30 @@ df_res = gcate_lfc_batch(
 See the [Replogle-E-K562 tutorial](https://causarray.readthedocs.io/en/latest/tutorial/replogle/replogle-py.html)
 for a demonstration on 200 perturbations from a genome-wide CRISPRi screen.
 
+### Diagnostic masks without refitting effects
+
+Treatment-by-gene support or quality-control rules can be aligned to an
+existing causarray result table by label, even when its rows are reordered:
+
+```python
+from causarray import align_test_mask
+
+keep = align_test_mask(
+    df_res,
+    support_mask,                 # treatments × genes, Boolean
+    treatment_names=perturbations,
+    gene_names=genes,
+)
+df_res_flagged = df_res.assign(support_keep=keep)
+```
+
+This operation only annotates or subsets existing results. It does not refit
+the causarray LFC or change standard errors and p-values. If a diagnostic rule
+was selected after inspecting the outcomes, retain the original adjusted
+p-values rather than redefining the multiple-testing family post hoc. The
+[Replogle tutorial](https://causarray.readthedocs.io/en/latest/tutorial/replogle/replogle-py.html)
+compares several expression-support rules with marginal Wilcoxon results.
+
 ## Changelog
 
 See [CHANGELOG](https://causarray.readthedocs.io/en/latest/changelog.html) for a full version history.

{causarray-0.0.8 → causarray-0.0.9}/README.md

@@ -73,6 +73,30 @@ df_res = gcate_lfc_batch(
 See the [Replogle-E-K562 tutorial](https://causarray.readthedocs.io/en/latest/tutorial/replogle/replogle-py.html)
 for a demonstration on 200 perturbations from a genome-wide CRISPRi screen.
 
+### Diagnostic masks without refitting effects
+
+Treatment-by-gene support or quality-control rules can be aligned to an
+existing causarray result table by label, even when its rows are reordered:
+
+```python
+from causarray import align_test_mask
+
+keep = align_test_mask(
+    df_res,
+    support_mask,                 # treatments × genes, Boolean
+    treatment_names=perturbations,
+    gene_names=genes,
+)
+df_res_flagged = df_res.assign(support_keep=keep)
+```
+
+This operation only annotates or subsets existing results. It does not refit
+the causarray LFC or change standard errors and p-values. If a diagnostic rule
+was selected after inspecting the outcomes, retain the original adjusted
+p-values rather than redefining the multiple-testing family post hoc. The
+[Replogle tutorial](https://causarray.readthedocs.io/en/latest/tutorial/replogle/replogle-py.html)
+compares several expression-support rules with marginal Wilcoxon results.
+
 ## Changelog
 
 See [CHANGELOG](https://causarray.readthedocs.io/en/latest/changelog.html) for a full version history.

{causarray-0.0.8 → causarray-0.0.9}/causarray/DR_estimation.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pandas as pd
 from sklearn.linear_model import LogisticRegression
 from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
 from sklearn_ensemble_cv import reset_random_seeds, Ensemble, ECV
@@ -10,6 +11,7 @@ from joblib import Parallel, delayed
 from tqdm import tqdm
 import pprint
 import warnings
+from collections.abc import Mapping
 
 from sklearn.model_selection import KFold, ShuffleSplit
 
@@ -40,6 +42,27 @@ def _get_func_ps(ps_model, **kwargs):
     return func_ps, params_ps
 
 
+def _validate_clip(clip):
+    """Validate a propensity clipping bound and return it as ``(lower, upper)``.
+
+    Raises
+    ------
+    ValueError
+        If ``clip`` is not a pair of numbers satisfying
+        ``0 <= lower < upper <= 1``.
+    """
+    message = 'clip must be None or a pair 0 <= lower < upper <= 1'
+    if isinstance(clip, (str, bytes)) or np.ndim(clip) != 1:
+        raise ValueError(message)
+    try:
+        lower, upper = (float(bound) for bound in clip)
+    except (TypeError, ValueError):
+        raise ValueError(message) from None
+    if not 0 <= lower < upper <= 1:
+        raise ValueError(message)
+    return lower, upper
+
+
 def estimate_propensity_scores(
     A, X_A, K=1, ps_model='logistic', mask=None, clip=None,
     random_state=0, verbose=False, class_weight='balanced', **kwargs,
@@ -172,12 +195,280 @@ def estimate_propensity_scores(
                 pi_hat[test_index, j] = func_ps(x_train, y_train, XA_test)
 
     if clip is not None:
-        if len(clip) != 2 or not 0 <= clip[0] < clip[1] <= 1:
-            raise ValueError('clip must be None or a pair 0 <= lower < upper <= 1')
-        pi_hat = np.clip(pi_hat, clip[0], clip[1])
+        lower, upper = _validate_clip(clip)
+        pi_hat = np.clip(pi_hat, lower, upper)
     return pi_hat
 
 
+def refit_propensity_scores(
+    A, X_A, drop_by_treatment=None, pi_hat=None, treatment_names=None,
+    covariate_names=None, penalty_factors_by_treatment=None, K=1,
+    ps_model='logistic', mask=None, clip=None, random_state=0, verbose=False,
+    class_weight='balanced', **kwargs,
+):
+    """Refit propensity scores with treatment-specific covariate filtering.
+
+    This helper supports sensitivity analyses in which different treatments
+    omit different observed covariates or latent factors, or apply stronger L2
+    regularization to selected covariates. When existing scores are supplied,
+    only treatments named in either treatment-specific mapping are refitted;
+    the remaining columns are carried over from ``pi_hat`` unchanged, up to the
+    shared ``clip`` described below. Outcome models are not fit by this function
+    and cached ``Y_hat`` values can be reused in :func:`LFC`.
+
+    Parameters
+    ----------
+    A : array-like, shape (n, a)
+        Binary treatment indicator matrix.
+    X_A : array-like, shape (n, d_A)
+        Full propensity-model design before treatment-specific filtering.
+    drop_by_treatment : mapping or None
+        Treatment names or indices mapped to covariate names or indices to
+        remove for that treatment. Defaults to no removals.
+    pi_hat : array-like or None, shape (n, a)
+        Existing raw propensity scores. If supplied, treatments absent from
+        ``drop_by_treatment`` are not refitted.
+    treatment_names, covariate_names : sequence, optional
+        Column labels, inferred from DataFrames when possible.
+    penalty_factors_by_treatment : mapping or None
+        Treatment names or indices mapped to ``{covariate: factor}`` mappings.
+        A factor greater than one applies that multiple of the ordinary L2
+        penalty to the named coefficient. This is implemented by dividing the
+        covariate by ``sqrt(factor)`` during both fitting and prediction and is
+        available only for ``ps_model='logistic'`` with an L2 penalty. A factor
+        of one leaves the covariate unchanged. Interpret relative penalties on
+        a common scale; ``prep_causarray_data`` standardizes log-library size.
+    K, ps_model, mask, random_state, verbose, class_weight, **kwargs
+        Passed to :func:`estimate_propensity_scores` for each refitted model.
+    clip : tuple(float, float) or None, optional
+        Bounds applied to the **whole** returned matrix, refitted columns and
+        carried-over columns alike, so that a single consistent bound reaches
+        :func:`LFC`. Pass ``None`` to leave carried-over scores exactly as
+        supplied.
+
+    Returns
+    -------
+    pi_updated : ndarray, shape (n, a)
+        Updated propensity scores.
+    report : DataFrame
+        Audit table of refitted treatments, retained/dropped covariates, and
+        the resulting score spread. ``degenerate_design`` flags a treatment
+        whose retained design is constant, and ``score_std`` reports the
+        standard deviation of its refitted scores on eligible rows; both make a
+        collapsed propensity model visible without reading the warning stream.
+
+    Warns
+    -----
+    RuntimeWarning
+        If filtering leaves a constant design for some treatment, in which case
+        its scores carry no covariate information.
+
+    Notes
+    -----
+    A very large penalty factor shrinks a coefficient towards zero without
+    making the design constant. That case leaves ``degenerate_design`` False
+    but drives ``score_std`` towards zero.
+
+    .. versionadded:: 0.0.9
+    """
+    if drop_by_treatment is None:
+        drop_by_treatment = {}
+    if penalty_factors_by_treatment is None:
+        penalty_factors_by_treatment = {}
+    if not isinstance(drop_by_treatment, Mapping):
+        raise ValueError('drop_by_treatment must be a mapping')
+    if not isinstance(penalty_factors_by_treatment, Mapping):
+        raise ValueError('penalty_factors_by_treatment must be a mapping')
+    if penalty_factors_by_treatment:
+        if ps_model != 'logistic':
+            raise ValueError(
+                'penalty_factors_by_treatment is supported only for logistic models'
+            )
+        if kwargs.get('penalty', 'l2') != 'l2':
+            raise ValueError(
+                'penalty_factors_by_treatment requires logistic penalty="l2"'
+            )
+
+    if isinstance(A, pd.DataFrame):
+        if treatment_names is None:
+            treatment_names = list(A.columns)
+        A_array = A.to_numpy()
+    else:
+        A_array = np.asarray(A)
+    if A_array.ndim == 1:
+        A_array = A_array[:, None]
+    if A_array.ndim != 2 or not np.all(np.isin(A_array, (0, 1))):
+        raise ValueError('A must be a one- or two-dimensional binary matrix')
+    if treatment_names is None:
+        treatment_names = list(range(A_array.shape[1]))
+    treatment_names = list(treatment_names)
+    if len(treatment_names) != A_array.shape[1]:
+        raise ValueError('treatment_names must match the number of treatments')
+    if len(set(treatment_names)) != len(treatment_names):
+        raise ValueError('treatment_names must be unique')
+
+    if isinstance(X_A, pd.DataFrame):
+        if covariate_names is None:
+            covariate_names = list(X_A.columns)
+        X_array = X_A.to_numpy()
+    else:
+        X_array = np.asarray(X_A)
+    if X_array.ndim != 2 or X_array.shape[0] != A_array.shape[0]:
+        raise ValueError('X_A must be two-dimensional with the same rows as A')
+    try:
+        X_array = np.asarray(X_array, dtype=float)
+    except (TypeError, ValueError) as exc:
+        raise ValueError('X_A must contain numeric covariates') from exc
+    if not np.all(np.isfinite(X_array)):
+        raise ValueError('X_A must contain only finite values')
+    if covariate_names is None:
+        covariate_names = [f'covariate_{j + 1}' for j in range(X_array.shape[1])]
+    covariate_names = list(covariate_names)
+    if len(covariate_names) != X_array.shape[1]:
+        raise ValueError('covariate_names must match the number of covariates')
+    if len(set(covariate_names)) != len(covariate_names):
+        raise ValueError('covariate_names must be unique')
+
+    treatment_lookup = {name: j for j, name in enumerate(treatment_names)}
+    covariate_lookup = {name: j for j, name in enumerate(covariate_names)}
+
+    def resolve_treatment(value):
+        if value in treatment_lookup:
+            return treatment_lookup[value]
+        if isinstance(value, (int, np.integer)) and 0 <= int(value) < A_array.shape[1]:
+            return int(value)
+        raise ValueError(f'Unknown treatment: {value}')
+
+    def resolve_covariate(value):
+        if value in covariate_lookup:
+            return covariate_lookup[value]
+        if isinstance(value, (int, np.integer)) and 0 <= int(value) < X_array.shape[1]:
+            return int(value)
+        raise ValueError(f'Unknown covariate: {value}')
+
+    drops = {}
+    for treatment, covariates in drop_by_treatment.items():
+        j = resolve_treatment(treatment)
+        if j in drops:
+            raise ValueError(f'Treatment {treatment_names[j]} is specified more than once')
+        if isinstance(covariates, (str, bytes)):
+            covariates = [covariates]
+        indices = [resolve_covariate(value) for value in covariates]
+        if len(set(indices)) != len(indices):
+            raise ValueError(
+                f'Duplicate dropped covariates for treatment {treatment_names[j]}'
+            )
+        drops[j] = set(indices)
+
+    penalty_factors = {}
+    for treatment, factors in penalty_factors_by_treatment.items():
+        j = resolve_treatment(treatment)
+        if j in penalty_factors:
+            raise ValueError(f'Treatment {treatment_names[j]} is specified more than once')
+        if not isinstance(factors, Mapping):
+            raise ValueError(
+                f'Penalty factors for treatment {treatment_names[j]} must be a mapping'
+            )
+        resolved = {}
+        for covariate, factor in factors.items():
+            k = resolve_covariate(covariate)
+            try:
+                factor = float(factor)
+            except (TypeError, ValueError) as exc:
+                raise ValueError('Penalty factors must be finite numbers at least 1') from exc
+            if not np.isfinite(factor) or factor < 1:
+                raise ValueError('Penalty factors must be finite numbers at least 1')
+            if k in resolved:
+                raise ValueError(
+                    f'Duplicate penalty factors for covariate {covariate_names[k]}'
+                )
+            resolved[k] = factor
+        penalty_factors[j] = resolved
+
+    for j in set(drops).intersection(penalty_factors):
+        conflict = drops[j].intersection(penalty_factors[j])
+        if conflict:
+            names = [covariate_names[k] for k in sorted(conflict)]
+            raise ValueError(
+                f'Covariates cannot be both dropped and penalized for '
+                f'{treatment_names[j]}: {names}'
+            )
+
+    if pi_hat is None:
+        pi_updated = np.empty(A_array.shape, dtype=float)
+        refit_indices = range(A_array.shape[1])
+    else:
+        pi_updated = np.asarray(pi_hat, dtype=float).copy()
+        if pi_updated.shape != A_array.shape:
+            raise ValueError('pi_hat must have the same shape as A')
+        if (not np.all(np.isfinite(pi_updated))
+                or np.any((pi_updated < 0) | (pi_updated > 1))):
+            raise ValueError('pi_hat must contain finite probabilities in [0, 1]')
+        refit_indices = sorted(set(drops).union(penalty_factors))
+
+    mask_array = None
+    if mask is not None:
+        mask_array = np.asarray(mask, dtype=bool)
+        if mask_array.ndim == 1:
+            mask_array = mask_array[:, None]
+        if mask_array.shape != A_array.shape:
+            raise ValueError('Mask must have the same shape as the treatment matrix')
+
+    ctrl = np.sum(A_array, axis=1) == 0
+    rows = []
+    for j in refit_indices:
+        dropped = drops.get(j, set())
+        retained = [k for k in range(X_array.shape[1]) if k not in dropped]
+        if not retained:
+            raise ValueError(
+                f'Filtering removes every covariate for treatment {treatment_names[j]}'
+            )
+        eligible = ctrl | (A_array[:, j] == 1)
+        if mask_array is not None:
+            eligible &= mask_array[:, j]
+        X_treatment = X_array[:, retained].copy()
+        treatment_penalties = penalty_factors.get(j, {})
+        for position, k in enumerate(retained):
+            factor = treatment_penalties.get(k, 1.0)
+            X_treatment[:, position] /= np.sqrt(factor)
+        degenerate = bool(np.all(np.ptp(X_treatment[eligible], axis=0) == 0))
+        if degenerate:
+            warnings.warn(
+                f'The propensity design for treatment {treatment_names[j]} is '
+                'constant after filtering; its scores fall back to a '
+                'class-weighted prevalence and carry no covariate information.',
+                RuntimeWarning, stacklevel=2,
+            )
+        scores = estimate_propensity_scores(
+            A_array[:, [j]], X_treatment, K=K,
+            ps_model=ps_model, mask=eligible[:, None], clip=None,
+            random_state=random_state, verbose=verbose,
+            class_weight=class_weight, **kwargs,
+        )
+        pi_updated[:, j] = scores[:, 0]
+        rows.append({
+            'treatment': treatment_names[j],
+            'dropped_covariates': [covariate_names[k] for k in sorted(dropped)],
+            'retained_covariates': [covariate_names[k] for k in retained],
+            'penalty_factors': {
+                covariate_names[k]: treatment_penalties[k]
+                for k in retained if treatment_penalties.get(k, 1.0) != 1.0
+            },
+            'n_retained': len(retained),
+            'degenerate_design': degenerate,
+            'score_std': float(np.std(scores[eligible, 0])),
+        })
+
+    if clip is not None:
+        lower, upper = _validate_clip(clip)
+        pi_updated = np.clip(pi_updated, lower, upper)
+    report = pd.DataFrame(rows, columns=[
+        'treatment', 'dropped_covariates', 'retained_covariates',
+        'penalty_factors', 'n_retained', 'degenerate_design', 'score_std',
+    ])
+    return pi_updated, report
+
+
 def cross_fitting(
     Y, A, X, X_A, family='poisson', K=1, glm_alpha=1e-4,
     ps_model='logistic', ps_class_weight='balanced',

causarray-0.0.9/causarray/__about__.py

@@ -0,0 +1 @@
+__version__ = "0.0.9"

{causarray-0.0.8 → causarray-0.0.9}/causarray/__init__.py

@@ -11,13 +11,19 @@ __all__ = [
     'fit_glm', 'fit_glm_fast', 'fit_glm_ondisk',
     'reset_random_seeds', 'fit_gcate', 'fit_gcate_batch',
     'estimate_propensity_scores', 'summarize_propensity_scores',
-    'plot_propensity_scores',
+    'plot_propensity_scores', 'refit_propensity_scores',
+    'summarize_treatment_associations', 'plot_treatment_associations',
+    'align_test_mask',
     ]
 
 
 from causarray.DR_learner import LFC, gcate_lfc_batch, LFC_batch  # ATE, SATE, FC
-from causarray.DR_estimation import estimate_propensity_scores
-from causarray.diagnostics import summarize_propensity_scores, plot_propensity_scores
+from causarray.DR_estimation import estimate_propensity_scores, refit_propensity_scores
+from causarray.diagnostics import (
+    summarize_propensity_scores, plot_propensity_scores,
+    summarize_treatment_associations, plot_treatment_associations,
+    align_test_mask,
+)
 from causarray.gcate_glm import fit_glm
 from causarray.nb_glm_fast import fit_glm_fast, fit_glm_ondisk
 from causarray.utils import prep_causarray_data, reset_random_seeds, comp_size_factor

causarray-0.0.9/causarray/diagnostics.py

@@ -0,0 +1,628 @@
+"""Diagnostics for treatment overlap, propensity scores, and result masks."""
+
+import math
+from typing import Hashable, Optional, Sequence, Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+from scipy import stats
+from sklearn.metrics import brier_score_loss, roc_auc_score
+from statsmodels.stats.multitest import multipletests
+
+
+_TestMask = Union[np.ndarray, pd.DataFrame, pd.Series]
+
+
+def _normalize_test_labels(values, name):
+    raw_labels = pd.Index(values)
+    if raw_labels.hasnans:
+        raise ValueError(f'{name} names must not be missing')
+    labels = pd.Index([str(value) for value in raw_labels], dtype=object)
+    if labels.has_duplicates:
+        duplicates = labels[labels.duplicated()].unique().tolist()
+        raise ValueError(f'{name} names must be unique; duplicates: {duplicates}')
+    return labels
+
+
+def _require_boolean_mask(values):
+    values = np.asarray(values)
+    if values.dtype.kind != 'b':
+        raise ValueError('test_mask must contain only Boolean values')
+    return values
+
+
+def align_test_mask(
+    results: pd.DataFrame,
+    test_mask: _TestMask,
+    treatment_names: Optional[Sequence[Hashable]] = None,
+    gene_names: Optional[Sequence[Hashable]] = None,
+    treatment_col: str = 'trt',
+    gene_col: str = 'gene_names',
+) -> np.ndarray:
+    """Align a treatment-by-gene diagnostic mask to an LFC result table.
+
+    Alignment uses treatment and gene labels rather than row position. This
+    makes it possible to annotate or subset existing causarray results with a
+    post-hoc diagnostic rule without refitting effects or changing their
+    standard errors and p-values. Labels are compared after conversion to
+    strings, matching the labeling convention of batch LFC results.
+
+    Parameters
+    ----------
+    results : DataFrame
+        Result table containing one row per treatment-gene test.
+    test_mask : ndarray, DataFrame, or Series
+        Boolean diagnostic mask. A two-dimensional array requires
+        ``treatment_names`` and ``gene_names``. A DataFrame uses its index as
+        treatment names and columns as gene names unless names are supplied
+        explicitly. A Series must have a two-level ``(treatment, gene)``
+        MultiIndex.
+    treatment_names, gene_names : sequence, optional
+        Labels for the rows and columns of a two-dimensional mask.
+    treatment_col, gene_col : str, optional
+        Columns containing treatment and gene labels in ``results``.
+
+    Returns
+    -------
+    aligned : ndarray of bool, shape (n_tests,)
+        Mask values in the original row order of ``results``.
+
+    Raises
+    ------
+    ValueError
+        If labels are missing or duplicated, the mask is malformed or
+        non-Boolean, or a result test is absent from the mask.
+
+    Notes
+    -----
+    This function only aligns a diagnostic mask. It does not modify the
+    multiple-testing family or recompute adjusted p-values.
+
+    .. versionadded:: 0.0.9
+    """
+    if not isinstance(results, pd.DataFrame):
+        raise ValueError('results must be a pandas DataFrame')
+    missing_columns = {treatment_col, gene_col}.difference(results.columns)
+    if missing_columns:
+        raise ValueError(
+            f'results is missing required columns: {sorted(missing_columns)}'
+        )
+    if results[[treatment_col, gene_col]].isna().any().any():
+        raise ValueError('result treatment and gene labels must not be missing')
+
+    result_keys = pd.MultiIndex.from_arrays([
+        results[treatment_col].astype(str),
+        results[gene_col].astype(str),
+    ])
+    if result_keys.has_duplicates:
+        duplicates = result_keys[result_keys.duplicated()].unique().tolist()
+        raise ValueError(
+            'results contains duplicate treatment-gene tests: '
+            f'{duplicates[:5]}'
+        )
+
+    if isinstance(test_mask, pd.Series):
+        if treatment_names is not None or gene_names is not None:
+            raise ValueError(
+                'treatment_names and gene_names must be omitted for a Series mask'
+            )
+        if (
+            not isinstance(test_mask.index, pd.MultiIndex)
+            or test_mask.index.nlevels != 2
+        ):
+            raise ValueError(
+                'a Series test_mask must have a two-level treatment-gene MultiIndex'
+            )
+        values = _require_boolean_mask(test_mask.to_numpy())
+        raw_treatments = pd.Index(test_mask.index.get_level_values(0))
+        raw_genes = pd.Index(test_mask.index.get_level_values(1))
+        if raw_treatments.hasnans or raw_genes.hasnans:
+            raise ValueError(
+                'test_mask treatment and gene labels must not be missing'
+            )
+        treatment_index = pd.Index(
+            [str(value) for value in raw_treatments], dtype=object,
+        )
+        gene_index = pd.Index(
+            [str(value) for value in raw_genes], dtype=object,
+        )
+        mask_index = pd.MultiIndex.from_arrays([treatment_index, gene_index])
+        if mask_index.has_duplicates:
+            duplicates = mask_index[mask_index.duplicated()].unique().tolist()
+            raise ValueError(
+                'test_mask contains duplicate treatment-gene tests: '
+                f'{duplicates[:5]}'
+            )
+        lookup = pd.Series(values, index=mask_index)
+    else:
+        if isinstance(test_mask, pd.DataFrame):
+            values = _require_boolean_mask(test_mask.to_numpy())
+            if treatment_names is None:
+                treatment_names = list(test_mask.index)
+            if gene_names is None:
+                gene_names = list(test_mask.columns)
+        else:
+            values = _require_boolean_mask(test_mask)
+        if values.ndim != 2:
+            raise ValueError('test_mask must be two-dimensional')
+        if treatment_names is None or gene_names is None:
+            raise ValueError(
+                'treatment_names and gene_names are required for an array mask'
+            )
+        treatments = _normalize_test_labels(treatment_names, 'treatment')
+        genes = _normalize_test_labels(gene_names, 'gene')
+        expected_shape = (len(treatments), len(genes))
+        if values.shape != expected_shape:
+            raise ValueError(
+                f'test_mask has shape {values.shape}; expected {expected_shape}'
+            )
+        mask_index = pd.MultiIndex.from_product([treatments, genes])
+        lookup = pd.Series(values.ravel(), index=mask_index)
+
+    aligned = lookup.reindex(result_keys)
+    if aligned.isna().any():
+        missing = result_keys[aligned.isna().to_numpy()].unique().tolist()
+        raise ValueError(
+            'test_mask is missing treatment-gene tests required by results: '
+            f'{missing[:5]}'
+        )
+    return aligned.to_numpy(dtype=bool)
+
+
+def _coerce_named_matrix(values, names, name, generated_prefix):
+    if isinstance(values, pd.DataFrame):
+        if names is None:
+            names = list(values.columns)
+        values = values.to_numpy()
+    else:
+        values = np.asarray(values)
+    if values.ndim == 1:
+        values = values[:, None]
+    if values.ndim != 2:
+        raise ValueError(f'{name} must be one- or two-dimensional')
+    if names is None:
+        names = [f'{generated_prefix}{j + 1}' for j in range(values.shape[1])]
+    names = list(names)
+    if len(names) != values.shape[1]:
+        raise ValueError(f'{name} names must match the number of columns')
+    if len(set(names)) != len(names):
+        raise ValueError(f'{name} names must be unique')
+    return values, names
+
+
+def summarize_treatment_associations(
+    A, Z, treatment_names=None, covariate_names=None, covariate_types=None,
+    bh_scope='global',
+):
+    """Summarize covariate association with each treatment.
+
+    Each treatment is compared with shared all-zero controls; rows assigned to
+    other treatments are excluded. Spearman p-values are adjusted with the
+    Benjamini-Hochberg procedure over the family selected by ``bh_scope``. The
+    returned statistics are diagnostics and do not automatically identify
+    variables that should be removed from an adjustment set.
+
+    Parameters
+    ----------
+    A : array-like, shape (n, a)
+        Binary treatment indicator matrix. DataFrame column names are retained.
+    Z : array-like, shape (n, q)
+        Observed covariates and/or estimated latent factors to diagnose.
+    treatment_names : sequence, optional
+        Treatment labels. Inferred from a DataFrame or generated when omitted.
+    covariate_names : sequence, optional
+        Covariate labels. Inferred from a DataFrame or generated when omitted.
+    covariate_types : sequence, optional
+        Labels such as ``'observed'`` and ``'latent'``. Defaults to
+        ``'observed'`` for every column.
+    bh_scope : {'global', 'per_treatment'}, optional
+        Multiple-testing family. ``'global'`` (the default) adjusts across every
+        finite treatment-by-covariate test at once, which with many treatments
+        is a large and correspondingly conservative family.
+        ``'per_treatment'`` adjusts within each treatment's own block, which is
+        the right choice when each treatment is screened on its own.
+
+    Returns
+    -------
+    summary : DataFrame
+        Long-form table containing ``spearman_rho``, ``pvalue``, BH-adjusted
+        ``padj``, ``n_tests_in_family``, and ``standardized_mean_difference``
+        for every treatment-by-covariate pair. ``n_tests_in_family`` records how
+        many tests entered that row's adjustment, so the correction is
+        self-describing.
+
+    .. versionadded:: 0.0.9
+    """
+    if bh_scope not in ('global', 'per_treatment'):
+        raise ValueError("bh_scope must be 'global' or 'per_treatment'")
+    A, treatment_names = _coerce_named_matrix(
+        A, treatment_names, 'treatment', 'treatment_',
+    )
+    Z, covariate_names = _coerce_named_matrix(
+        Z, covariate_names, 'covariate', 'covariate_',
+    )
+    if A.shape[0] != Z.shape[0]:
+        raise ValueError('A and Z must have the same number of rows')
+    if not np.all(np.isin(A, (0, 1))):
+        raise ValueError('A must contain only binary treatment indicators')
+    try:
+        Z = np.asarray(Z, dtype=float)
+    except (TypeError, ValueError) as exc:
+        raise ValueError('Z must contain numeric covariates') from exc
+    if not np.all(np.isfinite(Z)):
+        raise ValueError('Z must contain only finite values')
+
+    if covariate_types is None:
+        covariate_types = ['observed'] * Z.shape[1]
+    covariate_types = list(covariate_types)
+    if len(covariate_types) != Z.shape[1]:
+        raise ValueError('covariate_types must match the number of covariates')
+
+    ctrl = np.sum(A, axis=1) == 0
+    if not np.any(ctrl):
+        raise ValueError('At least one all-zero control row is required')
+
+    rows = []
+    for j, treatment in enumerate(treatment_names):
+        case = A[:, j] == 1
+        if not np.any(case):
+            raise ValueError(f'Treatment {treatment} has no treated rows')
+        eligible = ctrl | case
+        y = A[eligible, j]
+        for k, covariate in enumerate(covariate_names):
+            values = Z[eligible, k]
+            is_constant = bool(np.ptp(values) == 0)
+            if is_constant:
+                rho, pvalue = np.nan, np.nan
+            else:
+                rho, pvalue = stats.spearmanr(values, y)
+
+            control_values = Z[ctrl, k]
+            treated_values = Z[case, k]
+            if len(control_values) < 2 or len(treated_values) < 2:
+                smd = np.nan
+            else:
+                pooled_sd = np.sqrt(
+                    (np.var(control_values, ddof=1)
+                     + np.var(treated_values, ddof=1)) / 2
+                )
+                if pooled_sd == 0:
+                    smd = np.nan
+                else:
+                    smd = (
+                        np.mean(treated_values) - np.mean(control_values)
+                    ) / pooled_sd
+            rows.append({
+                'treatment': treatment,
+                'covariate': covariate,
+                'covariate_type': covariate_types[k],
+                'n_control': int(np.sum(ctrl)),
+                'n_treated': int(np.sum(case)),
+                'spearman_rho': float(rho),
+                'pvalue': float(pvalue),
+                'padj': np.nan,
+                'standardized_mean_difference': float(smd),
+                'constant': is_constant,
+            })
+
+    summary = pd.DataFrame(rows)
+    summary['n_tests_in_family'] = 0
+    finite = np.isfinite(summary['pvalue'].to_numpy())
+    if bh_scope == 'global':
+        families = [np.ones(len(summary), dtype=bool)]
+    else:
+        families = [
+            (summary['treatment'] == treatment).to_numpy()
+            for treatment in treatment_names
+        ]
+    for family in families:
+        adjust = family & finite
+        summary.loc[family, 'n_tests_in_family'] = int(adjust.sum())
+        if np.any(adjust):
+            summary.loc[adjust, 'padj'] = multipletests(
+                summary.loc[adjust, 'pvalue'], method='fdr_bh',
+            )[1]
+    return summary
+
+
+def plot_treatment_associations(
+    summary, value='spearman_rho', treatments=None, covariates=None,
+    alpha=None, vmax=None, ax=None,
+):
+    """Plot a heatmap of treatment-by-covariate associations.
+
+    Parameters
+    ----------
+    summary : DataFrame
+        Output from :func:`summarize_treatment_associations`.
+    value : {'spearman_rho', 'standardized_mean_difference'}, optional
+        Statistic displayed in the heatmap.
+    treatments, covariates : sequence, optional
+        Ordered subsets to display.
+    alpha : float or None, optional
+        If provided, mark cells whose adjusted p-value is at most ``alpha``. No
+        significance threshold is applied by default.
+    vmax : float or None, optional
+        Symmetric colour limit, so the heatmap spans ``(-vmax, vmax)``. When
+        omitted, ``'spearman_rho'`` uses the fixed range ``(-1, 1)`` given by
+        its natural bounds, which keeps panels comparable across subsets, while
+        the unbounded ``'standardized_mean_difference'`` scales to the largest
+        finite magnitude on display.
+    ax : matplotlib Axes, optional
+        Axes on which to draw the heatmap.
+
+    Returns
+    -------
+    fig, ax : matplotlib Figure and Axes
+        The populated figure and axes.
+
+    .. versionadded:: 0.0.9
+    """
+    if value not in ('spearman_rho', 'standardized_mean_difference'):
+        raise ValueError(
+            "value must be 'spearman_rho' or 'standardized_mean_difference'"
+        )
+    if not isinstance(summary, pd.DataFrame):
+        raise ValueError('summary must be a pandas DataFrame')
+    required = {'treatment', 'covariate', value, 'padj'}
+    missing = required.difference(summary.columns)
+    if missing:
+        raise ValueError(f'summary is missing required columns: {sorted(missing)}')
+    if alpha is not None and not 0 < alpha <= 1:
+        raise ValueError('alpha must be None or lie in (0, 1]')
+    if vmax is not None and not (np.isfinite(vmax) and vmax > 0):
+        raise ValueError('vmax must be None or a finite positive number')
+
+    all_treatments = list(dict.fromkeys(summary['treatment']))
+    all_covariates = list(dict.fromkeys(summary['covariate']))
+    treatments = all_treatments if treatments is None else list(treatments)
+    covariates = all_covariates if covariates is None else list(covariates)
+    unknown_treatments = set(treatments).difference(all_treatments)
+    unknown_covariates = set(covariates).difference(all_covariates)
+    if unknown_treatments:
+        raise ValueError(f'Unknown treatments: {sorted(unknown_treatments)}')
+    if unknown_covariates:
+        raise ValueError(f'Unknown covariates: {sorted(unknown_covariates)}')
+    if not treatments or not covariates:
+        raise ValueError('At least one treatment and covariate must be selected')
+
+    subset = summary[
+        summary['treatment'].isin(treatments)
+        & summary['covariate'].isin(covariates)
+    ]
+    if subset.duplicated(['treatment', 'covariate']).any():
+        raise ValueError('summary contains duplicate treatment-covariate pairs')
+    matrix = subset.pivot(
+        index='treatment', columns='covariate', values=value,
+    ).reindex(index=treatments, columns=covariates)
+    adjusted = subset.pivot(
+        index='treatment', columns='covariate', values='padj',
+    ).reindex(index=treatments, columns=covariates)
+
+    if ax is None:
+        fig, ax = plt.subplots(
+            figsize=(max(6, 0.7 * len(covariates)),
+                     max(3, 0.3 * len(treatments))),
+        )
+    else:
+        fig = ax.figure
+    values = matrix.to_numpy(dtype=float)
+    if vmax is not None:
+        limit = float(vmax)
+    elif value == 'spearman_rho':
+        limit = 1.0
+    else:
+        finite_values = np.abs(values[np.isfinite(values)])
+        limit = float(np.max(finite_values)) if finite_values.size else 1.0
+        if limit == 0:
+            limit = 1.0
+    image = ax.imshow(values, aspect='auto', cmap='coolwarm',
+                      vmin=-limit, vmax=limit)
+    ax.set_xticks(np.arange(len(covariates)))
+    ax.set_xticklabels(covariates, rotation=45, ha='right')
+    ax.set_yticks(np.arange(len(treatments)))
+    ax.set_yticklabels(treatments)
+    ax.set_xlabel('Covariate or latent factor')
+    ax.set_ylabel('Treatment')
+    colorbar = fig.colorbar(image, ax=ax)
+    colorbar.set_label(
+        'Spearman correlation' if value == 'spearman_rho'
+        else 'Standardized mean difference'
+    )
+    if alpha is not None:
+        padj = adjusted.to_numpy(dtype=float)
+        for row, col in np.argwhere(np.isfinite(padj) & (padj <= alpha)):
+            ax.text(col, row, '*', ha='center', va='center', color='black')
+    fig.tight_layout()
+    return fig, ax
+
+
+def _coerce_treatment_inputs(A, pi_hat, treatment_names=None):
+    if isinstance(A, pd.DataFrame):
+        if treatment_names is None:
+            treatment_names = list(A.columns)
+        A = A.values
+    else:
+        A = np.asarray(A)
+    if A.ndim == 1:
+        A = A[:, None]
+
+    pi_hat = np.asarray(pi_hat, dtype=float)
+    if pi_hat.ndim == 1:
+        pi_hat = pi_hat[:, None]
+    if A.shape != pi_hat.shape:
+        raise ValueError('A and pi_hat must have the same shape')
+    if not np.all(np.isin(A, (0, 1))):
+        raise ValueError('A must contain only binary treatment indicators')
+    if np.any(~np.isfinite(pi_hat)) or np.any((pi_hat < 0) | (pi_hat > 1)):
+        raise ValueError('pi_hat must contain finite probabilities in [0, 1]')
+
+    if treatment_names is None:
+        treatment_names = list(range(A.shape[1]))
+    if len(treatment_names) != A.shape[1]:
+        raise ValueError('treatment_names must match the number of treatments')
+    return A, pi_hat, list(treatment_names)
+
+
+def _effective_sample_size(weights):
+    weights = np.asarray(weights, dtype=float)
+    denom = np.sum(weights ** 2)
+    return float(np.sum(weights) ** 2 / denom) if denom > 0 else np.nan
+
+
+def summarize_propensity_scores(
+    A, pi_hat, treatment_names=None, overlap_bounds=(0.05, 0.95),
+    clip_bounds=(0.01, 0.99), bins=40,
+):
+    """Summarize overlap and inverse-weight stability for each treatment.
+
+    Other perturbations are excluded from a treatment's diagnostic comparison;
+    each row compares that treatment with shared all-zero controls.
+
+    Parameters
+    ----------
+    A, pi_hat, treatment_names, overlap_bounds, bins
+        Treatments, scores, labels, the reference overlap window, and the
+        histogram resolution used for ``overlap_ratio``.
+    clip_bounds : tuple(float, float) or None, optional
+        The bounds the caller **already applied** to ``pi_hat``, defaulting to
+        the ``LFC(ps_clip=(0.01, 0.99))`` default. ``clipped_fraction`` counts
+        scores sitting on those bounds. Pass ``None`` for raw, unclipped scores;
+        ``clipped_fraction`` is then ``NaN`` rather than a misleading ``0.0``,
+        since clipping cannot be inferred from the values alone.
+
+    .. versionchanged:: 0.0.9
+       ``clip_bounds=None`` reports ``NaN`` instead of ``0.0``.
+    """
+    A, pi_hat, treatment_names = _coerce_treatment_inputs(
+        A, pi_hat, treatment_names)
+    lo, hi = overlap_bounds
+    if not 0 <= lo < hi <= 1:
+        raise ValueError('overlap_bounds must satisfy 0 <= lower < upper <= 1')
+    if bins < 2:
+        raise ValueError('bins must be at least 2')
+
+    ctrl = np.sum(A, axis=1) == 0
+    if not np.any(ctrl):
+        raise ValueError('At least one all-zero control row is required')
+    rows = []
+    for j, name in enumerate(treatment_names):
+        case = A[:, j] == 1
+        if not np.any(case):
+            raise ValueError(f'Treatment {name} has no treated rows')
+        eligible = ctrl | case
+        y = A[eligible, j].astype(int)
+        p = pi_hat[eligible, j]
+        p_ctrl, p_case = p[y == 0], p[y == 1]
+
+        h_ctrl, edges = np.histogram(p_ctrl, bins=bins, range=(0, 1))
+        h_case, _ = np.histogram(p_case, bins=edges)
+        h_ctrl = h_ctrl / h_ctrl.sum() if h_ctrl.sum() else h_ctrl
+        h_case = h_case / h_case.sum() if h_case.sum() else h_case
+        overlap_ratio = float(np.minimum(h_ctrl, h_case).sum())
+
+        eps = np.finfo(float).eps
+        ess_ctrl = _effective_sample_size(1 / np.clip(1 - p_ctrl, eps, None))
+        ess_case = _effective_sample_size(1 / np.clip(p_case, eps, None))
+        if clip_bounds is None:
+            clipped_fraction = np.nan
+        else:
+            clipped = (
+                np.isclose(p, clip_bounds[0]) | np.isclose(p, clip_bounds[1])
+            )
+            clipped_fraction = float(clipped.mean())
+
+        rows.append({
+            'treatment': name,
+            'n_control': int((y == 0).sum()),
+            'n_treated': int((y == 1).sum()),
+            'prevalence': float(y.mean()),
+            'overlap_ratio': overlap_ratio,
+            'auc': float(roc_auc_score(y, p)),
+            'brier_score': float(brier_score_loss(y, p)),
+            'outside_overlap_fraction': float(np.mean((p < lo) | (p > hi))),
+            'clipped_fraction': clipped_fraction,
+            'ess_control': ess_ctrl,
+            'ess_treated': ess_case,
+            'ess_control_fraction': ess_ctrl / max(len(p_ctrl), 1),
+            'ess_treated_fraction': ess_case / max(len(p_case), 1),
+            'score_q01': float(np.quantile(p, 0.01)),
+            'score_median': float(np.median(p)),
+            'score_q99': float(np.quantile(p, 0.99)),
+        })
+    return pd.DataFrame(rows)
+
+
+def plot_propensity_scores(
+    A, pi_hat, treatments=None, treatment_names=None, overlap_bounds=(0.05, 0.95),
+    bins=40, max_panels=4, axes=None, clip_bounds=(0.01, 0.99),
+):
+    """Plot propensity distributions for treatment and control cells.
+
+    ``clip_bounds`` is forwarded to :func:`summarize_propensity_scores` for the
+    returned summary; pass ``None`` when ``pi_hat`` holds raw, unclipped scores.
+
+    .. versionchanged:: 0.0.9
+       ``clip_bounds`` is no longer fixed at its default.
+    """
+    A, pi_hat, treatment_names = _coerce_treatment_inputs(
+        A, pi_hat, treatment_names)
+    summary = summarize_propensity_scores(
+        A, pi_hat, treatment_names=treatment_names,
+        overlap_bounds=overlap_bounds, clip_bounds=clip_bounds, bins=bins,
+    )
+
+    if treatments is None:
+        indices = list(range(min(A.shape[1], max_panels)))
+    else:
+        indices = []
+        for treatment in treatments:
+            if treatment in treatment_names:
+                indices.append(treatment_names.index(treatment))
+            else:
+                index = int(treatment)
+                if index < 0 or index >= A.shape[1]:
+                    raise ValueError(f'Unknown treatment: {treatment}')
+                indices.append(index)
+    if not indices:
+        raise ValueError('At least one treatment must be selected')
+
+    if axes is None:
+        ncols = min(2, len(indices))
+        nrows = math.ceil(len(indices) / ncols)
+        fig, axes = plt.subplots(
+            nrows, ncols, figsize=(5 * ncols, 3.5 * nrows), squeeze=False,
+        )
+        axes_flat = axes.ravel()
+    else:
+        axes_flat = np.asarray(axes, dtype=object).ravel()
+        if len(axes_flat) < len(indices):
+            raise ValueError('Not enough axes for the selected treatments')
+        fig = axes_flat[0].figure
+
+    ctrl = np.sum(A, axis=1) == 0
+    for ax, j in zip(axes_flat, indices):
+        case = A[:, j] == 1
+        ax.hist(
+            pi_hat[ctrl, j], bins=bins, range=(0, 1), density=True,
+            histtype='step', linewidth=1.6, label='control', color='#4c78a8',
+        )
+        ax.hist(
+            pi_hat[case, j], bins=bins, range=(0, 1), density=True,
+            histtype='step', linewidth=1.6, label=str(treatment_names[j]),
+            color='#e45756',
+        )
+        ax.axvline(overlap_bounds[0], color='#777777', linestyle='--', linewidth=0.8)
+        ax.axvline(overlap_bounds[1], color='#777777', linestyle='--', linewidth=0.8)
+        overlap = summary.loc[j, 'overlap_ratio']
+        ess = summary.loc[j, 'ess_treated_fraction']
+        ax.set_title(f'{treatment_names[j]} | overlap={overlap:.2f}, ESS={ess:.2f}')
+        ax.set_xlim(0, 1)
+        ax.set_xlabel('Estimated propensity score')
+        ax.set_ylabel('Density')
+        ax.legend(frameon=False)
+    for ax in axes_flat[len(indices):]:
+        ax.set_visible(False)
+    fig.tight_layout()
+    return fig, axes, summary

causarray-0.0.8/causarray/__about__.py

@@ -1 +0,0 @@
-__version__ = "0.0.8"

causarray-0.0.8/causarray/diagnostics.py

@@ -1,172 +0,0 @@
-"""Diagnostics for treatment overlap and propensity-score quality."""
-
-import math
-
-import matplotlib.pyplot as plt
-import numpy as np
-import pandas as pd
-from sklearn.metrics import brier_score_loss, roc_auc_score
-
-
-def _coerce_treatment_inputs(A, pi_hat, treatment_names=None):
-    if isinstance(A, pd.DataFrame):
-        if treatment_names is None:
-            treatment_names = list(A.columns)
-        A = A.values
-    else:
-        A = np.asarray(A)
-    if A.ndim == 1:
-        A = A[:, None]
-
-    pi_hat = np.asarray(pi_hat, dtype=float)
-    if pi_hat.ndim == 1:
-        pi_hat = pi_hat[:, None]
-    if A.shape != pi_hat.shape:
-        raise ValueError('A and pi_hat must have the same shape')
-    if not np.all(np.isin(A, (0, 1))):
-        raise ValueError('A must contain only binary treatment indicators')
-    if np.any(~np.isfinite(pi_hat)) or np.any((pi_hat < 0) | (pi_hat > 1)):
-        raise ValueError('pi_hat must contain finite probabilities in [0, 1]')
-
-    if treatment_names is None:
-        treatment_names = list(range(A.shape[1]))
-    if len(treatment_names) != A.shape[1]:
-        raise ValueError('treatment_names must match the number of treatments')
-    return A, pi_hat, list(treatment_names)
-
-
-def _effective_sample_size(weights):
-    weights = np.asarray(weights, dtype=float)
-    denom = np.sum(weights ** 2)
-    return float(np.sum(weights) ** 2 / denom) if denom > 0 else np.nan
-
-
-def summarize_propensity_scores(
-    A, pi_hat, treatment_names=None, overlap_bounds=(0.05, 0.95),
-    clip_bounds=(0.01, 0.99), bins=40,
-):
-    """Summarize overlap and inverse-weight stability for each treatment.
-
-    Other perturbations are excluded from a treatment's diagnostic comparison;
-    each row compares that treatment with shared all-zero controls.
-    """
-    A, pi_hat, treatment_names = _coerce_treatment_inputs(
-        A, pi_hat, treatment_names)
-    lo, hi = overlap_bounds
-    if not 0 <= lo < hi <= 1:
-        raise ValueError('overlap_bounds must satisfy 0 <= lower < upper <= 1')
-    if bins < 2:
-        raise ValueError('bins must be at least 2')
-
-    ctrl = np.sum(A, axis=1) == 0
-    if not np.any(ctrl):
-        raise ValueError('At least one all-zero control row is required')
-    rows = []
-    for j, name in enumerate(treatment_names):
-        case = A[:, j] == 1
-        if not np.any(case):
-            raise ValueError(f'Treatment {name} has no treated rows')
-        eligible = ctrl | case
-        y = A[eligible, j].astype(int)
-        p = pi_hat[eligible, j]
-        p_ctrl, p_case = p[y == 0], p[y == 1]
-
-        h_ctrl, edges = np.histogram(p_ctrl, bins=bins, range=(0, 1))
-        h_case, _ = np.histogram(p_case, bins=edges)
-        h_ctrl = h_ctrl / h_ctrl.sum() if h_ctrl.sum() else h_ctrl
-        h_case = h_case / h_case.sum() if h_case.sum() else h_case
-        overlap_ratio = float(np.minimum(h_ctrl, h_case).sum())
-
-        eps = np.finfo(float).eps
-        ess_ctrl = _effective_sample_size(1 / np.clip(1 - p_ctrl, eps, None))
-        ess_case = _effective_sample_size(1 / np.clip(p_case, eps, None))
-        clipped = np.zeros(p.shape, dtype=bool)
-        if clip_bounds is not None:
-            clipped = np.isclose(p, clip_bounds[0]) | np.isclose(p, clip_bounds[1])
-
-        rows.append({
-            'treatment': name,
-            'n_control': int((y == 0).sum()),
-            'n_treated': int((y == 1).sum()),
-            'prevalence': float(y.mean()),
-            'overlap_ratio': overlap_ratio,
-            'auc': float(roc_auc_score(y, p)),
-            'brier_score': float(brier_score_loss(y, p)),
-            'outside_overlap_fraction': float(np.mean((p < lo) | (p > hi))),
-            'clipped_fraction': float(clipped.mean()),
-            'ess_control': ess_ctrl,
-            'ess_treated': ess_case,
-            'ess_control_fraction': ess_ctrl / max(len(p_ctrl), 1),
-            'ess_treated_fraction': ess_case / max(len(p_case), 1),
-            'score_q01': float(np.quantile(p, 0.01)),
-            'score_median': float(np.median(p)),
-            'score_q99': float(np.quantile(p, 0.99)),
-        })
-    return pd.DataFrame(rows)
-
-
-def plot_propensity_scores(
-    A, pi_hat, treatments=None, treatment_names=None, overlap_bounds=(0.05, 0.95),
-    bins=40, max_panels=4, axes=None,
-):
-    """Plot propensity distributions for treatment and control cells."""
-    A, pi_hat, treatment_names = _coerce_treatment_inputs(
-        A, pi_hat, treatment_names)
-    summary = summarize_propensity_scores(
-        A, pi_hat, treatment_names=treatment_names,
-        overlap_bounds=overlap_bounds, bins=bins,
-    )
-
-    if treatments is None:
-        indices = list(range(min(A.shape[1], max_panels)))
-    else:
-        indices = []
-        for treatment in treatments:
-            if treatment in treatment_names:
-                indices.append(treatment_names.index(treatment))
-            else:
-                index = int(treatment)
-                if index < 0 or index >= A.shape[1]:
-                    raise ValueError(f'Unknown treatment: {treatment}')
-                indices.append(index)
-    if not indices:
-        raise ValueError('At least one treatment must be selected')
-
-    if axes is None:
-        ncols = min(2, len(indices))
-        nrows = math.ceil(len(indices) / ncols)
-        fig, axes = plt.subplots(
-            nrows, ncols, figsize=(5 * ncols, 3.5 * nrows), squeeze=False,
-        )
-        axes_flat = axes.ravel()
-    else:
-        axes_flat = np.asarray(axes, dtype=object).ravel()
-        if len(axes_flat) < len(indices):
-            raise ValueError('Not enough axes for the selected treatments')
-        fig = axes_flat[0].figure
-
-    ctrl = np.sum(A, axis=1) == 0
-    for ax, j in zip(axes_flat, indices):
-        case = A[:, j] == 1
-        ax.hist(
-            pi_hat[ctrl, j], bins=bins, range=(0, 1), density=True,
-            histtype='step', linewidth=1.6, label='control', color='#4c78a8',
-        )
-        ax.hist(
-            pi_hat[case, j], bins=bins, range=(0, 1), density=True,
-            histtype='step', linewidth=1.6, label=str(treatment_names[j]),
-            color='#e45756',
-        )
-        ax.axvline(overlap_bounds[0], color='#777777', linestyle='--', linewidth=0.8)
-        ax.axvline(overlap_bounds[1], color='#777777', linestyle='--', linewidth=0.8)
-        overlap = summary.loc[j, 'overlap_ratio']
-        ess = summary.loc[j, 'ess_treated_fraction']
-        ax.set_title(f'{treatment_names[j]} | overlap={overlap:.2f}, ESS={ess:.2f}')
-        ax.set_xlim(0, 1)
-        ax.set_xlabel('Estimated propensity score')
-        ax.set_ylabel('Density')
-        ax.legend(frameon=False)
-    for ax in axes_flat[len(indices):]:
-        ax.set_visible(False)
-    fig.tight_layout()
-    return fig, axes, summary