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.
- causalscreen-0.2.0/LICENSE +21 -0
- causalscreen-0.2.0/MANIFEST.in +3 -0
- causalscreen-0.2.0/PKG-INFO +127 -0
- causalscreen-0.2.0/README.md +112 -0
- causalscreen-0.2.0/causalscreen/__init__.py +4 -0
- causalscreen-0.2.0/causalscreen/local.py +95 -0
- causalscreen-0.2.0/causalscreen/screening.py +159 -0
- causalscreen-0.2.0/causalscreen.egg-info/PKG-INFO +127 -0
- causalscreen-0.2.0/causalscreen.egg-info/SOURCES.txt +17 -0
- causalscreen-0.2.0/causalscreen.egg-info/dependency_links.txt +1 -0
- causalscreen-0.2.0/causalscreen.egg-info/requires.txt +3 -0
- causalscreen-0.2.0/causalscreen.egg-info/top_level.txt +1 -0
- causalscreen-0.2.0/examples/demo.py +39 -0
- causalscreen-0.2.0/pyproject.toml +19 -0
- causalscreen-0.2.0/setup.cfg +4 -0
- causalscreen-0.2.0/tests/reference_impl_local.py +42 -0
- causalscreen-0.2.0/tests/reference_impl_screening.py +89 -0
- causalscreen-0.2.0/tests/test_optimized.py +152 -0
- causalscreen-0.2.0/tests/test_screening.py +47 -0
|
@@ -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,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,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 @@
|
|
|
1
|
+
|
|
@@ -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,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
|