Python API¶

Every command is a thin layer over these modules, and they can be used directly, from a notebook or another pipeline. For example, the published models on a cohort, without the command line:

from dietnet import package, predict

packages = package.find("models/packages")             # every package below, in (seed, fold) order
matched = predict.matched_dataset(packages, "mycohort", "work/")
result = predict.predict(packages, matched, device="cpu")

result.table()      # one row per sample: prediction, probability, model_agreement
result.scores       # models x samples x classes
result.alignment    # the matching counts

Upstream's own code is reached through dietnet.upstream: run executes one of its scripts in a subprocess and modules imports its modules in this process. The modules below call it; they never reimplement what it does.

Prediction¶

dietnet.predict ¶

Predict a cohort with one or more model packages, at biobank scale.

The arithmetic is upstream's: its DietNetwork and its indep_test_step (imputation with the training means, standardisation, the D/(D-|C|) scale, forward pass, softmax), which is what upstream's test_independent_dataset.py runs; tests/parity checks the two give identical scores.

What differs is the I/O around it. Upstream's script loads the whole cohort into memory, matches SNPs with a Python loop, and does both again for every model. Here the cohort is matched once (allele-aware, :mod:dietnet.align) into an HDF5 file under work_dir, then memory-mapped -- possible because the file is contiguous -- and read once, a contiguous batch of samples at a time, with every model scoring each batch before the next is read.

Result `dataclass` ¶

Result(
    samples, model_names, classes, scores, alignment, scale
)

table ¶

table()

One row per sample: the ensemble's call and how much the models agree.

Source code in src/dietnet/predict.py

def table(self) -> pd.DataFrame:
    """One row per sample: the ensemble's call and how much the models agree."""
    out = pd.DataFrame({"sample": self.samples})
    if not self.classes:
        out["prediction"] = self.scores.mean(axis=0)
        out["sd_across_models"] = self.scores.std(axis=0)
        return out
    mean = self.scores.mean(axis=0)
    call = mean.argmax(axis=1)
    votes = self.scores.argmax(axis=2)
    out["prediction"] = np.asarray(self.classes)[call]
    out["probability"] = mean[np.arange(len(call)), call]
    out["model_agreement"] = (votes == call[None]).mean(axis=0)
    return out

matched_dataset ¶

matched_dataset(packages, cohort_prefix, work_dir)

The cohort on the models' SNPs: built once, reused while inputs are unchanged.

Source code in src/dietnet/predict.py

def matched_dataset(packages: list[Package], cohort_prefix, work_dir) -> Path:
    """The cohort on the models' SNPs: built once, reused while inputs are unchanged."""
    snps = packages[0].snps()
    for pkg in packages[1:]:
        if not pkg.snps()["name"].equals(snps["name"]):
            raise ValueError(f"{pkg.path} was trained on different SNPs than {packages[0].path}")

    out = (
        Path(work_dir)
        / f"{Path(str(cohort_prefix)).name}.{align.fingerprint(snps, cohort_prefix)}.matched.h5"
    )
    if not out.exists():
        align.write_matched(cohort_prefix, snps, out)
    return out

Model packages¶

dietnet.package ¶

Model packages: one trained (seed, fold) model, self-contained and portable.

A package is a directory holding everything upstream's model needs, so that prediction never needs the training data or the experiment directory::

model.pt                  upstream's best_model.pt, as train.py saved it
config.yaml               the training config (seed included)
embedding.npz             ``emb``: this fold's SNP embedding (SNPs x features);
                          absent for the MLP baseline, which has no auxiliary net
input_features_stats.npz  ``means_by_fold`` / ``sd_by_fold`` for this fold only
train_header.h5           ``snp_names``, ``class_label_names``, ``snps/``: the
                          parts of the training dataset upstream's classes read
metadata.json             task, normalisation, seed, fold, upstream commit, ...

The model is rebuilt with upstream's own classes (DietNetworkHandler) and its weights loaded strictly: a checkpoint that does not fit fails instead of leaving layers at their random initialisation.

Package `dataclass` ¶

Package(path)

handler ¶

handler(device)

Upstream's model handler with this package's weights, in eval mode.

Source code in src/dietnet/package.py

def handler(self, device):
    """Upstream's model handler with this package's weights, in eval mode."""
    m = upstream.modules()
    header = str(self.path / "train_header.h5")
    if self.task == "classification":
        task = m.task_handlers.ClassificationHandler(header, nn.CrossEntropyLoss())
    else:
        task = m.task_handlers.RegressionHandler(header, nn.MSELoss())
    if self.model == "Mlp":
        handler = m.model_handlers.MlpHandler(task, header, self.config)
    else:
        handler = m.model_handlers.DietNetworkHandler(
            task, 0, str(self.path / "embedding.npz"), device, header, self.config, None
        )
    checkpoint = torch.load(self.path / "model.pt", map_location=device, weights_only=False)
    handler.model.load_state_dict(checkpoint["model_state_dict"], strict=True)
    handler.model.to(device)
    handler.model.eval()
    return handler

build ¶

build(
    out,
    *,
    checkpoint,
    config,
    embedding,
    means,
    sds,
    train_dataset,
    task,
    normalize,
    seed,
    fold,
    model="Dietnet",
    extra=None
)

