How aleximmer/loci works
1. Repository purpose
- This is a research codebase for LOCI: a bivariate causal direction estimator for identifying whether x -> y or y -> x.
- The code is organized around two evaluation modes: pairwise causal benchmarks on curated datasets and a separate estimator benchmark for comparing estimator behavior.
2. Core library boundary: causa/
- causa is the reusable algorithm layer. It exposes the main causal scoring API in loci.py, plus supporting implementations for baselines and model variants.
- This package is the architectural center: scripts and experiments call into it, but it does not own orchestration or result generation.
3. Main inference path: LOCI
- The documented entry point is causa.loci.loci(x, y, ...).
- Conceptually, it fits a directional model in both orientations, computes a score, and optionally performs an independence test on scaled residuals.
- The score sign encodes the inferred direction: positive for x -> y, negative for y -> x.
4. Modeling subsystems inside causa/
- ml.py and utils.py provide the neural network-based regression/estimation path used by LOCI.
- het_ridge.py implements a convex location-scale / heteroskedastic regression estimator.
- iterative_fgls.py contains a baseline iterative feasible GLS estimator.
- Additional method modules such as heci.py, gnn_fast.py, and hsic.py support alternative estimators or the independence-testing machinery.
5. Data access layer: causa/datasets.py and data/
- Dataset classes in causa/datasets.py load pairwise cause-effect samples from the repository’s data/ directory.
- The data are organized by benchmark family, including ANLSMN pairs, Benchmark_simulated, Dataverse pairs, and Tübingen pairs.
- This forms a clear boundary: raw benchmark files live under data/, while the loader layer turns them into arrays/tensors for the algorithms.
6. Benchmark orchestration scripts
- run_individual.py runs a single causal pair benchmark instance from a chosen config.
- run_estimator_benchmark.py runs the estimator benchmark with a seed and knot heuristic selection.
- These scripts are the execution boundary between the library code and experiment runtime.
7. Config-driven experiment setup
- configs/*.gin contains per-dataset experiment configuration, suggesting the repo uses Gin config for dependency injection / parameterization.
- The configs are the mechanism for selecting preprocessing, model type, and evaluation settings without changing code.
- This makes the experiment pipeline reproducible and dataset-specific.
8. Bulk job generation
- commands/jobs_generate.py and commands/jobs_estimator_generate.py emit shell commands for running the full benchmark suite.
- The README describes these commands as safe to run in parallel, so the architecture assumes many independent jobs rather than a centralized scheduler.
- This is a thin orchestration layer whose output is meant to feed a batch runner or cluster queue.
9. Results and artifacts
- baseline_results/ stores precomputed outputs for comparison methods such as GRCI/QCCD-style baselines.
- The benchmark scripts appear to write fresh outputs into a results/ directory created by the user.
- This separates immutable reference results from generated experiment outputs.
10. Analysis and publication layer
- generate_figures_and_tables.py consumes benchmark outputs and baseline results to produce publication artifacts.
- It writes figure PDFs to paper_figures/ and prints LaTeX-ready tables to the console.
- This is the final reporting boundary: raw experiment outputs are transformed into figures and tables for the paper.
11. External runtime and dependencies
- The repo requires Python 3.8+ and is installed via requirements.txt.
- From the code structure and README, the important runtime dependencies materially shaping the architecture are likely NumPy/SciPy, scikit-learn preprocessing, and a deep learning stack for the neural estimator path.
- The entire system is file-based and local-execution oriented; there is no evidence of a service backend or database.
12. End-to-end data flow
- Benchmark pair files in data/ are loaded by causa.datasets.
- The selected estimator in causa/ fits both causal directions and computes a direction score.
- Experiment scripts aggregate results into results/.
- The analysis script merges those outputs with baseline_results/ to generate final figures and tables.
Main entry points: loci.py LOCI API · run_individual.py Pair runner · run_estimator_benchmark.py Estimator runner · generate_figures_and_tables.py Report gen
Components 19
Core algorithm
Reusable causal estimation and testing code.
- LOCI API · Main entry point that scores both causal directions and can run residual independence testing.
causa/loci.py - Neural model · Neural network estimation path used by LOCI for directional fitting.
causa/ml.py - Core utils · Shared utility functions for estimation and preprocessing.
causa/utils.py - Hetero ridge · Convex location-scale / heteroskedastic regression estimator.
causa/het_ridge.py - FGLS baseline · Iterative feasible generalized least squares baseline.
causa/iterative_fgls.py - HECI · Alternative causal estimator module.
causa/heci.py - GNN fast · Fast graph-neural estimator variant.
causa/gnn_fast.py - HSIC test · Independence-testing machinery for residuals and related comparisons.
causa/hsic.py
Data layer
Benchmark loaders and bundled pairwise datasets.
- Dataset loaders · Loads benchmark pair samples from the repository data directory into arrays or tensors.
causa/datasets.py - Benchmarks · Bundled cause-effect benchmark files organized by family and dataset.
data
Experiments
Benchmark runners, configs, and job generation.
- Gin configs · Per-dataset configuration files that parameterize preprocessing, models, and evaluation.
configs - Pair runner · Runs one causal-pair benchmark instance from a chosen config.
run_individual.py - Estimator runner · Runs the estimator benchmark with seed and knot-heuristic selection.
run_estimator_benchmark.py - Job gen · Emits shell commands for the main benchmark suite.
commands/jobs_generate.py - Estimator jobs · Emits shell commands for the estimator benchmark suite.
commands/jobs_estimator_generate.py
Reporting
Figure/table generation and stored comparison outputs.
- Baselines · Precomputed comparison results stored in tabular form.
baseline_results - Results · User-created benchmark outputs written by experiment scripts.
- Report gen · Consumes benchmark outputs and baseline results to produce figures and LaTeX-ready tables.
generate_figures_and_tables.py - Paper figs · Figure PDFs produced for publication output.
Connections 20
- Dataset loaders → Benchmarks: loads · Dataset loaders read benchmark files from the local data directory.
- Pair runner → Gin configs: uses · The pair benchmark runner is configured via Gin files.
- Estimator runner → Gin configs: uses · Estimator benchmark execution is parameterized by Gin configs.
- Pair runner → LOCI API: calls · Runs LOCI and related causal estimators on individual pairs.
- Estimator runner → Hetero ridge: calls · Benchmark compares estimator behavior using core estimators.
- Estimator runner → FGLS baseline: calls · Includes the iterative FGLS baseline in estimator comparisons.
- Pair runner → Neural model: uses · LOCI can use the neural regression path for directional fitting.
- LOCI API → Core utils: depends on · LOCI relies on shared utilities for fitting and scoring.
- LOCI API → HSIC test: tests · Optional residual independence testing is part of the LOCI flow.
- LOCI API → Neural model: fits · Main inference path fits models in both causal orientations.
- LOCI API → Hetero ridge: alternates · The library exposes a convex heteroskedastic regression variant.
- LOCI API → HECI: alternates · Alternative causal estimation implementations live in the same library boundary.
- LOCI API → GNN fast: alternates · Fast GNN-based estimator variant is available as a supporting method.
- Hetero ridge → Core utils: uses · Estimator implementations share utility routines.
- FGLS baseline → Core utils: uses · Baseline estimator uses shared helper code.
- Job gen → Pair runner: emits · Generates shell commands for the pair benchmark runner.
- Estimator jobs → Estimator runner: emits · Generates shell commands for the estimator benchmark runner.
- Report gen → Results: reads · Reporting consumes fresh benchmark outputs.
- Report gen → Baselines: reads · Reporting merges in precomputed baseline tables.
- Report gen → Paper figs: writes · Figure generation writes publication-ready PDFs.