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).
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).
mamba/condaavailable (this repo usesmamba; substitutecondaif that's what you have).- Read/write access to wherever your institution's UKB data lives.
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_py312This 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.
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-depsogm/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_demoThe 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_gafThis 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.)
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.shcd /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.
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.
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.yamlTraining is local-only. Each run receives a UUID and writes its logs, configuration,
checkpoints, and results under <logging.out_directory>/<run-id>/.
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_pathcan be used to inferrun.experiment_id. It points at the models.ckptfile (the inner single-quotes are needed because checkpoint filenames contain=).+offline_config=.../FLAGS.jsoncan be used to adjust the config, otherwise the training config is usedeval.keep_activations=accumulateis required for any interpretability analysis to run at all.- Toggle exactly the analyses you want via the
no_*/feature flags inevaluate_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=Falsefor the GO over-representation/GSEA baseline. Seeogm/config/evaluate_example.yamlfor the full, documented list of flags and their defaults.
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=FalseThe 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/geno2are just plugin names underdatamodule.plugins. Any additional genetic-like modality is added the same way, as a new+datamodule.plugins.<name>.*block withname: GeneticPluginandtags: [genetic]; the gene layer concatenates every plugin taggedgeneticbefore 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 viamodel_config.hierarchy.ontology: biological_process,chebi_ontologyplus a matchingobo_file(seeogm/core/agnostic/models/ogm.py'sontology/obo_filehandling for how the comma-separated namespace list is used).
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'sFLAGS.json.OGM_PRETRAINED_CHECKPOINT: path to that run's.ckptfile.DBGAP_CELLINK_DD_PATH: path to your single-cellcellinkDonorDatafile (.h5.dd).OGM_SC_DEFAULT_CONFIG: optional, overrides the default single-cell hyperparameters JSON (ogm/bulk_single_cell/config/default_config.json) that--configYAML files are merged on top of.
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 globalLike 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 pbmcThese 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.
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_hyperThis 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
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-genoThen 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"/plinkOnce everything looks as expected, remove the --dryrun to execute the workflow.
You can use the specific variant sets we used by:
- 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"- remove the
CHROM='[21,22]'subset from the config - set the
nfeatures=100k
This repo has been reset before publication. The contributors are juzb, LArnoldt, noahho, luisherrmann, khuenm, SedraAbouGhaloun, leoeckhoff
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).
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.