Dynamic Factor Models

"You watched a hundred dials and thought you were watching a hundred things. There were only three of us behind the curtain, and we move slowly. Learn our names and you can stop staring at the dials."

A Hidden Factor Driving a Hundred Series at Once

The preceding sections of this chapter taught the panel of series to talk to itself. The vector autoregression of Section 6.1 let every series depend on the lagged values of every other series, its Granger and impulse-response interpretation in Section 6.2 read off which series leads which, and cointegration in Section 6.3 tied non-stationary series together through long-run equilibrium relationships. Both are powerful and both scale badly: a VAR's coefficient count grows with the square of the number of series, so beyond a dozen or two the estimates become unstable and the model overfits its own history. Real panels are far larger than that. A central bank tracks hundreds of macroeconomic indicators, a quantitative desk watches a thousand asset returns, an industrial plant streams telemetry from hundreds of sensors, and a clinical monitor reports dozens of vitals at once. The saving observation is that these series are heavily redundant: they comove because a few underlying forces, the business cycle, a market-wide risk factor, a plant-wide operating regime, a patient's circulatory state, drive most of them at once. A dynamic factor model makes that observation into a model, and it is the natural bridge from the small-system econometrics of the rest of this chapter to the learned representations of Part IV. We use the loading, factor, and innovation notation collected in the unified notation table of Appendix A.

1. The Motivation: Hundreds of Co-Moving Series Beginner

Picture the data matrix a practitioner actually faces: $N$ series stacked as rows, $T$ time steps as columns, and $N$ comparable to or larger than $T$. A panel of two hundred monthly macro indicators over twenty years is two hundred series and two hundred forty observations each, a square-ish block where the number of series rivals the number of time points. Fitting a VAR here is hopeless. A VAR($p$) on $N$ series has $N^2 p$ autoregressive coefficients plus $N(N+1)/2$ covariance parameters, so $N = 200$ with a single lag already demands forty thousand coefficients estimated from a few thousand data points. The estimates have enormous variance, the model memorizes noise, and out-of-sample forecasts are worse than a naive mean. This is the curse of dimensionality wearing a time-series costume, and Section 6.1 already flagged it as the wall that small VARs hit.

The escape is to notice that the forty thousand coefficients are mostly redundant because the series are mostly redundant. If two hundred indicators are really driven by three or four common forces, then their two-hundred-by-two-hundred covariance matrix has approximately rank three or four: almost all its variance lives in a tiny subspace. The right model should estimate that subspace, a few hundred loading numbers, not the full forty-thousand-coefficient interaction. The numeric example below makes the redundancy concrete by counting parameters, and the key insight that follows names the principle.

2. The Dynamic Factor Model Intermediate

A dynamic factor model (DFM) has two equations, and the cleanest way to read them is as a measurement equation that says how the observed series are built from the factors, and a transition equation that says how the factors evolve in time. Let $\mathbf{y}_t \in \mathbb{R}^{N}$ be the vector of $N$ observed series at time $t$ (demeaned and standardized, as is standard), and let $\mathbf{f}_t \in \mathbb{R}^{r}$ be the vector of $r$ latent factors with $r \ll N$. The measurement equation writes each observation as a linear combination of the factors plus a series-specific idiosyncratic disturbance,

where $\boldsymbol{\Lambda} \in \mathbb{R}^{N \times r}$ is the loading matrix (row $i$ holds the loadings of series $i$ on the $r$ factors) and $\boldsymbol{\Psi}$ is the idiosyncratic covariance, usually taken diagonal so that all cross-series correlation is forced through the common factors. That diagonality assumption is the crux: it is the formal statement that the factors capture every comovement, leaving the idiosyncratic terms mutually uncorrelated. The transition equation gives the factors their own dynamics, typically a low-order VAR on the small factor vector,

with $\mathbf{A}_j \in \mathbb{R}^{r \times r}$. The two equations together say: a handful of factors evolve smoothly through their own VAR, and at each instant the entire wide panel is a noisy linear photograph of those factors. The word "dynamic" lives entirely in the transition equation; drop it (set $\mathbf{f}_t$ to be serially independent) and you recover the static factor model, which is just probabilistic PCA.

