causal-abstraction-eval 1.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (33) hide show
  1. causal_abstraction_eval-1.0.0/CONTRIBUTING.md +90 -0
  2. causal_abstraction_eval-1.0.0/LICENSE +21 -0
  3. causal_abstraction_eval-1.0.0/MANIFEST.in +23 -0
  4. causal_abstraction_eval-1.0.0/PKG-INFO +253 -0
  5. causal_abstraction_eval-1.0.0/README.md +215 -0
  6. causal_abstraction_eval-1.0.0/THIRD_PARTY +128 -0
  7. causal_abstraction_eval-1.0.0/causal_abstraction/__init__.py +177 -0
  8. causal_abstraction_eval-1.0.0/causal_abstraction/analytical_metrics.py +2146 -0
  9. causal_abstraction_eval-1.0.0/causal_abstraction/config.py +133 -0
  10. causal_abstraction_eval-1.0.0/causal_abstraction/engine.py +682 -0
  11. causal_abstraction_eval-1.0.0/causal_abstraction/experiment.py +174 -0
  12. causal_abstraction_eval-1.0.0/causal_abstraction/metrics.py +797 -0
  13. causal_abstraction_eval-1.0.0/causal_abstraction/models/__init__.py +25 -0
  14. causal_abstraction_eval-1.0.0/causal_abstraction/models/high_level.py +256 -0
  15. causal_abstraction_eval-1.0.0/causal_abstraction/models/low_level.py +128 -0
  16. causal_abstraction_eval-1.0.0/causal_abstraction/models/neural.py +418 -0
  17. causal_abstraction_eval-1.0.0/causal_abstraction/paths.py +454 -0
  18. causal_abstraction_eval-1.0.0/causal_abstraction/primitives.py +297 -0
  19. causal_abstraction_eval-1.0.0/causal_abstraction/py.typed +0 -0
  20. causal_abstraction_eval-1.0.0/causal_abstraction/sampling.py +330 -0
  21. causal_abstraction_eval-1.0.0/causal_abstraction/schema.py +249 -0
  22. causal_abstraction_eval-1.0.0/causal_abstraction/spaces/__init__.py +30 -0
  23. causal_abstraction_eval-1.0.0/causal_abstraction/spaces/base.py +384 -0
  24. causal_abstraction_eval-1.0.0/causal_abstraction/tasks.py +288 -0
  25. causal_abstraction_eval-1.0.0/causal_abstraction/valuemap.py +207 -0
  26. causal_abstraction_eval-1.0.0/causal_abstraction_eval.egg-info/PKG-INFO +253 -0
  27. causal_abstraction_eval-1.0.0/causal_abstraction_eval.egg-info/SOURCES.txt +31 -0
  28. causal_abstraction_eval-1.0.0/causal_abstraction_eval.egg-info/dependency_links.txt +1 -0
  29. causal_abstraction_eval-1.0.0/causal_abstraction_eval.egg-info/requires.txt +27 -0
  30. causal_abstraction_eval-1.0.0/causal_abstraction_eval.egg-info/top_level.txt +1 -0
  31. causal_abstraction_eval-1.0.0/pyproject.toml +79 -0
  32. causal_abstraction_eval-1.0.0/requirements.lock +256 -0
  33. causal_abstraction_eval-1.0.0/setup.cfg +4 -0
