Skip to content

chc.network_causal

Stability tier: evolving.

Network / interference-aware causal inference: direct + spillover effects when SUTVA fails.

On a network a unit's outcome depends on its neighbours' treatments too (interference), so an effect estimate that ignores the exposure is blind to the spillover. This adds a cross-fitted, Neyman-orthogonal Double ML that estimates both the direct and the spillover effect, using graph-aggregated (mean-neighbour) covariates as the nuisance features -- the lean, JAX-native counterpart of GNN-nuisance network DML (cf. arXiv 2509.18484, 2211.07823). Identifiability caveat: like confounding, the spillover is only recovered under the stated exposure model (mean neighbour treatment); real pilots need cluster/geo randomisation.

ConfoundedNetworkSystem dataclass

Units on a graph; y = a*x + direct*u + spillover*e + c*z, u confounded by z.

The confounder z is spatially smooth (neighbours share it), so the treatment u and the exposure e (mean neighbour treatment) correlate; an estimate that omits e misses it.

sample(key)

Draw n units as columns x, z, u, e, x_nb, z_nb, x_next (e = exposure).

DelayedNetworkPanel dataclass

Units on a graph observed over TIME, where a spillover takes lag steps per graph hop.

ConfoundedNetworkSystem draws one cross-section, so a propagation delay cannot even be written down there, let alone estimated. This is its panel counterpart: n_clusters disjoint cycles of cluster_size units, each observed at n_times times, with

e_i(t) = sum_d gamma_d * sum_{j : dist(i,j) = d} u_j(t - lag*d)

so a shell-d neighbour's treatment arrives lag*d steps late. That is exactly the generative model Result 51 is stated for, which is the point: it makes lag and the serial correlation phi ESTIMABLE from a log instead of assumed.

disturbance_scale adds an INDEPENDENT draw of the same propagated process to the outcome, and it is what decides whether Result 51 says anything about an estimator rather than only about a process. The delayed-network covariance the law is a functional of is the one the SCORE's noise carries; with the default i.i.d. outcome noise the score is white, the fold operator cannot move the estimator's variance at all, and the two fold schemes measure equal variance whatever the law predicts. This is the delayed-network analogue of the cluster random effect that carries the ICC in chc.regret.cluster_fold_leakage_certificate. It defaults to zero so the effect-recovery use of the panel is unchanged.

Construction detail that decides whether the law applies. The exogenous part of the treatment, eta, is i.i.d. ACROSS UNITS and AR(1) in time, while the confounder z carries the spatial smoothness. Result 51's covariance is the one the RESIDUALISED score has, and residualising on a nuisance that sees z leaves the eta-driven part -- so putting the spatial correlation in z and not in eta is what makes the law's premise hold after cross-fitting rather than before it. A spatially smooth eta would add a second correlation layer and the polynomial would be an approximation; that variant is reachable by raising eta_smoothing.

graph replaces the cycle with any regular adjacency, given as nested tuples so the frozen dataclass stays a value (torus_adjacency builds one). Regularity is required and checked, because the neighbours column is rectangular and a varying degree would have to be padded with a sentinel that every consumer would then have to know about. None keeps the cycle and the byte-identical draw. The topology matters to the design, not to the law: Result 52's closed form needs vertex transitivity, and on a torus Q(x) is not even circulant, so chc.regret.optimal_fold_partition takes the search route Result 61 priced.

sample takes a JAX key like its cross-sectional sibling and returns the same column-dict, so estimate_network_effects accepts a panel with no adapter. Internally the propagation runs through the NumPy routine chc.regret.delayed_network_certificate is verified against; the seed is derived from the key. One routine, one law -- a JAX re-implementation would be a second copy of the exact thing the certificate checks.

sample(key)

Draw the panel as flat columns, ordered cluster-major, then unit, then time.

PropagationEstimate dataclass

The two parameters Result 51's polynomial needs, read off a panel rather than assumed.

delay is the per-shell propagation lag delta in steps: a shell-d neighbour's treatment reaches the unit at lag delta * d. phi is the treatment's own AR(1) coefficient in time. Together they fix x = phi ** delta, the variable the delayed-network variance ratio is a polynomial in.

