CASTOR — Causal Temporal Regime Structure Learning

A complete, tested and documented Python implementation of

Abdellah Rahmani and Pascal Frossard, Causal Temporal Regime Structure Learning, AISTATS 2025, PMLR vol. 258.

Give CASTOR one multivariate time series made of an unknown number of unknown-length regimes, and it returns — jointly, with no prior knowledge —

1. the number of regimes K,
2. where each regime starts and ends, and
3. a full temporal causal graph per regime: instantaneous edges x_i(t) → x_j(t) and lagged edges x_i(t−τ) → x_j(t).

Most causal-discovery methods for time series assume one stationary regime. When the mechanism changes part-way through — a seizure begins, a market regime flips, a season turns — those methods return a single averaged graph that describes none of the actual regimes.

📖 Documentation

How to write this code, step by step	An 11-step guide for researchers implementing a method from a paper — conventions, the order to build in, the traps, and a debugging playbook. Written from this build.
API reference (syntax)	Every public function and parameter, with defaults and worked calls.
Paper → code map	Equation-by-equation correspondence, plus every verified disagreement between the paper and the authors' reference code.
Changelog	Release notes and known limitations.

---

Table of contents

• Install
• 60-second example
• What the algorithm does
• Choosing the three parameters that matter
• Working with real data
• Reproducing the paper
• Results
• Documentation
• Relationship to the authors' reference code
• Package layout
• Citing

---

Install

git clone https://github.com/merwanroudane/castor.git
cd castor
pip install -e .

Optional extras — comparison baselines and the developer tooling:

pip install -e ".[baselines,dev]"

Core requirements are NumPy, SciPy, pandas, PyTorch, networkx, scikit-learn and matplotlib. tigramite (PCMCI+), lingam (VARLiNGAM) and ruptures (KCP) are needed only for the comparison tables; without them those rows are skipped with a clear message rather than crashing.

---

60-second example

from castor import CASTOR, evaluate
from castor.datasets import simulate_regime_mts

# A 2-regime series: 400 samples under graph A, then 400 under graph B.
data = simulate_regime_mts(
    n_regimes=2, n_nodes=5, n_samples=[400, 400], lag=1, random_state=0
)

model = CASTOR(
    lag=1,                     # maximum lag L
    window=200,                # initial window length w
    min_regime_duration=100,   # zeta
    random_state=0,
).fit(data.X)

print(model.n_regimes_)            # 2
print(model.regime_intervals())    # {0: [(0, 401)], 1: [(402, 799)]}
print(model.get_graph(0).edge_list())
# [('x1', 'x4', 0, 1.0), ('x2', 'x2', 1, 1.0), ...]   (cause, effect, lag, weight)

print(evaluate(model, data))       # F1, SHD, regime accuracy -- permutation-matched

Every figure in the paper has a plotting counterpart:

from castor.plots import plot_regime_partition, plot_temporal_graph, save_figure

ax = plot_regime_partition(model.labels_, data.labels)
save_figure(ax, "figures/partition")

---

What the algorithm does

CASTOR maximises the data log-likelihood with an EM loop (Algorithm 1 of the paper). The difficulty is that regimes and graphs are entangled: you cannot segment the series without knowing the graphs, and you cannot fit a graph without knowing which samples belong to it.

Initialisation. Cut the series into N_w > K equal windows and call each a provisional regime. Some are pure (entirely inside one true regime), some are impure (they straddle a change point).

E-step — where does each sample belong? For each sample compute γ(t,u) ∝ π(α_u, t) · f^u(x_t), and assign it to the best u. The two factors do different jobs: f^u measures how well regime u's graph explains x_t, and π(α_u, t) is a smooth, time-indexed prior that keeps a sample from jumping to a distant regime. Pure regimes have meaningful graphs and win samples off impure ones, which shrink.

M-step — refit. Re-estimate π(α, t), then re-estimate one temporal graph per regime by γ-weighted DYNOTEARS (linear, Eq. 10) or a γ-weighted locally-connected MLP (non-linear, Eq. 11), each under the NOTEARS acyclicity constraint h(G_0) = tr(e^{G_0∘G_0}) − d = 0.

Pruning. Any regime left holding fewer than ζ samples is deleted and its samples returned to the pool. This is how N_w descends to K: K is never specified, it is discovered.

Under Gaussian noise with equal error variances, the regimes and their graphs are identifiable up to a permutation of the regime labels (Theorem 1) — which is why every metric in castor.metrics solves an assignment problem before scoring.

---

Choosing the three parameters that matter

Everything else has a sensible default taken from Appendix E.3 of the paper.

window (w) — the single most important knob

The initial window length. It must be shorter than your shortest true regime. If a window straddles a change point, its graph is fitted to a mixture and the two regimes can fuse. Too small is cheap (more EM iterations); too large is fatal.

Rule of thumb: if you believe no regime is shorter than m samples, set window ≈ m / 2 and n_windows will follow. Pass n_windows instead if you would rather fix the count. You need N_w > K, so err on the side of more windows.