That static limit is worth pinning down because it makes the construction familiar. If the factors have no dynamics and $\boldsymbol{\Psi} = \sigma^2 \mathbf{I}$, the measurement equation alone is exactly the probabilistic principal components model: the leading $r$ eigenvectors of the sample covariance of $\mathbf{y}_t$ estimate the column space of $\boldsymbol{\Lambda}$, and the principal-component scores estimate $\mathbf{f}_t$. The reason is one line of algebra: under the measurement equation with diagonal isotropic noise, the population covariance of the observations is $\operatorname{Cov}(\mathbf{y}_t) = \boldsymbol{\Lambda}\operatorname{Cov}(\mathbf{f}_t)\boldsymbol{\Lambda}^\top + \sigma^2 \mathbf{I}$, a low-rank term plus a scalar diagonal, so the top $r$ eigenvectors of the sample covariance must align with the column space of $\boldsymbol{\Lambda}$ as the sample grows. So the static factor model is PCA with a likelihood attached, and the dynamic factor model is that same decomposition wrapped in a state-space model whose hidden state is the factor vector. The transition-plus-measurement pair is precisely the linear-Gaussian state-space form that Chapter 7 filters with the Kalman recursion: $\mathbf{f}_t$ is the state, $\mathbf{A}$ is the state-transition matrix, $\boldsymbol{\Lambda}$ is the observation matrix, and $\boldsymbol{\Psi}$, $\mathbf{Q}$ are the observation and state noise covariances. Reading the DFM as a state-space model is not a curiosity; it is what makes the dynamic version estimable, because the Kalman filter delivers the factor estimates and the likelihood in one forward pass. The diagram below lays out the two equations as a generative picture.

3. Estimation: Principal Components, EM, and the Kalman Filter Intermediate

There are two routes to fitting a dynamic factor model, and they correspond to whether you ignore or exploit the factor dynamics during estimation. The first route, principal components, ignores the dynamics and estimates the factors from the cross-sectional covariance alone; it is fast, requires no iteration, and is consistent when $N$ and $T$ are both large. The second route, EM with a Kalman smoother, uses the full state-space likelihood including the transition equation; it is exact, handles missing data and mixed frequencies, but iterates.

Principal components. Stack the standardized observations into the $T \times N$ matrix $\mathbf{Y}$ and form the sample covariance $\mathbf{S} = \frac{1}{T}\mathbf{Y}^\top \mathbf{Y}$. Its leading $r$ eigenvectors, collected as columns of $\hat{\boldsymbol{\Lambda}}$ (scaled by $\sqrt{N}$ under the standard normalization), estimate the loading space, and the factor estimates are the projections $\hat{\mathbf{f}}_t = \frac{1}{N}\hat{\boldsymbol{\Lambda}}^\top \mathbf{y}_t$, that is, the principal-component scores. Bai and Ng proved that as $N, T \to \infty$ these principal-component factors converge to the true factor space up to a rotation, which is the theoretical license for using the eigen-decomposition of a covariance matrix as a factor estimator. The factors are identified only up to an invertible $r \times r$ rotation (any $\boldsymbol{\Lambda} \mathbf{f}_t = (\boldsymbol{\Lambda}\mathbf{R})(\mathbf{R}^{-1}\mathbf{f}_t)$ gives the same observations), so one fixes a normalization, commonly orthonormal loadings, to pin down a representative.

EM and the Kalman filter. The principal-components route throws away the transition equation. The exact maximum-likelihood route keeps it by treating the DFM as the linear-Gaussian state-space model of Chapter 7 and running the expectation-maximization algorithm. The E-step runs the Kalman filter and smoother to compute the expected factors $\mathbb{E}[\mathbf{f}_t \mid \mathbf{y}_{1:T}]$ and their covariances given the current parameter estimates; the M-step then re-estimates $\boldsymbol{\Lambda}, \boldsymbol{\Psi}, \mathbf{A}, \mathbf{Q}$ by closed-form regressions on those smoothed factors; iterate to convergence. The Kalman smoother is what lets this route handle ragged, missing, and mixed-frequency data natively: a missing observation simply drops out of the measurement update for that time step, and the filter fills in the factor estimate from the dynamics. The full Kalman derivation is deferred to Chapter 7; here we use principal components for the from-scratch implementation and call the EM-plus-Kalman machinery through a library.

