JiaruiLi/SpectraREML
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add arbitrary fixed-effect contrasts

Browse Source
This commit is contained in:
Jiarui Li
2026-08-13 13:39:33 +08:00
parent 8557daecac
commit 48ab77f9ab
8 changed files with 515 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

147
python/spectra_reml.py
View File

@@ -31,10 +31,10 @@ except ImportError as exc: # pragma: no cover
raise SystemExit("NumPy is required: {}".format(exc))
MANIFEST_FORMAT = "spectra-reml-manifest-v1"
BLOCK_FORMAT = "spectra-reml-block-v2"
RUN_SIGNATURE_FORMAT = "spectra-reml-run-signature-v1"
FINALIZE_FORMAT = "spectra-reml-finalize-v1"
MANIFEST_FORMAT = "spectra-reml-manifest-v2"
BLOCK_FORMAT = "spectra-reml-block-v3"
RUN_SIGNATURE_FORMAT = "spectra-reml-run-signature-v2"
FINALIZE_FORMAT = "spectra-reml-finalize-v2"
TASK_HEADER = (
"task_index",
"task_id",
@@ -59,6 +59,7 @@ SUMMARY_HEADER = (
"fixed_test_method",
"fixed_test_status",
"fixed_test_offset",
"contrast_test_offset",
"extra_joint_num_df",
"extra_joint_den_df",
"extra_joint_f",
@@ -75,6 +76,12 @@ BLOCK_PATTERNS = (
"block_*.fixed_stat.f64.bin",
"block_*.fixed_ddf.f64.bin",
"block_*.fixed_p.f64.bin",
"block_*.contrast_estimate.f64.bin",
"block_*.contrast_se.f64.bin",
"block_*.contrast_stat.f64.bin",
"block_*.contrast_numdf.f64.bin",
"block_*.contrast_ddf.f64.bin",
"block_*.contrast_p.f64.bin",
"block_*.complete",
)
@@ -180,6 +187,8 @@ def manifest_paths(manifest: Mapping[str, Any]) -> dict[str, Path | None]:
"extra_offsets_i64",
"extra_indices_i32",
"output_directory",
"contrast_matrix_f64",
"contrast_metadata_tsv",
)
missing = [name for name in required if name not in raw]
if missing:
@@ -199,6 +208,30 @@ def read_tasks(path: Path) -> list[dict[str, str]]:
return list(reader)
def read_contrast_metadata(path: Path) -> list[dict[str, str]]:
with path.open("r", encoding="utf-8-sig", newline="") as handle:
reader = csv.DictReader(handle, delimiter="\t")
expected = ("contrast_index", "contrast_id", "rhs")
if tuple(reader.fieldnames or ()) != expected:
fail("Contrast metadata header must be exactly: {}".format("\t".join(expected)))
rows = list(reader)
seen: set[str] = set()
for expected_index, row in enumerate(rows):
try:
index = int(row["contrast_index"])
rhs = float(row["rhs"])
except ValueError as exc:
fail("Invalid contrast metadata at row {}: {}".format(expected_index, exc))
if index != expected_index:
fail("Contrast indices must be consecutive and zero based.")
if not row["contrast_id"] or row["contrast_id"] in seen:
fail("Contrast IDs must be nonempty and unique.")
if not np.isfinite(rhs):
fail("Contrast rhs must be finite.")
seen.add(row["contrast_id"])
return rows
def validate_manifest(path: Path) -> dict[str, Any]:
path = resolved(path)
manifest = read_json(path)
@@ -214,6 +247,7 @@ def validate_manifest(path: Path) -> dict[str, Any]:
"extra_covariate_row_count",
"task_count",
"extra_index_count",
"contrast_count",
)
parsed: dict[str, int] = {}
for name in names:
@@ -240,6 +274,28 @@ def validate_manifest(path: Path) -> dict[str, Any]:
with paths["grm_id"].open("r", encoding="utf-8-sig") as handle:
if sum(bool(line.strip()) for line in handle) != n:
fail("GRM ID row count differs from sample_count.")
contrast_paths = (paths["contrast_matrix_f64"], paths["contrast_metadata_tsv"])
if parsed["contrast_count"]:
if any(path is None for path in contrast_paths):
fail("Both contrast_matrix_f64 and contrast_metadata_tsv are required.")
require_size(
paths["contrast_matrix_f64"],
parsed["contrast_count"] * parsed["base_covariate_count"] * 8,
"contrast matrix",
)
contrast_rows = read_contrast_metadata(paths["contrast_metadata_tsv"])
if len(contrast_rows) != parsed["contrast_count"]:
fail("Contrast metadata row count differs from contrast_count.")
contrast_matrix = np.fromfile(paths["contrast_matrix_f64"], dtype="<f8")
contrast_matrix = contrast_matrix.reshape(
parsed["contrast_count"], parsed["base_covariate_count"]
)
if not np.all(np.isfinite(contrast_matrix)):
fail("Contrast matrix contains a non-finite value.")
if np.any(np.all(contrast_matrix == 0.0, axis=1)):
fail("Every contrast row must contain at least one nonzero weight.")
elif any(path is not None for path in contrast_paths):
fail("Contrast paths must be omitted when contrast_count is zero.")
sources = manifest.get("source_files")
if not isinstance(sources, dict) or not sources:
@@ -263,6 +319,9 @@ def validate_manifest(path: Path) -> dict[str, Any]:
"extra_offsets_i64": paths["extra_offsets_i64"],
"extra_indices_i32": paths["extra_indices_i32"],
}
if parsed["contrast_count"]:
controls["contrast_matrix_f64"] = paths["contrast_matrix_f64"]
controls["contrast_metadata_tsv"] = paths["contrast_metadata_tsv"]
hashes = manifest.get("control_sha256")
if not isinstance(hashes, dict) or set(hashes) != set(controls):
fail("control_sha256 does not cover all control files.")
@@ -314,6 +373,8 @@ def make_manifest(args: argparse.Namespace) -> Path:
"tasks_tsv": resolved(args.tasks),
"extra_offsets_i64": resolved(args.extra_offsets),
"extra_indices_i32": resolved(args.extra_indices),
"contrast_matrix_f64": resolved(args.contrast_matrix) if args.contrast_matrix else None,
"contrast_metadata_tsv": resolved(args.contrast_metadata) if args.contrast_metadata else None,
"output_directory": resolved(args.output_dir),
}
tasks = read_tasks(paths["tasks_tsv"]) # type: ignore[arg-type]
@@ -323,7 +384,13 @@ def make_manifest(args: argparse.Namespace) -> Path:
for name in ("grm_id", "extra_covariate_f32"):
if paths[name] is not None:
sources[name] = file_identity(paths[name]) # type: ignore[arg-type]
controls = ("base_x_f64", "tasks_tsv", "extra_offsets_i64", "extra_indices_i32")
controls = ["base_x_f64", "tasks_tsv", "extra_offsets_i64", "extra_indices_i32"]
if args.n_contrasts:
if paths["contrast_matrix_f64"] is None or paths["contrast_metadata_tsv"] is None:
fail("--contrast-matrix and --contrast-metadata are required when --n-contrasts is positive.")
controls.extend(("contrast_matrix_f64", "contrast_metadata_tsv"))
elif paths["contrast_matrix_f64"] is not None or paths["contrast_metadata_tsv"] is not None:
fail("Contrast paths require a positive --n-contrasts.")
manifest = {
"format": MANIFEST_FORMAT,
"created_utc": utc_now(),
@@ -334,6 +401,7 @@ def make_manifest(args: argparse.Namespace) -> Path:
"extra_covariate_row_count": args.n_extra_covariate_rows,
"task_count": len(tasks),
"extra_index_count": extra_index_count,
"contrast_count": args.n_contrasts,
},
"paths": {name: str(value) if value is not None else None for name, value in paths.items()},
"source_files": sources,
@@ -371,6 +439,7 @@ def build_engine_command(engine: Path, manifest: Mapping[str, Any], args: argpar
"--n-base-covariates", str(dims["base_covariate_count"]),
"--n-phenotype-rows", str(dims["phenotype_row_count"]),
"--n-extra-covariate-rows", str(dims["extra_covariate_row_count"]),
"--n-contrasts", str(dims["contrast_count"]),
"--block-size", str(args.block_size),
"--threads", str(args.threads),
"--fixed-effect-test", args.fixed_effect_test,
@@ -379,6 +448,9 @@ def build_engine_command(engine: Path, manifest: Mapping[str, Any], args: argpar
command.extend(("--grm-id", str(paths["grm_id"])))
if int(dims["extra_index_count"]):
command.extend(("--extra-covariates", str(paths["extra_covariate_f32"])))
if int(dims["contrast_count"]):
command.extend(("--contrast-matrix", str(paths["contrast_matrix_f64"])))
command.extend(("--contrast-metadata", str(paths["contrast_metadata_tsv"])))
option_map = (
("max_iterations", "--max-iterations"),
("line_search_max_evals", "--line-search-max-evals"),
@@ -571,16 +643,25 @@ def read_completed(manifest: Mapping[str, Any]) -> list[dict[str, Any]]:
fixed_stat_path = output / (stem + ".fixed_stat.f64.bin")
fixed_ddf_path = output / (stem + ".fixed_ddf.f64.bin")
fixed_p_path = output / (stem + ".fixed_p.f64.bin")
contrast_estimate_path = output / (stem + ".contrast_estimate.f64.bin")
contrast_se_path = output / (stem + ".contrast_se.f64.bin")
contrast_stat_path = output / (stem + ".contrast_stat.f64.bin")
contrast_numdf_path = output / (stem + ".contrast_numdf.f64.bin")
contrast_ddf_path = output / (stem + ".contrast_ddf.f64.bin")
contrast_p_path = output / (stem + ".contrast_p.f64.bin")
marker_values = read_marker(marker)
block_format = marker_values.get("format")
try:
declared_tasks = int(marker_values["tasks"])
beta_elements = int(marker_values["beta_elements"])
cov_elements = int(marker_values["cov_elements"])
fixed_test_elements = int(marker_values["fixed_test_elements"])
contrast_test_elements = int(marker_values["contrast_test_elements"])
except (KeyError, ValueError) as exc:
fail("Malformed completion counts in {}: {}".format(marker, exc))
if marker_values.get("format") != BLOCK_FORMAT or min(
declared_tasks, beta_elements, cov_elements, fixed_test_elements
if block_format != BLOCK_FORMAT or min(
declared_tasks, beta_elements, cov_elements, fixed_test_elements,
contrast_test_elements,
) < 0:
fail("Invalid completion marker: {}".format(marker))
require_size(beta_path, beta_elements * 8, "block beta")
@@ -589,6 +670,15 @@ def read_completed(manifest: Mapping[str, Any]) -> list[dict[str, Any]]:
require_size(fixed_stat_path, fixed_test_elements * 8, "fixed-effect statistic")
require_size(fixed_ddf_path, fixed_test_elements * 8, "fixed-effect denominator df")
require_size(fixed_p_path, fixed_test_elements * 8, "fixed-effect p-value")
for path, label in (
(contrast_estimate_path, "contrast estimate"),
(contrast_se_path, "contrast SE"),
(contrast_stat_path, "contrast statistic"),
(contrast_numdf_path, "contrast numerator df"),
(contrast_ddf_path, "contrast denominator df"),
(contrast_p_path, "contrast p-value"),
):
require_size(path, contrast_test_elements * 8, label)
with summary_path.open("r", encoding="utf-8-sig", newline="") as handle:
reader = csv.DictReader(handle, delimiter="\t")
if tuple(reader.fieldnames or ()) != SUMMARY_HEADER:
@@ -602,9 +692,17 @@ def read_completed(manifest: Mapping[str, Any]) -> list[dict[str, Any]]:
fixed_stat = np.fromfile(fixed_stat_path, dtype="<f8")
fixed_ddf = np.fromfile(fixed_ddf_path, dtype="<f8")
fixed_p = np.fromfile(fixed_p_path, dtype="<f8")
contrast_estimate = np.fromfile(contrast_estimate_path, dtype="<f8")
contrast_se = np.fromfile(contrast_se_path, dtype="<f8")
contrast_stat = np.fromfile(contrast_stat_path, dtype="<f8")
contrast_numdf = np.fromfile(contrast_numdf_path, dtype="<f8")
contrast_ddf = np.fromfile(contrast_ddf_path, dtype="<f8")
contrast_p = np.fromfile(contrast_p_path, dtype="<f8")
next_beta_offset = 0
next_cov_offset = 0
next_fixed_test_offset = 0
next_contrast_test_offset = 0
q = int(manifest["dimensions"]["contrast_count"])
for row in rows:
index = int(row["task_index"])
if index in results or index < 0 or index >= task_count:
@@ -612,6 +710,7 @@ def read_completed(manifest: Mapping[str, Any]) -> list[dict[str, Any]]:
p = int(row["n_fixed"])
beta_offset, cov_offset = int(row["beta_offset"]), int(row["cov_offset"])
fixed_test_offset = int(row["fixed_test_offset"])
contrast_test_offset = int(row["contrast_test_offset"])
if beta_offset < 0 or cov_offset < 0:
if beta_offset != cov_offset:
fail("Only one output offset is negative for task {}".format(index))
@@ -641,6 +740,25 @@ def read_completed(manifest: Mapping[str, Any]) -> list[dict[str, Any]]:
fixed_ddf_values = fixed_ddf[fixed_test_offset:fixed_test_offset + p].tolist()
fixed_p_values = fixed_p[fixed_test_offset:fixed_test_offset + p].tolist()
next_fixed_test_offset += p
if contrast_test_offset < 0:
contrast_estimate_values: list[float] = []
contrast_se_values: list[float] = []
contrast_stat_values: list[float] = []
contrast_numdf_values: list[float] = []
contrast_ddf_values: list[float] = []
contrast_p_values: list[float] = []
else:
if (contrast_test_offset != next_contrast_test_offset or
contrast_test_offset + q > contrast_estimate.size):
fail("Invalid contrast test offset for task {}".format(index))
section = slice(contrast_test_offset, contrast_test_offset + q)
contrast_estimate_values = contrast_estimate[section].tolist()
contrast_se_values = contrast_se[section].tolist()
contrast_stat_values = contrast_stat[section].tolist()
contrast_numdf_values = contrast_numdf[section].tolist()
contrast_ddf_values = contrast_ddf[section].tolist()
contrast_p_values = contrast_p[section].tolist()
next_contrast_test_offset += q
result: dict[str, Any] = dict(row)
result["beta_json"] = json.dumps(beta_values, separators=(",", ":"))
result["covariance_packed_lower_json"] = json.dumps(cov_values, separators=(",", ":"))
@@ -648,9 +766,16 @@ def read_completed(manifest: Mapping[str, Any]) -> list[dict[str, Any]]:
result["fixed_effect_statistic_json"] = json.dumps(fixed_stat_values, separators=(",", ":"))
result["fixed_effect_denominator_df_json"] = json.dumps(fixed_ddf_values, separators=(",", ":"))
result["fixed_effect_p_value_json"] = json.dumps(fixed_p_values, separators=(",", ":"))
result["contrast_estimate_json"] = json.dumps(contrast_estimate_values, separators=(",", ":"))
result["contrast_standard_error_json"] = json.dumps(contrast_se_values, separators=(",", ":"))
result["contrast_statistic_json"] = json.dumps(contrast_stat_values, separators=(",", ":"))
result["contrast_numerator_df_json"] = json.dumps(contrast_numdf_values, separators=(",", ":"))
result["contrast_denominator_df_json"] = json.dumps(contrast_ddf_values, separators=(",", ":"))
result["contrast_p_value_json"] = json.dumps(contrast_p_values, separators=(",", ":"))
results[index] = result
if (next_beta_offset != beta_elements or next_cov_offset != cov_elements or
next_fixed_test_offset != fixed_test_elements):
next_fixed_test_offset != fixed_test_elements or
next_contrast_test_offset != contrast_test_elements):
fail("Block binary arrays contain unused elements: {}".format(stem))
missing = sorted(set(range(task_count)).difference(results))
if missing:
@@ -669,6 +794,9 @@ def finalize(manifest: Mapping[str, Any], output: Path) -> Path:
"beta_json", "covariance_packed_lower_json", "fixed_effect_se_json",
"fixed_effect_statistic_json", "fixed_effect_denominator_df_json",
"fixed_effect_p_value_json",
"contrast_estimate_json", "contrast_standard_error_json",
"contrast_statistic_json", "contrast_numerator_df_json",
"contrast_denominator_df_json", "contrast_p_value_json",
]
try:
with opener(temporary, "wt", encoding="utf-8", newline="") as handle:
@@ -771,11 +899,14 @@ def build_parser() -> argparse.ArgumentParser:
make.add_argument("--tasks", type=Path, required=True)
make.add_argument("--extra-offsets", type=Path, required=True)
make.add_argument("--extra-indices", type=Path, required=True)
make.add_argument("--contrast-matrix", type=Path)
make.add_argument("--contrast-metadata", type=Path)
make.add_argument("--output-dir", type=Path, required=True)
make.add_argument("--n-samples", type=positive_integer, required=True)
make.add_argument("--n-base-covariates", type=positive_integer, required=True)
make.add_argument("--n-phenotype-rows", type=positive_integer, required=True)
make.add_argument("--n-extra-covariate-rows", type=int, default=0)
make.add_argument("--n-contrasts", type=int, default=0)
for action in ("validate", "status"):
command = subparsers.add_parser(action)
command.add_argument("--manifest", type=Path, required=True)