# aleximmer/loci architecture

> Architecture diagram of the GitHub repository aleximmer/loci, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-19). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/aleximmer/loci
- Repository: https://github.com/aleximmer/loci

## Overview

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: [`causa/loci.py`](https://github.com/aleximmer/loci/blob/main/causa/loci.py) (LOCI API), [`run_individual.py`](https://github.com/aleximmer/loci/blob/main/run_individual.py) (Pair runner), [`run_estimator_benchmark.py`](https://github.com/aleximmer/loci/blob/main/run_estimator_benchmark.py) (Estimator runner), [`generate_figures_and_tables.py`](https://github.com/aleximmer/loci/blob/main/generate_figures_and_tables.py) (Report gen)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_core["Core algorithm"]
  node_causa_loci["LOCI API<br/>causal scorer<br/>[loci.py]"]
  node_causa_ml["Neural model<br/>regression path<br/>[ml.py]"]
  node_causa_utils["Core utils<br/>helpers<br/>[utils.py]"]
  node_causa_het_ridge["Hetero ridge<br/>convex estimator<br/>[het_ridge.py]"]
  node_causa_iterative_fgls["FGLS baseline<br/>baseline estimator<br/>[iterative_fgls.py]"]
  node_causa_heci["HECI<br/>estimator<br/>[heci.py]"]
  node_causa_gnn_fast["GNN fast<br/>estimator<br/>[gnn_fast.py]"]
  node_causa_hsic["HSIC test<br/>independence test<br/>[hsic.py]"]
end

subgraph group_data["Data layer"]
  node_causa_datasets["Dataset loaders<br/>data access<br/>[datasets.py]"]
  node_data_dir[("Benchmarks<br/>dataset files")]
end

subgraph group_experiments["Experiments"]
  node_configs_dir["Gin configs<br/>experiment config"]
  node_run_individual["Pair runner<br/>benchmark script<br/>[run_individual.py]"]
  node_run_estimator_benchmark["Estimator runner<br/>benchmark script"]
  node_jobs_generate["Job gen<br/>command generator<br/>[jobs_generate.py]"]
  node_jobs_estimator_generate["Estimator jobs<br/>command generator"]
end

subgraph group_reporting["Reporting"]
  node_baseline_results[("Baselines<br/>reference outputs")]
  node_results_dir[("Results<br/>generated outputs")]
  node_generate_figures_and_tables["Report gen<br/>analysis script"]
  node_paper_figures["Paper figs<br/>artifacts"]
end

node_causa_datasets -->|"loads"| node_data_dir
node_run_individual -->|"uses"| node_configs_dir
node_run_estimator_benchmark -->|"uses"| node_configs_dir
node_run_individual -->|"calls"| node_causa_loci
node_run_estimator_benchmark -->|"calls"| node_causa_het_ridge
node_run_estimator_benchmark -->|"calls"| node_causa_iterative_fgls
node_run_individual -->|"uses"| node_causa_ml
node_causa_loci -->|"depends on"| node_causa_utils
node_causa_loci -->|"tests"| node_causa_hsic
node_causa_loci -->|"fits"| node_causa_ml
node_causa_loci -->|"alternates"| node_causa_het_ridge
node_causa_loci -->|"alternates"| node_causa_heci
node_causa_loci -->|"alternates"| node_causa_gnn_fast
node_causa_het_ridge -->|"uses"| node_causa_utils
node_causa_iterative_fgls -->|"uses"| node_causa_utils
node_jobs_generate -->|"emits"| node_run_individual
node_jobs_estimator_generate -->|"emits"| node_run_estimator_benchmark
node_generate_figures_and_tables -->|"reads"| node_results_dir
node_generate_figures_and_tables -->|"reads"| node_baseline_results
node_generate_figures_and_tables -->|"writes"| node_paper_figures

click node_causa_loci "https://github.com/aleximmer/loci/blob/main/causa/loci.py"
click node_causa_ml "https://github.com/aleximmer/loci/blob/main/causa/ml.py"
click node_causa_utils "https://github.com/aleximmer/loci/blob/main/causa/utils.py"
click node_causa_het_ridge "https://github.com/aleximmer/loci/blob/main/causa/het_ridge.py"
click node_causa_iterative_fgls "https://github.com/aleximmer/loci/blob/main/causa/iterative_fgls.py"
click node_causa_heci "https://github.com/aleximmer/loci/blob/main/causa/heci.py"
click node_causa_gnn_fast "https://github.com/aleximmer/loci/blob/main/causa/gnn_fast.py"
click node_causa_hsic "https://github.com/aleximmer/loci/blob/main/causa/hsic.py"
click node_causa_datasets "https://github.com/aleximmer/loci/blob/main/causa/datasets.py"
click node_data_dir "https://github.com/aleximmer/loci/tree/main/data"
click node_configs_dir "https://github.com/aleximmer/loci/tree/main/configs"
click node_run_individual "https://github.com/aleximmer/loci/blob/main/run_individual.py"
click node_run_estimator_benchmark "https://github.com/aleximmer/loci/blob/main/run_estimator_benchmark.py"
click node_jobs_generate "https://github.com/aleximmer/loci/blob/main/commands/jobs_generate.py"
click node_jobs_estimator_generate "https://github.com/aleximmer/loci/blob/main/commands/jobs_estimator_generate.py"
click node_baseline_results "https://github.com/aleximmer/loci/tree/main/baseline_results"
click node_generate_figures_and_tables "https://github.com/aleximmer/loci/blob/main/generate_figures_and_tables.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_causa_loci,node_causa_ml,node_causa_utils,node_causa_het_ridge,node_causa_iterative_fgls,node_causa_heci,node_causa_gnn_fast,node_causa_hsic toneBlue
class node_causa_datasets,node_data_dir toneAmber
class node_configs_dir,node_run_individual,node_run_estimator_benchmark,node_jobs_generate,node_jobs_estimator_generate toneMint
class node_baseline_results,node_results_dir,node_generate_figures_and_tables,node_paper_figures toneRose
```

## Components

### Core algorithm

Reusable causal estimation and testing code.

- [LOCI API](https://github.com/aleximmer/loci/blob/main/causa/loci.py): causal scorer · Main entry point that scores both causal directions and can run residual independence testing. · `causa/loci.py`
- [Neural model](https://github.com/aleximmer/loci/blob/main/causa/ml.py): regression path · Neural network estimation path used by LOCI for directional fitting. · `causa/ml.py`
- [Core utils](https://github.com/aleximmer/loci/blob/main/causa/utils.py): helpers · Shared utility functions for estimation and preprocessing. · `causa/utils.py`
- [Hetero ridge](https://github.com/aleximmer/loci/blob/main/causa/het_ridge.py): convex estimator · Convex location-scale / heteroskedastic regression estimator. · `causa/het_ridge.py`
- [FGLS baseline](https://github.com/aleximmer/loci/blob/main/causa/iterative_fgls.py): baseline estimator · Iterative feasible generalized least squares baseline. · `causa/iterative_fgls.py`
- [HECI](https://github.com/aleximmer/loci/blob/main/causa/heci.py): estimator · Alternative causal estimator module. · `causa/heci.py`
- [GNN fast](https://github.com/aleximmer/loci/blob/main/causa/gnn_fast.py): estimator · Fast graph-neural estimator variant. · `causa/gnn_fast.py`
- [HSIC test](https://github.com/aleximmer/loci/blob/main/causa/hsic.py): independence test · Independence-testing machinery for residuals and related comparisons. · `causa/hsic.py`

### Data layer

Benchmark loaders and bundled pairwise datasets.

- [Dataset loaders](https://github.com/aleximmer/loci/blob/main/causa/datasets.py): data access · Loads benchmark pair samples from the repository data directory into arrays or tensors. · `causa/datasets.py`
- [Benchmarks](https://github.com/aleximmer/loci/tree/main/data): dataset files · Bundled cause-effect benchmark files organized by family and dataset. · `data`

### Experiments

Benchmark runners, configs, and job generation.

- [Gin configs](https://github.com/aleximmer/loci/tree/main/configs): experiment config · Per-dataset configuration files that parameterize preprocessing, models, and evaluation. · `configs`
- [Pair runner](https://github.com/aleximmer/loci/blob/main/run_individual.py): benchmark script · Runs one causal-pair benchmark instance from a chosen config. · `run_individual.py`
- [Estimator runner](https://github.com/aleximmer/loci/blob/main/run_estimator_benchmark.py): benchmark script · Runs the estimator benchmark with seed and knot-heuristic selection. · `run_estimator_benchmark.py`
- [Job gen](https://github.com/aleximmer/loci/blob/main/commands/jobs_generate.py): command generator · Emits shell commands for the main benchmark suite. · `commands/jobs_generate.py`
- [Estimator jobs](https://github.com/aleximmer/loci/blob/main/commands/jobs_estimator_generate.py): command generator · Emits shell commands for the estimator benchmark suite. · `commands/jobs_estimator_generate.py`

### Reporting

Figure/table generation and stored comparison outputs.

- [Baselines](https://github.com/aleximmer/loci/tree/main/baseline_results): reference outputs · Precomputed comparison results stored in tabular form. · `baseline_results`
- **Results**: generated outputs · User-created benchmark outputs written by experiment scripts.
- [Report gen](https://github.com/aleximmer/loci/blob/main/generate_figures_and_tables.py): analysis script · Consumes benchmark outputs and baseline results to produce figures and LaTeX-ready tables. · `generate_figures_and_tables.py`
- **Paper figs**: artifacts · Figure PDFs produced for publication output.

## Connections

- 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.