@@ -0,0 +1,90 @@
1
+ # Contributing
2
+
3
+ Thanks for your interest in the `causal_abstraction_eval` framework (import name
4
+ `causal_abstraction`). This repository accompanies the paper *"Validating Causal
5
+ Abstraction Metrics on Simulated Complex Systems"*.
6
+
7
+ Contributions of every kind are welcome: bug fixes, new benchmark systems, new or
8
+ revised metrics, documentation, and changes to terminology or notation. There are no
9
+ off-limits areas; the only expectation is that you explain the rationale for a
10
+ change in your pull request so reviewers can evaluate it (see
11
+ [Pull requests](#pull-requests)).
12
+
13
+ ## Development setup
14
+
15
+ ```bash
16
+ git clone https://github.com/MelouxM/CAE.git
17
+ cd CAE
18
+ pip install -e ".[dev]"
19
+ ```
20
+
21
+ Optional extras pull in per-system dependencies: `neural` (PyTorch), `physics`
22
+ (numba), `circuit` (PySpice), `figures` (matplotlib), and `all`. The Tracr system
23
+ needs `tracr` installed separately
24
+ (`pip install git+https://github.com/google-deepmind/tracr.git`).
25
+
26
+ ## Common tasks
27
+
28
+ A thin `Makefile` wraps the common commands:
29
+
30
+ | Command | What it does |
31
+ |---|---|
32
+ | `make test` | Run the pytest unit suite (`test/test_*.py`). |
33
+ | `make coverage` | Run the unit suite with coverage (`pytest --cov=causal_abstraction`). |
34
+ | `make lint` | Run `ruff check`. |
35
+ | `make format` | Run `ruff format`. |
36
+ | `make suite SYS=01_logic_circuit RUNS=10` | Run one evaluation suite for a bounded number of runs. |
37
+
38
+ Equivalently, by hand:
39
+
40
+ ```bash
41
+ pytest # unit tests
42
+ pytest --cov=causal_abstraction # with coverage
43
+ ruff check . # lint
44
+ python test/01_logic_circuit.py --runs 10 # one evaluation suite, bounded
45
+ ```
46
+
47
+ ## Tests and runnable suites
48
+
49
+ `test/` holds two different things:
50
+
51
+ - **Unit tests**: `test/test_*.py` (+ `conftest.py`), collected by `pytest`. Many of
52
+ them are **characterization tests** that pin the *current* output of a metric, path,
53
+ or baseline. Their job is to catch behavior that moves unexpectedly: if you are only
54
+ refactoring, a failing pin tells you something changed that shouldn't have. If you
55
+ *intentionally* change a metric's behavior, update its pinned values in the same PR
56
+ and explain the change.
57
+ - **Evaluation / power suites**: `test/NN_*.py`, `test/power_*.py`, built on
58
+ `runner.py` / `power_runner.py` / `utils.py`. They run the full metric battery and
59
+ write the paper's reported numbers to `test/results/*.json`; they are *not* collected
60
+ by `pytest`. The numeric-prefixed system modules (`01_logic_circuit.py`, …) are not
61
+ importable through normal package mechanics, so the suites load them by file path via
62
+ `utils.load_system`.
63
+
64
+ The result files under `test/results/` are generated locally and are not committed to the
65
+ repo. If a change affects them, regenerate the affected suites (e.g.
66
+ `python test/01_logic_circuit.py`) and describe what changed and why in the PR; those
67
+ numbers are cited in the paper, so reviewers need to weigh the impact.
68
+
69
+ ## Adding or changing a benchmark system
70
+
71
+ A benchmark system spans three places. A complete PR touches all of them so they stay
72
+ in sync:
73
+
74
+ - `systems/NN_*.py`: the system definition (the `(M, E, τ)` triple and its invalid
75
+ contrastive conditions);
76
+ - `test/NN_*.py`: the matching evaluation suite;
77
+ - `paper/`: the systems table (`sections/04_benchmark.tex`) and the appendix entry.
78
+
79
+ ## Pull requests
80
+
81
+ - Keep changes focused and explain why, referencing the relevant paper section where
82
+ applicable.
83
+ - Confirm `make test` and `make lint` pass.
84
+ - If your change affects metric implementations, benchmark systems, notation, or any
85
+ reported benchmark number, call it out explicitly and justify it: regenerate
86
+ the affected results and update the paper-facing artifacts (tables, definitions) so
87
+ the code, the numbers, and the paper stay consistent.
88
+
89
+ By contributing, you agree your contributions are licensed under the project's
90
+ MIT License.
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2025-2026 Maxime Méloux
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,23 @@
1
+ # Source-distribution contents.
2
+ #
3
+ # The wheel ships only the importable package. This file controls what the
4
+ # sdist additionally carries, so it stays small, predictable, and reproducible.
5
+
6
+ include LICENSE
7
+ include THIRD_PARTY
8
+ include README.md
9
+ include CONTRIBUTING.md
10
+ include requirements.lock
11
+ include pyproject.toml
12
+ include causal_abstraction/py.typed
13
+
14
+ # The benchmark systems, evaluation suites, paper sources, diagrams, legacy code,
15
+ # and build output are intentionally not part of the distribution.
16
+ prune systems
17
+ prune test
18
+ prune paper
19
+ prune svg
20
+ prune old
21
+ prune build
22
+
23
+ global-exclude __pycache__ *.py[cod]
@@ -0,0 +1,253 @@
1
+ Metadata-Version: 2.4
2
+ Name: causal_abstraction_eval
3
+ Version: 1.0.0
4
+ Summary: A framework for evaluating the validity of causal abstractions in complex systems.
5
+ Author-email: Maxime Méloux <maxime.meloux@univ-grenoble-alpes.fr>, Tiago Pimentel <tiago.pimentel@inf.ethz.ch>, François Portet <portetf@univ-grenoble-alpes.fr>, Maxime Peyrard <peyrardm@univ-grenoble-alpes.fr>
6
+ License-Expression: MIT
7
+ Project-URL: Homepage, https://github.com/MelouxM/CAE
8
+ Project-URL: Bug Tracker, https://github.com/MelouxM/CAE/issues
9
+ Classifier: Programming Language :: Python :: 3
10
+ Classifier: Programming Language :: Python :: 3.12
11
+ Classifier: Operating System :: OS Independent
12
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
13
+ Requires-Python: >=3.11
14
+ Description-Content-Type: text/markdown
15
+ License-File: LICENSE
16
+ Requires-Dist: numpy<3,>=2.2
17
+ Requires-Dist: scipy<2,>=1.16
18
+ Requires-Dist: tqdm<5,>=4.67
19
+ Requires-Dist: joblib<2,>=1.5
20
+ Provides-Extra: neural
21
+ Requires-Dist: torch<3,>=2.7; extra == "neural"
22
+ Provides-Extra: physics
23
+ Requires-Dist: numba<0.62,>=0.61.2; extra == "physics"
24
+ Provides-Extra: circuit
25
+ Requires-Dist: PySpice<2,>=1.5; extra == "circuit"
26
+ Provides-Extra: figures
27
+ Requires-Dist: matplotlib<4,>=3.10; extra == "figures"
28
+ Provides-Extra: dev
29
+ Requires-Dist: pytest<9,>=8.4; extra == "dev"
30
+ Requires-Dist: pytest-cov<7,>=5; extra == "dev"
31
+ Requires-Dist: ruff<1,>=0.6; extra == "dev"
32
+ Provides-Extra: all
33
+ Requires-Dist: torch<3,>=2.7; extra == "all"
34
+ Requires-Dist: numba<0.62,>=0.61.2; extra == "all"
35
+ Requires-Dist: PySpice<2,>=1.5; extra == "all"
36
+ Requires-Dist: matplotlib<4,>=3.10; extra == "all"
37
+ Dynamic: license-file
38
+
39
+ # Causal Abstraction Framework
40
+
41
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
42
+
43
+ A Python framework for **validating causal abstractions on simulated complex systems**. It accompanies
44
+ the paper *"Validating Causal Abstraction Metrics on Simulated Complex Systems"* and provides a unified
45
+ implementation of a benchmark of idealized complex systems, their ground-truth abstractions, and a
46
+ large battery of candidate validity metrics including the paper's contribution, the **Causal
47
+ Abstraction Error (CAE)**.
48
+
49
+ > The paper is the source of truth for all concepts, terminology, and notation. This
50
+ > README mirrors that notation; where the two differ, trust the paper.
51
+
52
+ ---
53
+
54
+ ### Key concepts
55
+
56
+ We reason about complex systems (a neural network, a physics simulation, a digital circuit, a CPU)
57
+ using simplified high-level models. The central question is: *when is a high-level causal model a
58
+ **valid** explanation of a low-level system?*
59
+
60
+ The framework formalizes this as **constructive causal abstraction**. A low-level structural causal
61
+ model `M` is abstracted by a high-level model `E` through a map `τ` (a variable coarse-graining `a`
62
+ plus per-variable value maps `τ_X`). Validity is tested on the **commuting diagram**: intervening on
63
+ `E` should yield the same result as grounding that intervention into `M`, running `M`, and abstracting
64
+ the result back up.
65
+
66
+ The paper's central metric is the **Causal Abstraction Error (CAE)**, a continuous validity score with
67
+ a built-in faithfulness test over unmapped variables `Φ` (variables claimed to be causally inert
68
+ for the abstraction). It comes in two directions and two forms, all exposed through `StandardTasks`:
69
+
70
+ * **`CAE↑` (`CAE_up`)** (bottom-up): micro-interventions are sampled in `M`'s space and lifted via
71
+ `τ_X` (`BottomUpSampler`).
72
+ * **`CAE↓` (`CAE_down`)** (top-down): macro-interventions are sampled in `E`'s domain and grounded
73
+ into `M` (`TopDownSampler`).
74
+ * Each runs in a faithful form (the canonical CAE, which also perturbs the `Φ` variables to
75
+ detect leakage) and a non-faithful ablation (`CAE_up_nf` / `CAE_down_nf`).
76
+
77
+ The paper benchmarks 30+ candidate metrics from observational, functional, information-theoretic, and
78
+ causal families, and shows that only causal-abstraction metrics *with* faithfulness testing reliably
79
+ discriminate valid from invalid abstractions.
80
+
81
+ ### Installation
82
+
83
+ Install the latest release from PyPI:
84
+
85
+ ```bash
86
+ pip install causal_abstraction_eval
87
+ ```
88
+
89
+ To install from source instead — for example to run the benchmark suites, which are not part of
90
+ the distribution (see [Distribution contents](#distribution-contents)):
91
+
92
+ ```bash
93
+ git clone https://github.com/MelouxM/CAE.git
94
+ cd CAE
95
+ pip install .
96
+ ```
97
+
98
+ For development (editable install with the test tooling):
99
+
100
+ ```bash
101
+ pip install -e ".[dev]"
102
+ ```
103
+
104
+ To reproduce the paper's exact numbers, install the pinned environment snapshot
105
+ rather than the loose ranges in `pyproject.toml`:
106
+
107
+ ```bash
108
+ pip install -r requirements.lock
109
+ ```
110
+
111
+ `requirements.lock` is a fully-resolved capture of the environment the benchmark
112
+ results were produced in (it excludes this package itself and `tracr`; install
113
+ those as above); `pyproject.toml` keeps deliberately loose ranges for downstream
114
+ consumers.
115
+
116
+ Optional dependency groups (extras): `neural` (PyTorch), `physics` (numba), `circuit` (PySpice),
117
+ `figures` (matplotlib, for the figure generators), `dev` (pytest), and `all` (every system
118
+ dependency installable from PyPI). Each benchmark system pulls in only the extras it needs.
119
+
120
+ The Tracr system (`systems/08_tracr.py`) additionally requires `tracr`, DeepMind's
121
+ RASP→transformer compiler. It is not published on PyPI (and a direct git reference would make this
122
+ package unpublishable to PyPI), so install it manually:
123
+
124
+ ```bash
125
+ pip install git+https://github.com/google-deepmind/tracr.git
126
+ ```
127
+
128
+ > Requires Python ≥ 3.11. The import name is `causal_abstraction`; the distribution name is
129
+ > `causal_abstraction_eval`.
130
+
131
+ ### Distribution contents
132
+
133
+ The published wheel (and sdist) ship only the importable `causal_abstraction` framework, not the
134
+ benchmark itself. Reproducing the benchmark requires a Git clone rather than a PyPI install.
135
+
136
+ ### Tutorial
137
+
138
+ New to the framework? Start with [`tutorial.ipynb`](tutorial.ipynb), a runnable notebook that builds a
139
+ minimal abstraction from scratch, evaluates one of the provided benchmark systems, and runs a selection
140
+ of baseline metrics, all in the same `M` / `E` / `τ` notation as the paper.
141
+
142
+ ### Benchmark systems and runnable suites
143
+
144
+ The benchmark comprises ten complex systems (rows 01–10 below) plus a suite of controlled
145
+ contrastive experiments (row 11). Each system is defined under `systems/` and evaluated by a matching
146
+ suite under `test/`; the `test/NN_*.py` script loads `systems/NN_*.py` by path and runs the full metric
147
+ battery across conditions (`valid`, plus invalid contrastive conditions such as `fail`, `inv_internal`,
148
+ `noise`), persisting results to `test/results/NN_*_results.json`.
149
+
150
+ | # | System | Definition | Suite |
151
+ |---|---|---|---|
152
+ | 01 | Logic circuit (2-bit adder) | `systems/01_logic_circuit.py` | `test/01_logic_circuit.py` |
153
+ | 02 | Transistor circuit (SPICE) | `systems/02_transistor_circuit.py` | `test/02_transistor_circuit.py` |
154
+ | 03 | Gas simulation (Lennard-Jones → ideal/Van der Waals gas) | `systems/03_gas_simulation.py` | `test/03_gas_simulation.py` |
155
+ | 04 | Predator–prey (agent-based → Lotka–Volterra) | `systems/04_predator_prey.py` | `test/04_predator_prey.py` |
156
+ | 05 | Heat equation (1D) | `systems/05_heat_equation.py` | `test/05_heat_equation.py` |
157
+ | 06 | Heat equation (2D) | `systems/06_heat_equation_2d.py` | `test/06_heat_equation_2d.py` |
158
+ | 07 | Ising model | `systems/07_ising_model.py` | `test/07_ising_model.py` |
159
+ | 08 | Tracr compiled transformer | `systems/08_tracr.py` | `test/08_tracr.py` |
160
+ | 09 | Gene regulatory network | `systems/09_grn/grn.py` | `test/09_grn.py` |
161
+ | 10 | MOS 6502 CPU | `systems/10_cpu_6502.py` | `test/10_cpu_6502.py` |
162
+ | 11 | Controlled contrastive experiments | `systems/11_controlled_experiments.py` | `test/11_controlled_experiments.py` |
163
+
164
+ > Note 1: The gene regulatory network does not follow the flat `systems/NN_<name>.py` convention
165
+ > and is stored at `systems/09_grn/grn.py` (its suite `test/09_grn.py` loads it specially).
166
+
167
+ > Note 2: system 10 (MOS 6502) requires native libraries built from three upstream emulators that
168
+ > are not vendored in this repo. Before running `test/10_cpu_6502.py`, follow
169
+ > [`systems/10_cpu_6502_libs/README.md`](systems/10_cpu_6502_libs/README.md) to fetch the pinned
170
+ > upstream sources and run `build_libs.sh`.
171
+
172
+ Run a suite forever (until `Ctrl-C`) or for a bounded number of runs:
173
+
174
+ ```bash
175
+ python test/01_logic_circuit.py # run until interrupted
176
+ python test/01_logic_circuit.py --runs 10 # bounded
177
+ ```
178
+
179
+ Companion scripts under `test/` cover statistical power / discrimination analysis
180
+ (`power_*.py`, `run_power.py`) and figure generation (`figures/get_figure_*.py`, needs the
181
+ `figures` extra). The unit-test suite (`test/test_*.py`) is run with `pytest`.
182
+
183
+ The `test/` directory is organized as follows:
184
+
185
+ * `test/test_*.py`, `test/conftest.py`: the pytest unit suite.
186
+ * `test/NN_*.py`, `test/power_*.py` and the shared infrastructure (`runner.py`, `utils.py`,
187
+ `power_runner.py`, `run_power.py`, `tracr_size.py`): the runnable evaluation and power suites.
188
+ * `test/results/`: where the suites write benchmark results (`*_results.json`, `power_*.json`)
189
+ when run; generated locally and not committed to the repo.
190
+ * `test/figures/`: figure-generation utilities (`get_figure_*.py`, `figure_utils.py`) and the
191
+ generated PDFs they produce.
192
+ * `test/artifacts/`: generated runtime artifacts (e.g. cached binaries).
193
+
194
+ ### Library architecture
195
+
196
+ The library (`causal_abstraction/`) is organized as follows; the canonical public API is whatever
197
+ `causal_abstraction/__init__.py` re-exports.
198
+
199
+ * **`primitives.py`**: core data types (`SystemState`, `AbstractVariable`, probability
200
+ distributions, and the `UNMAPPED` sentinel).
201
+ * **`schema.py`**: the low-level variable layout (`MicroVariableSchema`, `Variable`,
202
+ `MicroSelector`) and the coarse-graining map `CoarseGrainingMap` (the variable map `a`, which also
203
+ defines the unmapped set `Φ` and the internal variables).
204
+ * **`valuemap.py`**: the value map `τ_X` (`ValueMap`, `ContinuousValueMap`): `abstract()` lifts
205
+ micro-states to macro-labels, `ground()` does the reverse.
206
+ * **`models/`**: `high_level.py` defines the high-level SCM `E` (`CausalGraph`); `low_level.py`
207
+ defines the low-level SCM `M` (`LowLevelModel`, `NoisyLowLevelModel`); `neural.py` wraps PyTorch
208
+ models (`NeuralModel`).
209
+ * **`spaces/base.py`**: subspace geometry used to define variable domains (`RectSubspace`,
210
+ `SphereSubspace`, `UnionSubspace`, `ComplementSubspace`, `FullSubspace`, `UniformSubspace`,
211
+ `GaussianSubspace`).
212
+ * **`sampling.py`**: the intervention samplers `TopDownSampler` (`CAE↓`), `BottomUpSampler` (`CAE↑`),
213
+ `CombinedFaithfulnessSampler` (inline `Φ` faithfulness test), plus `PairedSampler` and
214
+ `NoisyMeasurementSampler`.
215
+ * **`paths.py`**: the commuting-diagram traversal (`DiagramBuilder`, `CausalPath`).
216
+ * **`metrics.py`**: the dissimilarity / divergence family `D` (`MSEMetric`, `L2Metric`,
217
+ `JSDivergenceMetric`, `KLDivergenceMetric`, `MMDMetric`, …) plus dynamical metrics
218
+ (`TrajectoryMSEMetric`, `DTWMetric`, `TemporalAutocorrelationMetric`, `SpectralMetric`) and
219
+ `get_metric`.
220
+ * **`analytical_metrics.py`**: analytical baseline metrics (`IIAMetric`, `BCCMetric`,
221
+ `DCCMetric`, `SobolSensitivityMetric`, `IBLagrangianMetric`, …) drawn from the four metric
222
+ families compared in the paper. With the dissimilarity metrics above and the `CAE↑`/`CAE↓`
223
+ tasks, these make up the 30+ candidate metrics benchmarked in the paper.
224
+ * **`tasks.py`**: `EvaluationTask`, `TaskResults`, and `StandardTasks`, the factory exposing the two
225
+ causal-abstraction tasks `CAE↑` / `CAE↓` (faithful and non-faithful) plus observational baselines.
226
+ * **`engine.py`**: the orchestrator `EvaluationEngine` (`run_tasks`, `run_analytical_metrics`).
227
+ * **`config.py`**: `EvaluationConfig`.
228
+ * **`experiment.py`**: `ExperimentResults`.
229
+
230
+ #### Auxiliary diagnostics
231
+
232
+ `PrecisionMetric` (in `metrics.py`, enabled via `EvaluationConfig(check_precision=True)`) is an auxiliary,
233
+ static diagnostic that measures how specific the value map `τ` is (how small the micro-space region each
234
+ abstract label maps to). It is not a validity score but a descriptive property of the mapping.
235
+
236
+ ### Typical wiring
237
+
238
+ Build a `MicroVariableSchema` → `CoarseGrainingMap` (`a`, defines `Φ`) → `ValueMap` (`τ_X`) → a
239
+ low-level model (`LowLevelModel`, `M`) and a high-level model (`CausalGraph`, `E`) → construct an
240
+ `EvaluationEngine` from them and an `EvaluationConfig` → run `StandardTasks` (CAE↑/CAE↓) and/or the
241
+ analytical metrics through the engine. See any `systems/NN_*.py` paired with its `test/NN_*.py` suite
242
+ for a complete worked example.
243
+
244
+ ### License
245
+
246
+ This project is licensed under the MIT License.
247
+
248
+ Some benchmark systems build on, link against, or redistribute third-party code
249
+ and data under their own licenses: notably, the MOS 6502 system (the `fake6502`,
250
+ `break6502`/M6502Core, and `perfect6502` emulator cores, plus the generated
251
+ `Decoder6502.bin`) and the gene-regulatory-network system (a GINsim
252
+ segment-polarity model). See [`THIRD_PARTY`](THIRD_PARTY) for the per-component
253
+ upstream URLs, pinned versions, and licenses.
@@ -0,0 +1,215 @@
1
+ # Causal Abstraction Framework
2
+
3
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
4
+
5
+ A Python framework for **validating causal abstractions on simulated complex systems**. It accompanies
6
+ the paper *"Validating Causal Abstraction Metrics on Simulated Complex Systems"* and provides a unified
7
+ implementation of a benchmark of idealized complex systems, their ground-truth abstractions, and a
8
+ large battery of candidate validity metrics including the paper's contribution, the **Causal
9
+ Abstraction Error (CAE)**.
10
+
11
+ > The paper is the source of truth for all concepts, terminology, and notation. This
12
+ > README mirrors that notation; where the two differ, trust the paper.
13
+
14
+ ---
15
+
16
+ ### Key concepts
17
+
18
+ We reason about complex systems (a neural network, a physics simulation, a digital circuit, a CPU)
19
+ using simplified high-level models. The central question is: *when is a high-level causal model a
20
+ **valid** explanation of a low-level system?*
21
+
22
+ The framework formalizes this as **constructive causal abstraction**. A low-level structural causal
23
+ model `M` is abstracted by a high-level model `E` through a map `τ` (a variable coarse-graining `a`
24
+ plus per-variable value maps `τ_X`). Validity is tested on the **commuting diagram**: intervening on
25
+ `E` should yield the same result as grounding that intervention into `M`, running `M`, and abstracting
26
+ the result back up.
27
+
28
+ The paper's central metric is the **Causal Abstraction Error (CAE)**, a continuous validity score with
29
+ a built-in faithfulness test over unmapped variables `Φ` (variables claimed to be causally inert
30
+ for the abstraction). It comes in two directions and two forms, all exposed through `StandardTasks`:
31
+
32
+ * **`CAE↑` (`CAE_up`)** (bottom-up): micro-interventions are sampled in `M`'s space and lifted via
33
+ `τ_X` (`BottomUpSampler`).
34
+ * **`CAE↓` (`CAE_down`)** (top-down): macro-interventions are sampled in `E`'s domain and grounded
35
+ into `M` (`TopDownSampler`).
36
+ * Each runs in a faithful form (the canonical CAE, which also perturbs the `Φ` variables to
37
+ detect leakage) and a non-faithful ablation (`CAE_up_nf` / `CAE_down_nf`).
38
+
39
+ The paper benchmarks 30+ candidate metrics from observational, functional, information-theoretic, and
40
+ causal families, and shows that only causal-abstraction metrics *with* faithfulness testing reliably
41
+ discriminate valid from invalid abstractions.
42
+
43
+ ### Installation
44
+
45
+ Install the latest release from PyPI:
46
+
47
+ ```bash
48
+ pip install causal_abstraction_eval
49
+ ```
50
+
51
+ To install from source instead — for example to run the benchmark suites, which are not part of
52
+ the distribution (see [Distribution contents](#distribution-contents)):
53
+
54
+ ```bash
55
+ git clone https://github.com/MelouxM/CAE.git
56
+ cd CAE
57
+ pip install .
58
+ ```
59
+
60
+ For development (editable install with the test tooling):
61
+
62
+ ```bash
63
+ pip install -e ".[dev]"
64
+ ```
65
+
66
+ To reproduce the paper's exact numbers, install the pinned environment snapshot
67
+ rather than the loose ranges in `pyproject.toml`:
68
+
69
+ ```bash
70
+ pip install -r requirements.lock
71
+ ```
72
+
73
+ `requirements.lock` is a fully-resolved capture of the environment the benchmark
74
+ results were produced in (it excludes this package itself and `tracr`; install
75
+ those as above); `pyproject.toml` keeps deliberately loose ranges for downstream
76
+ consumers.
77
+
78
+ Optional dependency groups (extras): `neural` (PyTorch), `physics` (numba), `circuit` (PySpice),
79
+ `figures` (matplotlib, for the figure generators), `dev` (pytest), and `all` (every system
80
+ dependency installable from PyPI). Each benchmark system pulls in only the extras it needs.
81
+
82
+ The Tracr system (`systems/08_tracr.py`) additionally requires `tracr`, DeepMind's
83
+ RASP→transformer compiler. It is not published on PyPI (and a direct git reference would make this
84
+ package unpublishable to PyPI), so install it manually:
85
+
86
+ ```bash
87
+ pip install git+https://github.com/google-deepmind/tracr.git
88
+ ```
89
+
90
+ > Requires Python ≥ 3.11. The import name is `causal_abstraction`; the distribution name is
91
+ > `causal_abstraction_eval`.
92
+
93
+ ### Distribution contents
94
+
95
+ The published wheel (and sdist) ship only the importable `causal_abstraction` framework, not the
96
+ benchmark itself. Reproducing the benchmark requires a Git clone rather than a PyPI install.
97
+
98
+ ### Tutorial
99
+
100
+ New to the framework? Start with [`tutorial.ipynb`](tutorial.ipynb), a runnable notebook that builds a
101
+ minimal abstraction from scratch, evaluates one of the provided benchmark systems, and runs a selection
102
+ of baseline metrics, all in the same `M` / `E` / `τ` notation as the paper.
103
+
104
+ ### Benchmark systems and runnable suites
105
+
106
+ The benchmark comprises ten complex systems (rows 01–10 below) plus a suite of controlled
107
+ contrastive experiments (row 11). Each system is defined under `systems/` and evaluated by a matching
108
+ suite under `test/`; the `test/NN_*.py` script loads `systems/NN_*.py` by path and runs the full metric
109
+ battery across conditions (`valid`, plus invalid contrastive conditions such as `fail`, `inv_internal`,
110
+ `noise`), persisting results to `test/results/NN_*_results.json`.
111
+
112
+ | # | System | Definition | Suite |
113
+ |---|---|---|---|
114
+ | 01 | Logic circuit (2-bit adder) | `systems/01_logic_circuit.py` | `test/01_logic_circuit.py` |
115
+ | 02 | Transistor circuit (SPICE) | `systems/02_transistor_circuit.py` | `test/02_transistor_circuit.py` |
116
+ | 03 | Gas simulation (Lennard-Jones → ideal/Van der Waals gas) | `systems/03_gas_simulation.py` | `test/03_gas_simulation.py` |
117
+ | 04 | Predator–prey (agent-based → Lotka–Volterra) | `systems/04_predator_prey.py` | `test/04_predator_prey.py` |
118
+ | 05 | Heat equation (1D) | `systems/05_heat_equation.py` | `test/05_heat_equation.py` |
119
+ | 06 | Heat equation (2D) | `systems/06_heat_equation_2d.py` | `test/06_heat_equation_2d.py` |
120
+ | 07 | Ising model | `systems/07_ising_model.py` | `test/07_ising_model.py` |
121
+ | 08 | Tracr compiled transformer | `systems/08_tracr.py` | `test/08_tracr.py` |
122
+ | 09 | Gene regulatory network | `systems/09_grn/grn.py` | `test/09_grn.py` |
123
+ | 10 | MOS 6502 CPU | `systems/10_cpu_6502.py` | `test/10_cpu_6502.py` |
124
+ | 11 | Controlled contrastive experiments | `systems/11_controlled_experiments.py` | `test/11_controlled_experiments.py` |
125
+
126
+ > Note 1: The gene regulatory network does not follow the flat `systems/NN_<name>.py` convention
127
+ > and is stored at `systems/09_grn/grn.py` (its suite `test/09_grn.py` loads it specially).
128
+
129
+ > Note 2: system 10 (MOS 6502) requires native libraries built from three upstream emulators that
130
+ > are not vendored in this repo. Before running `test/10_cpu_6502.py`, follow
131
+ > [`systems/10_cpu_6502_libs/README.md`](systems/10_cpu_6502_libs/README.md) to fetch the pinned
132
+ > upstream sources and run `build_libs.sh`.
133
+
134
+ Run a suite forever (until `Ctrl-C`) or for a bounded number of runs:
135
+
136
+ ```bash
137
+ python test/01_logic_circuit.py # run until interrupted
138
+ python test/01_logic_circuit.py --runs 10 # bounded
139
+ ```
140
+
141
+ Companion scripts under `test/` cover statistical power / discrimination analysis
142
+ (`power_*.py`, `run_power.py`) and figure generation (`figures/get_figure_*.py`, needs the
143
+ `figures` extra). The unit-test suite (`test/test_*.py`) is run with `pytest`.
144
+
145
+ The `test/` directory is organized as follows:
146
+
147
+ * `test/test_*.py`, `test/conftest.py`: the pytest unit suite.
148
+ * `test/NN_*.py`, `test/power_*.py` and the shared infrastructure (`runner.py`, `utils.py`,
149
+ `power_runner.py`, `run_power.py`, `tracr_size.py`): the runnable evaluation and power suites.
150
+ * `test/results/`: where the suites write benchmark results (`*_results.json`, `power_*.json`)
151
+ when run; generated locally and not committed to the repo.
152
+ * `test/figures/`: figure-generation utilities (`get_figure_*.py`, `figure_utils.py`) and the
153
+ generated PDFs they produce.
154
+ * `test/artifacts/`: generated runtime artifacts (e.g. cached binaries).
155
+
156
+ ### Library architecture
157
+
158
+ The library (`causal_abstraction/`) is organized as follows; the canonical public API is whatever
159
+ `causal_abstraction/__init__.py` re-exports.
160
+
161
+ * **`primitives.py`**: core data types (`SystemState`, `AbstractVariable`, probability
162
+ distributions, and the `UNMAPPED` sentinel).
163
+ * **`schema.py`**: the low-level variable layout (`MicroVariableSchema`, `Variable`,
164
+ `MicroSelector`) and the coarse-graining map `CoarseGrainingMap` (the variable map `a`, which also
165
+ defines the unmapped set `Φ` and the internal variables).
166
+ * **`valuemap.py`**: the value map `τ_X` (`ValueMap`, `ContinuousValueMap`): `abstract()` lifts
167
+ micro-states to macro-labels, `ground()` does the reverse.
168
+ * **`models/`**: `high_level.py` defines the high-level SCM `E` (`CausalGraph`); `low_level.py`
169
+ defines the low-level SCM `M` (`LowLevelModel`, `NoisyLowLevelModel`); `neural.py` wraps PyTorch
170
+ models (`NeuralModel`).
171
+ * **`spaces/base.py`**: subspace geometry used to define variable domains (`RectSubspace`,
172
+ `SphereSubspace`, `UnionSubspace`, `ComplementSubspace`, `FullSubspace`, `UniformSubspace`,
173
+ `GaussianSubspace`).
174
+ * **`sampling.py`**: the intervention samplers `TopDownSampler` (`CAE↓`), `BottomUpSampler` (`CAE↑`),
175
+ `CombinedFaithfulnessSampler` (inline `Φ` faithfulness test), plus `PairedSampler` and
176
+ `NoisyMeasurementSampler`.
177
+ * **`paths.py`**: the commuting-diagram traversal (`DiagramBuilder`, `CausalPath`).
178
+ * **`metrics.py`**: the dissimilarity / divergence family `D` (`MSEMetric`, `L2Metric`,
179
+ `JSDivergenceMetric`, `KLDivergenceMetric`, `MMDMetric`, …) plus dynamical metrics
180
+ (`TrajectoryMSEMetric`, `DTWMetric`, `TemporalAutocorrelationMetric`, `SpectralMetric`) and
181
+ `get_metric`.
182
+ * **`analytical_metrics.py`**: analytical baseline metrics (`IIAMetric`, `BCCMetric`,
183
+ `DCCMetric`, `SobolSensitivityMetric`, `IBLagrangianMetric`, …) drawn from the four metric
184
+ families compared in the paper. With the dissimilarity metrics above and the `CAE↑`/`CAE↓`
185
+ tasks, these make up the 30+ candidate metrics benchmarked in the paper.
186
+ * **`tasks.py`**: `EvaluationTask`, `TaskResults`, and `StandardTasks`, the factory exposing the two
187
+ causal-abstraction tasks `CAE↑` / `CAE↓` (faithful and non-faithful) plus observational baselines.
188
+ * **`engine.py`**: the orchestrator `EvaluationEngine` (`run_tasks`, `run_analytical_metrics`).
189
+ * **`config.py`**: `EvaluationConfig`.
190
+ * **`experiment.py`**: `ExperimentResults`.
191
+
192
+ #### Auxiliary diagnostics
193
+
194
+ `PrecisionMetric` (in `metrics.py`, enabled via `EvaluationConfig(check_precision=True)`) is an auxiliary,
195
+ static diagnostic that measures how specific the value map `τ` is (how small the micro-space region each
196
+ abstract label maps to). It is not a validity score but a descriptive property of the mapping.
197
+
198
+ ### Typical wiring
199
+
200
+ Build a `MicroVariableSchema` → `CoarseGrainingMap` (`a`, defines `Φ`) → `ValueMap` (`τ_X`) → a
201
+ low-level model (`LowLevelModel`, `M`) and a high-level model (`CausalGraph`, `E`) → construct an
202
+ `EvaluationEngine` from them and an `EvaluationConfig` → run `StandardTasks` (CAE↑/CAE↓) and/or the
203
+ analytical metrics through the engine. See any `systems/NN_*.py` paired with its `test/NN_*.py` suite
204
+ for a complete worked example.
205
+
206
+ ### License
207
+
208
+ This project is licensed under the MIT License.
209
+
210
+ Some benchmark systems build on, link against, or redistribute third-party code
211
+ and data under their own licenses: notably, the MOS 6502 system (the `fake6502`,
212
+ `break6502`/M6502Core, and `perfect6502` emulator cores, plus the generated
213
+ `Decoder6502.bin`) and the gene-regulatory-network system (a GINsim
214
+ segment-polarity model). See [`THIRD_PARTY`](THIRD_PARTY) for the per-component
215
+ upstream URLs, pinned versions, and licenses.