Skip to content

Repository files navigation

OGM

The omnigenic neural network repository

Welcome to the official OGM repository. This repository hosts the OGM project, a project for phenotype prediction from genotype data using Machine Learning models (notably the omnigenic neural network, OGM).

Getting Started


Dependencies are managed with a standard pyproject.toml (pip/uv-installable) plus a thin conda env per supported Python version (3.10, 3.11, 3.12).

0. Prerequisites

  • mamba/conda available (this repo uses mamba; substitute conda if that's what you have).
  • Read/write access to wherever your institution's UKB data lives.

1. Pick a Python version and create the conda env

Three Python env specs are provided; they only differ by interpreter version. Package versions are pinned identically (aside from numpy) in pyproject.toml. When in doubt, use 3.12.

mamba env create -f environment-py312.yml -n ogm_py312
conda activate ogm_py312

This installs ogm itself (editable) plus torch/torchvision (CUDA 12.6 build) and torch-scatter/torch-sparse from the matching wheel indexes, see the comments at the top of the environment-py3*.yml file for the exact pip install command if you want to run it by hand instead.

Optional extras (see [project.optional-dependencies] in pyproject.toml): preprocess (raw UKB .bgen parsing), singlecell, geno, interpret, and dev. Add them by editing the env yml's last pip line to -e .[interpret] etc.

2. Install the Unified DataModule

The real genotype/covariate dataloader (GeneralDatamodule, used by both train.py and evaluate.py) lives in the udm repository. Clone it somewhere persistent (a plain pip install from the git URL doesn't keep a checkout around, and the Demo below needs the fake datasets shipped under its res/ directory, which aren't packaged into the wheel), then install it as a no-deps package (this repo's pyproject.toml already provides compatible versions of everything udm needs):

mkdir -p ../udm/
git clone https://github.com/luisherrmann/udm.git ../udm/
pip install ../udm --no-deps

Demo

ogm/config/train_demo.yaml trains a genotype + covariates OGM end to end on fake data, no UKB access, no manual path edits beyond pointing at your own udm clone from step 2 above:

conda activate <your_ogm_env>
cd /path/to/your/ogm/checkout

TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 python -m ogm.core.genopheno.train --config-name train_demo

3. GO resources (only needed for real training, the Demo ships its own)

The default model config (ogm/config/lightning_module/model_config/default.yaml) maps genes to GO terms using a real ontology file plus two gene-annotation files. These are generic public bioinformatics resources, not private UKB data. Because they are downloaded/generated rather than maintained source files, they live under ../results/resources/agnostic/. Set OGM_RESULTS_DIR to use a different results root. Fetch the current release and map its UniProt IDs to Ensembl gene IDs with:

python -m ogm.pre.get_go_gaf

This downloads go-basic.obo plus goa_human.gaf/goa_human_isoform.gaf from EBI into ../results/resources/agnostic/, and writes the Ensembl-ID-mapped versions actually consumed by training, all stamped with today's date. Since the default config hardcodes a specific date (../results/resources/agnostic/go-basic_<date>.obo, goa_human{,_isoform}_ensembl_gene_mapping_<date>.gaf), point your own run at whatever date the script just produced:

python -m ogm.core.genopheno.train --config-name train_example \
    lightning_module.model_config.hierarchy.obo_file=../results/resources/agnostic/go-basic_<date>.obo \
    lightning_module.model_config.hierarchy.map_ensembl_go="[../results/resources/agnostic/goa_human_ensembl_gene_mapping_<date>.gaf, ../results/resources/agnostic/goa_human_isoform_ensembl_gene_mapping_<date>.gaf]" \
    ...

(hierarchy.gene_type_dict/transcript_type_dict are a separate, optional refinement, only used when remove_redundant_parents=True; there's no script for those in this checkout yet.)

4. optional R env for WebGestaltR (only if using go_enrichment_analysis_type=[webgestalt_*])

go_enrichment_analysis_type=[webgestalt_over_representation_analysis, webgestalt_gene_set_enrichment_analysis] shells out to a separate R environment via ogm/core/agnostic/enrichment_analysis.py. Create one from environment-r-interpretations.yml:

mamba env create -f environment-r-interpretations.yml -n <your_r_env_name>

Then point the code at it via two environment variables, no source edit needed:

export OGM_R_ENV=/path/to/your/<your_r_env_name>
export OGM_CONDA_SH=/path/to/your/conda/etc/profile.d/conda.sh

5. Verify

cd /path/to/ogm   # must run from the repo root, evaluate.py/train.py
                        # do sys.path.append(os.getcwd())
python -c "
import torch, torch_scatter, torch_sparse, pytorch_lightning as pl, torchmetrics, hydra
from udm.general_datamodule import GeneralDatamodule
from ogm.core.genopheno import evaluate, train
print('torch', torch.__version__, 'cuda avail:', torch.cuda.is_available())
print('OK')
"

If that prints OK with no traceback, the env is ready. For a real end-to-end check, run a training config with trainer_config.max_epochs=1 and confirm a checkpoint appears under logging.out_directory, or just run the Demo below.

Typical install time: steps 1-2 together take about 10-20 minutes on a normal desktop computer with a broadband connection, almost all of it downloading the CUDA-enabled torch/torchvision wheels in step 1 (a few GB); step 3 (only needed for WebGestaltR-based GO enrichment) adds a few more minutes for the R environment.

Overview: How to use


Both training and evaluation are driven by Hydra: every run is python -m ogm.core.genopheno.<train|evaluate> --config-name <config> [overrides...], executed from the repository root with the target conda env active (both entrypoints call os.chdir(hydra.utils.get_original_cwd()), so they need to be launched from here, but then run as if cwd were the repo root regardless of where you invoke them from).

Start from ogm/config/train_demo.yaml, which trains a genotype + covariates OGM end to end on fake data from this, and the udm repo. Adjust to your datasets and requirements.

conda activate <your_ogm_env>
cd /path/to/your/ogm/checkout

python -m ogm.core.genopheno.train --config-name train_demo \
    datamodule.plugins.geno.dataset_path=<udm_home>/res/datasets/fake/geno \
    datamodule.plugins.covariates_essential.data_paths=<udm_home>/res/datasets/fake/covariate/covariates_new_base_all.feather \
    datamodule.plugins.covariates_categorical.data_paths=<udm_home>/res/datasets/fake/covariate/covariates_new_extended_categorical_all.feather \
    datamodule.plugins.covariates_continuous.data_paths=<udm_home>/res/datasets/fake/covariate/covariates_new_extended_continous_all.feather \
    datamodule.plugins.labels.data_paths=<udm_home>/res/datasets/fake/pheno/cad_all.feather \
    datamodule.splits=<udm_home>/res/datasets/fake/splits.yaml \
    logging.out_directory=<scratch_dir>/

The genotype/covariate/label data is entirely synthetic (udm's own fake fincad100/covariates_new_*/cad datasets, 4096 samples, 10 SNPs), and the GO hierarchy uses a hand-written two-term .obo/.gaf pair (resources/demo/), so there's no biological signal to learn, this exists purely to prove your install works end to end. See the comments at the top of train_demo.yaml for what each override does.

Expected demo run time: about 1 minute on a normal CPU-only desktop computer (measured; the model is tiny and capped to a single epoch via trainer_config.max_epochs=1). On completion, look for Wrote extended scores to .../results/ukb_all_scores_extended.tsv and a final ukb_valid_genes_avg_roc line in the console output.

Training

Fill in the <...> dataset-path placeholders in that file with your own data first (see the comments at its top):

conda activate <your_ogm_env>
cd /path/to/your/ogm/checkout

python -m ogm.core.genopheno.train --config-name train_example \
    logging.out_directory=<results_dir>/

train_example.yaml composes its config from four groups (logging, lightning_module, trainer_config, callback_config), each under ogm/config/<group>/, and inlines its own datamodule.plugins block (rather than composing a datamodule group) so it stays self-contained. Override any leaf value with a dotted key, or swap an entire group for a different file in that directory:

# override individual leaf values
python -m ogm.core.genopheno.train --config-name train_example trainer_config.max_epochs=1 trainer_config.gpus=1

# swap a whole config group, e.g. the alternate trainer_config/old.yaml
python -m ogm.core.genopheno.train --config-name train_example trainer_config=old

# or point at a different top-level config file entirely (e.g. your own copy of train_example.yaml)
python -m ogm.core.genopheno.train --config-name my_config.yaml

Training is local-only. Each run receives a UUID and writes its logs, configuration, checkpoints, and results under <logging.out_directory>/<run-id>/.

Evaluation

Evaluation loads a trained checkpoint and runs one or more performance/interpretability analyses against it. Be sure to escape = in checkpoint paths as \=. Entrypoint: ogm/core/genopheno/evaluate.py; config: ogm/config/evaluate_example.yaml:

python -m ogm.core.genopheno.evaluate --config-name evaluate_example \
"overwrite_checkpoint_path=results/<model_id>/checkpoints/epoch\=<epoch>-step\=<step>.ckpt"

Notes on the key flags:

  • overwrite_checkpoint_path can be used to infer run.experiment_id. It points at the models .ckpt file (the inner single-quotes are needed because checkpoint filenames contain =).
  • +offline_config=.../FLAGS.json can be used to adjust the config, otherwise the training config is used
  • eval.keep_activations=accumulate is required for any interpretability analysis to run at all.
  • Toggle exactly the analyses you want via the no_*/feature flags in evaluate_example.yaml, e.g. snp_importance_analysis=True snp_importance_analysis_methods=[IntegratedGradients] for SNP-level Integrated Gradients, or [IntegratedHessians] for pairwise SNP interactions (snp_interaction_analysis_interaction_indices=[...]); no_go_enrichment_analysis=False for the GO over-representation/GSEA baseline. See ogm/config/evaluate_example.yaml for the full, documented list of flags and their defaults.

Quick sanity check

The Demo above already proves your install works with zero data access; this section is the equivalent check against your own real dataset, before committing to a full run. To confirm your environment/checkout works end-to-end before a real run, cap both entrypoints to a tiny epoch count and a small batch fraction (Lightning's native limit_*_batches). Also cap trainer_config.limit_final_inference_batches. Training always ends with a full-cohort trainer.predict() pass (run_final_inference, used to write final scores) which ignores limit_*_batches and is not capped by default; on a real dataset this step alone can run for hours and OOM once evaluate reloads the checkpoint and repeats it:

# training
python -m ogm.core.genopheno.train --config-name train_example \
    trainer_config.max_epochs=1 \
    trainer_config.limit_train_batches=0.01 \
    trainer_config.limit_val_batches=0.01 \
    trainer_config.limit_final_inference_batches=0.001 \
    logging.out_directory=<scratch_dir>/

The printed results path and console log share one freshly generated UUID per run.

# evaluation, against the checkpoint the smoke-training run above produced
python -m ogm.core.genopheno.evaluate --config-name evaluate_example \
    run.experiment_id=<run-id-from-above> \
    +logging.out_directory=<scratch_dir>/<run-id>/ \
    +offline_config=<scratch_dir>/<run-id>/FLAGS.json \
    overwrite_checkpoint_path="'<scratch_dir>/<run-id>/checkpoints/<checkpoint-file>.ckpt'" \
    evaluate_selector=test \
    no_performance_plots=False

Training with genotype + metabolomics

The gene layer can be fed more than one "genetic-like" input at once: metabolomics NMR panels (e.g. the UKB Nightingale panel) are ingested via a second GeneticPlugin (geno2) alongside the genotype plugin (geno) that train_example.yaml already inlines, and mapped into the same GO hierarchy as an additional block of "genes". Add it as a new plugin the same way geno is already defined in train_example.yaml:

python -m ogm.core.genopheno.train --config-name train_example \
    +datamodule.plugins.geno2.name=GeneticPlugin \
    +datamodule.plugins.geno2.__init__=__init__ \
    +datamodule.plugins.geno2.src=['<your_metabolomics_dataset_prefix>'] \
    +datamodule.plugins.geno2.src_names=['all'] \
    +datamodule.plugins.geno2.dataset_path=<path/to/your/metabolomics/dataset> \
    +datamodule.plugins.geno2.gene_transcript=gene \
    +datamodule.plugins.geno2.eid_map_from=EID.49966 \
    +datamodule.plugins.geno2.eid_map_to=EID.49966 \
    +datamodule.plugins.geno2.in_mem=False \
    +datamodule.plugins.geno2.tags=[genetic] \
    logging.out_directory=<results_dir>/

Notes:

  • geno/geno2 are just plugin names under datamodule.plugins. Any additional genetic-like modality is added the same way, as a new +datamodule.plugins.<name>.* block with name: GeneticPlugin and tags: [genetic]; the gene layer concatenates every plugin tagged genetic before building the GO hierarchy.
  • The metabolomics dataset's own var/annotation files need to carry a metabolite→gene(/GO) mapping for this "plain GO" combination to make sense; a CHEBI-based metabolite ontology is also supported via model_config.hierarchy.ontology: biological_process,chebi_ontology plus a matching obo_file (see ogm/core/agnostic/models/ogm.py's ontology/obo_file handling for how the comma-separated namespace list is used).

Transfer learning on single-cell data (bulk_single_cell)

ogm/bulk_single_cell/ is a separate pipeline that transfers a pretrained, frozen OGM (trained on bulk genotype data, as above) onto a single-cell PBMC dataset, fusing its genetic latent with a scRNAseq encoder. Entrypoint: ogm/bulk_single_cell/core/ogm_train.py.

_get_ogm_setup() in that file loads a real OGM training run's saved config (the FLAGS.json that ogm/core/genopheno/train.py writes for every run, the same file evaluate.py's +offline_config= flag consumes) to rebuild that run's frozen architecture for checkpoint loading -- no hand-written config, so it works with any of your own trained OGM runs, including the Demo checkpoint. Every path is environment-variable-overridable, with no source edit needed (see the os.environ.get(...) calls at the top of ogm_train.py for the full list):

  • OGM_PRETRAINED_RUN_DIR: directory containing that run's FLAGS.json.
  • OGM_PRETRAINED_CHECKPOINT: path to that run's .ckpt file.
  • DBGAP_CELLINK_DD_PATH: path to your single-cell cellink DonorData file (.h5.dd).
  • OGM_SC_DEFAULT_CONFIG: optional, overrides the default single-cell hyperparameters JSON (ogm/bulk_single_cell/config/default_config.json) that --config YAML files are merged on top of.

Single-cell demo

ogm/bulk_single_cell/scripts/config_demo.yaml runs the whole single-cell pipeline end to end on a tiny synthetic DonorData fixture, no private dbGaP access needed, using a demo OGM checkpoint you train yourself:

conda activate <your_ogm_env>
cd /path/to/your/ogm/checkout
export PYTHONPATH=$(pwd)

# 1. Train the demo OGM (see the genotype Demo section above) to get a checkpoint + FLAGS.json
TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 python -m ogm.core.genopheno.train --config-name train_demo
# note the printed run-id/results path, e.g. <out_directory>/<run-id>/

# 2. Point the single-cell pipeline at that run
export OGM_PRETRAINED_RUN_DIR=<out_directory>/<run-id>/
export OGM_PRETRAINED_CHECKPOINT=<out_directory>/<run-id>/checkpoints/<checkpoint-file>.ckpt

# 3. Point at the shipped synthetic single-cell fixture (10 SNPs, matching
#    train_demo.yaml's fake genotype dataset)
export DBGAP_CELLINK_DD_PATH=$(pwd)/resources/demo/bulk_single_cell/singlecell_demo.h5.dd

# 4. Run all three modes; mode 3 also runs Integrated Hessians and Integrated Gradients
#    (config_demo.yaml sets evaluate.run_ih=true and evaluate.run_ig=true)
python ogm/bulk_single_cell/core/ogm_train.py \
    --config ogm/bulk_single_cell/scripts/config_demo.yaml \
    --mode pseudobulk_sklearn --cell_type global
python ogm/bulk_single_cell/core/ogm_train.py \
    --config ogm/bulk_single_cell/scripts/config_demo.yaml --mode ogm_standalone
TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 python ogm/bulk_single_cell/core/ogm_train.py \
    --config ogm/bulk_single_cell/scripts/config_demo.yaml \
    --mode ogm_joint_pseudobulk_per_celltype --cell_type global

Like the genotype demo, resources/demo/bulk_single_cell/singlecell_demo.h5.dd is entirely synthetic data (built on cellink's own dummy-data generators), purely to prove the install and pipeline work end to end -- no biological signal to learn from.

Real production runs (ogm/bulk_single_cell/scripts/submit_stage1_pbmc.sh) drive it in three steps, via --mode:

conda activate <your_ogm_env>
cd /path/to/your/ogm/checkout
export PYTHONPATH=$(pwd)

# 1. Pseudobulk sklearn baseline (mean expression per donor -> logistic regression/SVM/RF)
python ogm/bulk_single_cell/core/ogm_train.py \
    --config ogm/bulk_single_cell/scripts/config_batch_split_stage1.yaml \
    --mode pseudobulk_sklearn --cell_type pbmc

# 2. Genetics-alone baseline (frozen OGM only, no single-cell branch)
python ogm/bulk_single_cell/core/ogm_train.py \
    --config ogm/bulk_single_cell/scripts/config_batch_split_stage1.yaml \
    --mode ogm_standalone

# 3. OGM + single-cell late-fusion transfer learning (the main model), per cell type
python ogm/bulk_single_cell/core/ogm_train.py \
    --config ogm/bulk_single_cell/scripts/config_batch_split_stage1.yaml \
    --mode ogm_joint_pseudobulk_per_celltype --cell_type pbmc

These are the only three modes run() supports (pseudobulk_sklearn, ogm_standalone, ogm_joint_pseudobulk_per_celltype); ogm_joint_pseudobulk_per_celltype, the default if --mode is omitted, is the real transfer-learning model: it loads the frozen OGM, fuses its latent with a per-donor MIL-pooled single-cell encoding (OGMLateFusionModel), trains, evaluates, and runs zero-/random-/permuted-genetics ablations to check the genetics branch is actually load-bearing. Passing --cell_type <name> trains one model for that cell type only; omitting it loops over all CANONICAL_CELL_TYPES plus global.

A YAML file passed via --config overrides only the keys it sets on top of ogm/bulk_single_cell/config/default_config.json (dataset paths, batch split, model hyperparameters, evaluate.* toggles); see config_batch_split_stage1.yaml for a real, documented example, including the use_genetics_aux_head/genetics_aux_weight covariate fix needed to stop the single-cell branch from learning to ignore the genetics signal entirely. Integrated Hessians (SNP-gene interaction attribution) is expensive and split into a separate stage, see config_batch_split_stage2.yaml and submit_stage2_global.sh, which reuse stage 1's checkpoint via --mode ogm_joint_pseudobulk_per_celltype with evaluate.run_ih=true. Integrated Gradients (whole-feature attribution over donor covariates + pooled cells, no SNP-anchor selection needed) runs the same way via evaluate.run_ig=true, and can be enabled alongside run_ih in the same run.

Nothing in the loading path (_prep_dd()) is hardcoded to the dbGaP/SLE dataset: donor/cell label column names (data.donor_labels_key/data.cell_labels_key) and their string-to-float mappings (data.donor_labels_map/data.cell_labels_map) are all --config-driven, defaulting to the dbGaP dataset's own values only because that's the dataset these defaults were validated against -- point a differently-labeled DonorData file at it by overriding those four keys (plus data.donor_batch_key if your batch/cohort column has a different name), no source changes needed.

Hyperparameter search

ogm/core/genopheno/hyper.py (delegating to hyper_funcs.HyperManager) runs hyperparameter search via the ax framework (Gaussian Process + Expected Improvement).

Caution: This is training multiple models and thus takes a lot more resources and time.

python -m ogm.core.genopheno.hyper --config_name=train_demo --parameters_json resources/hyper_params.json --pheno cad --out_directory ../results/hyper/ --name demo_hyper --gpus_per_script=0 --max_trials 4 --tags demo_hyper

This produces a table with the run details and performance that can inform your hyperparameter choide. With the demo code it will appear as

head -n3 ../results/hyper/overview.csv | cut -f1,3,5,6,7,8
trial_id        results_dir     finished        final_score     lightning_module.model_config.gomodel.standard_go_size  lightning_module.model_config.genemodel.standard_gene_size
0       ../results/hyper/4f07cb95-b975-4dd8-af35-9928863741c8/results   True    0.49463337659835815     6.0     7.0
1       ../results/hyper/21482567-d5a1-445b-8cdc-bf8175f8e61f/results   True    0.46925848722457886     8.0     6.0
2       ../results/hyper/f79c44d3-7154-4ab3-b3a2-8ffb28d04991/results   True    0.4871289134025574      9.0     9.0

Genotype datasets

This demo requires bgen files (and assumes they are from UKB, if not you may have to adjust variant-naming etc.). If available add --slurm to parallelize across nodes in a slurm cluster. Install the genotype-preprocessing dependencies into the active environment first:

mamba env create -f environment_ogm_geno.yaml
conda activate ogm-geno

Then run the demo as follows:

# Paths to external genotype resources

bgen_dir=path/to/dir/with/bgens/
snpeff_jar=path/to/snpEff.jar
plink2=path/to/plink2
out_dir=../results/geno_ds_demo


conda activate ogm-geno

mkdir -p \
  "$out_dir"/subsets_pgs \
  "$out_dir"/annotated \
  "$out_dir"/annotated_df \
  "$out_dir"/rsids \
  "$out_dir"/splits \
  "$out_dir"/plink \
  "$out_dir"/logs/{annotated,annotated_df,rsids,plink}

cp \
  resources/demo/variant_list_demo.tsv \
  "$out_dir/subsets_pgs/variant_subset.1k.tsv"

SLURM_JOB_ID=no_slurm_01 snakemake \
  -s ogm/pre/geno/genotype_selection.smk \
  --jobs 100 \
  --show-failed-logs \
  --latency-wait 60 \
  --keep-going \
  --use-conda \
  --configfile ogm/pre/geno/genotype_selection_config.yaml \
  --config \
    CHROM='[21,22]' \
    genetic_data="$bgen_dir" \
    snpeff_jar="$snpeff_jar" \
    plink2="$plink2" \
    subset_suffix=_pgs \
    variant_list="$out_dir/subsets_pgs/variant_subset.1k.tsv" \
    out="$out_dir" \
    nfeatures=1k \
    dsname=demo \
    pgsid=demo \
    pvar_all=none \
  --until require_nfeatures \
  -p  --dryrun

# generally, you can delete the leftover plink files if you do not want to run further analyses with them
rm -r   "$out_dir"/plink

Once everything looks as expected, remove the --dryrun to execute the workflow.

You can use the specific variant sets we used by:

  1. copying them from the resources directory into the subsets pgs:
cp resources/variant_sets/all_genes_all_prs_100k.tsv "$out_dir/variant_subset.100k.tsv"
  1. remove the CHROM='[21,22]' subset from the config
  2. set the nfeatures=100k

Original contributors

This repo has been reset before publication. The contributors are juzb, LArnoldt, noahho, luisherrmann, khuenm, SedraAbouGhaloun, leoeckhoff

Related work on single cell sequencing data

If you are interested in multimodal single cell data modelling with a similar architecture, check out our networkvi repository and preprint:

Arnoldt, L., Upmeier Zu Belzen, J., Herrmann L., Ishaque, N., Theis, F.J., Wild, B., Eils, R. “Biologically Informed Variational Inference Enables Interpretable Cell Phenotyping and Discovery.” Preprint, June 12, 2025. [doi.org/10.1101/2025.06.10.657924](https://doi.org/10.1101/2025.06.10.657924).

Contact

If you found a bug, please use the issue tracker. If you use the OGM in your research, please consider citing the preprint:

Upmeier zu Belzen, J., Arnoldt, L., Hollmann, N., Hermann, L., Nguyen, K., Eckhoff, L., Kohleick, L., Ghaloun, S.A., Schmidt, H., Hegselmann, S., Theis, F.J., Bürgel, T., Steinfeldt, J., Wild, B., Eils, R., "An interpretable omnigenic neural network architecture for the human genome", medRxiv, July 2026.

About

An interpretable omnigenic neural network architecture for the human genome

Resources

Contributing

Stars

2 stars

Watchers

2 watching

Forks

Releases

Packages

Used by

Contributors

Languages