Tutorial: the whole pipeline on a small 1KGP panel¶

This notebook trains Diet Network models and predicts with them, end to end, on kgp_mini: 480 samples (20 from each of the 24 populations) × 2,000 SNPs of public 1000 Genomes data, committed in tests/fixtures/. Every step is a dietnet command, run here from Python so the files between them can be inspected. It takes under a minute on a laptop CPU and needs no network.

With 2,000 SNPs and about 290 training samples per fold the models are far weaker than the published ones (229,986 SNPs); the point is the pipeline, not the accuracy.

It runs from a clone of the repository: uv sync --frozen --extra docs, then open docs/tutorial.ipynb.

In [1]:

import os
import subprocess
import sys
import tempfile
from pathlib import Path

import h5py
import numpy as np
import pandas as pd

# The fixtures are in the repository, next to docs/.
REPO = next(p for p in [Path.cwd(), *Path.cwd().parents] if (p / "tests" / "fixtures").is_dir())
FIXTURES = REPO / "tests" / "fixtures"
PANEL = FIXTURES / "kgp_mini"
LABELS = FIXTURES / "kgp_mini_labels.tsv"

# Everything this notebook writes goes to a scratch directory.
os.chdir(tempfile.mkdtemp(prefix="dietnet-tutorial-"))


def dietnet(*args, tail=None):
    """Run a dietnet command; print its output (only the last `tail` lines if given)."""
    result = subprocess.run(
        [sys.executable, "-m", "dietnet.cli", *map(str, args)],
        capture_output=True,
        text=True,
    )
    here = os.path.realpath(os.getcwd())
    output = (result.stdout + result.stderr).replace(here + os.sep, "")
    lines = [line for line in output.splitlines() if line.strip(" -")]
    print("\n".join(lines[-tail:] if tail else lines))
    if result.returncode != 0:
        raise RuntimeError(f"dietnet {args[0]} exited with status {result.returncode}")

import os import subprocess import sys import tempfile from pathlib import Path import h5py import numpy as np import pandas as pd # The fixtures are in the repository, next to docs/. REPO = next(p for p in [Path.cwd(), *Path.cwd().parents] if (p / "tests" / "fixtures").is_dir()) FIXTURES = REPO / "tests" / "fixtures" PANEL = FIXTURES / "kgp_mini" LABELS = FIXTURES / "kgp_mini_labels.tsv" # Everything this notebook writes goes to a scratch directory. os.chdir(tempfile.mkdtemp(prefix="dietnet-tutorial-")) def dietnet(*args, tail=None): """Run a dietnet command; print its output (only the last `tail` lines if given).""" result = subprocess.run( [sys.executable, "-m", "dietnet.cli", *map(str, args)], capture_output=True, text=True, ) here = os.path.realpath(os.getcwd()) output = (result.stdout + result.stderr).replace(here + os.sep, "") lines = [line for line in output.splitlines() if line.strip(" -")] print("\n".join(lines[-tail:] if tail else lines)) if result.returncode != 0: raise RuntimeError(f"dietnet {args[0]} exited with status {result.returncode}")

In [2]:

dietnet("info")

dietnet("info")

dietnet 0.3.0
upstream HussinLab/DIETNETWORK 5481556e472140c6d3a67363c1286f22109fb1ce at /home/runner/work/DIETNETWORK/DIETNETWORK
upstream files: identical to the pinned commit
torch 2.14.1+cpu, cuda available: False

The inputs¶

A PLINK fileset and a label table: a header row, then sample id and label.

In [3]:

labels = pd.read_csv(LABELS, sep="\t")
print(labels.head(3), "\n")
print(labels.iloc[:, 1].value_counts().describe()[["count", "min", "max"]])

labels = pd.read_csv(LABELS, sep="\t") print(labels.head(3), "\n") print(labels.iloc[:, 1].value_counts().describe()[["count", "min", "max"]])

    Sample Population
0  HG00096     CEUGBR
1  HG00130     CEUGBR
2  HG00132     CEUGBR 

count    24.0
min      20.0
max      20.0
Name: count, dtype: float64

In [4]:

pd.read_csv(f"{PANEL}.bim", sep="\t", header=None,
            names=["chrom", "snp", "cm", "pos", "a1", "a2"]).head(3)

