Add Satterthwaite and Kenward-Roger fixed-effect tests

This commit is contained in:
2026-08-13 07:02:10 +08:00
parent 76dd8b1379
commit 8557daecac
14 changed files with 1330 additions and 20 deletions

View File

@@ -32,7 +32,7 @@ except ImportError as exc: # pragma: no cover
MANIFEST_FORMAT = "spectra-reml-manifest-v1"
BLOCK_FORMAT = "spectra-reml-block-v1"
BLOCK_FORMAT = "spectra-reml-block-v2"
RUN_SIGNATURE_FORMAT = "spectra-reml-run-signature-v1"
FINALIZE_FORMAT = "spectra-reml-finalize-v1"
TASK_HEADER = (
@@ -56,6 +56,14 @@ SUMMARY_HEADER = (
"iterations",
"line_search_steps",
"grad_inf",
"fixed_test_method",
"fixed_test_status",
"fixed_test_offset",
"extra_joint_num_df",
"extra_joint_den_df",
"extra_joint_f",
"extra_joint_p",
"fixed_test_error",
"error",
)
SUCCESS_STATUSES = frozenset(("converged", "converged_boundary"))
@@ -63,6 +71,10 @@ BLOCK_PATTERNS = (
"block_*.summary.tsv",
"block_*.beta.f64.bin",
"block_*.cov.f64.bin",
"block_*.fixed_se.f64.bin",
"block_*.fixed_stat.f64.bin",
"block_*.fixed_ddf.f64.bin",
"block_*.fixed_p.f64.bin",
"block_*.complete",
)
@@ -361,6 +373,7 @@ def build_engine_command(engine: Path, manifest: Mapping[str, Any], args: argpar
"--n-extra-covariate-rows", str(dims["extra_covariate_row_count"]),
"--block-size", str(args.block_size),
"--threads", str(args.threads),
"--fixed-effect-test", args.fixed_effect_test,
]
if paths["grm_id"] is not None:
command.extend(("--grm-id", str(paths["grm_id"])))
@@ -554,19 +567,28 @@ def read_completed(manifest: Mapping[str, Any]) -> list[dict[str, Any]]:
summary_path = output / (stem + ".summary.tsv")
beta_path = output / (stem + ".beta.f64.bin")
cov_path = output / (stem + ".cov.f64.bin")
fixed_se_path = output / (stem + ".fixed_se.f64.bin")
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")
marker_values = read_marker(marker)
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"])
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
declared_tasks, beta_elements, cov_elements, fixed_test_elements
) < 0:
fail("Invalid completion marker: {}".format(marker))
require_size(beta_path, beta_elements * 8, "block beta")
require_size(cov_path, cov_elements * 8, "block covariance")
require_size(fixed_se_path, fixed_test_elements * 8, "fixed-effect SE")
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")
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:
@@ -576,14 +598,20 @@ def read_completed(manifest: Mapping[str, Any]) -> list[dict[str, Any]]:
fail("Summary row count disagrees with marker: {}".format(marker))
beta = np.fromfile(beta_path, dtype="<f8")
cov = np.fromfile(cov_path, dtype="<f8")
fixed_se = np.fromfile(fixed_se_path, dtype="<f8")
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")
next_beta_offset = 0
next_cov_offset = 0
next_fixed_test_offset = 0
for row in rows:
index = int(row["task_index"])
if index in results or index < 0 or index >= task_count:
fail("Duplicate/out-of-range task index: {}".format(index))
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"])
if beta_offset < 0 or cov_offset < 0:
if beta_offset != cov_offset:
fail("Only one output offset is negative for task {}".format(index))
@@ -599,11 +627,30 @@ def read_completed(manifest: Mapping[str, Any]) -> list[dict[str, Any]]:
cov_values = cov[cov_offset:cov_offset + packed].tolist()
next_beta_offset += p
next_cov_offset += packed
if fixed_test_offset < 0:
fixed_se_values: list[float] = []
fixed_stat_values: list[float] = []
fixed_ddf_values: list[float] = []
fixed_p_values: list[float] = []
else:
if (fixed_test_offset != next_fixed_test_offset or
fixed_test_offset + p > fixed_se.size):
fail("Invalid fixed-effect test offset for task {}".format(index))
fixed_se_values = fixed_se[fixed_test_offset:fixed_test_offset + p].tolist()
fixed_stat_values = fixed_stat[fixed_test_offset:fixed_test_offset + p].tolist()
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
result: dict[str, Any] = dict(row)
result["beta_json"] = json.dumps(beta_values, separators=(",", ":"))
result["covariance_packed_lower_json"] = json.dumps(cov_values, separators=(",", ":"))
result["fixed_effect_se_json"] = json.dumps(fixed_se_values, separators=(",", ":"))
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=(",", ":"))
results[index] = result
if next_beta_offset != beta_elements or next_cov_offset != cov_elements:
if (next_beta_offset != beta_elements or next_cov_offset != cov_elements or
next_fixed_test_offset != fixed_test_elements):
fail("Block binary arrays contain unused elements: {}".format(stem))
missing = sorted(set(range(task_count)).difference(results))
if missing:
@@ -618,7 +665,11 @@ def finalize(manifest: Mapping[str, Any], output: Path) -> Path:
output.parent.mkdir(parents=True, exist_ok=True)
temporary = output.with_name(output.name + ".tmp.{}".format(os.getpid()))
opener = gzip.open if output.suffix == ".gz" else open
fields = list(SUMMARY_HEADER) + ["beta_json", "covariance_packed_lower_json"]
fields = list(SUMMARY_HEADER) + [
"beta_json", "covariance_packed_lower_json", "fixed_effect_se_json",
"fixed_effect_statistic_json", "fixed_effect_denominator_df_json",
"fixed_effect_p_value_json",
]
try:
with opener(temporary, "wt", encoding="utf-8", newline="") as handle:
writer = csv.DictWriter(handle, fieldnames=fields, delimiter="\t", lineterminator="\n")
@@ -675,6 +726,11 @@ def add_run_options(parser: argparse.ArgumentParser) -> None:
parser.add_argument("--engine", default="spectra_reml")
parser.add_argument("--threads", type=positive_integer, default=1)
parser.add_argument("--blas-threads", type=positive_integer, default=1)
parser.add_argument(
"--fixed-effect-test",
choices=("satterthwaite", "kenward-roger", "none"),
default="satterthwaite",
)
parser.add_argument("--block-size", type=positive_integer, default=256)
parser.add_argument("--max-iterations", type=positive_integer, default=100)
parser.add_argument("--line-search-max-evals", type=positive_integer, default=48)