shell_lags is the peak lag measured at each shell separately and n_shells is the spillover truncation those lags imply -- the polynomial's degree, not the graph's diameter. lo and hi bound delay by a cluster bootstrap. censored propagates chc.irf.peak_lag's flag: a peak on the end of the horizon bounds the delay from below and is not an estimate of it.

NeighbourMessagePassing

Bases: Module

A mean-aggregation message-passing regressor that learns what to aggregate over neighbours.

Each layer is h <- tanh(W_self h + W_neigh mean_{j in N(i)} h_j), so it can represent mean_j f(h_j) for a learned f -- richer than a fixed mean-neighbour covariate.

cycle_shells(m, dmax)

0/1 distance-d indicator matrices of the cycle C_m, for d = 0..dmax.

S_d[i,j] = 1 iff j is exactly d hops from i. The shells partition the graph, so tr(S_d S_e) = 0 for d != e -- a vertex cannot sit at two distinct distances from i. That is the structural lemma Result 51 rests on.

graph_shells(adjacency, dmax)

0/1 distance-d indicator matrices of an arbitrary undirected graph, d = 0..dmax.

The generalisation of cycle_shells off the vertex-transitive case, and the piece Result 52 needs before its design law can be applied to a real topology: the law itself is closed-form only on vertex-transitive graphs, but the objective it minimises -- sum_f 1_f' Q(x) 1_f -- is defined for any shell decomposition, so chc.regret.optimal_fold_partition runs on whatever this returns.

Distances come from repeated boolean matrix products rather than a queue: dmax is the spillover truncation and is small, so dmax sparse-ish multiplies beat a Python BFS and keep the whole thing in NumPy. Vertices farther than dmax, and vertices in another component, simply appear in no shell -- which is the right reading, since a spillover model truncated at dmax says nothing about them.

The shells partition the reachable set, so tr(S_d S_e) = 0 for d != e -- the structural lemma Result 51 rests on -- and that holds here for the same reason it does on a cycle: a vertex sits at exactly one distance from i.

Raises:

Type Description
ValueError

if adjacency is not square, not symmetric, or carries a self-loop. Each would break the partition property silently rather than loudly.

torus_adjacency(rows, cols)

Adjacency rows of the rows x cols 2-D torus C_rows x C_cols, as nested tuples.

The immutable form is deliberate: it is what DelayedNetworkPanel stores, and a frozen dataclass carrying a NumPy array compares elementwise and stops being usable as a value.

The torus is the second topology Result 52's design law is tested on and the first one where the law has no closed form: it is vertex-transitive, so every unit sees the same neighbourhood, but Q(x) is NOT circulant for rows, cols > 1, so chc.regret.optimal_fold_partition cannot take the exact banded route and falls back to the search whose price Result 61 measured.

Raises:

Type Description
ValueError

if either side is below 3, where the wrap would create a double edge and the adjacency would stop being 0/1.

propagate_shells(innovations, shells, gammas, lag, n_times)

Sum shell-d neighbours' innovations arriving lag*d steps late.

x_i(t) = sum_d gamma_d * sum_{j : dist(i,j) = d} innovations_j(t - lag*d), read off a burnt-in window so the result is stationary. innovations is (..., m, span) with span >= n_times + lag*(len(gammas)-1); the returned window is the LAST n_times columns reachable by every shell, so every lag indexes real history rather than the burn-in edge.

Shared by DelayedNetworkPanel and chc.regret.delayed_network_certificate on purpose: the law and the design it is measured on must not drift apart.

panel_covariance(shells, gammas, phi, lag, n_times)

Space-time covariance of the panel propagate_shells draws, in kron(space, time).

From the generative law itself, with no separability assumed::

Cov(x_i(t), x_k(s)) = sum_{d,e} g_d g_e [S_d S_e]_{ik} phi**|(t-s) - lag*(d-e)|