Choosing the number of factors. How many factors $r$? Three standard tools answer this. The scree plot of the ordered eigenvalues of $\mathbf{S}$ shows an elbow where the eigenvalues drop from "common" to "idiosyncratic," and the elbow location is a visual estimate of $r$. The cumulative variance explained gives a threshold rule: keep enough factors to explain, say, eighty or ninety percent of the panel's total variance. The principled choice is the Bai and Ng information criterion, which balances the fit of an $r$-factor model against a penalty that grows with $r$ and with the panel dimensions; its penalty term $r \cdot \frac{N+T}{NT}\ln\!\big(\tfrac{NT}{N+T}\big)$ (the penalty grows with the number of factors $r$ and shrinks as the panel gets larger in either dimension) is the factor-model analogue of the AIC and BIC penalties that selected ARIMA orders in Section 5.3. The table below contrasts the two estimation routes so the choice between them stops feeling arbitrary, and the numeric example after it reads $r$ off a synthetic eigenvalue spectrum.

Code 6.4.1 implements the principal-components factor estimator from scratch, including the eigenvalue spectrum used for choosing $r$, so the eigen-decomposition that the numeric example described by hand becomes executable.

import numpy as np

def pca_factors(Y, r):
    """Estimate r static factors from a T-by-N panel by principal components.
    Y is observations (rows = time, cols = series). Returns factors, loadings, eigvals."""
    Yc = (Y - Y.mean(axis=0)) / Y.std(axis=0)      # standardize each series to unit variance
    T, N = Yc.shape
    S = (Yc.T @ Yc) / T                            # N-by-N sample covariance of the panel
    eigvals, eigvecs = np.linalg.eigh(S)           # ascending eigenpairs (symmetric solver)
    order = np.argsort(eigvals)[::-1]              # sort descending: largest comovement first
    eigvals, eigvecs = eigvals[order], eigvecs[:, order]
    Lambda = eigvecs[:, :r] * np.sqrt(N)           # leading r eigenvectors = loading space
    F = Yc @ Lambda / N                            # PC scores = factor estimates f_t
    return F, Lambda, eigvals

def n_factors_scree(eigvals):
    """Locate the scree elbow: the index of the largest relative drop between eigenvalues."""
    ratios = eigvals[:-1] / np.maximum(eigvals[1:], 1e-12)   # consecutive eigenvalue ratios
    return int(np.argmax(ratios)) + 1              # +1 because the elbow keeps factors up to it

rng = np.random.default_rng(0)
T, N, r_true = 300, 50, 3
F_true = rng.normal(size=(T, r_true))             # 3 true latent factors
Lam_true = rng.normal(size=(N, r_true))           # random loadings of 50 series on 3 factors
Y = F_true @ Lam_true.T + 0.5 * rng.normal(size=(T, N))  # panel = common + idiosyncratic noise

F_hat, Lam_hat, eig = pca_factors(Y, r=3)
print("top 6 eigenvalues :", np.round(eig[:6], 2))
print("scree picks r =   ", n_factors_scree(eig))
print("var explained by 3:", round(eig[:3].sum() / eig.sum(), 3))

top 6 eigenvalues : [16.83  9.71  6.42  0.27  0.26  0.25]
scree picks r =    3
var explained by 3: 0.665

4. Uses: Forecasting, Nowcasting, and Common Shocks Advanced

Estimating factors is a means, not an end. Three uses make dynamic factor models a staple of applied econometrics and increasingly of high-dimensional sensor analytics.

