Files
SpectraREML/examples/example.py

84 lines
3.0 KiB
Python

#!/usr/bin/env python3
"""Create a small, fully synthetic SpectraREML input bundle."""
from __future__ import annotations
import sys
from pathlib import Path
import numpy as np
PROJECT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT / "python"))
import spectra_reml # noqa: E402
def main() -> None:
output = Path(__file__).resolve().parent / "synthetic_work"
output.mkdir(parents=True, exist_ok=True)
rng = np.random.default_rng(20260809)
n = 12
z = rng.normal(size=(n, 5))
grm = z @ z.T / z.shape[1]
diagonal = np.sqrt(np.diag(grm))
grm /= np.outer(diagonal, diagonal)
packed = np.asarray(
[grm[row, col] for row in range(n) for col in range(row + 1)],
dtype="<f4",
)
packed.tofile(output / "example.grm.bin")
(output / "example.grm.id").write_text(
"".join("sample{0}\tsample{0}\n".format(i + 1) for i in range(n)),
encoding="utf-8",
)
base_x = np.column_stack((np.ones(n), np.linspace(-1.0, 1.0, n)))
extra = rng.normal(size=(2, n)).astype("<f4")
phenotypes = np.vstack(
(
0.3 + 0.5 * base_x[:, 1] + rng.normal(scale=0.2, size=n),
-0.2 + 0.8 * extra[0] + rng.normal(scale=0.2, size=n),
0.1 + 0.4 * extra[0] - 0.3 * extra[1] + rng.normal(scale=0.2, size=n),
)
)
np.asarray(base_x, dtype="<f8").tofile(output / "base_x.f64.bin")
np.asarray(phenotypes, dtype="<f8").tofile(output / "phenotypes.f64.bin")
extra.tofile(output / "extra_covariates.f32.bin")
(output / "tasks.tsv").write_text(
"task_index\ttask_id\tphenotype_row\tn_extra_covariates\n"
"0\ttrait_0\t0\t0\n"
"1\ttrait_1\t1\t1\n"
"2\ttrait_2\t2\t2\n",
encoding="utf-8",
)
np.asarray([0, 0, 1, 3], dtype="<i8").tofile(output / "extra_offsets.i64.bin")
np.asarray([0, 0, 1], dtype="<i4").tofile(output / "extra_indices.i32.bin")
manifest = output / "manifest.json"
spectra_reml.main(
[
"make-manifest",
"--manifest", str(manifest),
"--grm-bin", str(output / "example.grm.bin"),
"--grm-id", str(output / "example.grm.id"),
"--base-x", str(output / "base_x.f64.bin"),
"--phenotypes", str(output / "phenotypes.f64.bin"),
"--extra-covariates", str(output / "extra_covariates.f32.bin"),
"--tasks", str(output / "tasks.tsv"),
"--extra-offsets", str(output / "extra_offsets.i64.bin"),
"--extra-indices", str(output / "extra_indices.i32.bin"),
"--output-dir", str(output / "blocks"),
"--n-samples", str(n),
"--n-base-covariates", str(base_x.shape[1]),
"--n-phenotype-rows", str(phenotypes.shape[0]),
"--n-extra-covariate-rows", str(extra.shape[0]),
]
)
print("\nNext:")
print("python python/spectra_reml.py run --manifest {} --engine /path/to/spectra_reml --threads 4".format(manifest))
if __name__ == "__main__":
main()