pd.read_csv(f"{PANEL}.bim", sep="\t", header=None, names=["chrom", "snp", "cm", "pos", "a1", "a2"]).head(3)

Out[4]:

	chrom	snp	cm	pos	a1	a2
0	1	chr1:3377101:A:G	0	3377101	A	G
1	1	chr1:4109988:C:T	0	4109988	T	C
2	1	chr1:6107134:T:C	0	6107134	C	T

1. The HDF5 dataset¶

create-dataset reads the .bed directly and writes the HDF5 file upstream's scripts read. Dosages count the .bim A1 allele, and SNP names carry it (<id>_<A1>), as plink --recode A would name them.

In [5]:

dietnet("create-dataset", "--plink", PANEL, "--labels", LABELS, "--out", "dataset.h5")

with h5py.File("dataset.h5") as f:
    for key in ("inputs", "samples", "snp_names", "class_labels", "class_label_names"):
        print(f"{key:18s} {f[key].dtype}  {f[key].shape}")
    print("first SNP names:", np.array(f["snp_names"][:2]).astype(str))

dietnet("create-dataset", "--plink", PANEL, "--labels", LABELS, "--out", "dataset.h5") with h5py.File("dataset.h5") as f: for key in ("inputs", "samples", "snp_names", "class_labels", "class_label_names"): print(f"{key:18s} {f[key].dtype} {f[key].shape}") print("first SNP names:", np.array(f["snp_names"][:2]).astype(str))

wrote dataset.h5
inputs             int8  (480, 2000)
samples            |S7  (480,)
snp_names          |S21  (2000,)
class_labels       int64  (480,)
class_label_names  |S6  (24,)
first SNP names: ['chr1:3377101:A:G_A' 'chr1:4109988:C:T_T']

2. Folds¶

Upstream's partition_data.py: five folds, each with train, validation and test splits, class labels balanced across them.

In [6]:

dietnet("partition", "--dataset", "dataset.h5", "--out", "partition.npz", "--seed", 3, tail=1)

folds = np.load("partition.npz", allow_pickle=True)["folds_indexes"]
pd.DataFrame([[len(s) for s in fold] for fold in folds],
             columns=["train", "valid", "test"]).rename_axis("fold")

dietnet("partition", "--dataset", "dataset.h5", "--out", "partition.npz", "--seed", 3, tail=1) folds = np.load("partition.npz", allow_pickle=True)["folds_indexes"] pd.DataFrame([[len(s) for s in fold] for fold in folds], columns=["train", "valid", "test"]).rename_axis("fold")

Dataset partition was saved to partition.npz

Out[6]:

	train	valid	test
fold
0	288	96	96
1	288	96	96
2	288	96	96
3	288	96	96
4	288	96	96

3. Per-fold inputs¶

The auxiliary network's input (genotype frequencies per population, from each fold's training samples) and each SNP's mean and SD, both from upstream's scripts.

In [7]:

dietnet("generate-embedding", "--dataset", "dataset.h5", "--partition", "partition.npz",
        "--out", "embedding.npz", tail=1)
dietnet("compute-input-stats", "--dataset", "dataset.h5", "--partition", "partition.npz",
        "--out", "stats.npz", tail=1)

print("emb:", np.load("embedding.npz")["emb"].shape, "(folds x SNPs x 3 genotypes x 24 populations)")
stats = np.load("stats.npz")
print({key: stats[key].shape for key in stats.files})

dietnet("generate-embedding", "--dataset", "dataset.h5", "--partition", "partition.npz", "--out", "embedding.npz", tail=1) dietnet("compute-input-stats", "--dataset", "dataset.h5", "--partition", "partition.npz", "--out", "stats.npz", tail=1) print("emb:", np.load("embedding.npz")["emb"].shape, "(folds x SNPs x 3 genotypes x 24 populations)") stats = np.load("stats.npz") print({key: stats[key].shape for key in stats.files})

Saving results to embedding

Execution time 0.7771859169006348 seconds
emb: (5, 2000, 72) (folds x SNPs x 3 genotypes x 24 populations)
{'means_by_fold': (5, 2000), 'sd_by_fold': (5, 2000)}

4. Training¶