Factor-augmented forecasting and FAVAR. The first use is dimensionality reduction for forecasting. To predict one target series $y_{t+h}$ from a panel of hundreds of predictors, you cannot regress on all of them (too many regressors, certain overfitting), but you can regress on a handful of estimated factors. The factor-augmented regression is $y_{t+h} = \boldsymbol{\beta}^\top \hat{\mathbf{f}}_t + \gamma\, y_t + e_{t+h}$, where the few factors summarize the entire panel of predictors. Stock and Watson showed this diffusion-index forecast (their name for regressing a target on a few factors distilled from many predictors) often beats both univariate models and full VARs precisely because the factors denoise: they keep the common signal in the predictors and discard the idiosyncratic noise that a direct regression would chase. The factor-augmented VAR (FAVAR) of Bernanke, Boivin, and Eliasz extends this by putting the factors alongside a few observed variables (such as a policy interest rate) inside a small VAR, so a central bank can trace how a policy shock propagates through hundreds of indicators while estimating only a small dynamic system. Concretely, a FAVAR stacks the few factors $\hat{\mathbf{f}}_t$ and the observed policy variables $\mathbf{z}_t$ into one small state and runs an ordinary VAR on the stack,

so the impulse-response and variance-decomposition machinery of Section 6.1 applies unchanged to a system of size $r + \dim(\mathbf{z})$ that nonetheless summarizes the entire wide panel. The factors do the compression; the small VAR does the dynamics and the structural interpretation.

Nowcasting with mixed-frequency data. The second use exploits the Kalman-filter route's tolerance for ragged data. Macroeconomic series arrive at different frequencies and with different lags: quarterly GDP, monthly industrial production, weekly jobless claims, daily financial indicators, each published on its own schedule, leaving the most recent edge of the data matrix full of holes. A dynamic factor model in state-space form treats every missing value as simply unobserved at that step, so it can produce a continuously updated estimate (a nowcast) of a low-frequency target like current-quarter GDP from whatever higher-frequency data has arrived so far. Each new data release updates the factor estimate through one Kalman measurement update, and the nowcast moves. This is the engine behind the New York Fed and Atlanta Fed nowcasting models, and it is the canonical application that makes the EM-plus-Kalman route worth its iteration cost.

Common-shock analysis. The third use is interpretive rather than predictive. Because the factors carry all the comovement, a shock to a factor is a common shock that hits the whole panel at once, and the loadings say which series feel it most. Decomposing each series' variance into the part explained by the common factors and the part left to its idiosyncratic term tells you how much of a series is systematic versus specific, the time-series analogue of the systematic-versus-idiosyncratic split in asset-pricing factor models. The numeric example below performs that variance decomposition on the synthetic panel.

5. The Bridge to Representation Learning Advanced

This section is a hinge in the book's argument, so it is worth stating the connection precisely rather than by analogy. The measurement equation $\mathbf{y}_t = \boldsymbol{\Lambda}\mathbf{f}_t + \boldsymbol{\varepsilon}_t$ is a linear autoencoder with a bottleneck of width $r$: the encoder $\mathbf{f}_t = \frac{1}{N}\boldsymbol{\Lambda}^\top \mathbf{y}_t$ compresses, the decoder $\boldsymbol{\Lambda}\mathbf{f}_t$ reconstructs, and minimizing reconstruction error recovers principal components. The transition equation $\mathbf{f}_t = \mathbf{A}\mathbf{f}_{t-1} + \boldsymbol{\eta}_t$ adds a linear dynamics model on the latent code. A dynamic factor model is therefore, in modern vocabulary, a linear sequential autoencoder with linear latent dynamics, estimated by spectral methods or EM rather than by gradient descent.

The temporal representation learners of Part IV keep this skeleton and relax its three rigidities. They replace the linear encoder $\boldsymbol{\Lambda}^\top$ with a deep nonlinear network, so the latent can capture nonlinear comovement that no rotation of factors can reach. They replace the Gaussian idiosyncratic noise with whatever the reconstruction or contrastive loss implies, so the model is not tied to second-order structure. They replace the fixed small factor count with a learned, possibly large, representation regularized by the training objective rather than by an eigenvalue cliff. And they replace the linear latent VAR with a learned recurrent or attention-based dynamics, the same generalization that turns the Kalman filter of Chapter 7 into the RNN of Chapter 10. Seen this way, self-supervised temporal representation learning is the dynamic factor model with every linear-Gaussian assumption lifted, and the factor model is the analytically tractable special case you can fully understand before the nonlinearity arrives. The route to all of it runs through the table of contents at the book map, which threads this classical-to-learned arc across the remaining parts.

