Causal World Model for Genomic Prediction

Research code for the AAGI project Improving genomic prediction accuracy using causal machine learning approaches. This causal world model pipeline targets the AGT leave-one-environment-out benchmark. This repository contains the causal model, LD pruning, minimum-change transition construction, nested training, evaluation, and yield-blind checkpoint inference.

Installation

git clone https://github.com/AAGI-AUS/causalml.git
cd causalml

Python 3.11 is recommended.

python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements.txt

For tests, install requirements-dev.txt instead. requirements-lock.txt records the dependency versions used for local verification; GPU installations may need a platform-specific PyTorch distribution.

Input data

Supply your own processed data directory:

data/
  managed_outputs/
    phenotypes_location_year.csv
    genotypes_qc.npz

The phenotype CSV requires:

Column	Meaning
name	Genotype/sample identifier matching the NPZ
location_year, location	Environment and location identifiers
expt	Experiment identifier
bloc, prange, prow	Numeric block, range, and row
yieldkgha	Observed yield; missing values are permitted

Optional included_in_logo_splits controls row eligibility, and source_row retains an original row identifier. Inputs should already have completed QC. The NPZ contains numeric dosage (samples × markers, finite imputed diploid 0/1/2 dosages), string sample_ids, and string marker_ids. Identifiers must be unique. Alleles must be harmonized before use. Do not concatenate incompatible marker panels. The loader uses allow_pickle=False.

The LD allele-frequency calculation assumes 0/1/2 dosage. Some model tests use signed calls solely to exercise transition encoding; those are not a substitute for the documented LD input encoding.

Train the causal model benchmark

python run_benchmark.py \
  --data-dir data \
  --output-dir outputs/causal_world_model \
  --device cpu

Select a subset of outer folds or run a short development experiment:

python run_benchmark.py \
  --data-dir data \
  --output-dir outputs/site01_smoke \
  --environments SITE-01__2022 \
  --max-epochs 1

Use at least three environments for nested selection. Entirely unlabelled test environments are recorded as unscorable. Use a new output directory per run. config.json holds the model/training settings; --config loads an alternative. The default cap is 45 epochs, with patience 5 and minimum 3 epochs. A cap is not the actual training duration. Fold seeds are 42 plus the environment's position in the complete sorted environment list, even when only a subset is selected.

Outputs include selection/refit histories, encoder and LD metadata, sequence summaries, a checkpoint per fold, plot predictions, environment metrics, and macro/pooled summaries. Undefined correlations are null in fold JSON; aggregate metrics follow pandas' missing-value exclusion and should be interpreted alongside per-fold results. The normalized accuracy metric is 100 * (1 - MAE / mean(abs(y))), not classification accuracy.

Inference from a checkpoint

python predict.py \
  --checkpoint outputs/causal_world_model/folds/SITE-01__2022/causal_world_model.pt \
  --genotypes data/managed_outputs/genotypes_qc.npz \
  --records candidate_records.csv \
  --output outputs/candidate_predictions.csv

Candidate records need genotype, environment, experiment, and layout columns; yield labels are not needed or consumed. Marker columns are aligned by identifier, and sample row order is taken from the supplied NPZ. Missing checkpoint markers raise an error. Unknown environments/experiments map to the saved unknown token. Only load trusted PyTorch checkpoints: historical metadata uses pickle.

Learning objective

loss = latent_prediction_MSE
     + current_yield_MSE + future_yield_MSE
     + 0.1 * feature_balanced_reconstruction
     + 0.0001 * mean_absolute_latent_activation
     + 0.05 * SIGReg

Both yield targets use training-only mean/std. Genotype and context each receive half of the reconstruction weight. Actions include all signed observed changes, including layout changes, and a multi-hot factor mask. One coordinate changing by two dosage units still counts as one changed marker.

Interpretation

Constructed paths are comparisons of observed records, not necessarily real time. Layout covariates are not measured fertilizer/sowing interventions. The V10 benchmark evaluates current yield; it does not validate multi-step forecasting or causal intervention effects. LD windows follow input marker order, not physical positions. For historical compatibility, inner-validation genotypes participate in LD fitting, but validation labels never train the network and outer-test environments do not enter the fit. ReLU states cannot exactly follow a centered Gaussian; the regularizer is retained for compatibility.

Repository map

• agt_world_model.py: SAE, predictor, yield head, masked losses, SIGReg.
• agt_data.py: encoder, LD pruning, neighboring genotypes and sequence dataset.
• agt_logo.py: nested splits, mixed transition dataset, selection and refitting.
• training.py: transition-only optimization and device selection.
• run_benchmark.py: causal-model-only CLI and artifacts.
• predict.py: direct current-state inference, including saved V10 checkpoints.
• metrics.py: regression metrics.
• tests/: model, data, and end-to-end synthetic checks.

Tests

python -m pip install -r requirements-dev.txt
python -m pytest -q

The tests generate synthetic data in temporary directories; no AGT data are needed.

Contributing

See CONTRIBUTING.md for development setup, issue reporting, and the pull request workflow.

Citation

When using this implementation, cite this repository and the exact commit used: Causal World Model for Genomic Prediction, AAGI-AUS/causalml, https://github.com/AAGI-AUS/causalml. Describe it as research software; no released version or associated peer-reviewed model paper is claimed here.

See extraction provenance and verification results for reproducibility limits.

Licence and acknowledgements

The source project carries the MIT license reproduced in LICENSE (Lucas Maes, 2026). This extraction retains that notice. The tabular adaptation uses JEPA and SIGReg ideas; see I-JEPA and LeJEPA. No dataset rights are granted by the code license. AGT benchmark data must be obtained separately with the appropriate permissions. The upstream copyright notice is retained unchanged.