Upstream's train.py, once per (seed, fold). This config is small enough to train in seconds; the published one is in the training guide. Two folds of one seed give a two-model ensemble.

In [8]:

Path("config.yaml").write_text("""\
batch_size: 32
epochs: 60
input_dropout: 0.5
dropout_main: 0.5
lr_aux: 0.001
lr_main: 0.001
learning_rate_annealing: 0.99
nb_hidden_u_aux: [64, 64]
nb_hidden_u_main: [64]
patience: 60
seed: 1
uniform_init_limit: 0.02
""")

dietnet("train", "--dataset", "dataset.h5", "--partition", "partition.npz",
        "--embedding", "embedding.npz", "--input-stats", "stats.npz",
        "--config", "config.yaml", "--exp-dir", "exp", "--seeds", 1, "--folds", "0,1")
sorted(str(p) for p in Path("exp/packages").rglob("*") if p.is_file())

Path("config.yaml").write_text("""\ batch_size: 32 epochs: 60 input_dropout: 0.5 dropout_main: 0.5 lr_aux: 0.001 lr_main: 0.001 learning_rate_annealing: 0.99 nb_hidden_u_aux: [64, 64] nb_hidden_u_main: [64] patience: 60 seed: 1 uniform_init_limit: 0.02 """) dietnet("train", "--dataset", "dataset.h5", "--partition", "partition.npz", "--embedding", "embedding.npz", "--input-stats", "stats.npz", "--config", "config.yaml", "--exp-dir", "exp", "--seeds", 1, "--folds", "0,1") sorted(str(p) for p in Path("exp/packages").rglob("*") if p.is_file())

packaged exp/packages/seed_1/fold_0
packaged exp/packages/seed_1/fold_1

Out[8]:

['exp/packages/seed_1/fold_0/config.yaml',
 'exp/packages/seed_1/fold_0/embedding.npz',
 'exp/packages/seed_1/fold_0/input_features_stats.npz',
 'exp/packages/seed_1/fold_0/metadata.json',
 'exp/packages/seed_1/fold_0/model.pt',
 'exp/packages/seed_1/fold_0/test_predictions.tsv',
 'exp/packages/seed_1/fold_0/train_header.h5',
 'exp/packages/seed_1/fold_1/config.yaml',
 'exp/packages/seed_1/fold_1/embedding.npz',
 'exp/packages/seed_1/fold_1/input_features_stats.npz',
 'exp/packages/seed_1/fold_1/metadata.json',
 'exp/packages/seed_1/fold_1/model.pt',
 'exp/packages/seed_1/fold_1/test_predictions.tsv',
 'exp/packages/seed_1/fold_1/train_header.h5']

Each package is self-contained, and records how it was trained:

In [9]:

from dietnet import package

packages = package.find("exp/packages")
{k: v for k, v in packages[0].metadata.items() if k not in ("classes", "results_dir")}

from dietnet import package packages = package.find("exp/packages") {k: v for k, v in packages[0].metadata.items() if k not in ("classes", "results_dir")}

Out[9]:

{'format_version': 2,
 'model': 'Dietnet',
 'task': 'classification',
 'preprocessing': {'normalize': True, 'missing_snps': 'scale D/(D-|C|)'},
 'seed': 1,
 'fold': 0,
 'n_snps': 2000,
 'upstream_commit': '5481556e472140c6d3a67363c1286f22109fb1ce',
 'dietnet_version': '0.3.0',
 'created': '2026-10-28T18:00:33+0000',
 'embedding': None}

Each model's accuracy on its own held-out test fold (chance is 1/24):

In [10]:

dietnet("check", "--packages", "exp/packages")

dietnet("check", "--packages", "exp/packages")

 seed  fold  n  accuracy
    1     0 96  0.406250
    1     1 96  0.489583
mean held-out accuracy: 0.4479 (sd 0.0589)

5. Prediction¶

predict matches the cohort to the models' SNPs by position and allele, then scores it with every model and combines them. Here the "cohort" is the panel itself, so every SNP is found; on another cohort the counts show how many were copied, flipped or missing.

In [11]:

dietnet("predict", "--model", "exp/packages", "--plink", PANEL, "--out", "predictions.tsv",
        "--device", "cpu")
predictions = pd.read_csv("predictions.tsv", sep="\t")
predictions.head()