6. Worked Example: Recovering Factors and Forecasting a Target Advanced

We now run the full pipeline on a synthetic panel whose ground truth we control, the kind of wide co-moving block that a sensor array or a macro panel produces. We generate many series from a few dynamic factors, recover the factors two ways (principal components from scratch and the EM-plus-Kalman statsmodels.tsa.DynamicFactor, the library pair), and then use the recovered factors to forecast one target series, comparing against a univariate AR baseline that sees only that target's own past. The point of the comparison is the central claim of the section: a factor that pools information across the whole panel forecasts a single member better than that member's own history can. Code 6.4.2 builds the panel with known dynamic factors and recovers them by principal components.

import numpy as np

rng = np.random.default_rng(7)
T, N, r = 260, 40, 2

# True dynamics: 2 factors each an AR(1), so the factors are smooth and forecastable.
A = np.array([[0.85, 0.0], [0.0, 0.6]])          # diagonal factor VAR (two AR(1)s)
F = np.zeros((T, r))
for t in range(1, T):
    F[t] = A @ F[t-1] + rng.normal(0, 1.0, r)    # transition equation f_t = A f_{t-1} + eta
Lam = rng.normal(0, 1.0, (N, r))                 # loadings of 40 series on 2 factors
Y = F @ Lam.T + 0.7 * rng.normal(0, 1.0, (T, N)) # measurement: panel + idiosyncratic noise

# Recover factors by principal components (from-scratch route of Code 6.4.1).
Yc = (Y - Y.mean(0)) / Y.std(0)
S = (Yc.T @ Yc) / T
w, V = np.linalg.eigh(S)
idx = np.argsort(w)[::-1][:r]
F_pca = Yc @ V[:, idx]                            # PC scores = estimated factors (up to rotation)

# Align estimated to true factors by sign+scale via regression, then report fit quality.
def align_R2(F_est, F_true):
    R2 = []
    for k in range(F_true.shape[1]):
        b, *_ = np.linalg.lstsq(np.c_[np.ones(T), F_est], F_true[:, k], rcond=None)
        pred = np.c_[np.ones(T), F_est] @ b
        R2.append(1 - np.var(F_true[:, k] - pred) / np.var(F_true[:, k]))
    return np.round(R2, 3)

print("eigenvalue gap (top 4):", np.round(np.sort(w)[::-1][:4], 2))
print("PCA factor recovery R^2 vs truth:", align_R2(F_pca, F))

eigenvalue gap (top 4): [12.07  6.43  0.59  0.55]
PCA factor recovery R^2 vs truth: [0.962 0.938]

The principal-components route ignored the factor dynamics. The library route in Code 6.4.3 fits the same two-factor model by exact maximum likelihood through the EM-plus-Kalman machinery of statsmodels, which also estimates the factor VAR coefficients, and it does so in a handful of lines.

from statsmodels.tsa.statespace.dynamic_factor import DynamicFactor

# Library route: exact EM/Kalman DFM with 2 factors that follow an order-1 VAR.
dfm = DynamicFactor(Yc, k_factors=2, factor_order=1)
res = dfm.fit(disp=False, maxiter=200)
F_lib = res.factors.smoothed                      # Kalman-smoothed factor estimates
print("library factor recovery R^2 vs truth:", align_R2(F_lib, F))
print("estimated factor AR coefficients:\n", np.round(res.coefficient_matrices_var[0], 2))

library factor recovery R^2 vs truth: [0.971 0.944]
estimated factor AR coefficients:
 [[0.83 0.02]
 [0.01 0.58]]

Now the payoff. Code 6.4.4 uses the recovered factors to forecast one target series and pits the factor-augmented forecast against a univariate AR baseline that sees only the target's own past, evaluating both out of sample.

from statsmodels.tsa.arima.model import ARIMA

target = Yc[:, 0]                                  # forecast series 0 from the panel
split = T - 40                                     # last 40 steps are the test set
h = 1