min_regime_duration (ζ)

Regimes smaller than this are deleted. It encodes "a regime shorter than this is not a regime, it is noise". The paper uses 100 (linear) and 200 (non-linear). Must be smaller than window.

lag (L)

The maximum lag, in the ordinary sense: lag=1 means x(t−1) influences x(t).

Note. The authors' reference code calls this lags and expects L + 1 (a slice count). This package uses the true lag. If you are porting a script, lags=2 there means lag=1 here. See docs/PAPER_MAP.md.

Linear or non-linear?

functional_form="linear" (default) is fast and exact when relationships are linear. functional_form="nonlinear" fits a small neural network per component per regime — much slower, and worth it only when you expect genuine non-linearity. Start linear.

---

Working with real data

Two real datasets ship with the package, plus a reader for a third.

from castor.datasets import load_web_activity, load_us_macro

load_web_activity() — the IT-monitoring data behind the paper's Section 5.2: two 1106-sample, 7-node blocks from a web server, stacked into one 2212 sample series with a change point in the middle, with expert-annotated causal edges for each block. Ships with the package.

load_us_macro() — real US quarterly macroeconomic series (BEA/Federal Reserve, via statsmodels). No ground-truth graph, so this is an interpretive example in the spirit of the paper's Section 5.3: does the discovered partition line up with known economic history?

load_fluxnet(path) — reader for the biosphere–atmosphere data of Section 5.3. FLUXNET forbids redistribution, so you download it yourself; the loader raises with step-by-step instructions if you call it without a file.

Worked end-to-end analysis: examples/03_real_data_web_activity.py.

Three things to do before trusting a result on your own data

1. Standardise. The identifiability theorem assumes equal error variances. Columns spanning orders of magnitude break that assumption outright. Both bundled loaders z-score by default.
2. Make each regime plausibly stationary. CASTOR assumes stationarity within a regime (Assumption 1). Difference or log-difference trending series first — load_us_macro(transform="auto") shows the pattern.
3. Check model.history_. If label_changes has not settled, the EM has not converged; raise max_iter. plot_convergence(model.history_) shows it at a glance.

---

Reproducing the paper

python examples/06_reproduce_paper_tables.py          # tables -> results/
python examples/02_synthetic_benchmark.py             # synthetic grid
python examples/03_real_data_web_activity.py          # Section 5.2
python examples/04_nonlinear_regimes.py               # Section 3.5 / Figure 3

Programmatically:

from castor.experiments import run_synthetic_benchmark
from castor.tables import comparison_table, write_table

records = run_synthetic_benchmark(
    n_regimes=[2, 3], n_nodes=[5, 10], seeds=[0, 1, 2], models=["CASTOR", "DYNOTEARS-oracle"]
)
write_table(comparison_table(records, fmt="latex"), "results/table1.tex")

Tables come out as Markdown or LaTeX (booktabs), with mean ± std cells and best-in-column bolding, matching the paper's layout.

---

Results

Everything below was produced by python examples/05_full_benchmark.py --quick on a laptop CPU. Raw per-run records are in results/, rendered tables in results/tables/. These are what this code actually produces — not numbers copied from the paper. See docs/PAPER_MAP.md for why several are not directly comparable to the published ones.

Synthetic, linear, K=2, d=5, 500 samples per regime, seed 0. -oracle rows are handed the true regime partition; CASTOR has to discover it.

Model	Regime acc.	F1 inst.	F1 lag	SHD inst.	SHD lag	K̂	Time (s)
CASTOR	100.0	100.0	50.0	0.0	5.0	2	809
DYNOTEARS (regime-blind)	50.0	65.7	16.7	4.0	6.0	1	17
DYNOTEARS-oracle	100.0	100.0	50.0	0.0	2.5	2	25
PCMCI+ (regime-blind)	50.0	52.7	12.5	6.5	8.0	1	1.3
PCMCI+-oracle	100.0	75.0	59.5	2.5	2.5	2	1.0
VARLiNGAM-oracle	100.0	70.0	61.0	3.5	4.0	2	0.3
KCP	67.1	—	—	—	—	2	0.1

Reading it honestly:

• The paper's central claim holds. CASTOR matches DYNOTEARS-oracle on instantaneous edges (100 vs 100 F1) and on lagged edges (50 vs 50) while discovering K and the partition by itself — 100% regime accuracy against an oracle's free lunch.
• Regimes matter enormously. The same DYNOTEARS run regime-blind drops from 100 to 65.7 F1 on instantaneous edges and from 50 to 16.7 on lagged ones: it fits one averaged graph that describes neither regime.
• CASTOR is not uniformly best. On lagged edges PCMCI+-oracle (59.5) and VARLiNGAM-oracle (61.0) beat it (50.0), and its lagged SHD is twice the oracle's. The paper reports the same ordering on lagged links, attributing it to CASTOR having more to learn. Worth stating plainly rather than burying.
• KCP confirms the motivation. At 67.1% regime accuracy a state-of-the-art change-point detector is far behind the causal methods: a change of mechanism need not change any marginal distribution.
• Cost is the real weakness. 809 s versus 25 s for DYNOTEARS-oracle. Time is dominated by the augmented-Lagrangian graph fit, once per regime per EM iteration. graph_max_iter is the knob — dropping it from the default 100 to 20–30 was 4× faster with an identical recovered graph in our profiling.

