|
| 1 | +# Assimilative Causal Inference |
| 2 | + |
| 3 | +> *Nature Communications* 2026: bridging Bayesian data assimilation |
| 4 | +> (weather forecasting, oceanography) with streaming causal inference. |
| 5 | +
|
| 6 | +## 1. What problem does this solve? |
| 7 | + |
| 8 | +Most causal-inference pipelines assume the data arrive in one shot. |
| 9 | +But in A/B testing, pharmacovigilance, and policy evaluation you |
| 10 | +typically get **a batch at a time**, and the treatment effect itself |
| 11 | +may drift with seasons, user segments, or policy regimes. |
| 12 | + |
| 13 | +Assimilative causal inference treats the causal effect as a |
| 14 | +time-varying latent state and updates its posterior every time a new |
| 15 | +batch arrives — exactly the same mathematics that a weather model |
| 16 | +uses to fuse satellite observations into a global atmospheric state. |
| 17 | + |
| 18 | +## 2. State-space formulation |
| 19 | + |
| 20 | +$$ |
| 21 | +\begin{aligned} |
| 22 | +\theta_t &= \theta_{t-1} + w_t, |
| 23 | + \qquad w_t \sim \mathcal{N}(0, Q_t) \quad \text{(dynamics)}\\[2pt] |
| 24 | +\widehat{\theta}_t &= \theta_t + e_t, |
| 25 | + \qquad e_t \sim \mathcal{N}(0, \sigma_t^2) \quad \text{(observation)} |
| 26 | +\end{aligned} |
| 27 | +$$ |
| 28 | + |
| 29 | +Under Gaussianity the Kalman filter gives closed-form posteriors |
| 30 | +`θ_t | y_{1:t} ~ N(m_t, P_t)`. For non-Gaussian / nonlinear settings |
| 31 | +StatsPAI ships a bootstrap-SIR particle filter with systematic |
| 32 | +resampling. |
| 33 | + |
| 34 | +## 3. Two-level API |
| 35 | + |
| 36 | +### Low-level: fuse a pre-computed stream |
| 37 | + |
| 38 | +If you have already produced a stream of `(θ̂_t, σ_t)` from, say, |
| 39 | +running `sp.did` or `sp.dml` on each batch: |
| 40 | + |
| 41 | +```python |
| 42 | +import statspai as sp |
| 43 | + |
| 44 | +res = sp.causal_kalman( |
| 45 | + estimates=[θ̂_1, θ̂_2, ...], |
| 46 | + standard_errors=[σ_1, σ_2, ...], |
| 47 | + prior_mean=0.0, |
| 48 | + prior_var=1.0, |
| 49 | + process_var=0.0, # 0 = static effect, >0 = random-walk drift |
| 50 | + alpha=0.05, |
| 51 | +) |
| 52 | +print(res.summary()) |
| 53 | +``` |
| 54 | + |
| 55 | +For heavy-tailed observation noise or non-Gaussian priors swap in the |
| 56 | +particle filter: |
| 57 | + |
| 58 | +```python |
| 59 | +res = sp.assimilation.particle_filter( |
| 60 | + estimates, standard_errors, |
| 61 | + n_particles=3000, |
| 62 | + process_sd=0.05, # random-walk SD |
| 63 | + random_state=0, |
| 64 | +) |
| 65 | +``` |
| 66 | + |
| 67 | +### High-level: fuse directly from raw batches |
| 68 | + |
| 69 | +Pass the batches and a `(df → (θ̂, σ))` estimator callback and the |
| 70 | +pipeline handles both per-batch estimation and assimilation: |
| 71 | + |
| 72 | +```python |
| 73 | +def est(df): |
| 74 | + r = sp.regress('y ~ d + x', data=df) |
| 75 | + return float(r.params['d']), float(r.std_errors['d']) |
| 76 | + |
| 77 | +res = sp.assimilative_causal( |
| 78 | + batches=[batch_jan, batch_feb, ...], |
| 79 | + estimator=est, |
| 80 | + prior_mean=0.0, |
| 81 | + prior_var=1.0, |
| 82 | + process_var=0.01, # allow small month-to-month drift |
| 83 | + backend='kalman', # or 'particle' for non-Gaussian |
| 84 | +) |
| 85 | +``` |
| 86 | + |
| 87 | +## 4. Outputs |
| 88 | + |
| 89 | +Every backend returns the same `AssimilationResult` dataclass: |
| 90 | + |
| 91 | +| Field | Shape | Meaning | |
| 92 | +|--------------------|--------|-------------------------------------| |
| 93 | +| `posterior_mean` | (T,) | Running posterior mean `m_t` | |
| 94 | +| `posterior_sd` | (T,) | Running posterior SD `sqrt(P_t)` | |
| 95 | +| `posterior_ci` | (T, 2) | Per-step CI at level `alpha` | |
| 96 | +| `innovations` | (T,) | Observation surprise `θ̂_t − m_{t|t-1}` | |
| 97 | +| `ess` | (T,) | Effective sample size | |
| 98 | +| `final_mean`, `final_sd`, `final_ci` | scalars | End-of-stream summary | |
| 99 | +| `trajectory()` | method | Tidy DataFrame of the above | |
| 100 | + |
| 101 | +## 5. End-to-end example |
| 102 | + |
| 103 | +```python |
| 104 | +import numpy as np, pandas as pd, statspai as sp |
| 105 | + |
| 106 | +# Generate 12 monthly A/B-test batches with a slowly drifting effect. |
| 107 | +rng = np.random.default_rng(0) |
| 108 | +def make_batch(n, tau, seed): |
| 109 | + r = np.random.default_rng(seed) |
| 110 | + d = r.integers(0, 2, n); x = r.normal(size=n) |
| 111 | + y = tau * d + 0.2 * x + r.normal(scale=0.3, size=n) |
| 112 | + return pd.DataFrame({'y': y, 'd': d, 'x': x}) |
| 113 | + |
| 114 | +tau_path = 0.5 + np.linspace(0, 0.1, 12) # drifts from 0.5 to 0.6 |
| 115 | +batches = [make_batch(300, tau_path[t], seed=t) for t in range(12)] |
| 116 | + |
| 117 | +def est(df): |
| 118 | + r = sp.regress('y ~ d + x', data=df) |
| 119 | + return float(r.params['d']), float(r.std_errors['d']) |
| 120 | + |
| 121 | +res = sp.assimilative_causal( |
| 122 | + batches, est, prior_mean=0.0, prior_var=1.0, |
| 123 | + process_var=0.001, # mild drift allowance |
| 124 | +) |
| 125 | +print(res.summary()) |
| 126 | +print(res.trajectory().tail()) |
| 127 | +``` |
| 128 | + |
| 129 | +## 6. Choosing between the backends |
| 130 | + |
| 131 | +| Symptom | Use | |
| 132 | +|-------------------------------------------|----------------------------------------| |
| 133 | +| Per-batch estimates look roughly Gaussian | `backend='kalman'` (default, fast) | |
| 134 | +| Heavy tails / outliers per batch | `backend='particle'` with Student-t obs model | |
| 135 | +| Non-Gaussian prior (log-normal, bounded) | `backend='particle'` + `prior_sampler=...` | |
| 136 | +| You need CI bands in real time | Either — both return `posterior_ci` per step | |
| 137 | + |
| 138 | +## 7. Gotchas |
| 139 | + |
| 140 | +- **Garbage-in, garbage-out.** The filter trusts that your per-batch |
| 141 | + estimator is well-calibrated (nominal CIs cover at their stated |
| 142 | + rate). Run `sp.smart.assumption_audit` on the estimator first. |
| 143 | +- **`process_var = 0` means a static effect.** If the real effect |
| 144 | + drifts, you'll get overconfident CIs because the filter has no room |
| 145 | + for state innovation. Start with `process_var ≈ σ_t² / 10` and |
| 146 | + tune upwards if the innovations look persistently one-sided. |
| 147 | +- **Particle-filter degeneracy.** If you see |
| 148 | + `ESS / N < 0.2` persistently, increase `n_particles` or loosen the |
| 149 | + observation model (heavier tails). |
| 150 | + |
| 151 | +## 8. References |
| 152 | + |
| 153 | +- *Assimilative Causal Inference* — Nature Communications (2026). |
| 154 | +- Gordon, Salmond & Smith (1993). "Novel approach to nonlinear/non- |
| 155 | + Gaussian Bayesian state estimation." *IEE Proc. F.* |
| 156 | +- Douc & Cappé (2005). "Comparison of resampling schemes for |
| 157 | + particle filtering." *ISPA*. |
0 commit comments