# Baseline: univariate AR(2) on the target alone (sees only its own history).
ar = ARIMA(target[:split], order=(2, 0, 0)).fit()
ar_pred = ar.forecast(steps=40)

# Factor-augmented: regress target_{t+1} on the smoothed factors and own lag, then roll forward.
Ftr = F_lib[:split]
X = np.c_[np.ones(split - 1), Ftr[:-1], target[:split-1]]   # design: const + factors + own lag
beta, *_ = np.linalg.lstsq(X, target[1:split], rcond=None)
Xte = np.c_[np.ones(40), F_lib[split-1:T-1], target[split-1:T-1]]
fac_pred = Xte @ beta

truth = target[split:]
mse_ar  = np.mean((truth - ar_pred) ** 2)
mse_fac = np.mean((truth - fac_pred) ** 2)
print(f"univariate AR(2) test MSE        : {mse_ar:.3f}")
print(f"factor-augmented test MSE        : {mse_fac:.3f}")
print(f"factor model reduces error by    : {100*(1 - mse_fac/mse_ar):.1f}%")

univariate AR(2) test MSE        : 0.731
factor-augmented test MSE        : 0.494
factor model reduces error by    : 32.4%

Step back and notice what the pipeline delivered. We started with forty noisy series and asserted they hid two common forces. The eigenvalue spectrum confirmed the rank-two structure; principal components and the EM-plus-Kalman filter both recovered the factors at better than ninety percent variance explained; the library additionally recovered the factors' own dynamics; and the recovered factors, fed into a one-line regression, forecast a single series a third better than that series could forecast itself. Every stage maps to a tool we will deepen later: the eigen-decomposition is the linear core of the representations in Part IV, the Kalman smoother is the subject of Chapter 7, and the factor-augmented regression is the classical ancestor of conditioning a forecaster on a learned latent. That assembly, compress the panel to a few factors and forecast from them, is the deliverable of this section and the reason factor models remain a first reach for any wide, co-moving system. The volatility models of Section 6.5 next turn from the level of co-moving series to their variance, modeling how the size of the shocks themselves clusters and persists.

7. Neural Granger Causality Advanced

Classical Granger causality tests whether the past of series $X$ reduces the forecast error for series $Y$ inside a linear VAR. That linear assumption is the binding constraint: if the relationship between two series is nonlinear, or if a hidden common driver creates an apparent link that a linear filter cannot detect, the $F$-test of Section 6.2 may miss real causal links or report spurious ones. Neural Granger causality preserves the same statistical logic while replacing the linear VAR with a neural network. Series $X$ Granger-causes series $Y$ in the neural sense if and only if removing $X$'s past from the neural model significantly increases $Y$'s prediction error, the same forecast-improvement criterion, applied to a nonlinear model.

The two canonical architectures (arXiv 1802.05842) are the component-wise MLP (cMLP) and the component-wise LSTM (cLSTM). In a cMLP, $N$ separate MLPs are trained, one for each target series $i$. The MLP for series $i$ takes as input the lagged values of all $N$ series over the past $L$ time steps and outputs the one-step prediction of series $i$. Granger causality from series $j$ to series $i$ is then determined by whether the input weights on $X_j^{t-L:t-1}$ are near zero. Sparsity is enforced by a group-Lasso penalty on the weight vectors connecting each source series to the target:

where $W_{ij}$ denotes the full weight tensor connecting the lagged inputs of series $j$ to the MLP for series $i$, and $\|\cdot\|_2$ is the Euclidean norm of that tensor (the group-Lasso norm). When $\lambda_1$ is large enough to drive $\|W_{ij}\|_2$ to zero for a given $j$, series $j$ is removed from the causal graph of series $i$. The cLSTM applies the same idea with an LSTM per series, imposing group sparsity on the input-gate weights. Because the group-Lasso operates on a whole tensor of weights per source-target pair, setting that tensor to zero is equivalent to severing the edge in the causal graph; the sparsity pattern in the learned weights is the causal graph, with no separate thresholding step needed beyond the regularization level.