Real data (web activity, Section 5.2), window ablation and a scalability sweep are in results/tables/. On the real data CASTOR reaches 74.0% regime accuracy with F1 43.8 against the expert annotation; note that regime-blind DYNOTEARS scores higher on graph F1 there (63.6) — that dataset is hard for every method, for reasons documented in load_web_activity.

These are single-seed numbers from the --quick preset. Run python examples/05_full_benchmark.py (no flag) for the 3-setting, 3-seed grid with mean ± std; budget several hours.

---

Documentation

Document	What it is for
docs/SYNTAX.md	Complete API reference — every public function, every parameter, with defaults and worked calls
docs/GUIDE_STEP_BY_STEP.md	How to write this algorithm from the paper, equation by equation, for researchers implementing a method from a PDF
docs/PAPER_MAP.md	Equation-by-equation paper → code map, plus every verified disagreement between the paper and the authors' code

Every public function also carries a NumPy-style docstring with runnable examples (python -m pytest --doctest-modules castor).

---

Relationship to the authors' reference code

This is an independent reimplementation, checked line by line against the authors' repository (github.com/arahmani/CASTOR). Where the paper and that code disagree, castor.CASTOR follows the paper and the difference is documented in docs/PAPER_MAP.md with the file and line that settles it. The main ones:

• Per-lag graphs. The paper defines one G_τ per lag; the reference code's reshape collapses the hidden and lag dimensions, so it returns a single lag-aggregated matrix and cannot express G_1 ≠ G_2.
• Self-lag edges. Definition 1 forbids self-loops only at τ = 0. The reference code's non-linear branch also forbids x_i(t−1) → x_i(t), which caps recall on lagged links — in the paper's own IT-monitoring benchmark, 7 of the 15 annotated edges per regime are exactly those.
• Acyclicity. The paper specifies tr(e^{G∘G}) − d; the reference non-linear branch uses the Yu et al. (2019) polynomial surrogate instead (its trace_expm line is commented out).
• Numerical stability. The E-step is computed in log space; evaluating the Gaussian density directly underflows to exactly zero for wide series, after which every sample is silently assigned to regime 0.
• Recurrence. Appendix E.10 advertises recurring-regime support, but an affine π(α_u, t) provably cannot give one regime two separated intervals. CASTOR.merge_recurring_regimes() supplies it as an explicit post-processing step.
• Reproducibility. Every entry point takes random_state; the reference code seeds nothing.

If you need the reference behaviour exactly — quirks included — use castor.compat, which replicates its API and its arithmetic:

from castor.compat import CASTOR as ReferenceCASTOR
ref = ReferenceCASTOR(data, X, Xlags, lags=2, random_state=0)   # lags = L + 1
models, graphss, gamma, L = ref.run_linear(5, 1.0, 0.4, 150, 100)

---

Package layout

castor/
├── castor.py        the CASTOR estimator -- Algorithm 1, the EM loop
├── graph.py         TemporalGraph: the object of Definition 1
├── linear.py        Eq. (10) -- gamma-weighted DYNOTEARS
├── nonlinear.py     Eq. (11) -- gamma-weighted NOTEARS-MLP
├── mlp.py           locally-connected network, per-lag graph read-out
├── regime.py        Eq. (8)  -- the pi(alpha, t) alignment sub-problem
├── acyclicity.py    h(G) and its gradient, expm and polynomial forms
├── metrics.py       F1 / SHD / regime accuracy, permutation-matched
├── datasets/        Appendix E.1 generators + real data loaders
├── baselines.py     DYNOTEARS, PCMCI+, VARLiNGAM, KCP wrappers
├── experiments.py   benchmark drivers
├── plots.py         publication-quality figures
├── tables.py        Markdown / LaTeX table rendering
└── compat.py        exact replication of the authors' reference code

---

Citing

Cite the original paper:

@inproceedings{rahmani2025castor,
  title     = {Causal Temporal Regime Structure Learning},
  author    = {Rahmani, Abdellah and Frossard, Pascal},
  booktitle = {Proceedings of the 28th International Conference on
               Artificial Intelligence and Statistics (AISTATS)},
  series    = {PMLR},
  volume    = {258},
  year      = {2025}
}

If this implementation itself was useful, see CITATION.cff.

License

MIT — see LICENSE. castor/structure.py and the DYNOTEARS solver in castor/linear.py are adapted from causalnex (Apache-2.0, QuantumBlack Visual Analytics Limited); the non-linear model follows NOTEARS (Apache-2.0). Both retain their notices in-file.