so Sigma = sum_q P_q (x) T_q over shifts q = d-e, with P_q = sum_{d-e=q} g_d g_e S_d S_e and [T_q]_{ts} = phi**|(t-s) - lag*q|. Result 51 says this is separable only at lag = 0; Result 59 bounds how far from separable it gets. Since P_{-q} = P_q^T and T_{-q} = T_q^T, each +-q pair splits into a symmetric and an antisymmetric channel, so the KRONECKER RANK is at most 2*dmax -- set by the spillover truncation, never by the graph diameter or by n_times -- and exactly dmax+1 when the shell operators commute, which holds on any vertex-transitive graph. Measured on a cycle: 2, 3, 4, 5 at dmax = 1..4; on a path or an Erdos-Renyi graph: 2, 4, 6, 8. Never 2*dmax+1, because the q = 0 and q = dmax factors are symmetric on every graph.

Shares the shell convention with propagate_shells on purpose: this is the covariance of exactly that simulator's output, verified against 2e5 draws.

kronecker_spectrum(sigma, n_units, n_times)

Singular values of the space/time rearrangement -- the Kronecker rank of sigma.

Van Loan-Pitsianis: sum_r A_r (x) B_r becomes a sum of outer products once the entries are rearranged so that each (space, space) pair indexes a row and each (time, time) pair a column. The number of nonzero singular values is the smallest number of Kronecker terms that can represent sigma, so a value of 1 IS separability rather than evidence for it.

ar1_innovations(rng, shape, span, phi)

Unit-variance AR(1) paths of length span: z_t = phi z_{t-1} + sqrt(1-phi^2) e_t.

Drawn at the stationary variance from the first column, so no burn-in is needed for the MARGINAL scale; propagate_shells still needs slack for the lags.

within_ar1(series, corrections=3)

AR(1) coefficient of a panel series, corrected for the bias the within transform induces.

Sweeping out unit and time means to isolate the serial dependence also correlates the demeaned regressor with the demeaned innovation, biasing the lag-1 autocorrelation down by (1 + phi) / (n_times - 1) (Nickell 1981). The raw number is not reported: it is a prediction of a quantity smaller than the one Result 51 needs. Inverting that relation by fixed-point iteration recovered phi to within 0.007 for phi <= 0.6 over n_times in 20-160, and to 0.022 in the worst corner measured (phi = 0.85, n_times = 20), where the bias being removed is itself 0.097.

estimate_propagation(treatments, outcomes, shells, horizon, n_resamples=200, level=0.95, seed=0)

Estimate the propagation lag and persistence of a delayed network exposure from a panel.

treatments and outcomes are (clusters, units, times). shells[d] is the adjacency of the units exactly d apart, as cycle_shells returns it; shell 0 is ignored, since a unit's own treatment carries no propagation lag.

Each shell gets its own panel local projection of the outcome on that shell's mean treatment, controlling for the unit's own contemporaneous outcome, and the peak of the resulting response locates delta * d. Regressing those peaks on d through the origin gives delta. Doing it shell by shell rather than pooling is what makes the estimate falsifiable: if the peaks are not proportional to distance, the propagation model is wrong and the flat tail shows it, which is also how the spillover truncation is read off.

chc.irf.delay_estimate is the single-trajectory version of this and is the right tool for one series. It is not used here: its projection rows are cut by contiguous slicing, so on a flattened panel a fraction horizon / n_times of them span two different units. Measured on DelayedNetworkPanel the flattened estimate happened to agree down to n_times = 8, where three rows in four straddle a join -- the peak is a location statistic and the contaminated rows attenuate every horizon at once -- but agreement that depends on the contamination staying symmetric is not a guarantee, and cutting the rows inside a trajectory costs nothing.

The interval resamples whole clusters with replacement. Clusters are the independent unit here, so this preserves both the serial dependence within a unit and the spatial dependence within a cluster exactly, which a moving block over time would break. It needs enough clusters to be meaningful: below roughly 20 the percentile interval is reporting the resampling grid.

estimate_network_effects(data, covariates=('x', 'z', 'x_nb', 'z_nb'), exposure='e', degree=2, folds=5, ridge=1.0, seed=0, fold_groups=None, exclude_neighbours=False)

