causalscreen 0.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Rahul Kumar Mandal
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,3 @@
1
+ recursive-include tests *.py
2
+ include LICENSE
3
+ recursive-include examples *.py
@@ -0,0 +1,127 @@
1
+ Metadata-Version: 2.4
2
+ Name: causalscreen
3
+ Version: 0.2.0
4
+ Summary: Residual-correlation causal screening with iterative peeling, localized modeling, and drift early-warning
5
+ Author-email: Rahul Kumar Mandal <rahulkm.jobs@gmail.com>
6
+ License: MIT
7
+ Project-URL: Homepage, https://github.com/rahulkm3/causalscreen
8
+ Requires-Python: >=3.9
9
+ Description-Content-Type: text/markdown
10
+ License-File: LICENSE
11
+ Requires-Dist: numpy
12
+ Requires-Dist: scipy
13
+ Requires-Dist: scikit-learn
14
+ Dynamic: license-file
15
+
16
+ # causalscreen
17
+
18
+ Residual-correlation causal screening with iterative peeling, causal-factor
19
+ partitioning with inference-time local models, and a missing-causal-power
20
+ drift alarm.
21
+
22
+ Companion code for: **R. K. Mandal, "Residual-Correlation Causal Screening with
23
+ Localized Supervised Modeling: A Practical Framework for Confounder
24
+ Identification, Error Reduction, and Drift Early-Warning"** (working paper,
25
+ 2026; research 2022–present). Paper PDF in `/paper` · Preprint: https://papers.ssrn.com/sol3/papers.cfm?abstract_id=7135078
26
+
27
+ ## What it does
28
+
29
+ Supervised models fail in deployment because they lean on **confounded
30
+ proxies** — features whose correlation with the target is inherited from other
31
+ variables. `causalscreen`:
32
+
33
+ 1. **Screens** every candidate feature via the two-regression residual
34
+ construction (Frisch–Waugh–Lovell partial correlation): regress `y` on the
35
+ other features, regress `x_j` on the other features, correlate the
36
+ residuals. High residual correlation = genuine independent effect.
37
+ 2. **Peels** iteratively: each discovered factor's contribution is stripped
38
+ out of the response before the next round, so its influence cannot re-enter
39
+ the screen through the variables it confounds.
40
+ 3. **Ranks** features by causal power and reports the **missing causal
41
+ power** — the share of response variance the discovered factors cannot
42
+ explain — a deployable early-warning statistic for data/model drift.
43
+ 4. **Exploits** the ranking: `PartitionedRegressor` partitions data on top
44
+ causal factors and trains one small local model per query at inference
45
+ time (interpolation), falling back to the global model out of support
46
+ (extrapolation) — so the error reduction of local modeling never costs
47
+ out-of-sample generalization.
48
+
49
+ ## Quick result (synthetic ground truth, `examples/demo.py`)
50
+
51
+ | feature | naive \|corr\| with y | causal screen verdict |
52
+ |---|---|---|
53
+ | x1 (true cause) | 0.92 | **ranked** (power 0.78) |
54
+ | x2 (confounded proxy of x1) | 0.87 | **rejected** |
55
+ | x3 (noise) | 0.01 | rejected |
56
+ | x4 (true cause, weak) | 0.33 | **ranked** (power 0.83) |
57
+
58
+ Missing causal power: 5.1% — the exact noise floor of the generating process.
59
+ Partitioned local modeling: **99.0% out-of-sample MSE reduction** vs. a global
60
+ model on regime-structured data.
61
+
62
+ ## Install & use
63
+
64
+ ```bash
65
+ pip install -e .
66
+ python examples/demo.py
67
+ pytest tests/
68
+ ```
69
+
70
+ ```python
71
+ from causalscreen import CausalScreen, PartitionedRegressor
72
+ res = CausalScreen(alpha=0.05).fit(X, y, feature_names)
73
+ res.ranking # features in decreasing causal power
74
+ res.causal_power # partial correlation at selection round
75
+ res.missing_causal_power # drift-alarm statistic
76
+
77
+ pr = PartitionedRegressor(partition_features=res.ranking[:2], min_cell=50)
78
+ pr.fit(X, y, feature_names)
79
+ pr.predict(X_new) # local model in-support, global model out-of-support
80
+ ```
81
+
82
+ ## Assumptions (stated plainly)
83
+
84
+ Stationary cross-sectional data · causal sufficiency (all relevant variables
85
+ observed) · no reverse causation · features measured before the response.
86
+ The screen addresses the *confounding* component of causal analysis; it does
87
+ not prove causality in full generality. See paper §3, §7.
88
+
89
+ ## Deployed results
90
+
91
+ Layered over Random Forest, OLS and Prophet models in commercial deployments
92
+ (credit risk, segmentation, forecasting), the framework reduced out-of-sample
93
+ error by 65–85%. Client data is not releasable; this repo's synthetic
94
+ benchmarks are the reproducible counterpart. DoWhy/EconML comparison suite:
95
+ in progress.
96
+
97
+ ## Performance (v0.2.0)
98
+
99
+ Screening uses a single precision-matrix inversion per round (FWL-equivalent,
100
+ with an exact residual-regression fallback for ill-conditioned designs) —
101
+ ~77x faster than the reference implementation at n=5000, p=40. Partitioned
102
+ prediction indexes cells at fit time and caches local models — ~409x faster
103
+ on batch prediction (20k train / 5k queries). Equivalence to the reference
104
+ implementation is enforced in the test suite to ~1e-8 (38 tests; reference
105
+ implementation vendored in `tests/`).
106
+
107
+ ## License
108
+
109
+ MIT
110
+
111
+ ## Pipeline
112
+
113
+ ```mermaid
114
+ flowchart LR
115
+ A[Training data] --> B["CausalScreen<br/>residual-correlation test<br/>+ iterative peeling"]
116
+ B -->|"ranked causal factors<br/>(decreasing causal power)"| C["Partition data on<br/>top causal factors<br/>(strength x significance rule)"]
117
+ C --> G1[(Subgroup 1)]
118
+ C --> G2[(Subgroup 2)]
119
+ C --> Gn[(Subgroup n)]
120
+ Q[Query point] --> R{inside a<br/>subgroup?}
121
+ R -->|"yes (interpolation)"| L["Train ONE local model<br/>in that subgroup only<br/>(Optimisation 2)"]
122
+ R -->|"no (extrapolation)"| F["Global model fallback<br/>(retains significance &<br/>generalization)"]
123
+ G1 -.-> L
124
+ G2 -.-> L
125
+ Gn -.-> L
126
+ B --> M["Missing causal power<br/>= drift early-warning alarm"]
127
+ ```
@@ -0,0 +1,112 @@
1
+ # causalscreen
2
+
3
+ Residual-correlation causal screening with iterative peeling, causal-factor
4
+ partitioning with inference-time local models, and a missing-causal-power
5
+ drift alarm.
6
+
7
+ Companion code for: **R. K. Mandal, "Residual-Correlation Causal Screening with
8
+ Localized Supervised Modeling: A Practical Framework for Confounder
9
+ Identification, Error Reduction, and Drift Early-Warning"** (working paper,
10
+ 2026; research 2022–present). Paper PDF in `/paper` · Preprint: https://papers.ssrn.com/sol3/papers.cfm?abstract_id=7135078
11
+
12
+ ## What it does
13
+
14
+ Supervised models fail in deployment because they lean on **confounded
15
+ proxies** — features whose correlation with the target is inherited from other
16
+ variables. `causalscreen`:
17
+
18
+ 1. **Screens** every candidate feature via the two-regression residual
19
+ construction (Frisch–Waugh–Lovell partial correlation): regress `y` on the
20
+ other features, regress `x_j` on the other features, correlate the
21
+ residuals. High residual correlation = genuine independent effect.
22
+ 2. **Peels** iteratively: each discovered factor's contribution is stripped
23
+ out of the response before the next round, so its influence cannot re-enter
24
+ the screen through the variables it confounds.
25
+ 3. **Ranks** features by causal power and reports the **missing causal
26
+ power** — the share of response variance the discovered factors cannot
27
+ explain — a deployable early-warning statistic for data/model drift.
28
+ 4. **Exploits** the ranking: `PartitionedRegressor` partitions data on top
29
+ causal factors and trains one small local model per query at inference
30
+ time (interpolation), falling back to the global model out of support
31
+ (extrapolation) — so the error reduction of local modeling never costs
32
+ out-of-sample generalization.
33
+
34
+ ## Quick result (synthetic ground truth, `examples/demo.py`)
35
+
36
+ | feature | naive \|corr\| with y | causal screen verdict |
37
+ |---|---|---|
38
+ | x1 (true cause) | 0.92 | **ranked** (power 0.78) |
39
+ | x2 (confounded proxy of x1) | 0.87 | **rejected** |
40
+ | x3 (noise) | 0.01 | rejected |
41
+ | x4 (true cause, weak) | 0.33 | **ranked** (power 0.83) |
42
+
43
+ Missing causal power: 5.1% — the exact noise floor of the generating process.
44
+ Partitioned local modeling: **99.0% out-of-sample MSE reduction** vs. a global
45
+ model on regime-structured data.
46
+
47
+ ## Install & use
48
+
49
+ ```bash
50
+ pip install -e .
51
+ python examples/demo.py
52
+ pytest tests/
53
+ ```
54
+
55
+ ```python
56
+ from causalscreen import CausalScreen, PartitionedRegressor
57
+ res = CausalScreen(alpha=0.05).fit(X, y, feature_names)
58
+ res.ranking # features in decreasing causal power
59
+ res.causal_power # partial correlation at selection round
60
+ res.missing_causal_power # drift-alarm statistic
61
+
62
+ pr = PartitionedRegressor(partition_features=res.ranking[:2], min_cell=50)
63
+ pr.fit(X, y, feature_names)
64
+ pr.predict(X_new) # local model in-support, global model out-of-support
65
+ ```
66
+
67
+ ## Assumptions (stated plainly)
68
+
69
+ Stationary cross-sectional data · causal sufficiency (all relevant variables
70
+ observed) · no reverse causation · features measured before the response.
71
+ The screen addresses the *confounding* component of causal analysis; it does
72
+ not prove causality in full generality. See paper §3, §7.
73
+
74
+ ## Deployed results
75
+
76
+ Layered over Random Forest, OLS and Prophet models in commercial deployments
77
+ (credit risk, segmentation, forecasting), the framework reduced out-of-sample
78
+ error by 65–85%. Client data is not releasable; this repo's synthetic
79
+ benchmarks are the reproducible counterpart. DoWhy/EconML comparison suite:
80
+ in progress.
81
+
82
+ ## Performance (v0.2.0)
83
+
84
+ Screening uses a single precision-matrix inversion per round (FWL-equivalent,
85
+ with an exact residual-regression fallback for ill-conditioned designs) —
86
+ ~77x faster than the reference implementation at n=5000, p=40. Partitioned
87
+ prediction indexes cells at fit time and caches local models — ~409x faster
88
+ on batch prediction (20k train / 5k queries). Equivalence to the reference
89
+ implementation is enforced in the test suite to ~1e-8 (38 tests; reference
90
+ implementation vendored in `tests/`).
91
+
92
+ ## License
93
+
94
+ MIT
95
+
96
+ ## Pipeline
97
+
98
+ ```mermaid
99
+ flowchart LR
100
+ A[Training data] --> B["CausalScreen<br/>residual-correlation test<br/>+ iterative peeling"]
101
+ B -->|"ranked causal factors<br/>(decreasing causal power)"| C["Partition data on<br/>top causal factors<br/>(strength x significance rule)"]
102
+ C --> G1[(Subgroup 1)]
103
+ C --> G2[(Subgroup 2)]
104
+ C --> Gn[(Subgroup n)]
105
+ Q[Query point] --> R{inside a<br/>subgroup?}
106
+ R -->|"yes (interpolation)"| L["Train ONE local model<br/>in that subgroup only<br/>(Optimisation 2)"]
107
+ R -->|"no (extrapolation)"| F["Global model fallback<br/>(retains significance &<br/>generalization)"]
108
+ G1 -.-> L
109
+ G2 -.-> L
110
+ Gn -.-> L
111
+ B --> M["Missing causal power<br/>= drift early-warning alarm"]
112
+ ```
@@ -0,0 +1,4 @@
1
+ from .screening import CausalScreen, ScreenResult
2
+ from .local import PartitionedRegressor
3
+ __version__ = "0.1.0"
4
+ __all__ = ["CausalScreen", "ScreenResult", "PartitionedRegressor"]
@@ -0,0 +1,95 @@
1
+ """Causal-factor partitioning with inference-time localized training.
2
+
3
+ Optimized implementation: cells are indexed ONCE at fit time (dict of
4
+ cell-key -> row indices), local models are fitted at most once per cell
5
+ and cached, and batch prediction groups queries by cell so each cell's
6
+ model predicts its queries in a single vectorized call. Semantics are
7
+ identical to the per-query construction: a query gets a local model iff
8
+ its exact cell holds >= min_cell training rows, else the global model.
9
+ """
10
+ from __future__ import annotations
11
+
12
+ import numpy as np
13
+ from sklearn.base import clone
14
+ from sklearn.linear_model import LinearRegression
15
+
16
+
17
+ class PartitionedRegressor:
18
+ """Route each query to its causal-factor cell; train ONE local model there
19
+ (interpolation). Out-of-support queries fall back to the global model
20
+ (extrapolation), preserving generalization -- see paper Sec. 5.1.
21
+
22
+ Note: cell membership is exact-match on the partition features, so
23
+ partition features should be discrete/categorical. Continuous partition
24
+ features will rarely match and queries will fall back to the global model.
25
+ """
26
+
27
+ def __init__(self, base_estimator=None, partition_features=None, min_cell: int = 30):
28
+ if min_cell < 2:
29
+ raise ValueError("min_cell must be >= 2")
30
+ self.base = base_estimator if base_estimator is not None else LinearRegression()
31
+ self.partition_features = partition_features or []
32
+ self.min_cell = min_cell
33
+
34
+ def fit(self, X, y, feature_names):
35
+ self.names_ = list(feature_names)
36
+ self.X_, self.y_ = np.asarray(X, float), np.asarray(y, float).ravel()
37
+ if self.X_.ndim != 2 or self.X_.shape[0] != self.y_.shape[0]:
38
+ raise ValueError("X must be 2-D with rows matching y")
39
+ if not np.isfinite(self.X_).all() or not np.isfinite(self.y_).all():
40
+ raise ValueError("X or y contains NaN or inf")
41
+ missing = [f for f in self.partition_features if f not in self.names_]
42
+ if missing:
43
+ raise ValueError(f"partition features not in feature_names: {missing}")
44
+ self.idx_ = [self.names_.index(f) for f in self.partition_features]
45
+ self.global_ = clone(self.base).fit(self.X_, self.y_)
46
+ # Index cells once: cell key -> training-row indices.
47
+ self._cells: dict[tuple, np.ndarray] = {}
48
+ if self.idx_:
49
+ keys = self.X_[:, self.idx_]
50
+ order = np.lexsort(keys.T[::-1])
51
+ sorted_keys = keys[order]
52
+ boundaries = np.flatnonzero(
53
+ np.r_[True, (sorted_keys[1:] != sorted_keys[:-1]).any(axis=1)]
54
+ )
55
+ for s, e in zip(boundaries, np.r_[boundaries[1:], len(order)]):
56
+ if e - s >= self.min_cell:
57
+ self._cells[tuple(sorted_keys[s])] = order[s:e]
58
+ self._models: dict[tuple, object] = {} # lazy per-cell model cache
59
+ return self
60
+
61
+ def _cell_model(self, key: tuple):
62
+ """Fit (once) and return the cached local model for a qualifying cell."""
63
+ mdl = self._models.get(key)
64
+ if mdl is None:
65
+ rows = self._cells[key]
66
+ mdl = clone(self.base).fit(self.X_[rows], self.y_[rows])
67
+ self._models[key] = mdl
68
+ return mdl
69
+
70
+ def predict_one(self, xq) -> float:
71
+ xq = np.asarray(xq, float)
72
+ if self.idx_:
73
+ key = tuple(xq[self.idx_])
74
+ if key in self._cells:
75
+ return float(self._cell_model(key).predict(xq.reshape(1, -1))[0])
76
+ return float(self.global_.predict(xq.reshape(1, -1))[0])
77
+
78
+ def predict(self, Xq) -> np.ndarray:
79
+ Xq = np.asarray(Xq, float)
80
+ if Xq.ndim == 1:
81
+ Xq = Xq.reshape(1, -1)
82
+ out = np.empty(len(Xq), float)
83
+ if not self.idx_:
84
+ return self.global_.predict(Xq).astype(float)
85
+ # Group queries by cell so each model predicts once, vectorized.
86
+ keys = [tuple(row) for row in Xq[:, self.idx_]]
87
+ groups: dict[tuple, list[int]] = {}
88
+ fallback: list[int] = []
89
+ for i, key in enumerate(keys):
90
+ (groups.setdefault(key, []) if key in self._cells else fallback).append(i)
91
+ for key, idxs in groups.items():
92
+ out[idxs] = self._cell_model(key).predict(Xq[idxs])
93
+ if fallback:
94
+ out[fallback] = self.global_.predict(Xq[fallback])
95
+ return out
@@ -0,0 +1,159 @@
1
+ """Residual-correlation causal screening with iterative peeling.
2
+
3
+ Implements the two-model residual construction (Frisch-Waugh-Lovell partial
4
+ correlation) with an iterative peeling schedule, as described in:
5
+ Mandal, R.K. (2026). "Residual-Correlation Causal Screening with Localized
6
+ Supervised Modeling" (working paper).
7
+
8
+ Optimized implementation: each screening round computes ALL candidate
9
+ partial correlations from a single precision (inverse-covariance) matrix
10
+ of the remaining design -- O(m^3) per round -- instead of 2m separate
11
+ least-squares solves -- O(m * n * m^2). Results are numerically identical
12
+ to the residual construction (FWL theorem); a per-candidate residual
13
+ fallback guards ill-conditioned designs.
14
+ """
15
+ from __future__ import annotations
16
+
17
+ import numpy as np
18
+ from dataclasses import dataclass, field
19
+ from scipy import stats
20
+
21
+ _EPS = 1e-12
22
+
23
+
24
+ @dataclass
25
+ class ScreenResult:
26
+ ranking: list = field(default_factory=list) # feature names, decreasing causal power
27
+ causal_power: dict = field(default_factory=dict) # name -> partial correlation at selection round
28
+ p_value: dict = field(default_factory=dict) # name -> p-value at selection round
29
+ missing_causal_power: float = np.nan # share of response variance unexplained by ranked factors
30
+
31
+
32
+ def _validate(X: np.ndarray, y: np.ndarray) -> None:
33
+ if X.ndim != 2:
34
+ raise ValueError(f"X must be 2-D, got shape {X.shape}")
35
+ if y.ndim != 1:
36
+ raise ValueError(f"y must be 1-D after ravel, got shape {y.shape}")
37
+ if X.shape[0] != y.shape[0]:
38
+ raise ValueError(f"X and y have mismatched rows: {X.shape[0]} vs {y.shape[0]}")
39
+ if X.shape[0] < 3:
40
+ raise ValueError("need at least 3 samples")
41
+ if not np.isfinite(X).all():
42
+ raise ValueError("X contains NaN or inf")
43
+ if not np.isfinite(y).all():
44
+ raise ValueError("y contains NaN or inf")
45
+
46
+
47
+ def _residuals(y: np.ndarray, X: np.ndarray) -> np.ndarray:
48
+ """OLS residuals of y on X (with intercept)."""
49
+ if X.shape[1] == 0:
50
+ return y - y.mean()
51
+ A = np.column_stack([np.ones(len(X)), X])
52
+ beta, *_ = np.linalg.lstsq(A, y, rcond=None)
53
+ return y - A @ beta
54
+
55
+
56
+ def _partial_corr(y: np.ndarray, xj: np.ndarray, Xrest: np.ndarray) -> tuple[float, float]:
57
+ """Reference path: corr(resid(y|Xrest), resid(xj|Xrest)) and its p-value."""
58
+ ey = _residuals(y, Xrest)
59
+ ej = _residuals(xj, Xrest)
60
+ if ey.std() < _EPS or ej.std() < _EPS:
61
+ return 0.0, 1.0
62
+ r = float(np.corrcoef(ey, ej)[0, 1])
63
+ return r, _pval(r, len(y), Xrest.shape[1])
64
+
65
+
66
+ def _pval(r: float, n: int, k: int) -> float:
67
+ df = max(n - k - 2, 1)
68
+ t = r * np.sqrt(df / max(_EPS, 1.0 - r * r))
69
+ return float(2 * stats.t.sf(abs(t), df))
70
+
71
+
72
+ def _round_partial_corrs(y: np.ndarray, Xr: np.ndarray) -> np.ndarray | None:
73
+ """Partial corr of y with each column of Xr, given the other columns.
74
+
75
+ Single precision-matrix computation for the whole round:
76
+ Z = [Xr, y] centered; P = inv(cov(Z));
77
+ pcorr(y, x_j | rest) = -P[j, m] / sqrt(P[j, j] * P[m, m]).
78
+ Returns None if the design is too ill-conditioned to trust (caller
79
+ falls back to the per-candidate residual construction).
80
+ """
81
+ n, m = Xr.shape
82
+ Z = np.column_stack([Xr, y])
83
+ Z = Z - Z.mean(axis=0)
84
+ sd = Z.std(axis=0)
85
+ if (sd < _EPS).any():
86
+ return None # constant column: use reference path
87
+ Z /= sd # correlation scale for conditioning
88
+ C = (Z.T @ Z) / n
89
+ # reciprocal condition number guard
90
+ if np.linalg.cond(C) > 1e10:
91
+ return None
92
+ P = np.linalg.inv(C)
93
+ d = np.sqrt(np.diag(P))
94
+ r = -P[:m, m] / (d[:m] * d[m])
95
+ if not np.isfinite(r).all() or (np.abs(r) > 1.0 + 1e-8).any():
96
+ return None
97
+ return np.clip(r, -1.0, 1.0)
98
+
99
+
100
+ class CausalScreen:
101
+ """Iterative residual-correlation screen ('peeling').
102
+
103
+ Each round computes, for every remaining candidate j, the partial
104
+ correlation of X_j with y given all other remaining candidates; the
105
+ strongest significant candidate is recorded and REMOVED from all
106
+ subsequent conditioning sets (so a discovered factor's influence cannot
107
+ re-enter through the variables it confounds).
108
+ """
109
+
110
+ def __init__(self, alpha: float = 0.05, max_factors: int | None = None):
111
+ if not 0.0 < alpha <= 1.0:
112
+ raise ValueError("alpha must be in (0, 1]")
113
+ self.alpha = alpha
114
+ self.max_factors = max_factors
115
+
116
+ def fit(self, X, y, feature_names=None) -> ScreenResult:
117
+ X = np.asarray(X, dtype=float)
118
+ y = np.asarray(y, dtype=float).ravel()
119
+ _validate(X, y)
120
+ n, p = X.shape
121
+ names = list(feature_names) if feature_names is not None else [f"x{j}" for j in range(p)]
122
+ if len(names) != p:
123
+ raise ValueError(f"feature_names has {len(names)} entries for {p} columns")
124
+ remaining = list(range(p))
125
+ res = ScreenResult()
126
+ budget = self.max_factors or p
127
+ y_cur = y.copy() # response, progressively peeled of discovered factors
128
+ for _ in range(budget):
129
+ if not remaining:
130
+ break
131
+ m = len(remaining)
132
+ k = m - 1 # conditioning-set size per candidate
133
+ rs = _round_partial_corrs(y_cur, X[:, remaining])
134
+ if rs is None:
135
+ # ill-conditioned round: exact per-candidate reference path
136
+ best = None
137
+ for j in remaining:
138
+ rest = [t for t in remaining if t != j]
139
+ r, pv = _partial_corr(y_cur, X[:, j], X[:, rest])
140
+ if best is None or abs(r) > abs(best[1]):
141
+ best = (j, r, pv)
142
+ j, r, pv = best
143
+ else:
144
+ i = int(np.argmax(np.abs(rs)))
145
+ j, r = remaining[i], float(rs[i])
146
+ pv = _pval(r, n, k)
147
+ if pv > self.alpha:
148
+ break
149
+ res.ranking.append(names[j])
150
+ res.causal_power[names[j]] = r
151
+ res.p_value[names[j]] = pv
152
+ remaining.remove(j)
153
+ # PEEL: strip the discovered factor's contribution out of the
154
+ # response so its effect cannot re-enter the screen through
155
+ # variables it confounds (paper Sec. 4.3).
156
+ y_cur = _residuals(y_cur, X[:, [j]])
157
+ res.missing_causal_power = float(np.var(y_cur) / np.var(y)) if np.var(y) > 0 else np.nan
158
+ self.result_ = res
159
+ return res
@@ -0,0 +1,127 @@
1
+ Metadata-Version: 2.4
2
+ Name: causalscreen
3
+ Version: 0.2.0
4
+ Summary: Residual-correlation causal screening with iterative peeling, localized modeling, and drift early-warning
5
+ Author-email: Rahul Kumar Mandal <rahulkm.jobs@gmail.com>
6
+ License: MIT
7
+ Project-URL: Homepage, https://github.com/rahulkm3/causalscreen
8
+ Requires-Python: >=3.9
9
+ Description-Content-Type: text/markdown
10
+ License-File: LICENSE
11
+ Requires-Dist: numpy
12
+ Requires-Dist: scipy
13
+ Requires-Dist: scikit-learn
14
+ Dynamic: license-file
15
+
16
+ # causalscreen
17
+
18
+ Residual-correlation causal screening with iterative peeling, causal-factor
19
+ partitioning with inference-time local models, and a missing-causal-power
20
+ drift alarm.
21
+
22
+ Companion code for: **R. K. Mandal, "Residual-Correlation Causal Screening with
23
+ Localized Supervised Modeling: A Practical Framework for Confounder
24
+ Identification, Error Reduction, and Drift Early-Warning"** (working paper,
25
+ 2026; research 2022–present). Paper PDF in `/paper` · Preprint: https://papers.ssrn.com/sol3/papers.cfm?abstract_id=7135078
26
+
27
+ ## What it does
28
+
29
+ Supervised models fail in deployment because they lean on **confounded
30
+ proxies** — features whose correlation with the target is inherited from other
31
+ variables. `causalscreen`:
32
+
33
+ 1. **Screens** every candidate feature via the two-regression residual
34
+ construction (Frisch–Waugh–Lovell partial correlation): regress `y` on the
35
+ other features, regress `x_j` on the other features, correlate the
36
+ residuals. High residual correlation = genuine independent effect.
37
+ 2. **Peels** iteratively: each discovered factor's contribution is stripped
38
+ out of the response before the next round, so its influence cannot re-enter
39
+ the screen through the variables it confounds.
40
+ 3. **Ranks** features by causal power and reports the **missing causal
41
+ power** — the share of response variance the discovered factors cannot
42
+ explain — a deployable early-warning statistic for data/model drift.
43
+ 4. **Exploits** the ranking: `PartitionedRegressor` partitions data on top
44
+ causal factors and trains one small local model per query at inference
45
+ time (interpolation), falling back to the global model out of support
46
+ (extrapolation) — so the error reduction of local modeling never costs
47
+ out-of-sample generalization.
48
+
49
+ ## Quick result (synthetic ground truth, `examples/demo.py`)
50
+
51
+ | feature | naive \|corr\| with y | causal screen verdict |
52
+ |---|---|---|
53
+ | x1 (true cause) | 0.92 | **ranked** (power 0.78) |
54
+ | x2 (confounded proxy of x1) | 0.87 | **rejected** |
55
+ | x3 (noise) | 0.01 | rejected |
56
+ | x4 (true cause, weak) | 0.33 | **ranked** (power 0.83) |
57
+
58
+ Missing causal power: 5.1% — the exact noise floor of the generating process.
59
+ Partitioned local modeling: **99.0% out-of-sample MSE reduction** vs. a global
60
+ model on regime-structured data.
61
+
62
+ ## Install & use
63
+
64
+ ```bash
65
+ pip install -e .
66
+ python examples/demo.py
67
+ pytest tests/
68
+ ```
69
+
70
+ ```python
71
+ from causalscreen import CausalScreen, PartitionedRegressor
72
+ res = CausalScreen(alpha=0.05).fit(X, y, feature_names)
73
+ res.ranking # features in decreasing causal power
74
+ res.causal_power # partial correlation at selection round
75
+ res.missing_causal_power # drift-alarm statistic
76
+
77
+ pr = PartitionedRegressor(partition_features=res.ranking[:2], min_cell=50)
78
+ pr.fit(X, y, feature_names)
79
+ pr.predict(X_new) # local model in-support, global model out-of-support
80
+ ```
81
+
82
+ ## Assumptions (stated plainly)
83
+
84
+ Stationary cross-sectional data · causal sufficiency (all relevant variables
85
+ observed) · no reverse causation · features measured before the response.
86
+ The screen addresses the *confounding* component of causal analysis; it does
87
+ not prove causality in full generality. See paper §3, §7.
88
+
89
+ ## Deployed results
90
+
91
+ Layered over Random Forest, OLS and Prophet models in commercial deployments
92
+ (credit risk, segmentation, forecasting), the framework reduced out-of-sample
93
+ error by 65–85%. Client data is not releasable; this repo's synthetic
94
+ benchmarks are the reproducible counterpart. DoWhy/EconML comparison suite:
95
+ in progress.
96
+
97
+ ## Performance (v0.2.0)
98
+
99
+ Screening uses a single precision-matrix inversion per round (FWL-equivalent,
100
+ with an exact residual-regression fallback for ill-conditioned designs) —
101
+ ~77x faster than the reference implementation at n=5000, p=40. Partitioned
102
+ prediction indexes cells at fit time and caches local models — ~409x faster
103
+ on batch prediction (20k train / 5k queries). Equivalence to the reference
104
+ implementation is enforced in the test suite to ~1e-8 (38 tests; reference
105
+ implementation vendored in `tests/`).
106
+
107
+ ## License
108
+
109
+ MIT
110
+
111
+ ## Pipeline
112
+
113
+ ```mermaid
114
+ flowchart LR
115
+ A[Training data] --> B["CausalScreen<br/>residual-correlation test<br/>+ iterative peeling"]
116
+ B -->|"ranked causal factors<br/>(decreasing causal power)"| C["Partition data on<br/>top causal factors<br/>(strength x significance rule)"]
117
+ C --> G1[(Subgroup 1)]
118
+ C --> G2[(Subgroup 2)]
119
+ C --> Gn[(Subgroup n)]
120
+ Q[Query point] --> R{inside a<br/>subgroup?}
121
+ R -->|"yes (interpolation)"| L["Train ONE local model<br/>in that subgroup only<br/>(Optimisation 2)"]
122
+ R -->|"no (extrapolation)"| F["Global model fallback<br/>(retains significance &<br/>generalization)"]
123
+ G1 -.-> L
124
+ G2 -.-> L
125
+ Gn -.-> L
126
+ B --> M["Missing causal power<br/>= drift early-warning alarm"]
127
+ ```
@@ -0,0 +1,17 @@
1
+ LICENSE
2
+ MANIFEST.in
3
+ README.md
4
+ pyproject.toml
5
+ causalscreen/__init__.py
6
+ causalscreen/local.py
7
+ causalscreen/screening.py
8
+ causalscreen.egg-info/PKG-INFO
9
+ causalscreen.egg-info/SOURCES.txt
10
+ causalscreen.egg-info/dependency_links.txt
11
+ causalscreen.egg-info/requires.txt
12
+ causalscreen.egg-info/top_level.txt
13
+ examples/demo.py
14
+ tests/reference_impl_local.py
15
+ tests/reference_impl_screening.py
16
+ tests/test_optimized.py
17
+ tests/test_screening.py
@@ -0,0 +1,3 @@
1
+ numpy
2
+ scipy
3
+ scikit-learn
@@ -0,0 +1 @@
1
+ causalscreen
@@ -0,0 +1,39 @@
1
+ """Demo: confounded-proxy rejection + partitioned local modeling error reduction.
2
+ Run: python examples/demo.py
3
+ """
4
+ import numpy as np
5
+ from causalscreen import CausalScreen, PartitionedRegressor
6
+
7
+ rng = np.random.default_rng(7)
8
+ n = 3000
9
+
10
+ # --- Part 1: screening on synthetic ground truth ---
11
+ x1 = rng.normal(size=n) # true cause (strong)
12
+ x2 = 0.9 * x1 + 0.3 * rng.normal(size=n) # confounded proxy of x1
13
+ x3 = rng.normal(size=n) # noise
14
+ x4 = rng.normal(size=n) # true cause (weak)
15
+ y = 2.0 * x1 + 0.7 * x4 + 0.5 * rng.normal(size=n)
16
+ X = np.column_stack([x1, x2, x3, x4]); names = ["x1", "x2", "x3", "x4"]
17
+
18
+ res = CausalScreen().fit(X, y, names)
19
+ print("== Screening ==")
20
+ print(f"naive |corr| with y : " + ", ".join(f"{nm}={abs(np.corrcoef(X[:,j],y)[0,1]):.2f}" for j, nm in enumerate(names)))
21
+ print(f"causal ranking : {res.ranking}")
22
+ print(f"causal powers : " + ", ".join(f"{k}={v:.2f}" for k, v in res.causal_power.items()))
23
+ print(f"missing causal power: {res.missing_causal_power:.1%} (noise floor of the DGP is ~5%)")
24
+
25
+ # --- Part 2: partition on a causal categorical factor; local vs global model ---
26
+ g = rng.integers(0, 5, n).astype(float) # causal regime variable
27
+ xr = rng.normal(size=n)
28
+ slopes = np.array([2.0, -1.0, 0.5, 3.0, -2.5])
29
+ yr = slopes[g.astype(int)] * xr + 0.2 * rng.normal(size=n)
30
+ Xr = np.column_stack([g, xr])
31
+ tr, te = slice(0, 2500), slice(2500, None)
32
+
33
+ pr = PartitionedRegressor(partition_features=["g"], min_cell=50) # linear base model
34
+ pr.fit(Xr[tr], yr[tr], ["g", "x"])
35
+ mse_local = np.mean((pr.predict(Xr[te]) - yr[te]) ** 2)
36
+ mse_global = np.mean((pr.global_.predict(Xr[te]) - yr[te]) ** 2)
37
+ print("\n== Partitioned local modeling (out-of-sample) ==")
38
+ print(f"global model MSE : {mse_global:.4f}")
39
+ print(f"local models MSE : {mse_local:.4f} (reduction: {1 - mse_local/mse_global:.1%})")
@@ -0,0 +1,19 @@
1
+ [build-system]
2
+ requires = ["setuptools>=61"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "causalscreen"
7
+ version = "0.2.0"
8
+ description = "Residual-correlation causal screening with iterative peeling, localized modeling, and drift early-warning"
9
+ authors = [{name = "Rahul Kumar Mandal", email = "rahulkm.jobs@gmail.com"}]
10
+ readme = "README.md"
11
+ license = {text = "MIT"}
12
+ requires-python = ">=3.9"
13
+ dependencies = ["numpy", "scipy", "scikit-learn"]
14
+
15
+ [project.urls]
16
+ Homepage = "https://github.com/rahulkm3/causalscreen"
17
+
18
+ [tool.setuptools.packages.find]
19
+ include = ["causalscreen*"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,42 @@
1
+ """Reference per-query local regressor. Test oracle only."""
2
+ import numpy as np
3
+ from sklearn.base import clone
4
+ from sklearn.linear_model import LinearRegression
5
+
6
+
7
+ class PartitionedRegressor:
8
+ """Route each query to its causal-factor cell; train ONE local model there
9
+ (interpolation). Out-of-support queries fall back to the global model
10
+ (extrapolation), preserving generalization -- see paper Sec. 5.1.
11
+ """
12
+
13
+ def __init__(self, base_estimator=None, partition_features=None, min_cell=30):
14
+ self.base = base_estimator if base_estimator is not None else LinearRegression()
15
+ self.partition_features = partition_features or []
16
+ self.min_cell = min_cell
17
+
18
+ def fit(self, X, y, feature_names):
19
+ self.names_ = list(feature_names)
20
+ self.X_, self.y_ = np.asarray(X, float), np.asarray(y, float).ravel()
21
+ self.idx_ = [self.names_.index(f) for f in self.partition_features]
22
+ self.global_ = clone(self.base).fit(self.X_, self.y_)
23
+ return self
24
+
25
+ def _cell_mask(self, xq):
26
+ mask = np.ones(len(self.X_), bool)
27
+ for i in self.idx_:
28
+ mask &= (self.X_[:, i] == xq[i])
29
+ if mask.sum() < self.min_cell: # stop before losing significance
30
+ return None
31
+ return mask
32
+
33
+ def predict_one(self, xq):
34
+ xq = np.asarray(xq, float)
35
+ mask = self._cell_mask(xq) if self.idx_ else None
36
+ if mask is not None and mask.sum() >= self.min_cell:
37
+ local = clone(self.base).fit(self.X_[mask], self.y_[mask])
38
+ return float(local.predict(xq.reshape(1, -1))[0])
39
+ return float(self.global_.predict(xq.reshape(1, -1))[0])
40
+
41
+ def predict(self, Xq):
42
+ return np.array([self.predict_one(x) for x in np.asarray(Xq, float)])
@@ -0,0 +1,89 @@
1
+ """Reference implementation (naive residual construction). Test oracle only.
2
+
3
+ Implements the two-model residual construction (Frisch-Waugh-Lovell partial
4
+ correlation) with an iterative peeling schedule, as described in:
5
+ Mandal, R.K. (2026). "Residual-Correlation Causal Screening with Localized
6
+ Supervised Modeling" (working paper).
7
+ """
8
+ import numpy as np
9
+ from dataclasses import dataclass, field
10
+ from scipy import stats
11
+
12
+
13
+ @dataclass
14
+ class ScreenResult:
15
+ ranking: list = field(default_factory=list) # feature names, decreasing causal power
16
+ causal_power: dict = field(default_factory=dict) # name -> partial correlation at selection round
17
+ p_value: dict = field(default_factory=dict) # name -> p-value at selection round
18
+ missing_causal_power: float = np.nan # share of response variance unexplained by ranked factors
19
+
20
+
21
+ def _residuals(y, X):
22
+ """OLS residuals of y on X (with intercept)."""
23
+ if X.shape[1] == 0:
24
+ return y - y.mean()
25
+ A = np.column_stack([np.ones(len(X)), X])
26
+ beta, *_ = np.linalg.lstsq(A, y, rcond=None)
27
+ return y - A @ beta
28
+
29
+
30
+ def _partial_corr(y, xj, Xrest):
31
+ """corr(resid(y|Xrest), resid(xj|Xrest)) and its p-value."""
32
+ ey = _residuals(y, Xrest)
33
+ ej = _residuals(xj, Xrest)
34
+ if ey.std() < 1e-12 or ej.std() < 1e-12:
35
+ return 0.0, 1.0
36
+ r = float(np.corrcoef(ey, ej)[0, 1])
37
+ n, k = len(y), Xrest.shape[1]
38
+ df = max(n - k - 2, 1)
39
+ t = r * np.sqrt(df / max(1e-12, 1 - r * r))
40
+ p = 2 * stats.t.sf(abs(t), df)
41
+ return r, float(p)
42
+
43
+
44
+ class CausalScreen:
45
+ """Iterative residual-correlation screen ('peeling').
46
+
47
+ Each round computes, for every remaining candidate j, the partial
48
+ correlation of X_j with y given all other remaining candidates; the
49
+ strongest significant candidate is recorded and REMOVED from all
50
+ subsequent conditioning sets (so a discovered factor's influence cannot
51
+ re-enter through the variables it confounds).
52
+ """
53
+
54
+ def __init__(self, alpha=0.05, max_factors=None):
55
+ self.alpha = alpha
56
+ self.max_factors = max_factors
57
+
58
+ def fit(self, X, y, feature_names=None):
59
+ X = np.asarray(X, dtype=float)
60
+ y = np.asarray(y, dtype=float).ravel()
61
+ p = X.shape[1]
62
+ names = list(feature_names) if feature_names is not None else [f"x{j}" for j in range(p)]
63
+ remaining = list(range(p))
64
+ res = ScreenResult()
65
+ budget = self.max_factors or p
66
+ y_cur = y.copy() # response, progressively peeled of discovered factors
67
+ for _ in range(budget):
68
+ if not remaining:
69
+ break
70
+ best = None
71
+ for j in remaining:
72
+ rest = [k for k in remaining if k != j]
73
+ r, pv = _partial_corr(y_cur, X[:, j], X[:, rest])
74
+ if best is None or abs(r) > abs(best[1]):
75
+ best = (j, r, pv)
76
+ j, r, pv = best
77
+ if pv > self.alpha:
78
+ break
79
+ res.ranking.append(names[j])
80
+ res.causal_power[names[j]] = r
81
+ res.p_value[names[j]] = pv
82
+ remaining.remove(j)
83
+ # PEEL: strip the discovered factor's contribution out of the
84
+ # response so its effect cannot re-enter the screen through
85
+ # variables it confounds (paper Sec. 4.3).
86
+ y_cur = _residuals(y_cur, X[:, [j]])
87
+ res.missing_causal_power = float(np.var(y_cur) / np.var(y)) if np.var(y) > 0 else np.nan
88
+ self.result_ = res
89
+ return res
@@ -0,0 +1,152 @@
1
+ """Robustness suite: equivalence with the reference residual construction,
2
+ edge cases, batch/loop consistency, and input validation."""
3
+ import sys, pathlib
4
+ import numpy as np
5
+ import pytest
6
+
7
+ import importlib.util
8
+
9
+ _HERE = pathlib.Path(__file__).parent
10
+
11
+ def _load(name, path):
12
+ spec = importlib.util.spec_from_file_location(name, str(path))
13
+ mod = importlib.util.module_from_spec(spec)
14
+ spec.loader.exec_module(mod)
15
+ return mod
16
+
17
+ # Reference implementations: the naive residual (FWL) construction and the
18
+ # per-query local regressor, kept as the mathematical spec the optimized
19
+ # code must reproduce exactly.
20
+ ref_screen = _load("reference_impl_screening", _HERE / "reference_impl_screening.py")
21
+ ref_local = _load("reference_impl_local", _HERE / "reference_impl_local.py")
22
+
23
+ from causalscreen import CausalScreen, PartitionedRegressor
24
+ from causalscreen.screening import _round_partial_corrs, _partial_corr
25
+
26
+
27
+ # ---------- screening: numerical equivalence with reference ----------
28
+
29
+ @pytest.mark.parametrize("seed", range(6))
30
+ @pytest.mark.parametrize("n,p", [(200, 5), (500, 10), (120, 8)])
31
+ def test_round_pcorrs_match_residual_construction(seed, n, p):
32
+ rng = np.random.default_rng(seed)
33
+ A = rng.normal(size=(p, p))
34
+ X = rng.normal(size=(n, p)) @ A # correlated design
35
+ y = X @ rng.normal(size=p) + rng.normal(size=n)
36
+ rs = _round_partial_corrs(y, X)
37
+ assert rs is not None
38
+ for j in range(p):
39
+ rest = np.delete(np.arange(p), j)
40
+ r_ref, _ = _partial_corr(y, X[:, j], X[:, rest])
41
+ assert rs[j] == pytest.approx(r_ref, abs=1e-8)
42
+
43
+ @pytest.mark.parametrize("seed", range(5))
44
+ def test_full_screen_matches_reference(seed):
45
+ rng = np.random.default_rng(seed)
46
+ n, p = 800, 7
47
+ X = rng.normal(size=(n, p))
48
+ X[:, 1] = 0.8 * X[:, 0] + 0.3 * rng.normal(size=n) # proxy
49
+ y = 1.5 * X[:, 0] - 0.9 * X[:, 4] + 0.5 * rng.normal(size=n)
50
+ names = [f"f{j}" for j in range(p)]
51
+ new = CausalScreen().fit(X, y, names)
52
+ old = ref_screen.CausalScreen().fit(X, y, names)
53
+ assert new.ranking == old.ranking
54
+ for f in new.ranking:
55
+ assert new.causal_power[f] == pytest.approx(old.causal_power[f], abs=1e-7)
56
+ assert new.p_value[f] == pytest.approx(old.p_value[f], rel=1e-4, abs=1e-12)
57
+ assert new.missing_causal_power == pytest.approx(old.missing_causal_power, abs=1e-9)
58
+
59
+ def test_constant_feature_falls_back_and_matches_reference():
60
+ rng = np.random.default_rng(3)
61
+ n = 300
62
+ X = rng.normal(size=(n, 4)); X[:, 2] = 5.0 # constant column
63
+ y = 2 * X[:, 0] + rng.normal(size=n)
64
+ names = list("abcd")
65
+ new = CausalScreen().fit(X, y, names)
66
+ old = ref_screen.CausalScreen().fit(X, y, names)
67
+ assert new.ranking == old.ranking
68
+ assert "c" not in new.ranking
69
+
70
+ def test_perfectly_collinear_features_do_not_crash():
71
+ rng = np.random.default_rng(4)
72
+ n = 300
73
+ x = rng.normal(size=n)
74
+ X = np.column_stack([x, 2 * x, rng.normal(size=n)]) # exact collinearity
75
+ y = x + 0.1 * rng.normal(size=n)
76
+ res = CausalScreen().fit(X, y, ["a", "a2", "b"])
77
+ assert isinstance(res.ranking, list) # no crash; sane output
78
+ assert np.isfinite(res.missing_causal_power)
79
+
80
+ def test_max_factors_budget_respected():
81
+ rng = np.random.default_rng(5)
82
+ X = rng.normal(size=(500, 6))
83
+ y = X @ np.ones(6) + 0.1 * rng.normal(size=500)
84
+ res = CausalScreen(max_factors=2).fit(X, y)
85
+ assert len(res.ranking) <= 2
86
+
87
+ def test_screen_input_validation():
88
+ with pytest.raises(ValueError):
89
+ CausalScreen(alpha=0.0)
90
+ with pytest.raises(ValueError):
91
+ CausalScreen().fit(np.ones((5, 2)), np.ones(4))
92
+ Xbad = np.ones((10, 2)); Xbad[0, 0] = np.nan
93
+ with pytest.raises(ValueError):
94
+ CausalScreen().fit(Xbad, np.ones(10))
95
+ with pytest.raises(ValueError):
96
+ CausalScreen().fit(np.random.default_rng(0).normal(size=(10, 3)),
97
+ np.zeros(10), ["a", "b"]) # name-count mismatch
98
+
99
+ # ---------- partitioned regressor: equivalence + batching ----------
100
+
101
+ def _cells_data(seed=1, n=2500):
102
+ rng = np.random.default_rng(seed)
103
+ g = rng.integers(0, 5, n).astype(float)
104
+ h = rng.integers(0, 3, n).astype(float)
105
+ x = rng.normal(size=n)
106
+ y = (g - 2) * x + 0.5 * h + 0.2 * rng.normal(size=n)
107
+ return np.column_stack([g, h, x]), y, ["g", "h", "x"]
108
+
109
+ def test_predictions_match_reference_implementation():
110
+ X, y, names = _cells_data()
111
+ Xq = X[:300]
112
+ new = PartitionedRegressor(partition_features=["g", "h"], min_cell=40).fit(X, y, names)
113
+ old = ref_local.PartitionedRegressor(partition_features=["g", "h"], min_cell=40).fit(X, y, names)
114
+ np.testing.assert_allclose(new.predict(Xq), old.predict(Xq), atol=1e-8)
115
+
116
+ def test_batch_predict_equals_loop_predict():
117
+ X, y, names = _cells_data(seed=2)
118
+ pr = PartitionedRegressor(partition_features=["g"], min_cell=40).fit(X, y, names)
119
+ Xq = np.vstack([X[:100], [[99.0, 0.0, 0.3]]]) # incl. unseen cell
120
+ loop = np.array([pr.predict_one(x) for x in Xq])
121
+ np.testing.assert_allclose(pr.predict(Xq), loop, atol=1e-10)
122
+
123
+ def test_out_of_support_uses_global_model():
124
+ X, y, names = _cells_data(seed=3)
125
+ pr = PartitionedRegressor(partition_features=["g"], min_cell=40).fit(X, y, names)
126
+ xq = np.array([42.0, 1.0, 0.5]) # cell never seen
127
+ assert pr.predict_one(xq) == pytest.approx(float(pr.global_.predict(xq.reshape(1, -1))[0]))
128
+
129
+ def test_small_cell_falls_back_to_global():
130
+ X, y, names = _cells_data(seed=4)
131
+ pr = PartitionedRegressor(partition_features=["g"], min_cell=10**6).fit(X, y, names)
132
+ np.testing.assert_allclose(pr.predict(X[:50]), pr.global_.predict(X[:50]), atol=1e-10)
133
+
134
+ def test_no_partition_features_is_pure_global():
135
+ X, y, names = _cells_data(seed=5)
136
+ pr = PartitionedRegressor().fit(X, y, names)
137
+ np.testing.assert_allclose(pr.predict(X[:50]), pr.global_.predict(X[:50]), atol=1e-10)
138
+
139
+ def test_local_models_cached_one_fit_per_cell():
140
+ X, y, names = _cells_data(seed=6)
141
+ pr = PartitionedRegressor(partition_features=["g"], min_cell=40).fit(X, y, names)
142
+ pr.predict(X[:500]); n_models = len(pr._models)
143
+ pr.predict(X[:500]) # second pass: no refits
144
+ assert len(pr._models) == n_models
145
+ assert n_models <= len(pr._cells)
146
+
147
+ def test_partition_validation():
148
+ X, y, names = _cells_data(seed=7)
149
+ with pytest.raises(ValueError):
150
+ PartitionedRegressor(partition_features=["nope"]).fit(X, y, names)
151
+ with pytest.raises(ValueError):
152
+ PartitionedRegressor(min_cell=1)
@@ -0,0 +1,47 @@
1
+ import numpy as np
2
+ from causalscreen import CausalScreen, PartitionedRegressor
3
+
4
+ def _synth(n=4000, seed=0):
5
+ """Ground truth: x1 true cause; x2 = confounded proxy of x1 (no own effect);
6
+ x3 pure noise; x4 weaker true cause."""
7
+ rng = np.random.default_rng(seed)
8
+ x1 = rng.normal(size=n)
9
+ x2 = 0.9 * x1 + 0.3 * rng.normal(size=n) # rides on x1
10
+ x3 = rng.normal(size=n)
11
+ x4 = rng.normal(size=n)
12
+ y = 2.0 * x1 + 0.7 * x4 + 0.5 * rng.normal(size=n)
13
+ X = np.column_stack([x1, x2, x3, x4])
14
+ return X, y
15
+
16
+ def test_ranking_recovers_true_causes():
17
+ X, y = _synth()
18
+ res = CausalScreen().fit(X, y, ["x1", "x2", "x3", "x4"])
19
+ assert set(res.ranking[:2]) == {"x1", "x4"} # both true causes found
20
+ assert "x2" not in res.ranking # confounded proxy rejected
21
+ assert "x3" not in res.ranking # noise rejected
22
+
23
+ def test_raw_corr_vs_causal_power_exposes_proxy():
24
+ X, y = _synth()
25
+ raw2 = abs(np.corrcoef(X[:, 1], y)[0, 1])
26
+ res = CausalScreen().fit(X, y, ["x1", "x2", "x3", "x4"])
27
+ assert raw2 > 0.8 # proxy looks great naively...
28
+ assert "x2" not in res.ranking # ...but is rejected by the screen
29
+
30
+ def test_missing_causal_power_bounds():
31
+ X, y = _synth()
32
+ res = CausalScreen().fit(X, y, ["x1", "x2", "x3", "x4"])
33
+ assert 0.0 <= res.missing_causal_power < 0.2
34
+
35
+ def test_partitioned_regressor_runs_and_beats_nothing():
36
+ rng = np.random.default_rng(1)
37
+ n = 3000
38
+ g = rng.integers(0, 4, n).astype(float) # categorical causal factor
39
+ x = rng.normal(size=n)
40
+ slopes = np.array([2.0, -1.0, 0.5, 3.0])
41
+ y = slopes[g.astype(int)] * x + 0.2 * rng.normal(size=n)
42
+ X = np.column_stack([g, x])
43
+ pr = PartitionedRegressor(partition_features=["g"], min_cell=50).fit(X, y, ["g", "x"])
44
+ Xq = X[:200]
45
+ err_local = np.mean((pr.predict(Xq) - y[:200]) ** 2)
46
+ err_global = np.mean((pr.global_.predict(Xq) - y[:200]) ** 2)
47
+ assert err_local < 0.25 * err_global # local models slash error in-cell