causal-automata

A blazing-fast, JAX-accelerated Cellular Automaton engine for causal discovery.

v0.3.0 — Corrected V-well / gravity / conditioning physics, JIT-safe run(ticks, k), and Maximum Accuracy Mode (--multihop, Directed F1 0.696 on N=100). See CHANGELOG.md.

Installation

pip install "causal-automata==0.3.0"

With CPU JAX:

pip install "causal-automata[jax]"
# or: uv pip install "causal-automata[jax]"

With NVIDIA CUDA JAX (and the Colab-style stack):

pip install "causal-automata[cuda]"

Or using uv in a project:

uv add "causal-automata[jax]"

Benchmark — N=100

Default stack (Upgrade 2): k=3 if N≤50, k=5 if N>50, binary edges, 1-hop gravity. On N=100 (seed=0, static thresholds, ticks=100) that baseline hits Skeleton 0.937 / Directed 0.619 in about 2.8s.

Maximum Accuracy Mode — pass --multihop (Upgrade 4): 2-hop elevation gravity via A @ A lifts Directed F1 to 0.696 with essentially flat wall time.

PC / GOLEM wall times are Colab Tesla T4 (num_iter=20000). CA times are CPU JAX for the Upgrade-2 / Upgrade-4 A/B; at historical k=3 the same CA finished in 0.323s on T4 with Directed F1 0.483.

Opt-in flags (defaults stay fast & binary):

• --multihop — Maximum Accuracy Mode (2-hop gravity, Directed F1 0.696 on this SEM).
• --fuzzy — continuous A ∈ [0, 1] soft birth/death (experimental; did not beat binary here).

Conditioning-set sweep (CA-only, 1-hop binary, same SEM):

Background & Origin

causal-automata evolved from graphnet-automata, a sister project exploring the application of 1D cellular automata rules to dynamic network graphs. While graphnet-automata focuses on generalized structural evolution and terminal-based visualization, this repository specializes those concepts into a high-performance, JAX-accelerated engine specifically designed for causal discovery and DAG generation.

Documentation

Full docs (theory, physics, API, CLI, JAX static ticks / k): https://causal-automata.readthedocs.io/en/latest/

How it works — CA physics + V-well

State is a binary DAG adjacency A, node elevations E ∈ [0,1], and a fixed empirical covariance Cov from continuous observations. Each tick is a local cellular update:

1. Death / birth — prune weak edges and add strong ones using partial correlations conditioned on the top-k union neighbors ranked by |Cov|, after masking candidates with elevation below both endpoints (blocks collider / M-bias in S_{i,j}). Birth orients by elevation after the |ρ| gate. Optional --fuzzy keeps continuous weights in [0, 1] and binarizes only after the final tick.
2. V-structure gravity well — unshielded triples (i, k, j) that show explaining-away (Berkson): marginally weak |ρ_{ij}| but stronger |ρ_{ij|k}| mark middle node k as a collider. Collider elevations are pushed below min(E_i, E_j) − 0.15; parents are gently buoyed — the classic V-structure i → k ← j.
3. Gravity — drift elevations toward the midpoint of parent/child means (roots anchor missing parents at 1.0, leaves missing children at 0.0, isolates hold E), then keep only downhill edges with a tiny index tie-break. With --multihop, that drift also sees 2-hop neighbors from A @ A (discounted pull) to help break Markov-equivalence ties.

The pure functional JAX port (jax_core.step / jax_core.run) is @jax.jit-compiled; run(..., ticks=..., k=...) marks ticks and k as static so jax.lax.scan length and (2+k)×(2+k) partial-corr blocks stay compile-time constants. See the docs page JAX static arguments for functools.partial patterns that avoid ConcretizationTypeError.

Run

# Synthetic DAG harness (NumPy default; use --backend jax for JIT)
python -m causal_automata.bayesian_synth --n 10
python -m causal_automata.bayesian_synth --n 50 --backend jax
python -m causal_automata.bayesian_synth --n 100

# Benchmark vs gCastle PC / GOLEM (default N=100, k=5, 1-hop)
python -m causal_automata.benchmark
python -m causal_automata.benchmark --n 100 --seed 0 --ticks 100

# Maximum Accuracy Mode (Upgrade 4) — Directed F1 0.696 on the N=100 SEM
python -m causal_automata.benchmark --n 100 --ca-only --k 5 --multihop
python -m causal_automata.benchmark --n 100 --ca-only --k 5 --compare-multihop

python -m causal_automata.benchmark --n 100 --ca-only --k-sweep 5,6,7,8
python -m causal_automata.benchmark --smoke   # device checks + GOLEM init only

Smoke the JAX compile:

python -c "from causal_automata import warmup, run; print(warmup())"

Package layout

causal-automata/
├── causal_automata/
│   ├── __init__.py
│   ├── bayesian.py          # NumPy Bayesian CA step
│   ├── bayesian_synth.py    # Synthetic linear-DAG harness
│   ├── jax_core.py          # JAX @jit step / run
│   └── benchmark.py         # vs gCastle PC / GOLEM
├── docs/                    # Sphinx / Read the Docs
├── tests/
├── .readthedocs.yaml        # uv-based docs install (matches CI)
├── pyproject.toml
├── README.md
└── NOTICE

License / notice

See NOTICE. Copyright 2026 Masakazu Yamagiwa.