164 lines
7.3 KiB
Python
164 lines
7.3 KiB
Python
|
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()
|