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

548 lines
23 KiB
Python
Executable File

#!/usr/bin/env python3
"""
Draw a selected region from an existing phylogeny MSA.
This is intentionally a post-processing helper: it does not change the
phylogeny stage or its default full-length MSA figures.
"""
import argparse
import json
import shutil
import subprocess
import sys
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple
PROJECT_ROOT = Path(__file__).resolve().parents[2]
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
from gfa_core.gene_family_scope import read_scope
from gfa_core.gene_family_stage_context import apply_stage_context, build_stage_context
DEFAULT_RSCRIPT = Path(shutil.which("Rscript") or "Rscript")
FastaRecords = List[Tuple[str, str]]
def read_fasta(path: Path) -> FastaRecords:
records: FastaRecords = []
header: Optional[str] = None
chunks: List[str] = []
with path.open("r", encoding="utf-8", errors="ignore") as handle:
for raw in handle:
line = raw.strip()
if not line:
continue
if line.startswith(">"):
if header is not None:
records.append((header, "".join(chunks)))
header = line[1:].strip()
chunks = []
else:
chunks.append(line)
if header is not None:
records.append((header, "".join(chunks)))
return records
def write_fasta(path: Path, records: FastaRecords, width: int = 80) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("w", encoding="utf-8") as handle:
for header, seq in records:
handle.write(f">{header}\n")
for idx in range(0, len(seq), width):
handle.write(seq[idx:idx + width] + "\n")
def validate_alignment(records: FastaRecords) -> int:
if not records:
raise ValueError("MSA fasta contains no records")
lengths = {len(seq) for _, seq in records}
if len(lengths) != 1:
raise ValueError(f"MSA sequence lengths are not equal: {sorted(lengths)[:5]}")
return lengths.pop()
def conservation_profile(records: FastaRecords) -> Tuple[List[float], List[float]]:
aln_len = validate_alignment(records)
seqs = [seq for _, seq in records]
scores: List[float] = []
nongaps: List[float] = []
for col_idx in range(aln_len):
col = [seq[col_idx] for seq in seqs]
residues = [char for char in col if char not in {"-", "."}]
nongap_frac = len(residues) / len(col)
nongaps.append(nongap_frac)
if not residues:
scores.append(0.0)
continue
counts: Dict[str, int] = {}
for residue in residues:
counts[residue] = counts.get(residue, 0) + 1
scores.append((max(counts.values()) / len(residues)) * nongap_frac)
return scores, nongaps
def reference_residue_range_to_alignment_cols(
records: FastaRecords,
reference_id: str,
start_residue: int,
end_residue: int,
) -> Tuple[int, int, str]:
if start_residue < 1:
raise ValueError("--reference-start-residue must be >= 1")
if end_residue < start_residue:
raise ValueError("--reference-end-residue must be >= --reference-start-residue")
needle = (reference_id or "").strip()
if not needle:
raise ValueError("--reference-id is required for reference-residue mode")
matched_header = None
matched_seq = None
for header, seq in records:
primary_id = header.split()[0]
if header == needle or primary_id == needle or needle in header:
matched_header = header
matched_seq = seq
break
if matched_seq is None or matched_header is None:
available = ", ".join(header.split()[0] for header, _ in records[:10])
raise ValueError(f"Reference sequence not found: {needle}. Available examples: {available}")
residue_to_col: Dict[int, int] = {}
residue_idx = 0
for col_idx, char in enumerate(matched_seq, start=1):
if char not in {"-", "."}:
residue_idx += 1
residue_to_col[residue_idx] = col_idx
if start_residue not in residue_to_col or end_residue not in residue_to_col:
raise ValueError(
f"Reference residue range {start_residue}-{end_residue} exceeds "
f"nongap length {residue_idx} for {matched_header.split()[0]}"
)
return residue_to_col[start_residue], residue_to_col[end_residue], matched_header
def best_window(
scores: Sequence[float],
nongaps: Sequence[float],
window_len: int,
min_nongap: float,
) -> Tuple[int, int, float, float]:
if window_len < 1:
raise ValueError("--window-len must be >= 1")
if not scores:
raise ValueError("cannot choose a conserved window from an empty alignment")
aln_len = len(scores)
win = min(window_len, aln_len)
best_idx: Optional[int] = None
best_score = -1.0
best_nongap = 0.0
for idx in range(0, aln_len - win + 1):
mean_nongap = sum(nongaps[idx:idx + win]) / win
if mean_nongap < min_nongap:
continue
mean_score = sum(scores[idx:idx + win]) / win
if mean_score > best_score:
best_idx = idx
best_score = mean_score
best_nongap = mean_nongap
if best_idx is None:
for idx in range(0, aln_len - win + 1):
mean_score = sum(scores[idx:idx + win]) / win
if mean_score > best_score:
best_idx = idx
best_score = mean_score
best_nongap = sum(nongaps[idx:idx + win]) / win
assert best_idx is not None
return best_idx + 1, best_idx + win, best_score, best_nongap
def extract_region(records: FastaRecords, start_col: int, end_col: int) -> FastaRecords:
aln_len = validate_alignment(records)
if start_col < 1:
raise ValueError("--start-col must be >= 1")
if end_col < start_col:
raise ValueError("--end-col must be >= --start-col")
if end_col > aln_len:
raise ValueError(f"--end-col {end_col} exceeds alignment length {aln_len}")
start_idx = start_col - 1
return [(header, seq[start_idx:end_col]) for header, seq in records]
def resolve_msa(workspace: Path, prefix: str, msa_file: Optional[str]) -> Path:
if msa_file:
path = Path(msa_file).expanduser().resolve()
else:
path = (workspace / "final" / f"{prefix}.pep.msa.fa").resolve()
if not path.exists() or path.stat().st_size == 0:
raise FileNotFoundError(f"MSA file not found or empty: {path}")
return path
def collect_outputs(output_dir: Path, outbase: str) -> List[str]:
return sorted(
str(path.resolve())
for path in output_dir.iterdir()
if path.is_file() and path.name.startswith(outbase)
)
def run_r_plot(
region_fasta: Path,
output_dir: Path,
outbase: str,
r_script: Path,
rscript: Path,
method: str,
) -> Dict[str, Any]:
output_dir.mkdir(parents=True, exist_ok=True)
cmd = [
str(rscript),
str(r_script),
str(region_fasta.resolve()),
outbase,
"--method",
method,
]
result = subprocess.run(
cmd,
cwd=str(output_dir),
capture_output=True,
text=True,
check=False,
timeout=900,
)
return {
"command": cmd,
"returncode": result.returncode,
"stdout_tail": result.stdout[-2000:],
"stderr_tail": result.stderr[-2000:],
"success": result.returncode == 0,
}
def write_profile(path: Path, scores: Sequence[float], nongaps: Sequence[float]) -> None:
with path.open("w", encoding="utf-8") as handle:
handle.write("col_1based\tscore\tnongap_frac\n")
for idx, (score, nongap) in enumerate(zip(scores, nongaps), start=1):
handle.write(f"{idx}\t{score:.6f}\t{nongap:.6f}\n")
def read_json_dict(path: Path) -> dict[str, Any]:
try:
data = json.loads(path.read_text(encoding="utf-8"))
except Exception:
return {}
return data if isinstance(data, dict) else {}
def route_from_prefix(prefix: str) -> str:
text = str(prefix or "")
if ".subfamily_evidence" in text:
return "subfamily_evidence"
if ".subfamily_prediction" in text:
return "subfamily_prediction"
return ""
def find_source_manifest(workspace: Path, prefix: str, msa_file: Path, explicit: str = "") -> tuple[dict[str, Any], str]:
candidates: list[Path] = []
if explicit:
path = Path(explicit).expanduser()
if path.is_absolute():
candidates.append(path)
else:
candidates.extend([path, PROJECT_ROOT / path, workspace / path])
final_dir = workspace / "final"
if prefix:
candidates.extend([
final_dir / f"{prefix}.phylogeny.manifest.json",
final_dir / f"{prefix}.reference_classification.manifest.json",
])
candidates.extend(sorted(final_dir.glob(f"{prefix}*.phylogeny.manifest.json"), key=lambda p: p.stat().st_mtime, reverse=True))
candidates.extend(sorted(final_dir.glob(f"{prefix}*.reference_classification.manifest.json"), key=lambda p: p.stat().st_mtime, reverse=True))
candidates.extend(sorted(final_dir.glob("*.phylogeny.manifest.json"), key=lambda p: p.stat().st_mtime, reverse=True))
seen: set[Path] = set()
msa_resolved = str(msa_file.resolve()) if msa_file.exists() else str(msa_file)
for candidate in candidates:
if candidate in seen or not candidate.exists():
continue
seen.add(candidate)
manifest = read_json_dict(candidate)
if not manifest:
continue
out_prefix = str(manifest.get("output_prefix") or (manifest.get("parameters") or {}).get("output_prefix") or "")
if prefix and out_prefix == prefix:
manifest["_manifest_path"] = str(candidate.resolve())
return manifest, str(candidate.resolve())
artifacts = manifest.get("artifacts") if isinstance(manifest.get("artifacts"), dict) else {}
paths = []
for value in artifacts.values():
values = value if isinstance(value, list) else [value]
for item in values:
if isinstance(item, dict):
paths.extend([str(item.get("full_path") or ""), str(item.get("path") or "")])
elif isinstance(item, str):
paths.append(item)
params = manifest.get("parameters") if isinstance(manifest.get("parameters"), dict) else {}
paths.extend([str(params.get("msa_file") or ""), str(params.get("scope_input_fasta") or "")])
if any(path and (path == msa_resolved or Path(path).name == msa_file.name) for path in paths):
manifest["_manifest_path"] = str(candidate.resolve())
return manifest, str(candidate.resolve())
if not prefix:
manifest["_manifest_path"] = str(candidate.resolve())
return manifest, str(candidate.resolve())
return {}, ""
def scope_from_manifest(manifest: dict[str, Any], prefix: str) -> dict[str, Any]:
if not manifest:
return {}
stage_context = manifest.get("stage_context") if isinstance(manifest.get("stage_context"), dict) else {}
scope_obj = manifest.get("scope") if isinstance(manifest.get("scope"), dict) else {}
params = manifest.get("parameters") if isinstance(manifest.get("parameters"), dict) else {}
def pick(key: str, default: Any = "") -> Any:
for source in (manifest, stage_context, scope_obj, params):
if isinstance(source, dict) and source.get(key) not in (None, "", []):
return source.get(key)
return default
scope_id = str(pick("scope_id", "") or "").strip()
if not scope_id:
return {}
family = str(pick("family", "") or "").strip()
if not family and ".target-" in scope_id:
family = scope_id.split(".target-", 1)[0]
if not family and "." in prefix:
family = prefix.split(".")[-1]
targets = [str(x) for x in (pick("target_species", []) or []) if str(x).strip()]
refs = [str(x) for x in (pick("reference_species", []) or []) if str(x).strip()]
comps = [str(x) for x in (pick("comparative_species", []) or []) if str(x).strip()]
selected = [str(x) for x in (pick("selected_species", []) or []) if str(x).strip()] or [*targets, *refs, *comps]
return {
"family": family,
"prefix": pick("prefix", prefix),
"scope_id": scope_id,
"scope_mode": str(pick("scope_mode", "") or ""),
"target_species": targets,
"reference_species": refs,
"comparative_species": comps,
"selected_species": selected,
}
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Draw a manually selected or auto-detected conserved region from final/{prefix}.pep.msa.fa"
)
parser.add_argument("--workspace", required=True, help="Workspace directory, e.g. workspace/job_xxx")
parser.add_argument("--prefix", required=True, help="Analysis prefix, e.g. Ft.SPL")
parser.add_argument("--msa-file", help="Override aligned FASTA; defaults to final/{prefix}.pep.msa.fa")
parser.add_argument("--start-col", type=int, help="1-based inclusive alignment start column")
parser.add_argument("--end-col", type=int, help="1-based inclusive alignment end column")
parser.add_argument("--auto-conserved", action="store_true", help="Automatically choose the most conserved window")
parser.add_argument("--reference-id", help="Reference sequence ID/header for residue-number mode")
parser.add_argument("--reference-start-residue", type=int, help="1-based nongap residue start in the reference sequence")
parser.add_argument("--reference-end-residue", type=int, help="1-based nongap residue end in the reference sequence")
parser.add_argument("--window-len", type=int, default=60, help="Auto mode conserved window length (default: 60)")
parser.add_argument("--min-nongap", type=float, default=0.6, help="Auto mode mean nongap threshold (default: 0.6)")
parser.add_argument("--output-dir", help="Output directory; defaults to workspace/final")
parser.add_argument("--output-prefix", help="Output prefix basename; defaults to {prefix}.msa_region_START_END")
parser.add_argument("--r-script", default=str(PROJECT_ROOT / "scripts" / "msa_figure_from_aligned.R"))
parser.add_argument("--rscript", default=str(DEFAULT_RSCRIPT if DEFAULT_RSCRIPT.exists() else shutil.which("Rscript") or "Rscript"))
parser.add_argument("--method", default="mafft", help="Method label passed to R plotting script (default: mafft)")
parser.add_argument("--scope-file", default="", help="Scope JSON used by the source phylogeny/reference-classification route")
parser.add_argument("--analysis-route", default="", help="Analysis route identity copied into the manifest")
parser.add_argument("--member-set-id", default="", help="Member-set identity copied into the manifest")
parser.add_argument("--source-manifest", default="", help="Source phylogeny/reference-classification manifest used to inherit scope identity")
parser.add_argument("--force", action="store_true", help="Overwrite existing region FASTA/metadata files")
return parser.parse_args()
def main() -> int:
args = parse_args()
workspace = Path(args.workspace).expanduser().resolve()
output_dir = Path(args.output_dir).expanduser().resolve() if args.output_dir else (workspace / "final").resolve()
msa_file = resolve_msa(workspace, args.prefix, args.msa_file)
r_script = Path(args.r_script).expanduser().resolve()
rscript = Path(args.rscript).expanduser()
source_manifest, source_manifest_path = find_source_manifest(workspace, args.prefix, msa_file, args.source_manifest)
scope_data: dict[str, Any] = {}
scope_file_path = Path(args.scope_file).expanduser() if args.scope_file else None
if scope_file_path and scope_file_path.exists():
scope_data = read_scope(scope_file_path)
if not scope_data:
scope_data = scope_from_manifest(source_manifest, args.prefix)
analysis_route = args.analysis_route or str(source_manifest.get("analysis_route") or (source_manifest.get("stage_context") or {}).get("analysis_route") or (source_manifest.get("parameters") or {}).get("analysis_route") or "") or route_from_prefix(args.prefix)
member_set_id = args.member_set_id or str(source_manifest.get("member_set_id") or (source_manifest.get("stage_context") or {}).get("member_set_id") or (source_manifest.get("parameters") or {}).get("member_set_id") or "")
if not r_script.exists():
raise FileNotFoundError(f"R plotting script not found: {r_script}")
if not rscript.exists() and shutil.which(str(rscript)) is None:
raise FileNotFoundError(f"Rscript executable not found: {rscript}")
records = read_fasta(msa_file)
aln_len = validate_alignment(records)
scores, nongaps = conservation_profile(records)
manual = args.start_col is not None or args.end_col is not None
reference_mode = args.reference_id or args.reference_start_residue is not None or args.reference_end_residue is not None
active_modes = sum([bool(manual), bool(args.auto_conserved), bool(reference_mode)])
if active_modes != 1:
raise ValueError("Use exactly one mode: --start-col/--end-col, --reference-id with reference residues, or --auto-conserved")
reference_header = None
reference_start_residue = None
reference_end_residue = None
if manual:
if args.start_col is None or args.end_col is None:
raise ValueError("Manual mode requires both --start-col and --end-col")
start_col, end_col = args.start_col, args.end_col
selection_mode = "manual"
elif reference_mode:
if not args.reference_id or args.reference_start_residue is None or args.reference_end_residue is None:
raise ValueError("Reference-residue mode requires --reference-id, --reference-start-residue, and --reference-end-residue")
reference_start_residue = args.reference_start_residue
reference_end_residue = args.reference_end_residue
start_col, end_col, reference_header = reference_residue_range_to_alignment_cols(
records=records,
reference_id=args.reference_id,
start_residue=reference_start_residue,
end_residue=reference_end_residue,
)
selection_mode = "reference_residue"
else:
start_col, end_col, mean_score, mean_nongap = best_window(
scores=scores,
nongaps=nongaps,
window_len=args.window_len,
min_nongap=args.min_nongap,
)
selection_mode = "auto_conserved"
if selection_mode != "auto_conserved":
extract_region(records, start_col, end_col)
mean_score = sum(scores[start_col - 1:end_col]) / (end_col - start_col + 1)
mean_nongap = sum(nongaps[start_col - 1:end_col]) / (end_col - start_col + 1)
region_records = extract_region(records, start_col, end_col)
default_outbase_kind = "msa_conserved_region" if selection_mode == "auto_conserved" else "msa_region"
outbase = args.output_prefix or f"{args.prefix}.{default_outbase_kind}_{start_col}_{end_col}"
region_fasta = output_dir / f"{outbase}.fa"
region_tsv = output_dir / f"{outbase}.region.tsv"
profile_tsv = output_dir / f"{outbase}.conservation_profile.tsv"
manifest_json = output_dir / f"{outbase}.manifest.json"
for path in [region_fasta, region_tsv, profile_tsv, manifest_json]:
if path.exists() and not args.force:
raise FileExistsError(f"Output exists, use --force to overwrite: {path}")
write_fasta(region_fasta, region_records)
write_profile(profile_tsv, scores, nongaps)
with region_tsv.open("w", encoding="utf-8") as handle:
handle.write("prefix\tselection_mode\tmsa_file\talignment_length\tsequence_count\tstart_col_1based\tend_col_1based_inclusive\tregion_len\tmean_score\tmean_nongap\n")
handle.write(
f"{args.prefix}\t{selection_mode}\t{msa_file}\t{aln_len}\t{len(records)}\t"
f"{start_col}\t{end_col}\t{end_col - start_col + 1}\t{mean_score:.6f}\t{mean_nongap:.6f}\n"
)
plot_result = run_r_plot(
region_fasta=region_fasta,
output_dir=output_dir,
outbase=outbase,
r_script=r_script,
rscript=rscript,
method=args.method,
)
output_files = collect_outputs(output_dir, outbase)
manifest = {
"manifest_version": "1.0",
"stage": "msa_region_plot",
"status": "success" if plot_result["success"] else "failed",
"prefix": args.prefix,
"output_prefix": outbase,
"workspace": str(workspace),
"created_at": datetime.now().isoformat(timespec="seconds"),
"parameters": {
"selection_mode": selection_mode,
"start_col_1based": start_col,
"end_col_1based_inclusive": end_col,
"window_len": args.window_len if selection_mode == "auto_conserved" else None,
"min_nongap": args.min_nongap if selection_mode == "auto_conserved" else None,
"reference_id": reference_header,
"reference_start_residue": reference_start_residue,
"reference_end_residue": reference_end_residue,
"scope_file": args.scope_file,
"source_manifest": source_manifest_path,
"source_stage": source_manifest.get("stage", "") if source_manifest else "",
"source_output_prefix": source_manifest.get("output_prefix", "") if source_manifest else "",
"analysis_route": analysis_route,
"member_set_id": member_set_id,
},
"metrics": {
"alignment_length": aln_len,
"sequence_count": len(records),
"region_length": end_col - start_col + 1,
"mean_score": mean_score,
"mean_nongap": mean_nongap,
},
"artifacts": {
"source_msa": str(msa_file),
"region_fasta": str(region_fasta),
"region_tsv": str(region_tsv),
"conservation_profile_tsv": str(profile_tsv),
"output_files": output_files,
},
"plot": plot_result,
}
context = build_stage_context(
stage="msa_region_plot",
prefix=args.prefix,
scope_data=scope_data,
scope_file=args.scope_file,
output_prefix=outbase,
analysis_route=analysis_route,
member_set_id=member_set_id,
)
apply_stage_context(manifest, context)
if source_manifest_path:
manifest["source_manifest"] = source_manifest_path
manifest_json.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
print("MSA region plot finished.")
print(f" status: {manifest['status']}")
print(f" source_msa: {msa_file}")
print(f" region: {start_col}-{end_col} ({selection_mode})")
print(f" mean_score: {mean_score:.6f}")
print(f" mean_nongap: {mean_nongap:.6f}")
print(f" region_fasta: {region_fasta}")
print(f" manifest: {manifest_json}")
print(" outputs:")
for path in output_files:
print(f" - {path}")
return 0 if plot_result["success"] else 1
if __name__ == "__main__":
try:
raise SystemExit(main())
except Exception as exc:
print(f"ERROR: {exc}", file=sys.stderr)
raise SystemExit(1)