Files
PlantAI-rag-new/bioflow_kernel/cli/run_gwas_plot.py

562 lines
25 KiB
Python

#!/usr/bin/env python3
"""Render atomic MLM and CMLM Manhattan and QQ plots with CMplot."""
from __future__ import annotations
import argparse
import fcntl
import hashlib
import json
import math
import os
import platform
import re
import shlex
import shutil
import struct
import subprocess
import sys
import time
import uuid
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple
from gfa_core.gwas_scientific_identity import build_identity_recipe, materialize_stage_identity, new_run_id
STAGE = "gwas_plot"
METHOD_VERSION = "cmplot_manhattan_qq.v1"
CONTRACT_VERSION = "v1"
PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
SCOPES = (
("genotype_scope_id", "genotype"),
("phenotype_scope_id", "phenotype"),
("covariate_scope_id", "covariate"),
("source_model_scope_id", "gwas-model"),
)
def sha256_file(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def _write_json(path: Path, value: Mapping[str, Any]) -> None:
path.write_text(json.dumps(value, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
def _safe_prefix(prefix: str) -> str:
return re.sub(r"[^A-Za-z0-9._-]+", "_", prefix).strip("._") or "cohort"
def _natural_contig_key(contig: str) -> Tuple[Tuple[Tuple[int, Any], ...], str]:
tokens: List[Tuple[int, Any]] = []
for token in re.split(r"(\d+)", contig.casefold()):
if not token:
continue
tokens.append((0, int(token)) if token.isdigit() else (1, token))
return tuple(tokens), contig
def _sanitized_error(error: Exception) -> str:
return re.sub(r"[\x00-\x1f\x7f]+", " ", str(error)).strip() or error.__class__.__name__
def _markdown_table_cell(value: str) -> str:
return value.replace("\\", "\\\\").replace("|", "\\|")
def _input(path: Path, role: str) -> Dict[str, Any]:
return {
"source": "external",
"artifact": role,
"path": str(path),
"full_path": str(path),
"size": path.stat().st_size,
"sha256": sha256_file(path),
}
def _staged_artifact(
workspace: Path, source: Path, destination: Path, role: str, media_type: str
) -> Dict[str, Any]:
size = source.stat().st_size
return {
"artifact": role,
"role": role,
"path": str(destination.relative_to(workspace)),
"full_path": str(destination.resolve()),
"type": media_type,
"media_type": media_type,
"size": size,
"size_bytes": size,
"sha256": sha256_file(source),
}
def _read_pvalues(path: Path, model: str) -> Tuple[List[Tuple[str, int, float]], Dict[str, int]]:
if not path.is_file():
raise FileNotFoundError(f"Authoritative {model}_pvalues does not exist")
with path.open("r", encoding="utf-8") as handle:
lines = handle.read().splitlines()
if not lines:
raise ValueError(f"{model} p-value table header is missing")
if lines[0] != "chr\tpos\tp":
raise ValueError(f"{model} p-value table header must be exactly chr\\tpos\\tp")
rows: List[Tuple[str, int, float]] = []
coordinates: Counter[Tuple[str, int]] = Counter()
for line_number, line in enumerate(lines[1:], start=2):
fields = line.split("\t")
if len(fields) != 3:
raise ValueError(f"{model} row {line_number} must contain exactly three tab-separated fields")
contig, position_text, p_text = fields
if not contig:
raise ValueError(f"{model} row {line_number} chr must be nonempty")
try:
position = int(position_text)
except ValueError as exc:
raise ValueError(f"{model} row {line_number} pos must be a positive integer") from exc
if position <= 0:
raise ValueError(f"{model} row {line_number} pos must be a positive integer")
try:
p_value = float(p_text)
except ValueError as exc:
raise ValueError(f"{model} row {line_number} p must be finite and in (0,1]") from exc
if not math.isfinite(p_value):
raise ValueError(f"{model} row {line_number} p must be finite")
if not 0.0 < p_value <= 1.0:
raise ValueError(f"{model} row {line_number} p must be in (0,1]")
rows.append((contig, position, p_value))
coordinates[(contig, position)] += 1
if not rows:
raise ValueError(f"{model} p-value table contains zero plottable rows")
duplicates = sum(count - 1 for count in coordinates.values() if count > 1)
return rows, {
"valid_p_rows": len(rows),
"plotted_rows": len(rows),
"duplicate_coordinate_count": duplicates,
}
def _r_version(executable: Path) -> str:
completed = subprocess.run(
[str(executable), "--version"], text=True, stdout=subprocess.PIPE,
stderr=subprocess.STDOUT, check=False,
)
if completed.returncode != 0:
raise RuntimeError("Rscript version probe failed")
first = (completed.stdout or "").splitlines()
if not first:
raise RuntimeError("Rscript version could not be proven")
return first[0].strip()
def _write_plot_input(
path: Path, rows: Sequence[Tuple[str, int, float]], contig_to_plot: Mapping[str, int]
) -> None:
lines = ["SNP\tChromosome\tPosition\tP"]
lines.extend(
f"hypothesis_{index}\t{contig_to_plot[contig]}\t{position}\t{p_value:.17g}"
for index, (contig, position, p_value) in enumerate(rows, start=1)
)
path.write_text("\n".join(lines) + "\n", encoding="utf-8")
def _r_source() -> str:
return """args <- commandArgs(trailingOnly=TRUE)
suppressPackageStartupMessages(library(CMplot))
cat(paste0("CMplot_version=", as.character(packageVersion("CMplot")), "\\n"))
mlm <- read.table(args[[1]], header=TRUE, sep="\\t", check.names=FALSE)
cmlm <- read.table(args[[2]], header=TRUE, sep="\\t", check.names=FALSE)
png("mlm.manhattan.png", width=2400, height=1400, res=200)
CMplot(mlm, plot.type="m", LOG10=TRUE, threshold=as.numeric(args[[3]]), file.output=FALSE, verbose=FALSE)
dev.off()
png("mlm.qq.png", width=1800, height=1800, res=200)
CMplot(mlm, plot.type="q", LOG10=TRUE, file.output=FALSE, verbose=FALSE)
dev.off()
png("cmlm.manhattan.png", width=2400, height=1400, res=200)
CMplot(cmlm, plot.type="m", LOG10=TRUE, threshold=as.numeric(args[[4]]), file.output=FALSE, verbose=FALSE)
dev.off()
png("cmlm.qq.png", width=1800, height=1800, res=200)
CMplot(cmlm, plot.type="q", LOG10=TRUE, file.output=FALSE, verbose=FALSE)
dev.off()
"""
def _png_dimensions(path: Path, role: str) -> Tuple[int, int]:
if not path.is_file():
raise RuntimeError(f"CMplot did not produce required image: {role}")
data = path.read_bytes()
if len(data) < 33 or data[:8] != PNG_SIGNATURE:
raise ValueError(f"CMplot image is not a valid PNG: {role}")
if struct.unpack(">I", data[8:12])[0] != 13 or data[12:16] != b"IHDR":
raise ValueError(f"CMplot image has an invalid IHDR: {role}")
width, height = struct.unpack(">II", data[16:24])
if width <= 0 or height <= 0:
raise ValueError(f"CMplot image dimensions must be positive: {role}")
return width, height
def _runner_command(args: argparse.Namespace, paths: Mapping[str, Path]) -> List[str]:
return [
"python", "-m", "cli.run_gwas_plot",
"--workspace", str(args.workspace.resolve()), "--job-id", args.job_id,
"--prefix", args.prefix, "--mlm-pvalues", str(paths["mlm_pvalues"]),
"--cmlm-pvalues", str(paths["cmlm_pvalues"]),
"--genotype-scope-id", args.genotype_scope_id,
"--phenotype-scope-id", args.phenotype_scope_id,
"--covariate-scope-id", args.covariate_scope_id,
"--source-model-scope-id", args.source_model_scope_id,
"--mlm-tested-hypotheses", str(args.mlm_tested_hypotheses),
"--cmlm-tested-hypotheses", str(args.cmlm_tested_hypotheses),
"--rscript", str(args.rscript.resolve()), "--log-file", str(args.log_file.resolve()),
]
def _write_reproducibility(
workspace: Path, identity12: str, run_id: str, command: Sequence[str], manifest_path: Path,
public_manifest: Path, manifest: Mapping[str, Any], environment: str,
artifact_sources: Optional[Mapping[str, Path]] = None,
) -> Path:
bundle = workspace / "reproducibility" / STAGE / f"{_safe_prefix(str(manifest['prefix']))}.{identity12}" / run_id
bundle.mkdir(parents=True, exist_ok=False)
canonical = shlex.join(command)
commands = bundle / "commands.sh"
commands.write_text("#!/usr/bin/env bash\nset -euo pipefail\n" + canonical + "\n", encoding="utf-8")
commands.chmod(0o755)
(bundle / "environment.txt").write_text(environment, encoding="utf-8")
_write_json(bundle / "progress.json", {
"stage": STAGE, "status": manifest["status"], "run_id": run_id,
"manifest_path": str(public_manifest.relative_to(workspace)),
})
_write_json(bundle / "provenance.json", {
"stage": STAGE, "run_id": run_id, "runner_command": canonical,
"software": manifest.get("software", {}), "inputs": manifest.get("inputs", []),
"artifacts": manifest.get("artifacts", {}),
})
entries = [(manifest_path, str(public_manifest.relative_to(workspace)))]
for role, artifact in manifest.get("artifacts", {}).items():
source = artifact_sources.get(role) if artifact_sources else workspace / artifact["path"]
if source and source.is_file():
entries.append((source, artifact["path"]))
(bundle / "checksums.sha256").write_text(
"".join(f"{sha256_file(path)} {label}\n" for path, label in sorted(entries, key=lambda item: item[1])),
encoding="utf-8",
)
return bundle
def _publish_atomically(
sources: Mapping[str, Path], destinations: Mapping[str, Path], manifest_source: Path,
manifest_destination: Path, lock_path: Path,
) -> None:
token = uuid.uuid4().hex
staged: Dict[Path, Path] = {}
for role, destination in destinations.items():
temporary = destination.with_name(f".{destination.name}.{token}.publishing")
shutil.copy2(sources[role], temporary)
staged[destination] = temporary
temporary_manifest = manifest_destination.with_name(f".{manifest_destination.name}.{token}.publishing")
shutil.copy2(manifest_source, temporary_manifest)
staged[manifest_destination] = temporary_manifest
lock_path.parent.mkdir(parents=True, exist_ok=True)
with lock_path.open("a+", encoding="utf-8") as lock:
fcntl.flock(lock.fileno(), fcntl.LOCK_EX)
backups: Dict[Path, Optional[Path]] = {}
try:
for destination in staged:
if destination.is_file():
backup = destination.with_name(f".{destination.name}.{token}.rollback")
shutil.copy2(destination, backup)
backups[destination] = backup
else:
backups[destination] = None
for destination, temporary in staged.items():
os.replace(temporary, destination)
except Exception:
for destination, backup in backups.items():
if backup is None:
destination.unlink(missing_ok=True)
elif backup.is_file():
os.replace(backup, destination)
for temporary in staged.values():
temporary.unlink(missing_ok=True)
raise
finally:
for backup in backups.values():
if backup is not None:
backup.unlink(missing_ok=True)
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--workspace", required=True, type=Path)
parser.add_argument("--job-id", required=True)
parser.add_argument("--prefix", required=True)
parser.add_argument("--mlm-pvalues", required=True, type=Path)
parser.add_argument("--cmlm-pvalues", required=True, type=Path)
parser.add_argument("--genotype-scope-id", required=True)
parser.add_argument("--phenotype-scope-id", required=True)
parser.add_argument("--covariate-scope-id", required=True)
parser.add_argument("--source-model-scope-id", required=True)
parser.add_argument("--mlm-tested-hypotheses", required=True, type=int)
parser.add_argument("--cmlm-tested-hypotheses", required=True, type=int)
parser.add_argument("--rscript", required=True, type=Path)
parser.add_argument("--log-file", required=True, type=Path)
return parser
def main(argv: Optional[List[str]] = None) -> int:
args = build_parser().parse_args(argv)
workspace = args.workspace.resolve()
final_dir = workspace / "final"
final_dir.mkdir(parents=True, exist_ok=True)
run_id = new_run_id(STAGE)
started = datetime.now(timezone.utc)
started_clock = time.monotonic()
safe_prefix = _safe_prefix(args.prefix)
failed_manifest = final_dir / f"{safe_prefix}.{STAGE}.failed.manifest.json"
paths = {
"mlm_pvalues": args.mlm_pvalues.resolve(),
"cmlm_pvalues": args.cmlm_pvalues.resolve(),
}
command = _runner_command(args, paths)
tested = {"mlm": args.mlm_tested_hypotheses, "cmlm": args.cmlm_tested_hypotheses}
thresholds = {model: 0.05 / count if count > 0 else 0.0 for model, count in tested.items()}
parameters: Dict[str, Any] = {
"tested_hypotheses": tested,
"bonferroni_thresholds": thresholds,
"plots": ["manhattan", "qq"],
"plot_devices": {
"manhattan": {"width_px": 2400, "height_px": 1400, "resolution_dpi": 200},
"qq": {"width_px": 1800, "height_px": 1800, "resolution_dpi": 200},
},
"cmplot": {"log10": True, "file_output": False},
"contig_mapping_policy": "natural_union_v1",
}
manifest: Dict[str, Any] = {
"manifest_version": "1.0", "stage": STAGE, "prefix": args.prefix,
"job_id": args.job_id, "run_id": run_id, "status": "failed",
"timestamp": started.isoformat(), "start_time": started.isoformat(),
"end_time": started.isoformat(), "elapsed_sec": 0.0, "cmd": command,
"workspace_root": str(workspace), "workspace": str(final_dir),
"input_dir": str(workspace / "input"), "parameters": parameters,
"metrics": {}, "artifacts": {},
}
environment = ""
log_is_safe = True
identity12 = "failed"
success_bundle: Optional[Path] = None
success_manifest: Optional[Path] = None
attempt = workspace / ".attempts" / run_id
attempt.mkdir(parents=True, exist_ok=False)
try:
scopes = {name: str(getattr(args, name)) for name, _prefix in SCOPES}
for name, prefix in SCOPES:
if not re.fullmatch(prefix + r":[0-9a-f]{64}", scopes[name]):
raise ValueError(f"{name} is invalid")
manifest.update(scopes)
log_path = args.log_file.resolve()
for role, input_path in paths.items():
if log_path == input_path:
log_is_safe = False
raise ValueError("log path must not alias an authoritative input")
if not input_path.is_file():
raise FileNotFoundError(f"Authoritative {role} does not exist")
if log_path.exists() and os.path.samefile(log_path, input_path):
log_is_safe = False
raise ValueError("log path must not alias an authoritative input")
input_records = [_input(paths[role], role) for role in ("mlm_pvalues", "cmlm_pvalues")]
manifest["inputs"] = input_records
expected_input_hashes = {item["artifact"]: item["sha256"] for item in input_records}
rows: Dict[str, List[Tuple[str, int, float]]] = {}
metrics: Dict[str, Any] = {}
for model in ("mlm", "cmlm"):
model_rows, model_metrics = _read_pvalues(paths[f"{model}_pvalues"], model)
if tested[model] < model_metrics["valid_p_rows"]:
raise ValueError(
f"{model} tested_hypotheses must be >= valid_p_rows >= plotted_rows"
)
rows[model] = model_rows
metrics[model] = model_metrics
contig_order = sorted(
{row[0] for model_rows in rows.values() for row in model_rows},
key=_natural_contig_key,
)
contig_to_plot = {contig: index for index, contig in enumerate(contig_order, start=1)}
contig_counts = {
model: {contig: sum(row[0] == contig for row in rows[model]) for contig in contig_order}
for model in ("mlm", "cmlm")
}
rscript = args.rscript.resolve()
if not rscript.is_file() or not os.access(str(rscript), os.X_OK):
raise FileNotFoundError("explicit Rscript executable was not found or is not executable")
r_version = _r_version(rscript)
for model in ("mlm", "cmlm"):
_write_plot_input(attempt / f"{model}.cmplot.tsv", rows[model], contig_to_plot)
script = attempt / "run_cmplot.R"
script.write_text(_r_source(), encoding="utf-8")
completed = subprocess.run(
[str(rscript), str(script), str(attempt / "mlm.cmplot.tsv"),
str(attempt / "cmlm.cmplot.tsv"), str(thresholds["mlm"]), str(thresholds["cmlm"])],
cwd=attempt, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, check=False,
)
log_text = completed.stdout or ""
attempt_log = attempt / "cmplot.log"
attempt_log.write_text(log_text, encoding="utf-8")
args.log_file.parent.mkdir(parents=True, exist_ok=True)
shutil.copy2(attempt_log, args.log_file)
if completed.returncode != 0:
raise RuntimeError(f"CMplot failed with exit code {completed.returncode}")
for role, expected_hash in expected_input_hashes.items():
if sha256_file(paths[role]) != expected_hash:
raise RuntimeError(f"Authoritative input changed during plotting: {role}")
match = re.search(r"^CMplot_version=([^\r\n]+)$", log_text, re.MULTILINE)
if not match:
raise RuntimeError("CMplot version could not be proven")
cmplot_version = match.group(1).strip()
image_sources = {
f"{model}_{plot}": attempt / f"{model}.{plot}.png"
for model in ("mlm", "cmlm") for plot in ("manhattan", "qq")
}
image_dimensions = {
role: _png_dimensions(path, role) for role, path in image_sources.items()
}
environment = "\n".join([
"runner=cli.run_gwas_plot.py",
f"runtime_environment={os.environ.get('CONDA_DEFAULT_ENV', 'unknown')}",
f"python_resolved={Path(sys.executable).resolve()}",
f"python_version={platform.python_version()}",
f"rscript_resolved={rscript}", f"r_version={r_version}",
f"cmplot_version={cmplot_version}", f"platform={platform.platform()}",
]) + "\n"
recipe = build_identity_recipe(
stage=STAGE, required_artifact_roles=paths.keys(), parameters=parameters,
method_version=METHOD_VERSION, contract_version=CONTRACT_VERSION,
tool_versions={"R": r_version, "CMplot": cmplot_version},
)
identity = materialize_stage_identity(
recipe, {item["artifact"]: item for item in input_records}, scopes
)
identity12 = identity["analysis_scope_id"].split(":", 1)[1][:12]
stem = f"{safe_prefix}.{STAGE}.{identity12}"
success_manifest = final_dir / f"{stem}.manifest.json"
manifest.update(identity)
manifest.update({
"identity_recipe": recipe, "inputs": input_records,
"software": {
"R": {"resolved_path": str(rscript), "version": r_version},
"CMplot": {"version": cmplot_version},
},
})
report_source = attempt / "report.md"
environment_source = attempt / "environment.txt"
mapping_rows = "".join(
f"| {_markdown_table_cell(contig)} | {contig_to_plot[contig]} | {contig_counts['mlm'][contig]} | "
f"{contig_counts['cmlm'][contig]} |\n"
for contig in contig_order
)
report_source.write_text(
"# GWAS MLM/CMLM plots\n\n"
f"- Analysis scope: `{identity['analysis_scope_id']}`\n"
f"- Source model scope: `{args.source_model_scope_id}`\n"
f"- Genotype scope: `{args.genotype_scope_id}`\n"
f"- Phenotype scope: `{args.phenotype_scope_id}`\n"
f"- Covariate scope: `{args.covariate_scope_id}`\n"
"- Contig classification: unavailable; not inferred from contig names.\n\n"
"## Multiple-testing thresholds\n\n"
"| Model | Tested hypotheses | Bonferroni threshold |\n"
"|---|---:|---:|\n"
f"| MLM | {tested['mlm']} | {thresholds['mlm']} |\n"
f"| CMLM | {tested['cmlm']} | {thresholds['cmlm']} |\n\n"
"## Reversible contig plotting map\n\n"
"| Contig | Plot integer | MLM rows | CMLM rows |\n"
"|---|---:|---:|---:|\n" + mapping_rows + "\n"
"Duplicate hypotheses were preserved. Duplicate coordinate row counts: "
f"MLM={metrics['mlm']['duplicate_coordinate_count']}, "
f"CMLM={metrics['cmlm']['duplicate_coordinate_count']}.\n",
encoding="utf-8",
)
environment_source.write_text(environment, encoding="utf-8")
sources = dict(image_sources)
sources.update({"report": report_source, "environment": environment_source})
destinations = {
role: final_dir / f"{stem}.{role}.png" for role in image_sources
}
destinations.update({
"report": final_dir / f"{stem}.report.md",
"environment": final_dir / f"{stem}.environment.txt",
})
media = {"report": "text/markdown", "environment": "text/plain"}
metrics["contig_count"] = len(contig_order)
metrics["contig_order"] = contig_order
metrics["contig_to_plot"] = contig_to_plot
metrics["plot_to_contig"] = {str(value): key for key, value in contig_to_plot.items()}
metrics["contig_row_counts"] = contig_counts
metrics["contig_classification"] = "unavailable; not inferred from contig names"
metrics["image_dimensions"] = {
role: {"width": value[0], "height": value[1]} for role, value in image_dimensions.items()
}
manifest.update({
"status": "success", "metrics": metrics,
"artifacts": {
role: _staged_artifact(
workspace, sources[role], destinations[role], role, media.get(role, "image/png")
) for role in sources
},
})
manifest["end_time"] = datetime.now(timezone.utc).isoformat()
manifest["elapsed_sec"] = round(time.monotonic() - started_clock, 6)
manifest_source = attempt / "manifest.json"
_write_json(manifest_source, manifest)
success_bundle = _write_reproducibility(
workspace, identity12, run_id, command, manifest_source, success_manifest,
manifest, environment, sources,
)
_publish_atomically(
sources, destinations, manifest_source, success_manifest,
workspace / ".locks" / f"{stem}.lock",
)
failed_manifest.unlink(missing_ok=True)
return 0
except Exception as exc:
error = _sanitized_error(exc)
if success_bundle is not None and success_bundle.is_dir():
shutil.rmtree(success_bundle)
manifest.update({
"status": "failed", "artifacts": {}, "error": error,
"end_time": datetime.now(timezone.utc).isoformat(),
"elapsed_sec": round(time.monotonic() - started_clock, 6),
})
temporary = failed_manifest.with_name(f".{failed_manifest.name}.{uuid.uuid4().hex}.publishing")
_write_json(temporary, manifest)
os.replace(temporary, failed_manifest)
if log_is_safe:
args.log_file.parent.mkdir(parents=True, exist_ok=True)
needs_separator = False
if args.log_file.is_file() and args.log_file.stat().st_size:
with args.log_file.open("rb") as existing_log:
existing_log.seek(-1, os.SEEK_END)
needs_separator = existing_log.read(1) not in {b"\n", b"\r"}
with args.log_file.open("a", encoding="utf-8") as log_handle:
if needs_separator:
log_handle.write("\n")
log_handle.write(f"ERROR: {error}\n")
failure_source = attempt / "failed.manifest.json"
_write_json(failure_source, manifest)
_write_reproducibility(
workspace, identity12, run_id, command, failure_source, failed_manifest,
manifest, environment,
)
print(f"ERROR: {error}", file=sys.stderr)
return 1
if __name__ == "__main__":
raise SystemExit(main())