290 lines
14 KiB
Python
290 lines
14 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, with_contrasts: bool = False) -> 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")
|
|
if with_contrasts:
|
|
np.asarray([[0, 1], [1, -1]], dtype="<f8").tofile(
|
|
root / "contrasts.bin"
|
|
)
|
|
(root / "contrasts.tsv").write_text(
|
|
"contrast_index\tcontrast_id\trhs\n"
|
|
"0\tslope\t0\n"
|
|
"1\tintercept_minus_slope\t0\n",
|
|
encoding="utf-8",
|
|
)
|
|
manifest_path = root / "manifest.json"
|
|
command = [
|
|
"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",
|
|
]
|
|
if with_contrasts:
|
|
command.extend(
|
|
[
|
|
"--contrast-matrix", str(root / "contrasts.bin"),
|
|
"--contrast-metadata", str(root / "contrasts.tsv"),
|
|
"--n-contrasts", "2",
|
|
]
|
|
)
|
|
args = cli.build_parser().parse_args(command)
|
|
cli.make_manifest(args)
|
|
return manifest_path, cli.validate_manifest(manifest_path)
|
|
|
|
@staticmethod
|
|
def write_blocks(root: Path, contrast_count: int = 0) -> None:
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
first_contrast_offset = 0 if contrast_count else -1
|
|
second_contrast_offset = contrast_count if contrast_count else -1
|
|
(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"
|
|
"\tsatterthwaite\tok\t0\t{}\t0\tnan\tnan\tnan\t\t\n"
|
|
"1\ttrait_b\tconverged_boundary\t3\t1\t2\t3\t0\t1\t0\t-2\t3\t6\t1e-9"
|
|
"\tordinary-least-squares\tboundary_ols\t2\t{}\t1\t8\t2\t0.2\t\t\n".format(
|
|
first_contrast_offset, second_contrast_offset
|
|
),
|
|
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")
|
|
np.asarray([0.1, 0.2, 0.3, 0.4, 0.5], dtype="<f8").tofile(
|
|
root / "block_000000.fixed_se.f64.bin"
|
|
)
|
|
np.asarray([1, 2, 3, 4, 5], dtype="<f8").tofile(
|
|
root / "block_000000.fixed_stat.f64.bin"
|
|
)
|
|
np.asarray([10, 10, 8, 8, 8], dtype="<f8").tofile(
|
|
root / "block_000000.fixed_ddf.f64.bin"
|
|
)
|
|
np.asarray([0.5, 0.2, 0.1, 0.05, 0.01], dtype="<f8").tofile(
|
|
root / "block_000000.fixed_p.f64.bin"
|
|
)
|
|
contrast_elements = 2 * contrast_count
|
|
contrast_arrays = {
|
|
"contrast_estimate": np.arange(1, contrast_elements + 1, dtype="<f8"),
|
|
"contrast_se": np.full(contrast_elements, 0.25, dtype="<f8"),
|
|
"contrast_stat": np.arange(2, contrast_elements + 2, dtype="<f8"),
|
|
"contrast_numdf": np.ones(contrast_elements, dtype="<f8"),
|
|
"contrast_ddf": np.full(contrast_elements, 12.0, dtype="<f8"),
|
|
"contrast_p": np.full(contrast_elements, 0.05, dtype="<f8"),
|
|
}
|
|
for suffix, values in contrast_arrays.items():
|
|
values.tofile(root / "block_000000.{}.f64.bin".format(suffix))
|
|
(root / "block_000000.complete").write_text(
|
|
"format\t{}\nblock\t0\ntasks\t2\nbeta_elements\t5\ncov_elements\t9\n"
|
|
"fixed_test_elements\t5\ncontrast_test_elements\t{}\n".format(
|
|
cli.BLOCK_FORMAT, contrast_elements
|
|
),
|
|
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",
|
|
"--line-search-expansion",
|
|
"--zoom-safeguard",
|
|
"--n-contrasts",
|
|
):
|
|
self.assertEqual(command.count(option), 1)
|
|
|
|
def test_contrast_manifest_and_engine_command(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
root = Path(temporary)
|
|
_, manifest = self.fixture(root, with_contrasts=True)
|
|
self.assertEqual(manifest["dimensions"]["contrast_count"], 2)
|
|
args = cli.build_parser().parse_args(
|
|
[
|
|
"run", "--manifest", str(root / "manifest.json"),
|
|
"--engine", "spectra_reml", "--fixed-effect-test",
|
|
"kenward-roger",
|
|
]
|
|
)
|
|
command = cli.build_engine_command(Path("spectra_reml"), manifest, args)
|
|
self.assertEqual(command.count("--contrast-matrix"), 1)
|
|
self.assertEqual(command.count("--contrast-metadata"), 1)
|
|
self.assertEqual(command[command.index("--n-contrasts") + 1], "2")
|
|
self.assertEqual(command[command.index("--fixed-effect-test") + 1], "kenward-roger")
|
|
|
|
def test_new_contract_rejects_missing_contrast_fields(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
root = Path(temporary)
|
|
manifest_path, _ = self.fixture(root)
|
|
manifest = cli.read_json(manifest_path)
|
|
del manifest["dimensions"]["contrast_count"]
|
|
del manifest["paths"]["contrast_matrix_f64"]
|
|
del manifest["paths"]["contrast_metadata_tsv"]
|
|
cli.atomic_json(manifest_path, manifest)
|
|
with self.assertRaises(RuntimeError):
|
|
cli.validate_manifest(manifest_path)
|
|
|
|
def test_both_fixed_effect_methods_are_forwarded(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
root = Path(temporary)
|
|
_, manifest = self.fixture(root, with_contrasts=True)
|
|
for method in ("satterthwaite", "kenward-roger"):
|
|
args = cli.build_parser().parse_args(
|
|
[
|
|
"run", "--manifest", str(root / "manifest.json"),
|
|
"--engine", "spectra_reml", "--fixed-effect-test", method,
|
|
]
|
|
)
|
|
command = cli.build_engine_command(Path("spectra_reml"), manifest, args)
|
|
self.assertEqual(
|
|
command[command.index("--fixed-effect-test") + 1], method
|
|
)
|
|
|
|
def test_contrast_finalize_contract(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
root = Path(temporary)
|
|
manifest_path, manifest = self.fixture(root, with_contrasts=True)
|
|
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", "--fixed-effect-test",
|
|
"satterthwaite",
|
|
]
|
|
)
|
|
with mock.patch.object(
|
|
cli.subprocess, "run", return_value=SimpleNamespace(returncode=0)
|
|
):
|
|
cli.run_engine(manifest, args)
|
|
output = Path(manifest["paths"]["output_directory"])
|
|
self.write_blocks(output, contrast_count=2)
|
|
result = root / "contrast_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(json.loads(rows[0]["contrast_estimate_json"]), [1.0, 2.0])
|
|
self.assertEqual(json.loads(rows[1]["contrast_estimate_json"]), [3.0, 4.0])
|
|
self.assertEqual(json.loads(rows[0]["contrast_numerator_df_json"]), [1.0, 1.0])
|
|
|
|
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)
|
|
self.assertEqual(len(json.loads(rows[1]["fixed_effect_p_value_json"])), 3)
|
|
|
|
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()
|