causarray 0.0.6__tar.gz → 0.0.7__tar.gz

{causarray-0.0.6 → causarray-0.0.7}/.gitignore

@@ -1,6 +1,9 @@
 # Planning documents
 plan/
 
+# Local Codex project guidance
+/AGENTS.md
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
@@ -170,11 +173,15 @@ cython_debug/
 tests/benchmark_adamson.py
 
 causarray/___*.py
-/docs/source/tutorial/adamson
+
+# Adamson tutorial is currently a local analysis and is not part of the
+# committed documentation.
+/docs/source/tutorial/adamson/
 
 # Replogle tutorial: large data files (regenerate with prep_tutorial_data.py or
 # download from the project data repository; replogle_subset.h5ad is ~2 GB)
 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
 

{causarray-0.0.6 → causarray-0.0.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: causarray
-Version: 0.0.6
+Version: 0.0.7
 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>
@@ -30,7 +30,7 @@ Description-Content-Type: text/markdown
 
 # causarray
 
-Advances in single-cell sequencing and CRISPR technologies have enabled detailed case-control comparisons and experimental perturbations at single-cell resolution. However, uncovering causal relationships in observational genomic data remains challenging due to selection bias and inadequate adjustment for unmeasured confounders, particularly in heterogeneous datasets. To address these challenges, we introduce `causarray` [Du25], a doubly robust causal inference framework for analyzing array-based genomic data at both bulk-cell and single-cell levels. `causarray` integrates a generalized confounder adjustment method to account for unmeasured confounders and employs semiparametric inference with flexible machine learning techniques to ensure robust statistical estimation of treatment effects.
+Advances in single-cell sequencing and CRISPR technologies have enabled detailed case-control comparisons and experimental perturbations at single-cell resolution. However, uncovering causal relationships in observational genomic data remains challenging due to selection bias and inadequate adjustment for unmeasured confounders, particularly in heterogeneous datasets. To address these challenges, we introduce `causarray` [Du26], a doubly robust causal inference framework for analyzing array-based genomic data at both bulk-cell and single-cell levels. `causarray` integrates a generalized confounder adjustment method to account for unmeasured confounders and employs semiparametric inference with flexible machine learning techniques to ensure robust statistical estimation of treatment effects.
 
 
 ## Usage
@@ -59,10 +59,19 @@ The documentation and tutorials using both `Python` and `R` are available at [ca
 
 
 
-## Batch fitting for large-scale screens
+## Tutorials
 
-For screens with hundreds to thousands of perturbations, use the batch API
-so that peak memory is bounded by one batch at a time:
+| Tutorial | Language | Description | Link |
+|----------|----------|-------------|------|
+| Perturb-seq [Jin20] | Python | CRISPR screen analysis on excitatory neurons | [Notebook](https://causarray.readthedocs.io/en/latest/tutorial/perturbseq/perturbseq-py.html) |
+| Perturb-seq [Jin20] | R | Same analysis using `reticulate` | [Notebook](https://causarray.readthedocs.io/en/latest/tutorial/perturbseq/perturbseq-r.html) |
+| Genome-wide CRISPRi screen [Replogle22] | Python | Batch fitting on 200 perturbations from a K562 genome-wide CRISPRi screen | [Notebook](https://causarray.readthedocs.io/en/latest/tutorial/replogle/replogle-py.html) |
+| Case-control: SEA-AD [Gabitto24] | Python | Causal inference on observational single-cell data (Alzheimer's disease) | [Notebook](https://causarray.readthedocs.io/en/latest/tutorial/case_control/sea_ad_case_control.html) |
+
+### Batch fitting API
+
+For screens with hundreds to thousands of perturbations, use `gcate_lfc_batch` so
+that peak memory is bounded by one batch at a time:
 
 ```python
 from causarray import gcate_lfc_batch
@@ -77,16 +86,12 @@ df_res = gcate_lfc_batch(
 )
 ```
 
-See the [Replogle-E-K562 tutorial](https://causarray.readthedocs.io/en/latest/)
+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.
 
 ## Changelog
 
-- [x] (2025-01-30) Python package released on PyPI
-- [x] (2025-02-01) Code for reproducing figures in paper
-- [x] (2025-02-02) Tutorial for Python and R
-- [x] (2026-05-31) Batch fitting API (`gcate_lfc_batch`) for large-scale screens
-- [x] (2026-05-31) Documentation at [causarray.readthedocs.io](https://causarray.readthedocs.io/en/latest/)
+See [CHANGELOG](https://causarray.readthedocs.io/en/latest/changelog.html) for a full version history.
 
 
 <!-- 
@@ -128,4 +133,10 @@ rmarkdown::render("perturbseq.Rmd", rmarkdown::md_document(variant = "markdown_g
 
 
 ## References
-[Du25] Jin-Hong Du, Maya Shen, Hansruedi Mathys, and Kathryn Roeder (2025). Causal differential expression analysis under unmeasured confounders with causarray. bioRxiv, 2025-01.
+[Du26] Jin-Hong Du, Maya Shen, Hansruedi Mathys, and Kathryn Roeder. "Uncovering causal relationships in single cell omic studies with causarray". In: Briefings in Bioinformatics (2026).
+
+[Gabitto24] Mariano I. Gabitto et al. "Integrated multimodal cell atlas of Alzheimer's disease". In: Nature Neuroscience (2024).
+
+[Jin20] Xin Jin et al. "In vivo Perturb-seq reveals neuronal and glial abnormalities associated with autism risk genes". In: Nature Neuroscience (2020).
+
+[Replogle22] Joseph M. Replogle et al. "Mapping information-rich genotype-phenotype landscapes with genome-scale Perturb-seq". In: Cell (2022).

{causarray-0.0.6 → causarray-0.0.7}/README.md

@@ -5,7 +5,7 @@
 
 # causarray
 
-Advances in single-cell sequencing and CRISPR technologies have enabled detailed case-control comparisons and experimental perturbations at single-cell resolution. However, uncovering causal relationships in observational genomic data remains challenging due to selection bias and inadequate adjustment for unmeasured confounders, particularly in heterogeneous datasets. To address these challenges, we introduce `causarray` [Du25], a doubly robust causal inference framework for analyzing array-based genomic data at both bulk-cell and single-cell levels. `causarray` integrates a generalized confounder adjustment method to account for unmeasured confounders and employs semiparametric inference with flexible machine learning techniques to ensure robust statistical estimation of treatment effects.
+Advances in single-cell sequencing and CRISPR technologies have enabled detailed case-control comparisons and experimental perturbations at single-cell resolution. However, uncovering causal relationships in observational genomic data remains challenging due to selection bias and inadequate adjustment for unmeasured confounders, particularly in heterogeneous datasets. To address these challenges, we introduce `causarray` [Du26], a doubly robust causal inference framework for analyzing array-based genomic data at both bulk-cell and single-cell levels. `causarray` integrates a generalized confounder adjustment method to account for unmeasured confounders and employs semiparametric inference with flexible machine learning techniques to ensure robust statistical estimation of treatment effects.
 
 
 ## Usage
@@ -34,10 +34,19 @@ The documentation and tutorials using both `Python` and `R` are available at [ca
 
 
 
-## Batch fitting for large-scale screens
+## Tutorials
 
-For screens with hundreds to thousands of perturbations, use the batch API
-so that peak memory is bounded by one batch at a time:
+| Tutorial | Language | Description | Link |
+|----------|----------|-------------|------|
+| Perturb-seq [Jin20] | Python | CRISPR screen analysis on excitatory neurons | [Notebook](https://causarray.readthedocs.io/en/latest/tutorial/perturbseq/perturbseq-py.html) |
+| Perturb-seq [Jin20] | R | Same analysis using `reticulate` | [Notebook](https://causarray.readthedocs.io/en/latest/tutorial/perturbseq/perturbseq-r.html) |
+| Genome-wide CRISPRi screen [Replogle22] | Python | Batch fitting on 200 perturbations from a K562 genome-wide CRISPRi screen | [Notebook](https://causarray.readthedocs.io/en/latest/tutorial/replogle/replogle-py.html) |
+| Case-control: SEA-AD [Gabitto24] | Python | Causal inference on observational single-cell data (Alzheimer's disease) | [Notebook](https://causarray.readthedocs.io/en/latest/tutorial/case_control/sea_ad_case_control.html) |
+
+### Batch fitting API
+
+For screens with hundreds to thousands of perturbations, use `gcate_lfc_batch` so
+that peak memory is bounded by one batch at a time:
 
 ```python
 from causarray import gcate_lfc_batch
@@ -52,16 +61,12 @@ df_res = gcate_lfc_batch(
 )
 ```
 
-See the [Replogle-E-K562 tutorial](https://causarray.readthedocs.io/en/latest/)
+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.
 
 ## Changelog
 
-- [x] (2025-01-30) Python package released on PyPI
-- [x] (2025-02-01) Code for reproducing figures in paper
-- [x] (2025-02-02) Tutorial for Python and R
-- [x] (2026-05-31) Batch fitting API (`gcate_lfc_batch`) for large-scale screens
-- [x] (2026-05-31) Documentation at [causarray.readthedocs.io](https://causarray.readthedocs.io/en/latest/)
+See [CHANGELOG](https://causarray.readthedocs.io/en/latest/changelog.html) for a full version history.
 
 
 <!-- 
@@ -103,4 +108,10 @@ rmarkdown::render("perturbseq.Rmd", rmarkdown::md_document(variant = "markdown_g
 
 
 ## References
-[Du25] Jin-Hong Du, Maya Shen, Hansruedi Mathys, and Kathryn Roeder (2025). Causal differential expression analysis under unmeasured confounders with causarray. bioRxiv, 2025-01.
+[Du26] Jin-Hong Du, Maya Shen, Hansruedi Mathys, and Kathryn Roeder. "Uncovering causal relationships in single cell omic studies with causarray". In: Briefings in Bioinformatics (2026).
+
+[Gabitto24] Mariano I. Gabitto et al. "Integrated multimodal cell atlas of Alzheimer's disease". In: Nature Neuroscience (2024).
+
+[Jin20] Xin Jin et al. "In vivo Perturb-seq reveals neuronal and glial abnormalities associated with autism risk genes". In: Nature Neuroscience (2020).
+
+[Replogle22] Joseph M. Replogle et al. "Mapping information-rich genotype-phenotype landscapes with genome-scale Perturb-seq". In: Cell (2022).

{causarray-0.0.6 → causarray-0.0.7}/causarray/DR_estimation.py

@@ -9,6 +9,7 @@ from causarray.utils import _filter_params
 from joblib import Parallel, delayed
 from tqdm import tqdm
 import pprint
+import warnings
 
 from sklearn.model_selection import KFold, ShuffleSplit
 
@@ -18,13 +19,18 @@ def _get_func_ps(ps_model, **kwargs):
         func_ps = lambda X, Y, X_test:fit_rf_ind_ps(X, Y[:,None], X_test=X_test, **params_ps)[:,0]
     elif ps_model=='logistic':
         clf_ps = LogisticRegression
-        kwargs = {**{'fit_intercept':False, 'C':1e0, 'class_weight':'balanced', 'random_state':0}, **kwargs}
+        kwargs = {**{'fit_intercept':False, 'C':1e0, 'class_weight':None, 'random_state':0}, **kwargs}
         params_ps = _filter_params(clf_ps().get_params(), kwargs)
         func_ps = lambda X, Y, X_test: clf_ps(**params_ps).fit(X, Y).predict_proba(X_test)[:,1]
     elif ps_model=='ensemble':
+        kwargs = dict(kwargs)
+        logistic_class_weight = kwargs.pop('class_weight', None)
         params_ps_rf = _filter_params(fit_rf, kwargs)
         clf_ps = LogisticRegression
-        kwargs = {**{'fit_intercept':False, 'C':1e0, 'class_weight':'balanced', 'random_state':0}, **kwargs}
+        kwargs = {**{
+            'fit_intercept': False, 'C': 1e0,
+            'class_weight': logistic_class_weight, 'random_state': 0,
+        }, **kwargs}
         params_ps_lr = _filter_params(clf_ps().get_params(), kwargs)
         params_ps = {'params_ps_rf':params_ps_rf, 'params_ps_lr':params_ps_lr}
         func_ps = lambda X, Y, X_test:(fit_rf_ind(X, Y[:,None], X_test=X_test, **params_ps_rf)[:,0] + clf_ps(**params_ps_lr).fit(X, Y).predict_proba(X_test)[:,1])/2
@@ -34,10 +40,149 @@ def _get_func_ps(ps_model, **kwargs):
     return func_ps, params_ps
 
 
+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,
+):
+    """Estimate per-treatment propensity scores.
+
+    Each treatment is compared with the shared all-zero control group.  With
+    ``K > 1``, every returned score is predicted by a model that did not train
+    on that cell. Logistic models use ``class_weight='balanced'`` by default,
+    matching :func:`LFC` and historical causarray fits. Pass
+    ``class_weight=None`` for calibrated treatment probabilities.
+
+    Parameters
+    ----------
+    A : array-like, shape (n,) or (n, a)
+        Binary treatment indicators.  Rows containing only zeros are controls.
+    X_A : array-like, shape (n, d_A)
+        Covariates used by the propensity model, including an intercept column
+        when ``fit_intercept=False``.
+    K : int, optional
+        Number of folds.  ``1`` fits and predicts on all eligible cells;
+        values greater than one produce out-of-fold predictions.
+    ps_model : {'logistic', 'random_forest_cv', 'ensemble'}, optional
+        Propensity model.
+    mask : array-like or None, shape (n,) or (n, a)
+        Optional per-treatment eligibility mask for model fitting.
+    clip : tuple(float, float) or None, optional
+        Bounds applied after prediction.  ``None`` returns raw probabilities.
+    random_state : int, optional
+        Random seed used for fold construction and supported estimators.
+    class_weight : str, dict or None, optional
+        Class weighting for logistic propensity estimation. The default
+        ``'balanced'`` matches :func:`LFC`; pass ``None`` for calibrated
+        probabilities.
+
+    Returns
+    -------
+    pi_hat : ndarray, shape (n, a)
+        Estimated probabilities ``P(A_j=1 | X_A)``.
+    """
+    A = np.asarray(A)
+    if A.ndim == 1:
+        A = A[:, None]
+    X_A = np.asarray(X_A)
+    if A.ndim != 2 or X_A.ndim != 2 or A.shape[0] != X_A.shape[0]:
+        raise ValueError('A and X_A must be two-dimensional with matching rows')
+    if not np.all(np.isin(A, (0, 1))):
+        raise ValueError('A must contain only binary treatment indicators')
+
+    try:
+        K_int = int(K)
+    except (TypeError, ValueError) as exc:
+        raise ValueError('K must be a positive integer') from exc
+    if K_int != K:
+        raise ValueError('K must be a positive integer')
+    K = K_int
+    if K < 1:
+        raise ValueError('K must be a positive integer')
+    if K > A.shape[0]:
+        raise ValueError('K cannot exceed the number of samples')
+
+    if mask is not None:
+        mask = np.asarray(mask, dtype=bool)
+        if mask.ndim == 1:
+            mask = mask[:, None]
+        if mask.shape != A.shape:
+            raise ValueError('Mask must have the same shape as the treatment matrix')
+
+    func_ps, params_ps = _get_func_ps(
+        ps_model, verbose=False, random_state=random_state,
+        class_weight=class_weight, **kwargs)
+    if verbose:
+        pprint.pprint(params_ps)
+
+    if ps_model == 'random_forest_cv':
+        info_ecv = run_ecv(X_A, A, **params_ps)
+        func_ps, params_ps = _get_func_ps(
+            ps_model, verbose=False, ecv=False,
+            kwargs_ensemble=info_ecv['best_params_ensemble'],
+            kwargs_regr=info_ecv['best_params_regr'],
+        )
+        if verbose:
+            pprint.pprint('Best parameters for the regression model:')
+            pprint.pprint(info_ecv['best_params_regr'])
+            pprint.pprint('Best parameters for the ensemble model:')
+            pprint.pprint(info_ecv['best_params_ensemble'])
+
+    n = A.shape[0]
+    if K == 1:
+        folds = [(np.arange(n), np.arange(n))]
+    elif K == n:
+        folds = [
+            (np.delete(np.arange(n), j), np.array([j])) for j in range(n)
+        ]
+    else:
+        folds = KFold(
+            n_splits=K, random_state=random_state, shuffle=True,
+        ).split(X_A)
+
+    pi_hat = np.zeros(A.shape, dtype=float)
+    for train_index, test_index in folds:
+        A_train = A[train_index]
+        XA_train, XA_test = X_A[train_index], X_A[test_index]
+        i_ctrl = np.sum(A_train, axis=1) == 0
+
+        for j in range(A.shape[1]):
+            i_case = A_train[:, j] == 1
+            eligible = mask[train_index, j] if mask is not None else (i_ctrl | i_case)
+            y_train = A_train[eligible, j]
+            if y_train.size == 0 or np.unique(y_train).size != 2:
+                raise ValueError(
+                    f'Treatment {j} needs at least one eligible control and case '
+                    'in every training fold'
+                )
+
+            x_train = XA_train[eligible]
+            constant_design = np.all(np.ptp(x_train, axis=0) == 0)
+            if ps_model == 'logistic' and constant_design:
+                if class_weight == 'balanced':
+                    pi_hat[test_index, j] = 0.5
+                elif isinstance(class_weight, dict):
+                    sample_weight = np.asarray([
+                        class_weight.get(int(value), 1.0) for value in y_train
+                    ])
+                    pi_hat[test_index, j] = np.average(
+                        y_train, weights=sample_weight)
+                else:
+                    pi_hat[test_index, j] = np.mean(y_train)
+            else:
+                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])
+    return pi_hat
+
+
 def cross_fitting(
     Y, A, X, X_A, family='poisson', K=1, glm_alpha=1e-4,
-    ps_model='logistic', 
-    Y_hat=None, pi_hat=None, mask=None, verbose=False, **kwargs):
+    ps_model='logistic', ps_class_weight='balanced',
+    Y_hat=None, pi_hat=None, mask=None, ps_clip=(0.01, 0.99),
+    return_raw_pi=False, verbose=False, **kwargs):
     '''
     Cross-fitting for causal estimands.
 
@@ -59,6 +204,10 @@ def cross_fitting(
         The regularization parameter for the generalized linear model. The default is 1e-4.
     ps_model : str, optional
         The propensity score model. The default is 'logistic'.
+    ps_class_weight : str, dict or None, optional
+        Class weighting used by the propensity model. ``'balanced'`` preserves
+        the established ``LFC`` nuisance fit; pass ``None`` for calibrated
+        treatment probabilities.
     
     Y_hat : array, optional
         Estimated potential outcome of shape (n, p, a, 2). The default is None.
@@ -66,8 +215,11 @@ def cross_fitting(
         Propensity score of shape (n, a). The default is None.
     mask : array, optional
         Boolean mask of shape (n, a) for the treatment, indicating which samples are used for 
-        the estimation of the estimand. This does not affect the estimation of pseudo-outcomes
-        and propensity scores.
+        propensity-model fitting and the downstream estimand.
+    ps_clip : tuple(float, float) or None, optional
+        Bounds applied to scores used by AIPW. ``None`` disables clipping.
+    return_raw_pi : bool, optional
+        Return raw scores as a third result when true.
 
     **kwargs : dict
         Additional arguments to pass to the model.
@@ -78,12 +230,22 @@ def cross_fitting(
         Estimated potential outcome under control.
     pi_hat : array
         Estimated propensity score.
+    pi_hat_raw : array
+        Unclipped propensity score, returned only when ``return_raw_pi=True``.
     '''
-    func_ps, params_ps = _get_func_ps(ps_model, verbose=False, **kwargs)
+    kwargs = dict(kwargs)
+    if 'class_weight' in kwargs:
+        legacy_class_weight = kwargs.pop('class_weight')
+        warnings.warn(
+            'Passing class_weight through LFC/cross_fitting is deprecated; '
+            'use ps_class_weight instead.',
+            FutureWarning, stacklevel=2,
+        )
+        ps_class_weight = legacy_class_weight
+
     params_glm = _filter_params(fit_glm, {**kwargs, 'verbose': verbose})
 
     if verbose:
-        pprint.pprint(params_ps)
         pprint.pprint(params_glm)
     
     if K > 1:
@@ -99,14 +261,31 @@ def cross_fitting(
         folds = [(np.arange(X.shape[0]), np.arange(X.shape[0]))]
 
     # Initialize lists to store results
-    fit_pi = True if pi_hat is None else False
-    pi_hat = np.zeros_like(A, dtype=float) if fit_pi else pi_hat
+    if pi_hat is None:
+        if verbose:
+            pprint.pprint('Fit propensity score models...')
+        ps_kwargs = {k: v for k, v in kwargs.items() if k != 'random_state'}
+        if ps_model in ('logistic', 'ensemble'):
+            ps_kwargs['class_weight'] = ps_class_weight
+        pi_hat_raw = estimate_propensity_scores(
+            A, X_A, K=K, ps_model=ps_model, mask=mask,
+            random_state=kwargs.get('random_state', 0), verbose=verbose,
+            **ps_kwargs,
+        )
+    else:
+        pi_hat_raw = np.asarray(pi_hat, dtype=float).reshape(A.shape)
+    if ps_clip is None:
+        pi_hat = pi_hat_raw.copy()
+    else:
+        if len(ps_clip) != 2 or not 0 <= ps_clip[0] < ps_clip[1] <= 1:
+            raise ValueError(
+                'ps_clip must be None or a pair 0 <= lower < upper <= 1')
+        pi_hat = np.clip(pi_hat_raw, ps_clip[0], ps_clip[1])
     fit_Y = True if Y_hat is None else False
     if fit_Y:
         _yhat_gb = Y.shape[0] * Y.shape[1] * A.shape[1] * 2 * 8 / 1e9
         _mem_limit_gb = kwargs.get('mem_limit_gb', None)
         if _mem_limit_gb is not None and _yhat_gb > _mem_limit_gb:
-            import warnings
             warnings.warn(
                 f"Y_hat allocation ({_yhat_gb:.1f} GB as float64) exceeds "
                 f"mem_limit_gb={_mem_limit_gb} GB; using float32 to halve peak memory.",
@@ -116,16 +295,6 @@ def cross_fitting(
         else:
             Y_hat = np.zeros((Y.shape[0], Y.shape[1], A.shape[1], 2), dtype=float)
 
-    # perform ECV at once
-    if fit_pi and ps_model == 'random_forest_cv':
-        info_ecv = run_ecv(X_A, A, **params_ps)
-        func_ps, params_ps = _get_func_ps(ps_model, verbose=False, ecv=False, 
-                kwargs_ensemble=info_ecv['best_params_ensemble'], kwargs_regr=info_ecv['best_params_regr'])
-        pprint.pprint('Best parameters for the regression model:')
-        pprint.pprint(info_ecv['best_params_regr'])
-        pprint.pprint('Best parameters for the ensemble model:')
-        pprint.pprint(info_ecv['best_params_ensemble'])
-
     # Perform cross-fitting
     for train_index, test_index in folds:
         # Split data
@@ -134,28 +303,6 @@ def cross_fitting(
         A_train, A_test = A[train_index], A[test_index]
         Y_train, Y_test = Y[train_index], Y[test_index]
 
-        if fit_pi:
-            if verbose: pprint.pprint('Fit propensity score models...')
-            i_ctrl = (np.sum(A_train, axis=1) == 0.)
-
-            pi = np.zeros_like(A_test, dtype=float)
-            for j in range(A.shape[1]):
-                i_case = (A_train[:,j] == 1.)
-
-                if mask is not None:
-                    i_cells = mask[:, j]
-                else:
-                    i_ctrl = (np.sum(A_train, axis=1) == 0.)
-                    i_cells = i_ctrl | i_case
-
-                if ps_model=='logistic' and XA_train.shape[1]==1 and np.all(XA_train==1):
-                    prob = np.sum(i_case)/np.sum(i_cells)
-                    pi[A_train[:,j] == 1., j] = prob
-                    pi[A_train[:,j] == 0., j] = 1 - prob
-                else:
-                    pi[:,j] = func_ps(XA_train[i_cells], A_train[i_cells][:,j], XA_test)
-            pi_hat[test_index] = pi
-
         if fit_Y:
             if verbose: pprint.pprint('Fit outcome models...')
             # Subset offset to training fold (for fitting) and test fold (for
@@ -174,15 +321,16 @@ def cross_fitting(
             Y_hat[test_index,:,:,0] = res[1][0]
             Y_hat[test_index,:,:,1] = res[1][1]
 
-    pi_hat = np.clip(pi_hat, 0.01, 0.99)
     Y_hat = np.clip(Y_hat, None, 1e5)
+    if return_raw_pi:
+        return Y_hat, pi_hat, pi_hat_raw
     return Y_hat, pi_hat
 
 
 
 
 
-def AIPW_mean(Y, A, mu, pi, positive=False):
+def AIPW_mean(Y, A, mu, pi):
     '''
     Augmented inverse probability weighted estimator (AIPW)
 
@@ -196,9 +344,6 @@ def AIPW_mean(Y, A, mu, pi, positive=False):
         Conditional outcome distribution estimate of shape (n, p, a, 2).
     pi : array
         Propensity score of shape (n, a, 2).
-    positive : bool, optional
-        Whether to restrict the pseudo-outcome to be positive.
-
     Returns
     -------
     tau : array
@@ -207,16 +352,18 @@ def AIPW_mean(Y, A, mu, pi, positive=False):
         Pseudo-outcome of shape (n, p, a, 2).
     '''
     
-    weight = A / pi
+    with np.errstate(divide='ignore', invalid='ignore', over='ignore'):
+        weight = A / pi
     weight = weight[:, None, ...]
     Y = Y[:, :, None, None]
 
+    # Influence-function values are intentionally left unconstrained.  Even
+    # for a nonnegative outcome, individual AIPW pseudo-outcomes may be
+    # negative; projecting them cell by cell changes their mean and biases the
+    # estimator.  Parameter-space constraints belong after aggregation.
     pseudo_y = weight * (Y - mu) + mu
-    
-    if positive:
-        pseudo_y = np.clip(pseudo_y, 0, None)
 
-    tau = np.mean(pseudo_y, axis=0)
+    tau = np.mean(pseudo_y, axis=0, dtype=np.float64)
 
     return tau, pseudo_y
 
@@ -350,4 +497,4 @@ def fit_rf_ind_outcome(W, Y, A, *args, **kwargs):
     Y_pred = Y_pred.reshape(X.shape[0],1+a,Y.shape[1])
     Yhat_1 = Y_pred[:,1:,:].transpose(0,2,1)
     Yhat_0 = np.tile(Y_pred[:,0,:][:,:,None], (1,1,a))
-    return Yhat_0, Yhat_1
+    return Yhat_0, Yhat_1

{causarray-0.0.6 → causarray-0.0.7}/causarray/DR_learner.py

@@ -1,6 +1,8 @@
 import numpy as np
 import contextlib
 import pandas as pd
+import warnings
+from typing import Literal
 from causarray.DR_estimation import AIPW_mean, cross_fitting
 from causarray.gcate_glm import loess_fit, ls_fit
 import causarray.gcate_glm as _gcate_glm  # for _backend_override
@@ -14,7 +16,9 @@ def compute_causal_estimand(
     Y, W, A, W_A=None, family='nb', offset=False,
     Y_hat=None, pi_hat=None, mask=None,
     fdx=False, fdx_B=1000, fdx_alpha=0.05, fdx_c=0.1,
-    verbose=False, random_state=0, backend: str = "auto", **kwargs):
+    verbose=False, random_state=0, backend: str = "auto", K=1,
+    ps_clip=(0.01, 0.99), ps_class_weight='balanced',
+    **kwargs):
     """Estimate causal treatment effects using AIPW with a user-supplied estimand.
 
     Parameters
@@ -42,9 +46,17 @@ def compute_causal_estimand(
     pi_hat : array or None, shape (n, a)
         Pre-computed propensity scores.  When provided, propensity fitting is
         skipped.
+    K : int
+        Number of folds used for nuisance estimation.  The default ``1``
+        preserves in-sample fitting.
+    ps_clip : tuple(float, float)
+        Bounds applied to propensity scores used by AIPW.
+    ps_class_weight : str, dict or None
+        Class weighting for the propensity model. ``'balanced'`` preserves the
+        established nuisance fit; pass ``None`` for calibrated probabilities.
     mask : array or None, shape (n, a)
-        Boolean mask indicating which cells are used for the final estimand
-        computation.  Does not affect cross-fitting or propensity estimation.
+        Boolean mask indicating eligible cells for each treatment. It limits
+        propensity-model fitting and final estimand computation.
     fdx : bool
         Whether to apply FDX control (``P(FDP > fdx_c) < fdx_alpha``).
     fdx_B : int
@@ -64,12 +76,17 @@ def compute_causal_estimand(
     Returns
     -------
     df_res : DataFrame
-        Test results with columns ``gene_names``, ``tau``, ``std``, ``stat``,
-        ``rej``, ``pvalue``, ``padj``, ``pvalue_emp_null_adj``,
-        ``padj_emp_null_adj`` (and ``trt`` when ``A`` has multiple columns).
+        Test results produced by ``estimand``. An estimand may optionally return
+        a fifth dictionary whose arrays are added as diagnostic columns.
     """
     reset_random_seeds(random_state)
 
+    if 'clip_pseudo_outcomes' in kwargs:
+        raise TypeError(
+            'clip_pseudo_outcomes has been removed; AIPW pseudo-outcomes are '
+            'always left unclipped'
+        )
+
     ctx = _gcate_glm._backend_override(backend) if backend != "auto" else contextlib.nullcontext()
 
     # check the input data
@@ -85,7 +102,6 @@ def compute_causal_estimand(
     _mem_limit = kwargs.get('mem_limit_gb', None)
     _use_f32 = _mem_limit is not None and _yhat_gb > _mem_limit
     if _use_f32:
-        import warnings
         warnings.warn(
             f"AIPW intermediate Y ({_yhat_gb:.1f} GB as float64) exceeds "
             f"mem_limit_gb={_mem_limit} GB; downcasting Y, A, and pi_hat to "
@@ -94,6 +110,8 @@ def compute_causal_estimand(
         )
     Y = Y.astype(np.float32 if _use_f32 else float)
     n, p = Y.shape
+    if not np.all(np.isfinite(Y)):
+        raise ValueError('Y must contain only finite values')
 
     if len(A.shape) == 1:
         A = A.reshape(-1,1)
@@ -135,11 +153,25 @@ def compute_causal_estimand(
     else:
         offset = None
         size_factors = np.ones(n)
+    size_factors = np.asarray(size_factors, dtype=float)
+    if size_factors.shape != (n,) or not np.all(np.isfinite(size_factors)) or np.any(size_factors <= 0):
+        raise ValueError('Size factors must be finite, positive, and have length n')
     
     with ctx:
-        Y_hat, pi_hat = cross_fitting(Y, A, W, W_A, family=family, offset=offset,
-            Y_hat=Y_hat, pi_hat=pi_hat, mask=mask, random_state=random_state, verbose=verbose, **kwargs)
+        Y_hat, pi_hat, pi_hat_raw = cross_fitting(
+            Y, A, W, W_A, family=family, K=K, offset=offset,
+            Y_hat=Y_hat, pi_hat=pi_hat, mask=mask, ps_clip=ps_clip,
+            ps_class_weight=ps_class_weight,
+            return_raw_pi=True, random_state=random_state, verbose=verbose,
+            **kwargs,
+        )
     pi_hat = pi_hat.reshape(*A.shape)
+    if not np.all(np.isfinite(Y_hat)):
+        raise ValueError('Outcome-model predictions must contain only finite values')
+    if not np.all(np.isfinite(pi_hat)) or np.any((pi_hat <= 0) | (pi_hat >= 1)):
+        raise ValueError(
+            'Propensity scores used by AIPW must be finite and strictly between 0 and 1; '
+            'pass ps_clip bounds or provide valid pi_hat values')
 
     if verbose: pprint.pprint('Estimating AIPW mean...')
     _aipw_dtype = Y_hat.dtype if _use_f32 else None
@@ -147,8 +179,12 @@ def compute_causal_estimand(
     pi_hat_aipw = pi_hat.astype(_aipw_dtype) if _aipw_dtype is not None else pi_hat
 
     # point estimation of the treatment effect
-    _, etas = AIPW_mean(Y, np.stack([1-A_aipw, A_aipw], axis=-1),
-        Y_hat, np.stack([1-pi_hat_aipw, pi_hat_aipw], axis=-1), positive=True)
+    _, etas = AIPW_mean(
+        Y,
+        np.stack([1-A_aipw, A_aipw], axis=-1),
+        Y_hat,
+        np.stack([1-pi_hat_aipw, pi_hat_aipw], axis=-1),
+    )
 
     # normalize the influence function values
     etas /= size_factors[:,None,None,None]
@@ -165,6 +201,7 @@ def compute_causal_estimand(
         _ret = estimand(etas[i_cells,:,j], A[i_cells,j], **kwargs)
         eta_est, tau_est, var_est = _ret[:3]
         df_est = _ret[3] if len(_ret) > 3 else None
+        estimand_info = _ret[4] if len(_ret) > 4 else None
 
         std_est = np.sqrt(var_est)
         tvalues_init = tau_est / std_est
@@ -187,20 +224,41 @@ def compute_causal_estimand(
             'pvalue_emp_null_adj': pvals_adj,
             'padj_emp_null_adj': qvals_adj,            
             })
+        if estimand_info is not None:
+            for name, values in estimand_info.items():
+                df_res[name] = values
+            n_nonestimable = int(np.sum(~np.asarray(estimand_info['estimable'], dtype=bool)))
+            if n_nonestimable:
+                treatment = trt_names[j] if A.shape[1] > 1 else 0
+                warnings.warn(
+                    f'{n_nonestimable} gene(s) for treatment {treatment!r} have '
+                    'nonfinite or nonpositive AIPW counterfactual means and are '
+                    'reported as non-estimable.',
+                    RuntimeWarning, stacklevel=2,
+                )
         if A.shape[1]>1:
             df_res['trt'] = trt_names[j]
         res.append(df_res)
     df_res = pd.concat(res, axis=0).reset_index(drop=True)
-    estimation = {**{'pi_hat':pi_hat, 'Y_hat':Y_hat, 'offset':offset, 'size_factors':size_factors}, **kwargs}
+    estimation = {**{
+        'pi_hat': pi_hat,
+        'pi_hat_raw': pi_hat_raw,
+        'Y_hat': Y_hat,
+        'offset': offset,
+        'size_factors': size_factors,
+        'ps_class_weight': ps_class_weight,
+    }, **kwargs}
     return df_res, estimation
 
 
 def LFC(
     Y, W, A, W_A=None, family='nb', offset=False,
-    Y_hat=None, pi_hat=None, cross_est=False, mask=None, usevar='unequal',
+    Y_hat=None, pi_hat=None, cross_est=False, K=None, mask=None,
+    usevar: Literal['unequal', 'pooled'] = 'unequal',
     thres_min=1e-2, thres_diff=1e-2, eps_var=1e-4,
     fdx=False, fdx_alpha=0.05, fdx_c=0.1,
-    verbose=False, backend: str = "auto", **kwargs):
+    verbose=False, backend: str = "auto", ps_clip=(0.01, 0.99),
+    ps_class_weight='balanced', **kwargs):
     """Estimate log-fold changes of treatment effects (LFCs) using AIPW.
 
     Fits a doubly-robust AIPW estimator for the log-ratio of counterfactual
@@ -230,17 +288,38 @@ def LFC(
         Pre-computed propensity scores.  When provided, propensity fitting is
         skipped.
     cross_est : bool
-        Whether to use cross-estimation for nuisance parameters.
+        Whether to use two-fold cross-estimation for nuisance parameters.
+        An explicit ``K`` takes precedence.
+    K : int or None
+        Number of nuisance-estimation folds.  ``None`` uses 2 when
+        ``cross_est=True`` and 1 otherwise.
     mask : array or None, shape (n, a)
-        Boolean mask indicating which cells are used for the final estimand
-        computation.  Does not affect cross-fitting or propensity estimation.
+        Boolean mask indicating eligible cells for each treatment. It limits
+        propensity-model fitting and final estimand computation.
     usevar : str
         Variance estimator for the AIPW pseudo-outcomes:
 
         * ``'unequal'`` (default, v0.0.6+): Welch variance
           ``s₀²/n₀ + s₁²/n₁`` with Welch-Satterthwaite degrees of
-          freedom; p-values use the t-distribution.
+          freedom; p-values use the t-distribution. This is the recommended
+          choice for perturbation screens and for case-control, bulk, and
+          donor-level pseudo-bulk analyses. Independence of the rows does not
+          imply equal treatment-arm variances, and biological heterogeneity or
+          unequal effective sample sizes can make pooled inference
+          anti-conservative.
         * ``'pooled'``: pooled-variance estimator ``(s² + eps_var) / n``.
+          Treat this as an opt-in sensitivity or legacy analysis, not as the
+          default for small case-control studies. Use it only when equal arm
+          variances have a strong scientific and empirical justification.
+          Pooled inference can produce substantially smaller standard errors
+          and many more discoveries; donor-level independence alone is not a
+          justification for pooling.
+
+        ``'unequal'`` accommodates arm-specific variance but does not model
+        within-donor or within-subject correlation. Repeated cells from the
+        same biological unit should still be pseudo-bulked or analyzed with a
+        cluster-aware method; changing ``usevar`` alone does not remove
+        pseudoreplication.
 
         .. versionchanged:: 0.0.6
             Default changed from ``'pooled'`` to ``'unequal'``.  The
@@ -267,25 +346,41 @@ def LFC(
     backend : str
         GLM backend: ``"auto"`` (default), ``"fast"`` (force crispyx),
         or ``"original"`` (force statsmodels).
+    ps_clip : tuple(float, float)
+        Bounds applied to propensity scores used by AIPW. Raw, unclipped scores
+        remain available as ``estimation['pi_hat_raw']``.
+    ps_class_weight : str, dict or None
+        Class weighting for the propensity model. ``'balanced'`` remains the
+        default to limit nuisance-model drift; pass ``None`` for calibrated
+        probabilities.
     **kwargs
         Additional arguments forwarded to the GLM fitting functions.
 
     Returns
     -------
     df_res : DataFrame
-        Test results with columns ``gene_names``, ``tau``, ``std``, ``stat``,
-        ``rej``, ``pvalue``, ``padj``, ``pvalue_emp_null_adj``,
-        ``padj_emp_null_adj`` (and ``trt`` when ``A`` has multiple columns).
+        Test results with effect estimates and inference columns, raw
+        ``mean_control`` and ``mean_treated`` counterfactual means, and an
+        ``estimable`` flag (plus ``trt`` for multiple treatments).
     """
 
     def estimand(etas, A, **kwargs):
         eta_0, eta_1 = etas[..., 0], etas[..., 1]
-        tau_0, tau_1 = np.mean(eta_0, axis=0), np.mean(eta_1, axis=0)
-
-        tau_1 = np.clip(tau_1, thres_diff, None)
-        tau_0 = np.clip(tau_0, thres_diff, None)
-        tau_est = np.log(tau_1/tau_0)
-        eta_est = eta_1 / tau_1[None,:] -  eta_0 / tau_0[None,:]
+        mean_0 = np.mean(eta_0, axis=0, dtype=np.float64)
+        mean_1 = np.mean(eta_1, axis=0, dtype=np.float64)
+        finite_means = np.isfinite(mean_0) & np.isfinite(mean_1)
+        estimable = finite_means & (mean_0 > 0) & (mean_1 > 0)
+
+        # Apply the count-mean parameter-space constraint only after averaging
+        # the unmodified AIPW pseudo-outcomes.  The floor makes the logarithm
+        # and its delta-method denominator numerically safe.  Nonpositive raw
+        # means remain non-estimable and are filtered below, so this safeguard
+        # cannot turn an invalid aggregate into a discovery.
+        tau_0 = np.where(estimable, np.maximum(mean_0, thres_diff), 1.0)
+        tau_1 = np.where(estimable, np.maximum(mean_1, thres_diff), 1.0)
+        with np.errstate(invalid='ignore', divide='ignore', over='ignore'):
+            tau_est = np.log(tau_1/tau_0)
+            eta_est = eta_1 / tau_1[None,:] - eta_0 / tau_0[None,:]
 
         df_eff = None
         if usevar == 'pooled':
@@ -316,8 +411,13 @@ def LFC(
         else:
             raise ValueError('usevar must be either "pooled" or "unequal"')
 
-        # filter out low-expressed genes
-        idx = (np.maximum(np.abs(tau_0),np.abs(tau_1))<thres_min) | (np.abs(tau_1-tau_0)<thres_diff)
+        # Filter on the raw aggregate estimates, not on cell-level projections
+        # or the numerical floor used for the log transform.
+        idx = (
+            ~estimable |
+            (np.maximum(mean_0, mean_1) < thres_min) |
+            (np.abs(mean_1 - mean_0) < thres_diff)
+        )
         tau_est[idx] = 0.; eta_est[:,idx] = 0.; var_est[idx] = np.inf
         if df_eff is not None:
             df_eff[idx] = np.nan
@@ -338,13 +438,22 @@ def LFC(
                     RuntimeWarning, stacklevel=3,
                 )
 
-        return eta_est, tau_est, var_est, df_eff
+        info = {
+            'mean_control': mean_0,
+            'mean_treated': mean_1,
+            'estimable': estimable,
+        }
+        return eta_est, tau_est, var_est, df_eff, info
+
+    if K is None:
+        K = 2 if cross_est else 1
 
     return compute_causal_estimand(
         estimand, Y, W, A, W_A, family, offset,    
         Y_hat=Y_hat, pi_hat=pi_hat, mask=mask,
         fdx=fdx, fdx_alpha=fdx_alpha, fdx_c=fdx_c,
-        verbose=verbose, backend=backend, **kwargs)
+        verbose=verbose, backend=backend, K=K, ps_clip=ps_clip,
+        ps_class_weight=ps_class_weight, **kwargs)
 
 
 
@@ -524,7 +633,11 @@ def gcate_lfc_batch(
 
     lfc_kwargs : dict or None
         Extra keyword arguments forwarded to :func:`LFC`
-        (e.g. ``usevar``, ``fdx``, ``thres_min``).
+        (e.g. ``usevar``, ``fdx``, ``thres_min``). Retain the default
+        ``usevar='unequal'`` for perturbation screens and for case-control,
+        bulk, or pseudo-bulk analyses. Pass
+        ``lfc_kwargs=dict(usevar='pooled')`` only for a deliberately justified
+        equal-variance sensitivity or legacy analysis.
     **kwargs
         Additional arguments forwarded to both :func:`fit_gcate_batch` and
         :func:`LFC`.  When a key collides with ``gcate_kwargs`` /
@@ -725,4 +838,4 @@ def LFC_batch(*args, **kwargs):
         DeprecationWarning,
         stacklevel=2,
     )
-    return gcate_lfc_batch(*args, **kwargs)
+    return gcate_lfc_batch(*args, **kwargs)

causarray-0.0.7/causarray/__about__.py

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

{causarray-0.0.6 → causarray-0.0.7}/causarray/__init__.py

@@ -10,10 +10,14 @@ __all__ = [
     'LFC', 'gcate_lfc_batch', 'LFC_batch',
     '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',
     ]
 
 
 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.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
@@ -28,4 +32,4 @@ __maintainer__ = "Jin-Hong Du"
 __maintainer_email__ = "jinhongd@andrew.cmu.edu"
 __description__ = ("Causarray: A Python package for simultaneous causal inference"
     " with an array of outcomes."
-    )
+    )

causarray-0.0.7/causarray/diagnostics.py

@@ -0,0 +1,172 @@
+"""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

causarray-0.0.6/causarray/__about__.py

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