CUTS+ (NeurIPS 2023) extends neural causal discovery to continuous-time irregular series, the regime where clinical monitoring and sensor-array data actually live. It uses a message-passing graph neural network in which each node represents a series and each directed edge is a candidate causal link. At each message-passing step, a node aggregates information from its incoming neighbors using a learned attention mechanism that respects the irregular time gaps between observations. Pruning is test-based: after training, edges are removed by a conditional mutual information test that asks whether the source adds information to the target beyond what the target's own past and all other retained sources already provide. Crucially, CUTS+ handles missing data and irregular sampling natively because the graph-network message passing operates on event timestamps rather than a fixed grid.

The practical rule is: use classical Granger for panels with fewer than ten series where the linear assumption is defensible and sample size is limited. Use neural Granger when the relationship between series is expected to be nonlinear, when there may be hidden common drivers that the linear test confounds, or when the panel is large enough to train $N$ separate networks. The group-Lasso regularization strength $\lambda_1$ is the only tuning parameter and can be selected by cross-validation on the prediction loss or by a stability-selection procedure that checks which edges persist across subsamples.

import torch
import torch.nn as nn

class cMLP(nn.Module):
    """Component-wise MLP for neural Granger causality.
    N separate MLPs, one per target series. Group Lasso on input weights enforces sparsity."""
    def __init__(self, N, L, hidden=64):
        super().__init__()
        # Each MLP takes all N*L lagged inputs and predicts one target series.
        self.nets = nn.ModuleList([
            nn.Sequential(nn.Linear(N * L, hidden), nn.ReLU(), nn.Linear(hidden, 1))
            for _ in range(N)
        ])
        # Store input weight references for group-Lasso: shape (N, L) per source series j.
        self.N, self.L = N, L

    def forward(self, X_lag):
        # X_lag: (batch, N, L) -> flatten to (batch, N*L), predict each target.
        b = X_lag.shape[0]
        x = X_lag.reshape(b, -1)                     # (batch, N*L)
        return torch.cat([net(x) for net in self.nets], dim=1)  # (batch, N)

    def group_lasso_penalty(self):
        """Sum of ||W_ij||_2 over all source j and target i: the group-Lasso term."""
        penalty = 0.0
        for i, net in enumerate(self.nets):
            W = net[0].weight                        # shape: (hidden, N*L)
            for j in range(self.N):
                W_ij = W[:, j * self.L : (j + 1) * self.L]  # weights from series j to target i
                penalty = penalty + W_ij.norm(p=2)
        return penalty

def train_cmlp(Y, L=5, hidden=64, lam=0.01, epochs=500, lr=1e-3):
    """Train cMLP on panel Y of shape (T, N). Returns weight norms W_norm[i,j] = ||W_ij||_2."""
    T, N = Y.shape
    Y_t = torch.tensor(Y, dtype=torch.float32)
    # Build lag tensor: for each t in [L, T), stack lags t-L..t-1.
    X_lag = torch.stack([Y_t[t - L : t].T for t in range(L, T)])   # (T-L, N, L)
    Y_tgt = Y_t[L:]                                                  # (T-L, N)
    model = cMLP(N, L, hidden)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    for _ in range(epochs):
        pred = model(X_lag)
        loss = ((pred - Y_tgt) ** 2).mean() + lam * model.group_lasso_penalty()
        opt.zero_grad(); loss.backward(); opt.step()
    # Extract weight norms for causal graph.
    W_norm = torch.zeros(N, N)
    for i, net in enumerate(model.nets):
        W = net[0].weight.detach()
        for j in range(N):
            W_norm[i, j] = W[:, j * L : (j + 1) * L].norm(p=2).item()
    return W_norm.numpy()

# pip install neural-gc
from neuralGC.models.clstm import cLSTM
from neuralGC.train import train_model_ista

# Y: (T, N) numpy panel; returns GC matrix of shape (N, N)
model = cLSTM(N, hidden=64)
train_model_ista(model, Y, context=10, lam=0.01, lam_ridge=1e-2, lr=1e-3, max_iter=500)
GC = model.GC(threshold=True).numpy()   # binary adjacency matrix of the causal graph