braintools.analysis.population.granger_causality

braintools.analysis.population.granger_causality#

braintools.analysis.population.granger_causality(x, order)#

Compute pairwise conditional Granger causality from a VAR model.

For every target \(i\) and source \(j \ne i\), two vector autoregressive (VAR) models of order \(p\) with an intercept are fitted to \(x_i(t)\) by ordinary least squares [1] [2]:

  • the full model regresses \(x_i(t)\) on the past \(x_k(t-1), \dots, x_k(t-p)\) of all \(N\) variables;

  • the restricted model drops the past of the source \(j\).

The conditional Granger causality is the log ratio of the residual variances,

\[\mathrm{GC}_{j \to i} = \ln \frac{\sigma^2_{i \mid \setminus j}}{\sigma^2_{i \mid \mathrm{all}}} \ge 0 .\]
Parameters:

  • x (Any) – Signals with time on axis 0, shape (T, N). The unit of a Quantity does not affect the (dimensionless) result.

  • order (int) – Model order \(p\) (number of lags), a static Python integer.

Returns:

Matrix of shape (N, N) with GC[i, j] the causal influence of j on i (j -> i) and a zero diagonal.

Return type:

Array

Raises:

  • ValueError – If x is not 2-D, or T <= order.

  • TypeError – If order is not a static integer.

See also

braintools.analysis.population.functional_connectivity

Undirected (zero-lag) correlations.

Notes

The fits use the normal equations. The data are centred, and the Gram matrix \(D^\top D\) of the lagged design (with \(P = N p + 1\) columns) and \(D^\top y\) are accumulated once from lagged cross-products, without materialising the \((T - p) \times P\) design matrix. Each restricted model is a static sub-block of the same Gram matrix, solved sequentially (lax.map), and the residual sum of squares is \(y^\top y - \hat\beta^\top D^\top y\). Time is \(O(T N^2 p + N P^3)\); memory is a few copies of x plus \(O(P^2)\).

Each small system is Jacobi-scaled and solved with a symmetric pseudo-inverse (jnp.linalg.pinv, relative cutoff \(10 P \epsilon\)), so collinear or duplicated signals give finite values instead of NaN. The Gram matrices are accumulated in float64 when jax_enable_x64 is on, otherwise in float32. Residual variances are floored at \(\epsilon\) times the mean square of the raw target plus \(P \epsilon\) times its variance (the round-off of an exact fit), so a perfectly predictable or constant target gives 0, and tiny negative round-off values are clipped to 0. T should be much larger than \(P\) for a meaningful estimate.

With N = 2 this is the bivariate (pairwise) Granger causality of Geweke [2], as computed by elephant.causality.granger.pairwise_granger (which fits the model with Yule-Walker equations instead of least squares, so finite-sample values differ slightly).

Works under jax.jit and jax.vmap (for example over trials) and is differentiable with respect to x for well-conditioned data (the gradient of the pseudo-inverse is undefined for collinear signals).

References

Examples

>>> import jax
>>> import jax.numpy as jnp
>>> from braintools.analysis.population import granger_causality
>>> # x drives y with one step of delay; y does not drive x.
>>> noise = jax.random.normal(jax.random.key(0), (2000, 2))
>>> x = noise[:, 0]
>>> y = jnp.concatenate([noise[:1, 1], 0.8 * x[:-1] + 0.5 * noise[1:, 1]])
>>> gc = granger_causality(jnp.stack([x, y], axis=1), order=1)
>>> bool(gc[1, 0] > 0.5), bool(gc[0, 1] < 0.01)
(True, True)
>>> gc.shape
(2, 2)