402 lines
18 KiB
Python
402 lines
18 KiB
Python
#!/usr/bin/env python3
|
|
"""Run LD-pruned PLINK PCA with fail-closed population-analysis publishing."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import gzip
|
|
import json
|
|
import os
|
|
import platform
|
|
import re
|
|
import shlex
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any, TextIO
|
|
|
|
from gfa_core.genotype_context import (
|
|
build_population_analysis_scope_id,
|
|
genotype_scope_id_for_vcf,
|
|
sha256_file,
|
|
)
|
|
from gfa_core.population_stage_runtime import acquire_identity_lock, population_manifest_base, release_identity_lock, write_reproducibility_bundle
|
|
|
|
|
|
STAGE = "pop_pca"
|
|
METHOD_VERSION = "plink_ld_pruned_pca.v1"
|
|
CONTRACT_VERSION = "v1"
|
|
|
|
|
|
def _open_text(path: Path) -> TextIO:
|
|
if path.suffix == ".gz":
|
|
return gzip.open(path, "rt", encoding="utf-8")
|
|
return path.open("r", encoding="utf-8")
|
|
|
|
|
|
def scan_vcf_marker_policy(path: Path) -> tuple[list[str], dict[str, int]]:
|
|
if not path.is_file() or path.stat().st_size <= 0:
|
|
raise FileNotFoundError(f"Population VCF is missing or empty: {path}")
|
|
samples: list[str] = []
|
|
metrics = {
|
|
"input_record_count": 0,
|
|
"eligible_biallelic_snp_count": 0,
|
|
"excluded_multiallelic_count": 0,
|
|
"excluded_non_snp_count": 0,
|
|
}
|
|
with _open_text(path) as handle:
|
|
for raw in handle:
|
|
if raw.startswith("#CHROM\t"):
|
|
samples = raw.rstrip("\n").split("\t")[9:]
|
|
continue
|
|
if raw.startswith("#") or not raw.strip():
|
|
continue
|
|
fields = raw.rstrip("\n").split("\t")
|
|
if len(fields) < 5:
|
|
raise ValueError("VCF record has fewer than five columns")
|
|
metrics["input_record_count"] += 1
|
|
ref, alt = fields[3].upper(), fields[4].upper()
|
|
if "," in alt:
|
|
metrics["excluded_multiallelic_count"] += 1
|
|
elif len(ref) != 1 or len(alt) != 1 or ref not in "ACGT" or alt not in "ACGT":
|
|
metrics["excluded_non_snp_count"] += 1
|
|
else:
|
|
metrics["eligible_biallelic_snp_count"] += 1
|
|
if len(samples) < 2:
|
|
raise ValueError("Population PCA requires at least two VCF samples")
|
|
if len(samples) != len(set(samples)):
|
|
raise ValueError("VCF sample IDs must be unique")
|
|
if metrics["eligible_biallelic_snp_count"] <= 0:
|
|
raise ValueError("No eligible biallelic single-nucleotide SNPs")
|
|
return samples, metrics
|
|
|
|
|
|
def validate_sample_group(path: Path, vcf_samples: list[str]) -> dict[str, str]:
|
|
groups: dict[str, str] = {}
|
|
allowed = set(vcf_samples)
|
|
with path.open("r", encoding="utf-8") as handle:
|
|
for line_number, raw in enumerate(handle, start=1):
|
|
line = raw.strip()
|
|
if not line or line.startswith("#"):
|
|
continue
|
|
fields = line.split()
|
|
if len(fields) != 2 or not fields[1]:
|
|
raise ValueError(f"sample_group row {line_number} must have exactly two columns")
|
|
sample, group = fields
|
|
if sample in groups:
|
|
raise ValueError("sample_group sample IDs must be unique")
|
|
if sample not in allowed:
|
|
raise ValueError(f"sample_group ID not present in VCF: {sample}")
|
|
groups[sample] = group
|
|
if not groups:
|
|
raise ValueError("sample_group is empty")
|
|
return groups
|
|
|
|
|
|
def _write_json(path: Path, value: dict[str, Any]) -> None:
|
|
path.write_text(json.dumps(value, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
|
|
|
|
|
def _artifact(workspace: Path, path: Path, role: str, media_type: str) -> dict[str, Any]:
|
|
if not path.is_file() or path.stat().st_size <= 0:
|
|
raise ValueError(f"Required artifact is missing or empty: {role}")
|
|
size = path.stat().st_size
|
|
return {
|
|
"artifact": role,
|
|
"role": role,
|
|
"path": str(path.relative_to(workspace)),
|
|
"full_path": str(path.resolve()),
|
|
"type": media_type,
|
|
"size": size,
|
|
"size_bytes": size,
|
|
"sha256": sha256_file(path),
|
|
}
|
|
|
|
|
|
def _publish(source: Path, destination: Path) -> None:
|
|
temporary = destination.with_name(f".{destination.name}.{os.getpid()}.publishing")
|
|
shutil.copy2(source, temporary)
|
|
os.replace(temporary, destination)
|
|
|
|
|
|
def _tool_version(executable: Path) -> str:
|
|
result = subprocess.run([str(executable), "--version"], text=True, capture_output=True, timeout=30)
|
|
if result.returncode != 0:
|
|
raise RuntimeError(f"PLINK version probe failed ({result.returncode}): {result.stderr.strip()}")
|
|
return (result.stdout or result.stderr).strip().splitlines()[0]
|
|
|
|
|
|
def _run(command: list[str], log_path: Path) -> None:
|
|
result = subprocess.run(command, text=True, capture_output=True)
|
|
log_path.write_text(
|
|
"$ " + shlex.join(command) + "\n\nSTDOUT\n" + result.stdout + "\nSTDERR\n" + result.stderr,
|
|
encoding="utf-8",
|
|
)
|
|
if result.returncode != 0:
|
|
raise RuntimeError(f"Command failed with exit code {result.returncode}: {shlex.join(command)}")
|
|
|
|
|
|
def _parse_pca(eigenvec: Path, eigenval: Path, expected_samples: list[str]) -> tuple[list[str], list[list[str]], list[float]]:
|
|
lines = [line.split() for line in eigenvec.read_text(encoding="utf-8").splitlines() if line.strip()]
|
|
if not lines:
|
|
raise ValueError("PLINK PCA eigenvector file is empty")
|
|
if lines[0][0].lstrip("#") == "FID":
|
|
header, rows = [lines[0][0].lstrip("#"), *lines[0][1:]], lines[1:]
|
|
else:
|
|
rows = lines
|
|
header = ["FID", "IID", *(f"PC{i}" for i in range(1, len(rows[0]) - 1))]
|
|
if any(len(row) != len(header) for row in rows):
|
|
raise ValueError("PLINK PCA rows have inconsistent columns")
|
|
sample_ids = [row[1] for row in rows]
|
|
if sample_ids != expected_samples:
|
|
raise ValueError("PCA sample order does not exactly match the VCF")
|
|
eigenvalues = [float(value) for value in eigenval.read_text(encoding="utf-8").split()]
|
|
if len(eigenvalues) != len(header) - 2 or any(value < 0 for value in eigenvalues):
|
|
raise ValueError("PCA eigenvalue count does not match eigenvectors")
|
|
return header, rows, eigenvalues
|
|
|
|
|
|
def _write_outputs(
|
|
attempt: Path,
|
|
header: list[str],
|
|
rows: list[list[str]],
|
|
eigenvalues: list[float],
|
|
groups: dict[str, str],
|
|
requested_plot_pcs: int,
|
|
tassel_covariate_pcs: int,
|
|
) -> dict[str, Path]:
|
|
scores = attempt / "pca_scores.tsv"
|
|
scores.write_text("\t".join(header) + "\n" + "\n".join("\t".join(row) for row in rows) + "\n", encoding="utf-8")
|
|
eigen = attempt / "pca_eigenvalues.tsv"
|
|
eigen.write_text("PC\teigenvalue\n" + "".join(f"PC{i}\t{value:.15g}\n" for i, value in enumerate(eigenvalues, 1)), encoding="utf-8")
|
|
total = sum(eigenvalues)
|
|
if total <= 0:
|
|
raise ValueError("PCA eigenvalues have no positive total")
|
|
variance = attempt / "pca_variance.tsv"
|
|
cumulative = 0.0
|
|
variance_lines = ["PC\teigenvalue\tvariance_percent\tcumulative_percent"]
|
|
for index, value in enumerate(eigenvalues, 1):
|
|
percent = 100.0 * value / total
|
|
cumulative += percent
|
|
variance_lines.append(f"PC{index}\t{value:.15g}\t{percent:.12g}\t{cumulative:.12g}")
|
|
variance.write_text("\n".join(variance_lines) + "\n", encoding="utf-8")
|
|
|
|
covariate_count = min(tassel_covariate_pcs, len(eigenvalues))
|
|
covariate = attempt / "pca_tassel_covariate.txt"
|
|
cov_lines = ["<Covariate>", "<Trait>\t" + "\t".join(f"PC{i}" for i in range(1, covariate_count + 1))]
|
|
for row in rows:
|
|
cov_lines.append("\t".join([row[1], *row[2 : 2 + covariate_count]]))
|
|
covariate.write_text("\n".join(cov_lines) + "\n", encoding="utf-8")
|
|
|
|
import matplotlib
|
|
matplotlib.use("Agg")
|
|
import matplotlib.pyplot as plt
|
|
|
|
if len(eigenvalues) < 2:
|
|
raise ValueError("At least two computed PCs are required for the PC1/PC2 plot")
|
|
pc1 = [float(row[2]) for row in rows]
|
|
pc2 = [float(row[3]) for row in rows]
|
|
scatter = attempt / "pca_pc1_pc2.png"
|
|
fig, ax = plt.subplots(figsize=(7, 6))
|
|
labels = [groups.get(row[1], "Unassigned") for row in rows]
|
|
for label in sorted(set(labels)):
|
|
indices = [i for i, value in enumerate(labels) if value == label]
|
|
ax.scatter([pc1[i] for i in indices], [pc2[i] for i in indices], label=label)
|
|
ax.set_xlabel("PC1")
|
|
ax.set_ylabel("PC2")
|
|
if groups:
|
|
ax.legend()
|
|
fig.tight_layout()
|
|
fig.savefig(scatter, dpi=180)
|
|
plt.close(fig)
|
|
|
|
scree = attempt / "pca_scree.png"
|
|
fig, ax = plt.subplots(figsize=(7, 5))
|
|
percentages = [100.0 * value / total for value in eigenvalues[:requested_plot_pcs]]
|
|
ax.bar(range(1, len(percentages) + 1), percentages)
|
|
ax.set_xlabel("Principal component")
|
|
ax.set_ylabel("Variance explained (%)")
|
|
fig.tight_layout()
|
|
fig.savefig(scree, dpi=180)
|
|
plt.close(fig)
|
|
return {
|
|
"pca_scores": scores,
|
|
"pca_eigenvalues": eigen,
|
|
"pca_variance_table": variance,
|
|
"pca_pc1_pc2_plot": scatter,
|
|
"pca_scree_plot": scree,
|
|
"pca_tassel_covariate": covariate,
|
|
}
|
|
|
|
|
|
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("--vcf", required=True, type=Path)
|
|
parser.add_argument("--sample-group", type=Path)
|
|
parser.add_argument("--plink", default="plink", type=Path)
|
|
parser.add_argument("--ld-window-variants", default=50, type=int)
|
|
parser.add_argument("--ld-step-variants", default=10, type=int)
|
|
parser.add_argument("--ld-r2", default=0.2, type=float)
|
|
parser.add_argument("--requested-plot-pcs", default=10, type=int)
|
|
parser.add_argument("--tassel-covariate-pcs", default=3, type=int)
|
|
return parser
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
args = build_parser().parse_args(argv)
|
|
workspace = args.workspace.resolve()
|
|
final = workspace / "final"
|
|
final.mkdir(parents=True, exist_ok=True)
|
|
started = datetime.now(timezone.utc)
|
|
start_clock = time.monotonic()
|
|
run_id = f"{STAGE}_{started.strftime('%Y%m%dT%H%M%SZ')}_{uuid.uuid4().hex[:8]}"
|
|
failed_manifest = final / f"{args.prefix}.{STAGE}.failed.{run_id}.manifest.json"
|
|
identity12 = "unresolved"
|
|
lock_handle = None
|
|
runner_command = [sys.executable, "-m", "apps.cli.run_pop_pca", *(argv if argv is not None else sys.argv[1:])]
|
|
parameters = {
|
|
"marker_policy": "biallelic_snp",
|
|
"ld_window_variants": args.ld_window_variants,
|
|
"ld_step_variants": args.ld_step_variants,
|
|
"ld_r2": args.ld_r2,
|
|
"requested_plot_pcs": args.requested_plot_pcs,
|
|
"tassel_covariate_pcs": args.tassel_covariate_pcs,
|
|
}
|
|
manifest = population_manifest_base(
|
|
workspace=workspace, stage=STAGE, prefix=args.prefix, job_id=args.job_id,
|
|
run_id=run_id, started=started, runner_command=runner_command,
|
|
parameters=parameters,
|
|
)
|
|
try:
|
|
vcf = args.vcf.resolve()
|
|
samples, marker_metrics = scan_vcf_marker_policy(vcf)
|
|
groups = validate_sample_group(args.sample_group.resolve(), samples) if args.sample_group else {}
|
|
if args.ld_window_variants <= 0 or args.ld_step_variants <= 0 or not 0 < args.ld_r2 <= 1:
|
|
raise ValueError("Invalid LD pruning parameters")
|
|
if args.tassel_covariate_pcs <= 0 or args.requested_plot_pcs <= 0:
|
|
raise ValueError("requested_plot_pcs and tassel_covariate_pcs must be positive")
|
|
plink = args.plink.resolve() if args.plink.is_absolute() or args.plink.exists() else Path(shutil.which(str(args.plink)) or "")
|
|
if not plink.is_file():
|
|
raise FileNotFoundError(f"PLINK executable not found: {args.plink}")
|
|
plink_version = _tool_version(plink)
|
|
genotype_scope = genotype_scope_id_for_vcf(vcf)
|
|
input_hashes = {"population_vcf": sha256_file(vcf)}
|
|
if args.sample_group:
|
|
input_hashes["sample_group"] = sha256_file(args.sample_group.resolve())
|
|
analysis_scope = build_population_analysis_scope_id(
|
|
stage=STAGE, genotype_scope_id=genotype_scope,
|
|
input_artifact_hashes=input_hashes, parameters=parameters,
|
|
method_version=METHOD_VERSION, contract_version=CONTRACT_VERSION,
|
|
tool_versions={"plink": plink_version, "python": platform.python_version()},
|
|
)
|
|
identity12 = analysis_scope.split(":", 1)[1][:12]
|
|
lock_handle = acquire_identity_lock(workspace, STAGE, args.prefix, identity12)
|
|
manifest.update({
|
|
"genotype_scope_id": genotype_scope, "analysis_scope_id": analysis_scope,
|
|
"method_version": METHOD_VERSION, "contract_version": CONTRACT_VERSION,
|
|
"tool_versions": {"plink": plink_version, "python": platform.python_version()},
|
|
"inputs": [{"role": role, "sha256": digest} for role, digest in input_hashes.items()],
|
|
})
|
|
attempt = workspace / ".attempts" / run_id
|
|
attempt.mkdir(parents=True, exist_ok=False)
|
|
prune_prefix = attempt / "prune"
|
|
common = [str(plink), "--vcf", str(vcf), "--double-id", "--allow-extra-chr", "--biallelic-only", "strict", "--snps-only", "just-acgt", "--set-missing-var-ids", "@:#:$1:$2"]
|
|
prune_command = [*common, "--indep-pairwise", str(args.ld_window_variants), str(args.ld_step_variants), str(args.ld_r2), "--out", str(prune_prefix)]
|
|
_run(prune_command, attempt / "plink_prune.log")
|
|
prune_in = prune_prefix.with_suffix(".prune.in")
|
|
markers = [line.strip() for line in prune_in.read_text(encoding="utf-8").splitlines() if line.strip()] if prune_in.is_file() else []
|
|
if not markers:
|
|
raise ValueError("PLINK LD pruning retained zero markers")
|
|
pc_count = min(len(samples) - 1, len(markers))
|
|
if pc_count < 2:
|
|
raise ValueError("Fewer than two PCs are available after LD pruning")
|
|
pca_prefix = attempt / "pca"
|
|
pca_command = [*common, "--extract", str(prune_in), "--pca", str(pc_count), "header", "tabs", "--out", str(pca_prefix)]
|
|
_run(pca_command, attempt / "plink_pca.log")
|
|
header, rows, eigenvalues = _parse_pca(pca_prefix.with_suffix(".eigenvec"), pca_prefix.with_suffix(".eigenval"), samples)
|
|
outputs = _write_outputs(
|
|
attempt, header, rows, eigenvalues, groups,
|
|
args.requested_plot_pcs, args.tassel_covariate_pcs,
|
|
)
|
|
outputs["pca_pruning_markers"] = prune_in
|
|
environment = attempt / "environment.txt"
|
|
environment.write_text(
|
|
f"runner=apps.cli.run_pop_pca.py\npython={sys.executable}\npython_version={platform.python_version()}\nplink={plink}\nplink_version={plink_version}\n",
|
|
encoding="utf-8",
|
|
)
|
|
outputs["environment"] = environment
|
|
report = attempt / "report.md"
|
|
report.write_text(
|
|
"# Population PCA\n\n"
|
|
f"- Samples: {len(samples)}\n- Eligible biallelic SNPs: {marker_metrics['eligible_biallelic_snp_count']}\n"
|
|
f"- LD-pruned markers: {len(markers)}\n- Computed PCs: {len(eigenvalues)}\n"
|
|
"- Validation scope: program/flow output; biological inference depends on cohort adequacy.\n",
|
|
encoding="utf-8",
|
|
)
|
|
outputs["report"] = report
|
|
stem = f"{args.prefix}.{STAGE}.{identity12}.{run_id}"
|
|
media_types = {
|
|
"pca_scores": "text/tab-separated-values", "pca_eigenvalues": "text/tab-separated-values",
|
|
"pca_variance_table": "text/tab-separated-values", "pca_pruning_markers": "text/plain",
|
|
"pca_pc1_pc2_plot": "image/png", "pca_scree_plot": "image/png",
|
|
"pca_tassel_covariate": "text/tab-separated-values", "report": "text/markdown",
|
|
"environment": "text/plain",
|
|
}
|
|
suffixes = {
|
|
"pca_scores": "scores.tsv", "pca_eigenvalues": "eigenvalues.tsv", "pca_variance_table": "variance.tsv",
|
|
"pca_pruning_markers": "prune.in", "pca_pc1_pc2_plot": "pc1_pc2.png", "pca_scree_plot": "scree.png",
|
|
"pca_tassel_covariate": "tassel_covariate.txt", "report": "report.md", "environment": "environment.txt",
|
|
}
|
|
published: dict[str, Path] = {}
|
|
for role, source in outputs.items():
|
|
destination = final / f"{stem}.{suffixes[role]}"
|
|
_publish(source, destination)
|
|
published[role] = destination
|
|
manifest.update({
|
|
"status": "success",
|
|
"metrics": {**marker_metrics, "sample_count": len(samples), "grouped_sample_count": len(groups), "pruned_marker_count": len(markers), "computed_pc_count": len(eigenvalues)},
|
|
"commands": [prune_command, pca_command],
|
|
"artifacts": {role: _artifact(workspace, path, role, media_types[role]) for role, path in published.items()},
|
|
})
|
|
manifest["end_time"] = datetime.now(timezone.utc).isoformat()
|
|
manifest["elapsed_sec"] = round(time.monotonic() - start_clock, 6)
|
|
manifest_work = attempt / "manifest.json"
|
|
_write_json(manifest_work, manifest)
|
|
manifest_path = final / f"{stem}.manifest.json"
|
|
write_reproducibility_bundle(
|
|
workspace=workspace, stage=STAGE, prefix=args.prefix, identity12=identity12,
|
|
run_id=run_id, runner_command=runner_command, manifest_source=manifest_work,
|
|
public_manifest=manifest_path, manifest=manifest,
|
|
environment=environment.read_text(encoding="utf-8"), artifact_sources=outputs,
|
|
)
|
|
_publish(manifest_work, manifest_path)
|
|
return 0
|
|
except Exception as exc:
|
|
manifest.update({"status": "failed", "artifacts": {}, "error": str(exc), "end_time": datetime.now(timezone.utc).isoformat(), "elapsed_sec": round(time.monotonic() - start_clock, 6)})
|
|
_write_json(failed_manifest, manifest)
|
|
try:
|
|
write_reproducibility_bundle(
|
|
workspace=workspace, stage=STAGE, prefix=args.prefix, identity12=identity12,
|
|
run_id=run_id, runner_command=runner_command, manifest_source=failed_manifest,
|
|
public_manifest=failed_manifest, manifest=manifest,
|
|
environment=json.dumps(manifest.get("tool_versions", {}), ensure_ascii=False),
|
|
)
|
|
except Exception as bundle_exc:
|
|
manifest["reproducibility_error"] = str(bundle_exc)
|
|
_write_json(failed_manifest, manifest)
|
|
print(f"ERROR: {exc}", file=sys.stderr)
|
|
return 1
|
|
finally:
|
|
release_identity_lock(lock_handle)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|