Cross-fitted DML for the direct and spillover effects (partials out graph-aware nuisances).

Residualises the outcome, treatment u, and exposure e on flexible predictions from the (own + mean-neighbour) covariates, then regresses the outcome residual on the two treatment residuals -- the coefficients are the direct and spillover effects. The covariates are centred and scaled before the polynomial basis is built, and the intercept takes no ridge, so ridge reads the same in any units of the covariates and at any level of them, of the treatment, the exposure and the outcome.

fold_groups labels each row with a unit that must not be split across folds; whole labels are permuted and chunked instead of rows. The default of None permutes rows, which is graph-blind: a unit's neighbour then usually lands in the training fold that predicts it, so the same-fold edge fraction sits near 1 / folds. Result 51 shows the delayed cross-fit variance ratio is affine in that fraction at spillover truncation D = 1, with a root at theta* = K^3 / (4K^3 - 6K^2 + 4K - 1), which exceeds 1 / K for every K -- so random rows always sit on the side where a propagation delay drives the ratio down. Passing the unit index recovers the graph structure the permutation destroys; passing a graph-aware grouping (contiguous blocks of a cycle, a partition of a real graph) moves theta toward the root. Folds become unbalanced when the labels are, which is the price of keeping them intact.

exclude_neighbours drops, from each fold's TRAINING set, every row whose unit is a graph neighbour of a unit in the test fold -- the Emmenegger-style rule, which buys validity by throwing data away rather than by choosing the split. It is a baseline to design against, not a default: it is hop-1 (the neighbours column is what a panel carries), and on a graph-blind split it removes most of the training set, which the estimator feels.

Returns the two effects and their cluster-robust influence-function standard errors, clustered on cid when the data carry it and heteroskedastic otherwise. Under interference an i.i.d. SE is not conservative-or-not, it is wrong: the score correlates across units that share a spillover, and the cluster is the only level at which independence is credible here. A 95% interval is effect +/- 1.96 * se.

Raises:

Type Description
ValueError

when a column is read in two roles -- the exposure or a covariate named x_next, u or cid, which this reads as the outcome, the treatment and the cluster, the exposure among the covariates, or a covariate named twice; or when the neighbour exclusion leaves a fold too few training rows.

estimate_network_effects_gnn(data, neighbours, *, features=('x', 'z'), treatment='u', exposure='e', folds=2, hidden=16, steps=400, lr=0.003, seed=0, fold_groups=None)

Cross-fitted network DML with learned message-passing GNN nuisances, not fixed features.

The GNN counterpart of estimate_network_effects: it residualises the outcome, treatment, and exposure on a GNN of the raw node features plus the graph, so neighbour aggregation is learned, not hand-specified (no x_nb/z_nb columns needed). Same orthogonal score and estimand (direct + spillover). Accuracy matches the lean version on smooth confounding -- DML orthogonality makes the effect robust to nuisance form -- so it earns its cost when the graph representation is better learned than crafted, not as a free accuracy win.

fold_groups behaves as in estimate_network_effects, but its payoff here is CHANNEL-SPECIFIC and smaller than Result 51's scalar law predicts. Measured on DelayedNetworkPanel at m=6, G=15, p=12, delta=1, phi=0.6, 400 paired replications (both designs share each draw, so the ratio is bootstrapped over seeds; arm correlation 0.88 and 0.97): parity-vs-block unit folds move the DIRECT effect's variance to 0.887 [0.803, 0.976] and the SPILLOVER effect's to 0.999 [0.950, 1.052]. The direct channel therefore gains a real ~11% variance reduction, the spillover channel is inert, and the law's 0.715 is outside BOTH intervals -- a scalar Psi cannot produce a channel-dependent answer at all, so the two-column sandwich is what governs here. An independent-arm standard error is not valid for this comparison; assuming one is what left an earlier 150-replication run unable to decide.

Raises:

Type Description
ValueError

when a column is read in two roles: the outcome is x_next, and the treatment, the exposure and each feature must be columns of their own.