Metadata-Version: 2.4
Name: castor-causal
Version: 0.1.0
Summary: Causal Temporal Regime Structure Learning (CASTOR): joint regime and temporal causal graph discovery
Author-email: Merwan Roudane <merwanroudane920@gmail.com>
Maintainer-email: Merwan Roudane <merwanroudane920@gmail.com>
License: MIT
Project-URL: Homepage, https://github.com/merwanroudane/castor
Project-URL: Repository, https://github.com/merwanroudane/castor
Project-URL: Step-by-step guide, https://github.com/merwanroudane/castor/blob/main/docs/GUIDE_STEP_BY_STEP.md
Project-URL: API reference, https://github.com/merwanroudane/castor/blob/main/docs/SYNTAX.md
Project-URL: Paper to code map, https://github.com/merwanroudane/castor/blob/main/docs/PAPER_MAP.md
Project-URL: Changelog, https://github.com/merwanroudane/castor/blob/main/CHANGELOG.md
Project-URL: Issues, https://github.com/merwanroudane/castor/issues
Project-URL: Paper, https://arxiv.org/abs/2311.01412
Keywords: causal discovery,causal inference,time series,regime switching,structure learning,change point detection,DAG,NOTEARS,DYNOTEARS,expectation maximization
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Mathematics
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy>=1.23
Requires-Dist: scipy>=1.10
Requires-Dist: pandas>=1.5
Requires-Dist: torch>=2.0
Requires-Dist: networkx>=3.0
Requires-Dist: scikit-learn>=1.2
Requires-Dist: matplotlib>=3.6
Requires-Dist: tabulate>=0.9
Provides-Extra: baselines
Requires-Dist: ruptures>=1.1; extra == "baselines"
Requires-Dist: tigramite>=5.2; extra == "baselines"
Requires-Dist: lingam>=1.8; extra == "baselines"
Provides-Extra: dev
Requires-Dist: pytest>=7.4; extra == "dev"
Requires-Dist: pytest-cov>=4.1; extra == "dev"
Requires-Dist: ruff>=0.5; extra == "dev"
Provides-Extra: docs
Requires-Dist: sphinx>=7.0; extra == "docs"
Requires-Dist: furo; extra == "docs"
Requires-Dist: myst-parser; extra == "docs"
Provides-Extra: all
Requires-Dist: castor-causal[baselines,dev,docs]; extra == "all"
Dynamic: license-file

# CASTOR — Causal Temporal Regime Structure Learning

[![Python](https://img.shields.io/badge/Python-3.10%2B-blue)](https://www.python.org/)
[![License: MIT](https://img.shields.io/badge/License-MIT-informational.svg)](LICENSE)
[![arXiv](https://img.shields.io/badge/arXiv-2311.01412-b31b1b.svg)](https://arxiv.org/abs/2311.01412)

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](https://github.com/merwanroudane/castor/blob/main/docs/GUIDE_STEP_BY_STEP.md)** | 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)](https://github.com/merwanroudane/castor/blob/main/docs/SYNTAX.md)** | Every public function and parameter, with defaults and worked calls. |
| **[Paper → code map](https://github.com/merwanroudane/castor/blob/main/docs/PAPER_MAP.md)** | Equation-by-equation correspondence, plus every verified disagreement between the paper and the authors' reference code. |
| **[Changelog](https://github.com/merwanroudane/castor/blob/main/CHANGELOG.md)** | Release notes and known limitations. |

---

## Table of contents

- [Install](#install)
- [60-second example](#60-second-example)
- [What the algorithm does](#what-the-algorithm-does)
- [Choosing the three parameters that matter](#choosing-the-three-parameters-that-matter)
- [Working with real data](#working-with-real-data)
- [Reproducing the paper](#reproducing-the-paper)
- [Results](#results)
- [Documentation](#documentation)
- [Relationship to the authors' reference code](#relationship-to-the-authors-reference-code)
- [Package layout](#package-layout)
- [Citing](#citing)

---

## Install

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

Optional extras — comparison baselines and the developer tooling:

```bash
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

```python
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:

```python
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`](https://github.com/merwanroudane/castor/blob/main/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.

```python
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`](https://github.com/merwanroudane/castor/blob/main/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

```bash
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:

```python
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/`](https://github.com/merwanroudane/castor/tree/main/results/), rendered
tables in [`results/tables/`](https://github.com/merwanroudane/castor/tree/main/results/tables/). These are what this code
actually produces — not numbers copied from the paper. See
[`docs/PAPER_MAP.md`](https://github.com/merwanroudane/castor/blob/main/docs/PAPER_MAP.md#5-empirical-section) 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/`](https://github.com/merwanroudane/castor/tree/main/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`](https://github.com/merwanroudane/castor/blob/main/docs/SYNTAX.md) | Complete API reference — every public function, every parameter, with defaults and worked calls |
| [`docs/GUIDE_STEP_BY_STEP.md`](https://github.com/merwanroudane/castor/blob/main/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`](https://github.com/merwanroudane/castor/blob/main/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`](https://github.com/merwanroudane/castor/blob/main/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:

```python
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:

```bibtex
@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`](https://github.com/merwanroudane/castor/blob/main/CITATION.cff).

## License

MIT — see [`LICENSE`](https://github.com/merwanroudane/castor/blob/main/LICENSE). `castor/structure.py` and the DYNOTEARS solver in
`castor/linear.py` are adapted from
[causalnex](https://github.com/mckinsey/causalnex) (Apache-2.0, QuantumBlack
Visual Analytics Limited); the non-linear model follows
[NOTEARS](https://github.com/xunzheng/notears) (Apache-2.0). Both retain their
notices in-file.