dietnet("predict", "--model", "exp/packages", "--plink", PANEL, "--out", "predictions.tsv", "--device", "cpu") predictions = pd.read_csv("predictions.tsv", sep="\t") predictions.head()

Embedding size: torch.Size([2000, 72])
Embedding size: torch.Size([2000, 72])
matching 0s, inference 1s (2 model(s))
480 samples, 2 model(s); 2000/2000 model SNPs found (scale 1.0000): {"same": 2000, "matched_by": "position", "model_snps": 2000, "found_with_other_alleles": 0, "ambiguous_kept": 4, "missing": 0}
wrote predictions.tsv

Out[11]:

	sample	prediction	probability	model_agreement
0	HG00096	IBS	0.276338	1.0
1	HG00130	IBS	0.264165	0.5
2	HG00132	FIN	0.254850	0.0
3	HG00138	CEUGBR	0.877572	1.0
4	HG00140	CEUGBR	0.876383	1.0

In [12]:

dietnet("check", "--predictions", "predictions.tsv", "--labels", LABELS, tail=1)

dietnet("check", "--predictions", "predictions.tsv", "--labels", LABELS, tail=1)

accuracy: 0.9000 over 480 samples

Most of these samples were in the models' training splits, so this accuracy flatters them; the held-out numbers above are the honest ones.

The same through the Python API, which also gives every model's scores:

In [13]:

from dietnet import predict

matched = predict.matched_dataset(packages, PANEL, "work")
result = predict.predict(packages, matched, device="cpu")
print("scores:", result.scores.shape, "(models x samples x classes)")
print("scale:", result.scale)
result.alignment

from dietnet import predict matched = predict.matched_dataset(packages, PANEL, "work") result = predict.predict(packages, matched, device="cpu") print("scores:", result.scores.shape, "(models x samples x classes)") print("scale:", result.scale) result.alignment

Embedding size: torch.Size([2000, 72])

Embedding size: torch.Size([2000, 72])
scores: (2, 480, 24) (models x samples x classes)
scale: 1.0

Out[13]:

{'same': 2000,
 'matched_by': 'position',
 'model_snps': 2000,
 'found_with_other_alleles': 0,
 'ambiguous_kept': 4,
 'missing': 0}

6. Other SNP features and the MLP baseline¶

build-embedding assembles the auxiliary network's input from blocks. Here: the published features plus a block of random columns, the control for "more columns" rather than "more information". See SNP features.

In [14]:

Path("features.yaml").write_text("""\
blocks:
  - provider: class_genotype_freq
  - provider: random
    columns: 16
""")
dietnet("build-embedding", "--dataset", "dataset.h5", "--partition", "partition.npz",
        "--spec", "features.yaml", "--out", "embedding_plus_random.npz")

Path("features.yaml").write_text("""\ blocks: - provider: class_genotype_freq - provider: random columns: 16 """) dietnet("build-embedding", "--dataset", "dataset.h5", "--partition", "partition.npz", "--spec", "features.yaml", "--out", "embedding_plus_random.npz")

  class_genotype_freq  columns 0-71
  random               columns 72-87
wrote embedding_plus_random.npz (5, 2000, 88) and embedding_plus_random.npz.json

Upstream's MLP has no auxiliary network, so it takes no embedding. It is the baseline any change to the SNP features should beat.

In [15]:

dietnet("train", "--model", "mlp", "--dataset", "dataset.h5", "--partition", "partition.npz",
        "--input-stats", "stats.npz", "--config", "config.yaml", "--exp-dir", "exp_mlp",
        "--seeds", 1, "--folds", "0,1")
dietnet("check", "--packages", "exp_mlp/packages")

dietnet("train", "--model", "mlp", "--dataset", "dataset.h5", "--partition", "partition.npz", "--input-stats", "stats.npz", "--config", "config.yaml", "--exp-dir", "exp_mlp", "--seeds", 1, "--folds", "0,1") dietnet("check", "--packages", "exp_mlp/packages")

packaged exp_mlp/packages/seed_1/fold_0
packaged exp_mlp/packages/seed_1/fold_1

 seed  fold  n  accuracy
    1     0 96  0.343750
    1     1 96  0.395833
mean held-out accuracy: 0.3698 (sd 0.0368)