Add reusable SpectraREML batch AI-REML engine
This commit is contained in:
163
tests/python/test_spectra_reml_cli.py
Normal file
163
tests/python/test_spectra_reml_cli.py
Normal file
@@ -0,0 +1,163 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import gzip
|
||||
import importlib.util
|
||||
import json
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
SCRIPT = Path(__file__).resolve().parents[2] / "python" / "spectra_reml.py"
|
||||
SPEC = importlib.util.spec_from_file_location("spectra_reml", SCRIPT)
|
||||
assert SPEC is not None and SPEC.loader is not None
|
||||
cli = importlib.util.module_from_spec(SPEC)
|
||||
SPEC.loader.exec_module(cli)
|
||||
|
||||
|
||||
class SpectraRemlCliTests(unittest.TestCase):
|
||||
def fixture(self, root: Path) -> tuple[Path, dict]:
|
||||
n, p0 = 3, 2
|
||||
np.zeros(n * (n + 1) // 2, dtype="<f4").tofile(root / "grm.bin")
|
||||
(root / "grm.id").write_text("F1 I1\nF2 I2\nF3 I3\n", encoding="utf-8")
|
||||
np.asarray([[1, -1], [1, 0], [1, 1]], dtype="<f8").tofile(root / "base.bin")
|
||||
np.zeros((2, n), dtype="<f8").tofile(root / "phenotypes.bin")
|
||||
np.asarray([[0, 1, 2]], dtype="<f4").tofile(root / "extra.bin")
|
||||
(root / "tasks.tsv").write_text(
|
||||
"task_index\ttask_id\tphenotype_row\tn_extra_covariates\n"
|
||||
"0\ttrait_a\t0\t0\n"
|
||||
"1\ttrait_b\t1\t1\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
np.asarray([0, 0, 1], dtype="<i8").tofile(root / "offsets.bin")
|
||||
np.asarray([0], dtype="<i4").tofile(root / "indices.bin")
|
||||
manifest_path = root / "manifest.json"
|
||||
args = cli.build_parser().parse_args(
|
||||
[
|
||||
"make-manifest",
|
||||
"--manifest", str(manifest_path),
|
||||
"--grm-bin", str(root / "grm.bin"),
|
||||
"--grm-id", str(root / "grm.id"),
|
||||
"--base-x", str(root / "base.bin"),
|
||||
"--phenotypes", str(root / "phenotypes.bin"),
|
||||
"--extra-covariates", str(root / "extra.bin"),
|
||||
"--tasks", str(root / "tasks.tsv"),
|
||||
"--extra-offsets", str(root / "offsets.bin"),
|
||||
"--extra-indices", str(root / "indices.bin"),
|
||||
"--output-dir", str(root / "blocks"),
|
||||
"--n-samples", "3",
|
||||
"--n-base-covariates", "2",
|
||||
"--n-phenotype-rows", "2",
|
||||
"--n-extra-covariate-rows", "1",
|
||||
]
|
||||
)
|
||||
cli.make_manifest(args)
|
||||
return manifest_path, cli.validate_manifest(manifest_path)
|
||||
|
||||
@staticmethod
|
||||
def write_blocks(root: Path) -> None:
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
(root / "block_000000.summary.tsv").write_text(
|
||||
"\t".join(cli.SUMMARY_HEADER) + "\n"
|
||||
"0\ttrait_a\tconverged\t2\t0\t0\t0\t1\t1\t0.5\t-1\t4\t8\t1e-8\t\n"
|
||||
"1\ttrait_b\tconverged_boundary\t3\t1\t2\t3\t0\t1\t0\t-2\t3\t6\t1e-9\t\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
np.asarray([1, 2, 3, 4, 5], dtype="<f8").tofile(root / "block_000000.beta.f64.bin")
|
||||
np.arange(1, 10, dtype="<f8").tofile(root / "block_000000.cov.f64.bin")
|
||||
(root / "block_000000.complete").write_text(
|
||||
"format\t{}\nblock\t0\ntasks\t2\nbeta_elements\t5\ncov_elements\t9\n".format(cli.BLOCK_FORMAT),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
def test_manifest_and_generic_command(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as temporary:
|
||||
root = Path(temporary)
|
||||
_, manifest = self.fixture(root)
|
||||
args = cli.build_parser().parse_args(
|
||||
["run", "--manifest", str(root / "manifest.json"), "--engine", "spectra_reml"]
|
||||
)
|
||||
command = cli.build_engine_command(Path("spectra_reml"), manifest, args)
|
||||
for option in (
|
||||
"--phenotypes",
|
||||
"--extra-covariates",
|
||||
"--extra-offsets",
|
||||
"--extra-indices",
|
||||
"--n-phenotype-rows",
|
||||
"--n-extra-covariate-rows",
|
||||
):
|
||||
self.assertEqual(command.count(option), 1)
|
||||
|
||||
def test_signature_recovery_and_finalize(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as temporary:
|
||||
root = Path(temporary)
|
||||
manifest_path, manifest = self.fixture(root)
|
||||
engine = root / "spectra_reml"
|
||||
engine.write_bytes(b"synthetic engine")
|
||||
args = cli.build_parser().parse_args(
|
||||
["run", "--manifest", str(manifest_path), "--engine", str(engine), "--resume"]
|
||||
)
|
||||
with mock.patch.object(cli.subprocess, "run", return_value=SimpleNamespace(returncode=0)):
|
||||
cli.run_engine(manifest, args)
|
||||
output = Path(manifest["paths"]["output_directory"])
|
||||
signature = cli.read_signature(output / "run.signature.json")
|
||||
self.assertEqual(signature["format"], cli.RUN_SIGNATURE_FORMAT)
|
||||
|
||||
self.write_blocks(output)
|
||||
changed = cli.build_parser().parse_args(
|
||||
[
|
||||
"run", "--manifest", str(manifest_path), "--engine", str(engine),
|
||||
"--resume", "--max-iterations", "101",
|
||||
]
|
||||
)
|
||||
with self.assertRaisesRegex(RuntimeError, "signature differs"):
|
||||
cli.run_engine(manifest, changed)
|
||||
|
||||
result = root / "results.tsv.gz"
|
||||
cli.finalize(manifest, result)
|
||||
with gzip.open(result, "rt", encoding="utf-8", newline="") as handle:
|
||||
rows = list(csv.DictReader(handle, delimiter="\t"))
|
||||
self.assertEqual([row["task_id"] for row in rows], ["trait_a", "trait_b"])
|
||||
self.assertEqual(json.loads(rows[0]["beta_json"]), [1.0, 2.0])
|
||||
self.assertEqual(len(json.loads(rows[1]["covariance_packed_lower_json"])), 6)
|
||||
|
||||
altered = dict(manifest)
|
||||
altered["created_utc"] = "changed"
|
||||
with self.assertRaisesRegex(RuntimeError, "does not belong"):
|
||||
cli.finalize(altered, root / "unsafe.tsv")
|
||||
|
||||
def test_force_invalidates_before_cleanup_and_dry_run_is_read_only(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as temporary:
|
||||
root = Path(temporary)
|
||||
manifest_path, manifest = self.fixture(root)
|
||||
engine = root / "spectra_reml"
|
||||
engine.write_bytes(b"engine")
|
||||
output = Path(manifest["paths"]["output_directory"])
|
||||
self.write_blocks(output)
|
||||
resume_args = cli.build_parser().parse_args(
|
||||
["run", "--manifest", str(manifest_path), "--engine", str(engine), "--resume"]
|
||||
)
|
||||
cli.atomic_json(output / "run.signature.json", cli.make_signature(manifest, engine, resume_args))
|
||||
before = sorted(path.name for path in output.iterdir())
|
||||
dry = cli.build_parser().parse_args(
|
||||
["run", "--manifest", str(manifest_path), "--engine", "missing", "--force", "--dry-run"]
|
||||
)
|
||||
cli.run_engine(manifest, dry)
|
||||
self.assertEqual(before, sorted(path.name for path in output.iterdir()))
|
||||
|
||||
force = cli.build_parser().parse_args(
|
||||
["run", "--manifest", str(manifest_path), "--engine", str(engine), "--force"]
|
||||
)
|
||||
with mock.patch.object(cli, "block_files", side_effect=RuntimeError("cleanup interrupted")):
|
||||
with self.assertRaisesRegex(RuntimeError, "cleanup interrupted"):
|
||||
cli.run_engine(manifest, force)
|
||||
self.assertFalse((output / "run.signature.json").exists())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user