Add arbitrary fixed-effect contrasts

tests/python/test_spectra_reml_cli.py

@@ -21,7 +21,7 @@ SPEC.loader.exec_module(cli)
 
 
 class SpectraRemlCliTests(unittest.TestCase):
-    def fixture(self, root: Path) -> tuple[Path, dict]:
+    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")
@@ -36,9 +36,18 @@ class SpectraRemlCliTests(unittest.TestCase):
         )
         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"
-        args = cli.build_parser().parse_args(
-            [
+        command = [
                 "make-manifest",
                 "--manifest", str(manifest_path),
                 "--grm-bin", str(root / "grm.bin"),
@@ -55,19 +64,31 @@ class SpectraRemlCliTests(unittest.TestCase):
                 "--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) -> None:
+    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\t0\tnan\tnan\tnan\t\t\n"
+            "\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"
-            "\tsatterthwaite\tboundary_conditional\t2\t1\t8\t2\t0.2\t\t\n",
+            "\tsatterthwaite\tboundary_conditional\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")
@@ -84,9 +105,22 @@ class SpectraRemlCliTests(unittest.TestCase):
         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\n".format(cli.BLOCK_FORMAT),
+            "fixed_test_elements\t5\ncontrast_test_elements\t{}\n".format(
+                cli.BLOCK_FORMAT, contrast_elements
+            ),
             encoding="utf-8",
         )
 
@@ -107,9 +141,83 @@ class SpectraRemlCliTests(unittest.TestCase):
                 "--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)