Assemble a package from a finished upstream training run.

embedding is the fold's auxiliary-network input; None for the MLP, which has no auxiliary network.

Source code in src/dietnet/package.py

def build(
    out,
    *,
    checkpoint,
    config: dict,
    embedding: np.ndarray | None,
    means: np.ndarray,
    sds: np.ndarray | None,
    train_dataset,
    task: str,
    normalize: bool,
    seed: int,
    fold: int,
    model: str = "Dietnet",
    extra: dict | None = None,
) -> "Package":
    """Assemble a package from a finished upstream training run.

    ``embedding`` is the fold's auxiliary-network input; None for the MLP,
    which has no auxiliary network.
    """
    out = Path(out)
    if out.exists():
        shutil.rmtree(out)
    out.mkdir(parents=True)

    shutil.copy(checkpoint, out / "model.pt")
    (out / "config.yaml").write_text(yaml.safe_dump(config, sort_keys=False))
    if embedding is not None:
        np.savez(out / "embedding.npz", emb=np.asarray(embedding))
    stats = {"means_by_fold": np.asarray(means)[None]}
    if sds is not None:
        stats["sd_by_fold"] = np.asarray(sds)[None]
    np.savez(out / "input_features_stats.npz", **stats)
    _write_header(train_dataset, out / "train_header.h5")

    with h5py.File(out / "train_header.h5", "r") as f:
        n_snps = len(f["snp_names"])
        classes = (
            list(np.array(f["class_label_names"]).astype(str)) if task == "classification" else []
        )
    if embedding is not None and np.asarray(embedding).shape[0] != n_snps:
        raise ValueError(f"embedding has {np.asarray(embedding).shape[0]} rows for {n_snps} SNPs")

    metadata = {
        "format_version": FORMAT_VERSION,
        "model": model,
        "task": task,
        "classes": classes,
        "preprocessing": {"normalize": bool(normalize), "missing_snps": "scale D/(D-|C|)"},
        "seed": int(seed),
        "fold": int(fold),
        "n_snps": n_snps,
        "upstream_commit": upstream.commit(),
        "dietnet_version": __version__,
        "created": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
        **(extra or {}),
    }
    (out / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")
    return Package(out)

find ¶

find(root)

Every package at or below root, in (seed, fold) order.

Source code in src/dietnet/package.py

def find(root) -> list[Package]:
    """Every package at or below ``root``, in (seed, fold) order."""
    root = Path(root)
    found = [Package(p.parent) for p in root.rglob("metadata.json")]
    if not found:
        raise FileNotFoundError(f"no model packages under {root}")
    return sorted(found, key=lambda p: (p.metadata["seed"], p.metadata["fold"]))

Matching a cohort¶

dietnet.align ¶

Match a new cohort's SNPs to a trained model's, by position and allele.

A dosage counts one allele, and which one PLINK 1.9 picks (A1, by default the minor allele) depends on the cohort's own frequencies. In the 1KGP training data a quarter of SNPs count the reference allele and the rest the alternate. So the same SNP can be coded x in one cohort and 2 - x in another; matching by id alone would feed the model mirror-coded genotypes for those SNPs, silently.

Per model SNP, against the cohort's .bim:

same A1/A2 -> copy
A1 and A2 swapped -> flip (2 - x; missing stays missing)
the same on the other strand -> copy or flip, except for A/T and C/G SNPs, whose strand cannot be told from the alleles (a swap and the other strand look the same): those are kept only when both alleles match outright, and never flipped
otherwise, or absent -> missing

Missing SNPs then follow the published method (Rochefort-Boulanger, Scicluna et al. 2026; upstream's match_input_features): they are set to -1, which imputation and standardisation turn into 0, and the observed ones are scaled by D / (D - |C|), stored as scale in the matched dataset.

The matched dataset has upstream's schema for it (inputs, samples, snp_names, scale), so upstream's test_independent_dataset.py reads it with --matched-test-dataset.

Alignment `dataclass` ¶

Alignment(source, flip, counts=dict())

Where each model SNP comes from in the cohort, and whether to flip it.

scale `property` ¶

scale

D / (D - |C|): the published rescaling of observed SNPs.

align ¶

align(model_snps, cohort_prefix)

Align a cohort's .bim to model_snps (from :func:dietnet.dataset.snp_table).

Source code in src/dietnet/align.py

def align(model_snps: pd.DataFrame, cohort_prefix) -> Alignment:
    """Align a cohort's .bim to ``model_snps`` (from :func:`dietnet.dataset.snp_table`)."""
    cohort = read_bim(cohort_prefix).reset_index(names="row")
    model = model_snps.reset_index(drop=True).reset_index(names="col")

    by_position = {"chrom", "pos"} <= set(model.columns) and model[
        ["chrom", "pos"]
    ].notna().all().all()
    if by_position:
        model["key"] = model["chrom"].astype(str) + ":" + model["pos"].astype(str)
        cohort["key"] = cohort["chrom"] + ":" + cohort["pos"].astype(str)
    else:
        model["key"] = model["snp"]
        cohort["key"] = cohort["snp"]

    if "a1" not in model.columns or model["a1"].isna().all():
        # Nothing records which allele the model counts (a dataset whose SNP
        # names lack plink's _A1 suffix): match on the key alone and say so.
        pairs = model[["col", "key"]].merge(cohort[["row", "key"]], on="key")
        pairs = pairs.sort_values(["col", "row"]).drop_duplicates("col")
        source = np.full(len(model), -1, dtype=np.int64)
        source[pairs["col"].to_numpy()] = pairs["row"].to_numpy()
        counts = dict(
            alleles_unchecked=len(pairs),
            matched_by="position" if by_position else "id",
            model_snps=len(model),
            missing=int((source < 0).sum()),
        )
        return Alignment(source=source, flip=np.zeros(len(model), dtype=bool), counts=counts)

    model["a2"] = model["a2"] if "a2" in model.columns else None
    pairs = model.merge(cohort, on="key", suffixes=("_m", "_c"))
    a1m, a2m, a1c, a2c = pairs["a1_m"], pairs["a2_m"], pairs["a1_c"], pairs["a2_c"]
    known_a2 = a2m.notna()
    ambiguous = known_a2 & (a1m == _complement(a2m))

    same = (a1c == a1m) & (~known_a2 | (a2c == a2m))
    # For an A/T or C/G SNP a swap looks exactly like the other strand, so
    # only an outright match is trusted.
    swapped = ~ambiguous & (a2c == a1m) & (~known_a2 | (a1c == a2m))
    strand_same = ~ambiguous & (_complement(a1c) == a1m) & (~known_a2 | (_complement(a2c) == a2m))
    strand_swapped = (
        ~ambiguous & (_complement(a2c) == a1m) & (~known_a2 | (_complement(a1c) == a2m))
    )

    pairs["kind"] = np.select(
        [same, swapped, strand_same, strand_swapped],
        ["same", "swapped", "strand", "strand_swapped"],
        default="",
    )
    usable = pairs[pairs["kind"] != ""]
    # A position can hold several cohort variants (multi-allelic sites split
    # into biallelic rows); keep the first whose alleles fit.
    usable = usable.sort_values(["col", "row"]).drop_duplicates("col")

    source = np.full(len(model), -1, dtype=np.int64)
    flip = np.zeros(len(model), dtype=bool)
    source[usable["col"].to_numpy()] = usable["row"].to_numpy()
    flip[usable["col"].to_numpy()] = usable["kind"].isin(["swapped", "strand_swapped"]).to_numpy()

    counts = usable["kind"].value_counts().to_dict()
    counts.update(
        matched_by="position" if by_position else "id",
        model_snps=len(model),
        found_with_other_alleles=int(pairs.loc[pairs["kind"] == "", "col"].nunique()),
        ambiguous_kept=int((ambiguous & same)[pairs["kind"] != ""].sum()),
        missing=int((source < 0).sum()),
    )
    return Alignment(source=source, flip=flip, counts=counts)

fingerprint ¶

fingerprint(model_snps, cohort_prefix)

Identifies an (model SNP set, cohort) pair, to reuse a matched dataset.

Source code in src/dietnet/align.py

def fingerprint(model_snps: pd.DataFrame, cohort_prefix) -> str:
    """Identifies an (model SNP set, cohort) pair, to reuse a matched dataset."""
    digest = hashlib.sha256()
    digest.update("\n".join(model_snps["name"]).encode())
    for ext in (".bim", ".fam"):
        digest.update(Path(f"{cohort_prefix}{ext}").read_bytes())
    bed = Path(f"{cohort_prefix}.bed")
    digest.update(str(bed.stat().st_size).encode())
    return digest.hexdigest()[:16]

write_matched ¶

write_matched(
    cohort_prefix,
    model_snps,
    out,
    alignment=None,
    block_bytes=BLOCK_BYTES,
)

Write the cohort on the model's SNPs, in upstream's matched-dataset schema.

Built a block of samples at a time and written as contiguous runs of rows: a biobank cohort is matched in bounded memory and with sequential writes.

Source code in src/dietnet/align.py

def write_matched(
    cohort_prefix, model_snps: pd.DataFrame, out, alignment=None, block_bytes: int = BLOCK_BYTES
) -> Path:
    """Write the cohort on the model's SNPs, in upstream's matched-dataset schema.

    Built a block of samples at a time and written as contiguous runs of rows:
    a biobank cohort is matched in bounded memory and with sequential writes.
    """
    alignment = alignment or align(model_snps, cohort_prefix)
    bed = Bed.open(cohort_prefix)
    samples = read_fam(cohort_prefix)["iid"].to_numpy()
    present = np.flatnonzero(alignment.source >= 0)
    flipped = alignment.flip[present]

    out = Path(out)
    out.parent.mkdir(parents=True, exist_ok=True)
    tmp = out.with_name(out.name + ".partial")
    with h5py.File(tmp, "w") as f:
        inputs = f.create_dataset("inputs", shape=(bed.n_samples, alignment.n_model), dtype=np.int8)
        blocks = bed.sample_blocks(alignment.source[present], block_bytes, width=alignment.n_model)
        for start, dosages in blocks:
            swapped = dosages[:, flipped]
            dosages[:, flipped] = np.where(swapped >= 0, 2 - swapped, swapped)
            block = np.full((len(dosages), alignment.n_model), -1, dtype=np.int8)
            block[:, present] = dosages
            inputs[start : start + len(block)] = block

        f.create_dataset("samples", data=samples.astype("S"))
        f.create_dataset("snp_names", data=model_snps["name"].to_numpy().astype("S"))
        f.create_dataset("scale", data=np.array([alignment.scale]))
        f.attrs["alignment"] = json.dumps(alignment.counts)
        f.attrs["source_plink_prefix"] = str(Path(bed.prefix).resolve())
    tmp.replace(out)
    return out

The HDF5 dataset¶

dietnet.dataset ¶

Build the HDF5 dataset upstream's scripts read, straight from PLINK.

Upstream makes it from a genotype text table (plink --recode A, then create_dataset.py). This writes the same file from the .bed directly: same datasets, same dtypes, same values (tests/parity), without the text round-trip, which does not scale to a biobank.

Upstream schema, read by partition_data.py, generate_embedding.py, compute_input_features_statistics.py and train.py:

inputs int8, samples x SNPs, contiguous; dosage of the .bim A1, -1 missing samples bytes, sample ids (.fam IID), in .fam order snp_names bytes, <.bim id>_<A1>: the plink --recode A column names, so the counted allele travels with the name (as in the published 1KGP dataset: chr1:858952:G:A_A) class_labels int64 index into class_label_names (classification) class_label_names bytes, sorted unique labels regression_labels float64 (regression)

Added here, ignored by upstream: snps/{chrom,pos,a1,a2} (what matching a new cohort by position and allele needs) and provenance attributes.

The layout is contiguous and uncompressed on purpose: chunked or compressed layouts make the random minibatch reads training does 100-300x slower (docs/refactor-plan.md, section 2).

read_labels ¶

read_labels(path)

Sample id -> label from a TSV with a header row (upstream's format).

Source code in src/dietnet/dataset.py

def read_labels(path) -> pd.Series:
    """Sample id -> label from a TSV with a header row (upstream's format)."""
    table = pd.read_csv(path, sep="\t", dtype=str)
    if table.shape[1] < 2:
        raise ValueError(f"{path}: expected two tab-separated columns, sample id and label")
    labels = table.set_index(table.columns[0])[table.columns[1]]
    duplicated = labels.index[labels.index.duplicated()]
    if len(duplicated):
        raise ValueError(f"{path}: duplicated sample ids, e.g. {list(duplicated[:3])}")
    return labels

build ¶

build(
    plink_prefix,
    out,
    class_labels=None,
    regression_labels=None,
    block_bytes=BLOCK_BYTES,
)

Write out (upstream's HDF5 schema) from a PLINK fileset.

Source code in src/dietnet/dataset.py

def build(
    plink_prefix,
    out,
    class_labels=None,
    regression_labels=None,
    block_bytes: int = BLOCK_BYTES,
) -> Path:
    """Write ``out`` (upstream's HDF5 schema) from a PLINK fileset."""
    bed = Bed.open(plink_prefix)
    fam = read_fam(bed.prefix)
    bim = read_bim(bed.prefix)
    samples = fam["iid"].to_numpy()

    out = Path(out)
    out.parent.mkdir(parents=True, exist_ok=True)
    tmp = out.with_name(out.name + ".partial")
    with h5py.File(tmp, "w") as f:
        inputs = f.create_dataset("inputs", shape=(bed.n_samples, bed.n_snps), dtype=np.int8)
        for start, block in bed.sample_blocks(block_bytes=block_bytes):
            inputs[start : start + len(block)] = block

        names = (bim["snp"] + "_" + bim["a1"]).to_numpy()
        f.create_dataset("snp_names", data=names.astype("S"))
        f.create_dataset("samples", data=samples.astype("S"))

        if class_labels is not None:
            ordered = _ordered(read_labels(class_labels), samples, class_labels)
            names = np.sort(np.unique(ordered))
            f.create_dataset("class_labels", data=np.searchsorted(names, ordered).astype(np.int64))
            f.create_dataset("class_label_names", data=names.astype("S"))

        if regression_labels is not None:
            ordered = _ordered(read_labels(regression_labels), samples, regression_labels)
            f.create_dataset("regression_labels", data=ordered.astype("float64"))

        snps = f.create_group("snps")
        snps.create_dataset("chrom", data=bim["chrom"].to_numpy().astype("S"))
        snps.create_dataset("pos", data=bim["pos"].to_numpy().astype(np.int64))
        snps.create_dataset("a1", data=bim["a1"].to_numpy().astype("S"))
        snps.create_dataset("a2", data=bim["a2"].to_numpy().astype("S"))

        f.attrs["source_plink_prefix"] = str(Path(bed.prefix).resolve())
        f.attrs["source_bed_fingerprint"] = _fingerprint(bed.prefix)
        f.attrs["built_by"] = f"dietnet {__version__}"
        f.attrs["built_at"] = time.strftime("%Y-%m-%dT%H:%M:%S%z")
    tmp.replace(out)
    return out

snp_table ¶

snp_table(path)

One row per SNP: name as stored, snp id, a1 counted allele, and chrom, pos, a2 where they can be known.

Files built here carry them in snps/. Files made by upstream's create_dataset.py, including the published 1KGP dataset, only have snp_names; those are parsed: <id>_<A1> gives the counted allele, and an id of the form chr:pos:ref:alt gives position and other allele.

Source code in src/dietnet/dataset.py

def snp_table(path) -> pd.DataFrame:
    """One row per SNP: ``name`` as stored, ``snp`` id, ``a1`` counted allele, and
    ``chrom``, ``pos``, ``a2`` where they can be known.

    Files built here carry them in ``snps/``. Files made by upstream's
    create_dataset.py, including the published 1KGP dataset, only have
    ``snp_names``; those are parsed: ``<id>_<A1>`` gives the counted allele, and
    an id of the form ``chr:pos:ref:alt`` gives position and other allele.
    """
    with h5py.File(path, "r") as f:
        names = np.array(f["snp_names"]).astype(str)
        stored = None
        if "snps" in f:
            g = f["snps"]
            stored = {k: np.array(g[k]) for k in ("chrom", "pos", "a1", "a2")}

    table = pd.DataFrame({"name": names})
    split = table["name"].str.rsplit("_", n=1, expand=True)
    has_allele = split.shape[1] == 2 and split[1].fillna("").str.fullmatch(r"[ACGT]+").all()
    table["snp"] = split[0] if has_allele else table["name"]
    if stored is not None:
        table["chrom"] = stored["chrom"].astype(str)
        table["pos"] = stored["pos"].astype(np.int64)
        table["a1"] = stored["a1"].astype(str)
        table["a2"] = stored["a2"].astype(str)
        return table

    table["a1"] = split[1] if has_allele else None
    parts = table["snp"].str.split(":", expand=True)
    if parts.shape[1] == 4 and parts[1].str.fullmatch(r"\d+").all():
        table["chrom"] = parts[0].str.removeprefix("chr")
        table["pos"] = parts[1].astype(np.int64)
        if has_allele:
            ref, alt = parts[2], parts[3]
            table["a2"] = np.where(table["a1"] == alt, ref, np.where(table["a1"] == ref, alt, None))
    return table

memmap_inputs ¶

memmap_inputs(path)

inputs as a read-only memory map: no copy into RAM, pages read on demand.

Possible because the dataset is contiguous and uncompressed; HDF5 then stores it as one run of bytes at a fixed offset.

Source code in src/dietnet/dataset.py

def memmap_inputs(path) -> np.ndarray:
    """``inputs`` as a read-only memory map: no copy into RAM, pages read on demand.

    Possible because the dataset is contiguous and uncompressed; HDF5 then
    stores it as one run of bytes at a fixed offset.
    """
    with h5py.File(path, "r") as f:
        ds = f["inputs"]
        offset = ds.id.get_offset()
        if offset is None or ds.chunks is not None or ds.compression is not None:
            raise ValueError(f"{path}: 'inputs' is not contiguous and uncompressed")
        shape, dtype = ds.shape, ds.dtype
    # Copy-on-write: rows read as writable arrays (torch warns on read-only ones)
    # while the file itself is never modified.
    return np.memmap(path, dtype=dtype, mode="c", offset=offset, shape=shape)

Training¶

dietnet.training ¶

Train with upstream's train.py and package what it produces.

Training is upstream's script, run unmodified: the published model, loop, early stopping and checkpoint selection. This module supplies its inputs (one config per seed, since upstream reads the seed from the config), finds its results directory by upstream's own naming, and turns each (seed, fold) into a model package with the fold's held-out predictions next to it.

InputError ¶

Bases: ValueError

Training inputs upstream's model cannot learn from.

default_workers ¶

default_workers()

Loader processes: the CPUs this process may use (a SLURM allocation's, not the node's), less one for training itself, at most 8. Past ~8 the disk, not the CPUs, sets the pace (docs/scaling.md).

Source code in src/dietnet/training.py

def default_workers() -> int:
    """Loader processes: the CPUs this process may use (a SLURM allocation's, not
    the node's), less one for training itself, at most 8. Past ~8 the disk, not
    the CPUs, sets the pace (docs/scaling.md)."""
    try:
        cpus = len(os.sched_getaffinity(0))
    except AttributeError:  # macOS
        cpus = os.cpu_count() or 1
    return max(0, min(8, cpus - 1))

fold_slice ¶

fold_slice(array, fold)

The fold's embedding from a per-fold (folds x SNPs x features) array, or a fold-independent (SNPs x features) one as it is -- the two shapes upstream's DietNetwork accepts.

Source code in src/dietnet/training.py

def fold_slice(array: np.ndarray, fold: int) -> np.ndarray:
    """The fold's embedding from a per-fold (folds x SNPs x features) array, or a
    fold-independent (SNPs x features) one as it is -- the two shapes upstream's
    DietNetwork accepts."""
    array = np.asarray(array)
    return array[fold] if array.ndim == 3 else array

check_inputs ¶

check_inputs(
    embedding, input_stats, dataset, partition, folds=None
)

Refuse per-fold inputs upstream's model cannot learn from, saying why.

input_stats may be None: generate-embedding checks before they exist; embedding may be None: the MLP baseline has none.

Upstream's generate_embedding.py divides each class's genotype counts by the number of its samples in the fold's training split, and sizes the embedding by the largest class index present. A class with no training samples in a fold -- a population of one or two people, split five ways -- therefore either gives 0/0 (NaN, which spreads to every weight: the model trains to NaN and predicts the first class for everyone) or, if it is the last class, silently drops its columns. Either way the class cannot be learned. Caught here, before any training time is spent.

Source code in src/dietnet/training.py

def check_inputs(embedding, input_stats, dataset, partition, folds=None) -> None:
    """Refuse per-fold inputs upstream's model cannot learn from, saying why.

    ``input_stats`` may be None: generate-embedding checks before they exist;
    ``embedding`` may be None: the MLP baseline has none.

    Upstream's generate_embedding.py divides each class's genotype counts by
    the number of its samples in the fold's training split, and sizes the
    embedding by the largest class index present. A class with no training
    samples in a fold -- a population of one or two people, split five ways --
    therefore either gives 0/0 (NaN, which spreads to every weight: the model
    trains to NaN and predicts the first class for everyone) or, if it is the
    last class, silently drops its columns. Either way the class cannot be
    learned. Caught here, before any training time is spent.
    """
    emb = np.load(embedding, allow_pickle=True)["emb"] if embedding else None
    stats = np.load(input_stats, allow_pickle=True) if input_stats else None
    parts = np.load(partition, allow_pickle=True)["folds_indexes"]
    with h5py.File(dataset, "r") as f:
        labels = np.array(f["class_labels"]) if "class_labels" in f else None
        names = np.array(f["class_label_names"]).astype(str) if "class_label_names" in f else None

    problems = []
    for fold in range(len(parts)) if folds is None else folds:
        if labels is not None:
            train = np.asarray(parts[fold][0], dtype=int)
            absent = names[np.bincount(labels[train], minlength=len(names)) == 0]
            if len(absent):
                problems.append(f"fold {fold}: no training samples for {', '.join(absent)}")
        nan = ["embedding"] if emb is not None and np.isnan(fold_slice(emb, fold)).any() else []
        for key in ("means_by_fold", "sd_by_fold"):
            if stats is not None and key in stats.files and np.isnan(stats[key][fold]).any():
                nan.append(key)
        if nan:
            problems.append(f"fold {fold}: NaN in {', '.join(nan)}")

    if problems:
        raise InputError(
            "the model cannot learn from these training inputs:\n  "
            + "\n  ".join(problems)
            + "\nMerge or drop classes too small to appear in every fold's training split "
            "(at least nb-folds samples each), then rebuild the dataset and rerun."
        )

train ¶

train(
    *,
    dataset,
    partition,
    embedding,
    input_stats,
    config,
    exp_dir,
    seeds,
    folds,
    task="classification",
    normalize=True,
    model="Dietnet",
    inputs="auto",
    loader="fast",
    workers=None,
    gpus=1,
    engine=None
)

Train every (seed, fold); return their packages.

model is upstream's: "Dietnet", or "Mlp", the baseline without an auxiliary network, which takes no embedding.

engine "upstream" (the default on one GPU) runs upstream's train.py, unmodified, fed so that the trained model is bit-identical to train.py's:

inputs: the genotype store it reads (:mod:dietnet.genotypes): "ram" (upstream's copy of the whole dataset), "memmap", "packed", or "auto".
loader: "fast" reads each batch in one call with workers loader processes (default :func:default_workers) and fills missing values in one operation; "upstream" leaves train.py's own loader alone.

engine "dietnet" (the default, and the only choice, for gpus > 1) trains with :mod:dietnet.parallel: upstream's model and epoch logic in a loop of its own, data-parallel over gpus, packed genotypes decoded on the GPU, without train.py's per-batch bookkeeping; on Mila's A100 its training pass ran ~20x upstream's (docs/scaling.md). Its weights differ from train.py's.

Source code in src/dietnet/training.py

def train(
    *,
    dataset,
    partition,
    embedding,
    input_stats,
    config,
    exp_dir,
    seeds,
    folds,
    task="classification",
    normalize=True,
    model="Dietnet",
    inputs="auto",
    loader="fast",
    workers: int | None = None,
    gpus: int = 1,
    engine: str | None = None,
) -> list[package.Package]:
    """Train every (seed, fold); return their packages.

    ``model`` is upstream's: "Dietnet", or "Mlp", the baseline without an
    auxiliary network, which takes no embedding.

    ``engine`` "upstream" (the default on one GPU) runs upstream's train.py,
    unmodified, fed so that the trained model is bit-identical to train.py's:

    - ``inputs``: the genotype store it reads (:mod:`dietnet.genotypes`): "ram"
      (upstream's copy of the whole dataset), "memmap", "packed", or "auto".
    - ``loader``: "fast" reads each batch in one call with ``workers`` loader
      processes (default :func:`default_workers`) and fills missing values in
      one operation; "upstream" leaves train.py's own loader alone.

    ``engine`` "dietnet" (the default, and the only choice, for ``gpus`` > 1)
    trains with :mod:`dietnet.parallel`: upstream's model and epoch logic in a
    loop of its own, data-parallel over ``gpus``, packed genotypes decoded on the
    GPU, without train.py's per-batch bookkeeping; on Mila's A100 its training
    pass ran ~20x upstream's (docs/scaling.md). Its weights differ from train.py's.
    """
    engine = engine or ("dietnet" if gpus > 1 else "upstream")
    if engine not in ("upstream", "dietnet"):
        raise ValueError(f"engine must be 'upstream' or 'dietnet', not {engine!r}")
    if loader not in ("fast", "upstream"):
        raise ValueError(f"loader must be 'fast' or 'upstream', not {loader!r}")
    if gpus < 1:
        raise ValueError("gpus must be at least 1")
    if gpus > 1 and engine == "upstream":
        raise ValueError("engine 'upstream' is train.py's own loop, on one GPU")
    if engine == "dietnet" and loader == "upstream":
        raise ValueError("--loader upstream applies to engine 'upstream' only")
    if model not in upstream.MODELS:
        raise ValueError(f"model must be one of {sorted(upstream.MODELS)}")
    if (embedding is None) != (model == "Mlp"):
        raise ValueError("an embedding is needed for the Diet Network, and unused by the MLP")
    dataset, partition = Path(dataset).resolve(), Path(partition).resolve()
    embedding = Path(embedding).resolve() if embedding else None
    input_stats = Path(input_stats).resolve()
    exp_dir = Path(exp_dir).resolve()
    base = yaml.safe_load(Path(config).read_text())
    check_inputs(embedding, input_stats, dataset, partition, folds)
    try:
        env = {_launch.INPUTS_ENV: genotypes.choose(dataset, inputs)}
    except ValueError as error:
        raise InputError(str(error)) from None
    if loader == "fast":
        workers = default_workers() if workers is None else workers
        env.update({_launch.LOADER_ENV: "fast", _launch.WORKERS_ENV: str(workers)})

    emb_by_fold = np.load(embedding, allow_pickle=True)["emb"] if embedding else None
    described = Path(f"{embedding}.json") if embedding else None
    feature_meta = json.loads(described.read_text()) if described and described.exists() else None
    stats = np.load(input_stats, allow_pickle=True)

    packages = []
    for seed in seeds:
        cfg = {**base, "seed": int(seed)}
        cfg_path = exp_dir / "configs" / f"config_seed{seed}.yaml"
        cfg_path.parent.mkdir(parents=True, exist_ok=True)
        cfg_path.write_text(yaml.safe_dump(cfg, sort_keys=False))

        for fold in folds:
            args = [
                "--exp-path", exp_dir, "--exp-name", "runs", "--config", cfg_path,
                "--model", model, "--dataset", dataset, "--partition", partition,
                # Required by train.py, read only for the Diet Network.
                "--embedding", embedding or "none", "--input-features-stats", input_stats,
                "--task", task, "--which-fold", fold,
            ]  # fmt: skip
            if normalize:
                args.append("--normalize")
            log = exp_dir / "logs" / f"train_seed{seed}_fold{fold}.log"
            if engine == "dietnet":
                args += ["--inputs", env[_launch.INPUTS_ENV], "--workers", workers]
                _run_parallel(args, gpus, log)
            else:
                upstream.run("train.py", args, log=log, env=env)

            results = (
                exp_dir / "runs" / f"RESULTS_{upstream.exp_identifier(cfg, fold, task, model)}"
            )
            pkg = package.build(
                exp_dir / "packages" / f"seed_{seed}" / f"fold_{fold}",
                checkpoint=results / "best_model.pt",
                config=cfg,
                embedding=None if emb_by_fold is None else fold_slice(emb_by_fold, fold),
                means=stats["means_by_fold"][fold],
                sds=stats["sd_by_fold"][fold] if "sd_by_fold" in stats.files else None,
                train_dataset=dataset,
                task=task,
                normalize=normalize,
                seed=seed,
                fold=fold,
                model=model,
                extra={"results_dir": str(results), "embedding": feature_meta},
            )
            _held_out_predictions(results, dataset, task).to_csv(
                pkg.path / "test_predictions.tsv", sep="\t", index=False
            )
            packages.append(pkg)
    return packages

SNP features¶

dietnet.features ¶

Per-SNP features for the auxiliary network, from several sources.

Upstream's DietNetwork reads its auxiliary-network input from embedding.npz: a SNPs x features matrix, one per fold (3-D) or shared by all folds (2-D), whose columns it L2-normalises itself. Anything that can be written as such a matrix can feed the model, without touching upstream's code. This module builds one from blocks, described in a YAML spec::

blocks:
  - provider: class_genotype_freq          # upstream's embedding (per fold)
  - provider: precomputed                  # any per-SNP table, e.g. AlphaGenome scores
    path: alphagenome_scores.tsv
    orient: effect                         # values are ALT - REF effects
    pca: 64                                # reduce to 64 components
  - provider: random                       # control: same width, no information
    columns: 64
    seed: 0

Providers:

class_genotype_freq Genotype frequencies per class and fold: upstream's generate_embedding.py, run unmodified. Alone, the result is exactly what dietnet generate-embedding writes. genotype_freq Genotype frequencies per fold over all its training samples, classes pooled: the same definition without the split by class, for regression tasks (which have no classes) or as a baseline. Read a block of samples at a time from a genotype store (inputs: as for train), so it scales to a biobank. precomputed A table with chrom, pos, ref, alt and numeric feature columns (TSV, optionally gzipped, or parquet), matched to the dataset's SNPs by position and alleles, on either strand. With orient: effect a value describes ALT relative to REF, so it is negated where the dataset counts REF (plink 1.9 makes the minor allele A1, which is REF for a quarter of the 1KGP SNPs). SNPs the table lacks get zeros and are marked in an added <name>:missing column. random Standard-normal columns: the control for "more columns" rather than "more information". shuffled Another block with its SNP rows permuted (of: that block's index): the control for "this information" rather than "information of this shape".

Transforms per block, applied after the provider: center: true subtracts each column's mean over SNPs; pca: k keeps the first k principal components over SNPs (computed per fold for a per-fold block). Upstream's per-column L2 normalisation undoes any per-column scaling, so these are the transforms that matter: centering, and controlling how many columns -- and so how much weight -- a block brings.

Every column must vary: an all-zero column would be divided by a zero norm inside the model and turn it to NaN, so building refuses one.

Block `dataclass` ¶

Block(values, columns, info=dict())

A provider's output: SNPs x columns, or folds x SNPs x columns.

build ¶

build(spec, dataset, partition, out)

Write out (upstream's embedding.npz) and <out>.json from a spec.

Source code in src/dietnet/features.py

def build(spec, dataset, partition, out) -> tuple[np.ndarray, dict]:
    """Write ``out`` (upstream's embedding.npz) and ``<out>.json`` from a spec."""
    spec = yaml.safe_load(Path(spec).read_text()) if not isinstance(spec, dict) else spec
    # Resolved: upstream's scripts run from their own directory, where a path
    # relative to the caller's would not be found.
    dataset, partition, out = Path(dataset).resolve(), Path(partition).resolve(), Path(out)
    n_folds = len(np.load(partition, allow_pickle=True)["folds_indexes"])
    ctx = Context(dataset, partition, ds.snp_table(dataset), n_folds, [])

    described = []
    for i, block_spec in enumerate(spec["blocks"]):
        provider = block_spec.get("provider")
        if provider not in PROVIDERS:
            raise ValueError(
                f"block {i}: unknown provider {provider!r}; one of {sorted(PROVIDERS)}"
            )
        block = _transform(PROVIDERS[provider](ctx, block_spec), block_spec)
        ctx.built.append(block)
        described.append({"provider": provider, "spec": block_spec, **block.info})

    if any(b.fold_dependent for b in ctx.built):
        values = np.concatenate(
            [b.values if b.fold_dependent else np.broadcast_to(b.values, (n_folds, *b.values.shape))
             for b in ctx.built],
            axis=-1,
        )  # fmt: skip
    else:
        values = np.concatenate([b.values for b in ctx.built], axis=-1)

    columns = [c for b in ctx.built for c in b.columns]
    flat = values.reshape(-1, values.shape[-2], values.shape[-1])
    constant = sorted(
        {columns[j] for v in flat for j in np.flatnonzero(np.abs(v).sum(axis=0) == 0)}
    )
    if constant:
        raise ValueError(
            f"all-zero feature columns would be NaN in the model: {constant[:5]}"
            + (" ..." if len(constant) > 5 else "")
        )

    start = 0
    for entry, block in zip(described, ctx.built):
        entry["columns"] = [start, start + len(block.columns)]
        start += len(block.columns)
    meta = {
        "blocks": described,
        "column_names": columns,
        "fold_dependent": values.ndim == 3,
        "shape": list(values.shape),
    }
    out.parent.mkdir(parents=True, exist_ok=True)
    np.savez(out, emb=values)
    Path(f"{out}.json").write_text(json.dumps(meta, indent=2) + "\n")
    return values, meta

dietnet.alphagenome ¶

AlphaGenome variant-effect scores for a dataset's SNPs, as an auxiliary-network feature table.

AlphaGenome (Avsec et al.) predicts many genomic tracks from DNA sequence, and scores a variant by how its ALT allele changes them relative to REF. Through the hosted API (pip install 'dietnetwork[alphagenome]', key in ALPHAGENOME_API_KEY), :func:score turns every SNP of a dataset into a row of such effects:

chrom  pos  ref  alt  <output_type>|<track_name>  ...

one signed ALT-minus-REF raw_score per track. Gene-level scorers report one score per nearby gene; the gene with the largest effect on the track is kept. The table is what build-embedding's precomputed provider reads, with orient: effect: it flips the sign wherever the dataset's dosage counts REF rather than ALT, and pca: k reduces the thousands of tracks to a block the size of the rest of the embedding.

Scoring is resumable: rows are appended a chunk at a time, and a rerun skips SNPs already in the file. Hundreds of thousands of API calls will be interrupted at some point.

REF and ALT come from SNP ids of the form chr:pos:REF:ALT (the 1KGP data's), since a .bim's A1/A2 say which allele is counted, not which is the reference. Coordinates must be GRCh38, which is what AlphaGenome uses.

Backend ¶

Bases: Protocol

score ¶

score(variants)

Tidy scores for variants: one row per (variant, track[, gene]) with chrom, pos, ref, alt, output_type, track_name, raw_score.