Files
PlantAI-dev/apps/cli/run_reference_classification.py
T

1811 lines
77 KiB
Python
Raw Normal View History

2026-08-03 18:42:47 +08:00
#!/usr/bin/env python3
"""Run reference-guided subfamily classification from a joint phylogenetic tree."""
from __future__ import annotations
import argparse
import copy
import csv
import json
import os
import re
import shutil
import subprocess
import sys
from collections import Counter
from datetime import datetime
from pathlib import Path
from typing import Iterable
PROJECT_ROOT = Path(__file__).resolve().parents[2]
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
CONDA_BIN = "/home/xzq/miniconda3/envs/test4py311/bin"
if CONDA_BIN not in os.environ.get("PATH", ""):
os.environ["PATH"] = CONDA_BIN + ":" + os.environ.get("PATH", "")
from gfa_core.gene_family_scope import read_scope, sanitize_token
from gfa_core.gene_family_stage_context import apply_stage_context, build_stage_context, user_output_prefix_for, write_legacy_manifest_copy
def resolve_path(path_text: str | None, base_dirs: list[Path]) -> Path | None:
if not path_text:
return None
path = Path(path_text)
if path.is_absolute() and path.exists():
return path
candidates = [path]
for base in base_dirs:
candidates.append(base / path)
candidates.append(base / path.name)
for candidate in candidates:
if candidate.exists():
return candidate.resolve()
return None
def has_two_column_data(path: Path) -> bool:
try:
with path.open(encoding="utf-8", errors="ignore") as handle:
for raw in handle:
line = raw.strip()
if not line or line.startswith("#"):
continue
return len(line.split()) >= 2
except OSError:
return False
return False
def read_reference_gene_ids(path: Path) -> set[str]:
genes: set[str] = set()
try:
with path.open(encoding="utf-8", errors="ignore") as handle:
for raw in handle:
line = raw.strip()
if not line or line.startswith("#"):
continue
parts = line.split()
if len(parts) >= 2:
genes.add(parts[0])
genes.add(parts[0].lower())
except OSError:
pass
return genes
def read_reference_gene_labels(path: Path, reference_prefix: str) -> set[str]:
labels: set[str] = set()
for gene in read_reference_gene_ids(path):
labels.add(gene)
labels.add(gene.lower())
labels.add(f"{reference_prefix}|{gene}")
labels.add(f"{reference_prefix}|{gene}".lower())
return labels
def fasta_records(path: Path) -> Iterable[tuple[str, str]]:
header = ""
seq_parts: list[str] = []
with path.open(encoding="utf-8", errors="ignore") as handle:
for raw in handle:
line = raw.strip()
if not line:
continue
if line.startswith(">"):
if header:
yield header, "".join(seq_parts)
header = line[1:].strip()
seq_parts = []
else:
seq_parts.append(line)
if header:
yield header, "".join(seq_parts)
def fasta_id(header: str) -> str:
return header.split(None, 1)[0].strip()
def read_fasta_ids(path: Path) -> set[str]:
ids: set[str] = set()
if not path.exists() or path.stat().st_size == 0:
return ids
for header, _seq in fasta_records(path):
gene = fasta_id(header)
if gene:
ids.add(gene)
return ids
def gene_id_aliases(gene_id: str) -> set[str]:
gene = str(gene_id or "").strip()
if not gene:
return set()
aliases = {gene}
aliases.add(gene.replace(":", "_"))
aliases.add(gene.replace("_", ":", 1))
aliases.add(gene.replace("|", "_"))
aliases.add(gene.replace("_", "|", 1))
return {item for item in aliases if item}
def clean_protein_sequence(seq: str) -> str:
return "".join(ch for ch in seq.replace("*", "").upper() if ch.isalpha() or ch == "-")
def find_reference_subfamily(
workspace: Path,
input_dir: Path,
prefix: str,
reference_prefix: str = "",
) -> Path | None:
final_dir = workspace / "final"
search_dirs = [input_dir, final_dir]
candidates: list[Path] = []
family = prefix.split(".")[-1] if "." in prefix else prefix
preferred_names = []
if reference_prefix and family:
preferred_names.extend([
f"{reference_prefix}.{family}.subfamily",
f"{reference_prefix}.{family}.subfamily.tsv",
f"{reference_prefix}.{family}.geneID.subfamily",
])
for directory in search_dirs:
for name in preferred_names:
candidate = directory / name
if candidate.exists() and has_two_column_data(candidate):
return candidate.resolve()
for directory in search_dirs:
if not directory.exists():
continue
for pattern in ("*.subfamily", "*.subfamily.tsv", "*subfamily*.tsv", "*subfamily*.txt"):
candidates.extend(directory.glob(pattern))
generated_names = {
f"{prefix}.geneID.subfamily",
f"{prefix}.reference_classification.subfamily.tsv",
}
filtered = []
for candidate in candidates:
if candidate.name in generated_names:
continue
if ".reference_classification." in candidate.name:
continue
if has_two_column_data(candidate):
filtered.append(candidate)
if not filtered:
return None
return sorted(filtered, key=lambda p: (p.stat().st_mtime, str(p)), reverse=True)[0].resolve()
def find_target_proteins(final_dir: Path, prefix: str, workspace: Path | None = None, scope_data: dict | None = None) -> Path | None:
candidates = [
final_dir / f"{prefix}.final_members.pep.clean.fa",
final_dir / f"{prefix}.final_members.pep.fa",
final_dir / f"{prefix}.final_members.pep.fasta",
]
if workspace is not None and scope_data:
family = str(scope_data.get("family") or (prefix.split(".")[-1] if "." in prefix else "")).strip()
targets = [str(x) for x in scope_data.get("target_species") or [] if str(x).strip()]
if family and len(targets) == 1:
species = targets[0]
candidates[0:0] = [
final_dir / f"Target.{species}.{family}.final_members.pep.clean.fa",
final_dir / f"Target.{species}.{family}.final_members.pep.fa",
final_dir / f"{species}.{family}.final_members.pep.clean.fa",
final_dir / f"{species}.{family}.final_members.pep.fa",
]
scope_id = str(scope_data.get("scope_id") or "").strip()
manifest_candidates = []
if scope_id:
manifest_candidates.append(final_dir / f"{prefix}.{sanitize_token(scope_id)}.phylogeny.manifest.json")
else:
manifest_candidates.append(final_dir / f"{prefix}.phylogeny.manifest.json")
for manifest_path in manifest_candidates:
if not manifest_path.exists() or manifest_path.stat().st_size == 0:
continue
try:
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
except Exception:
continue
scope_input = str(manifest.get("parameters", {}).get("scope_input_fasta") or "").strip()
if scope_input:
path = Path(scope_input)
if not path.is_absolute():
path = workspace / path
candidates.append(path)
for candidate in candidates:
if candidate.exists() and candidate.stat().st_size > 0:
return candidate
return None
def find_reference_proteins(workspace: Path, input_dir: Path, reference_prefix: str, family: str) -> Path | None:
final_dir = workspace / "final"
search_dirs = [final_dir, input_dir]
names = [
f"{reference_prefix}.{family}.pep.fa",
f"{reference_prefix}.{family}.pep.fasta",
f"{reference_prefix}.{family}.sequence_extract.pep.fa",
f"{reference_prefix}.{family}.sequence_extract.pep.fasta",
f"{reference_prefix}.{family}.ncbi_fetch.protein.fa",
f"{reference_prefix}.{family}.protein.fa",
f"{reference_prefix}.{family}.protein.fasta",
]
for directory in search_dirs:
for name in names:
candidate = directory / name
if candidate.exists() and candidate.stat().st_size > 0:
return candidate
for directory in search_dirs:
for pattern in (
f"{reference_prefix}.{family}*.sequence_extract.pep.fa",
f"{reference_prefix}.{family}*.protein.fa",
f"{reference_prefix}.{family}*.pep.fa",
f"{reference_prefix}.{family}*.fa",
f"{reference_prefix}.{family}*.fasta",
):
for candidate in sorted(directory.glob(pattern)):
if candidate.exists() and candidate.stat().st_size > 0:
return candidate
return None
def resolve_manifest_path(workspace: Path, relative_or_absolute: str) -> Path | None:
if not relative_or_absolute:
return None
path = Path(relative_or_absolute)
candidates = [path] if path.is_absolute() else [workspace / path, workspace / "final" / path, path]
for candidate in candidates:
if candidate.exists() and candidate.stat().st_size > 0:
return candidate.resolve()
return None
def collect_target_ids_from_scope(workspace: Path, prefix: str, scope_data: dict | None) -> tuple[set[str], dict[str, int], list[str]]:
if not scope_data:
return set(), {}, []
target_species = [str(item) for item in scope_data.get("target_species") or [] if str(item).strip()]
comparative_species = [str(item) for item in scope_data.get("comparative_species") or [] if str(item).strip()]
for item in comparative_species:
if item not in target_species:
target_species.append(item)
ids: set[str] = set()
counts: dict[str, int] = {}
warnings: list[str] = []
family = str(scope_data.get("family") or (prefix.split(".")[-1] if "." in prefix else "")).strip()
for species in target_species:
refined_paths = [
workspace / "final" / f"Target.{species}.{family}.final_members.pep.fa",
workspace / "final" / f"{species}.{family}.final_members.pep.fa",
workspace / "final" / f"{species}.final_members.pep.fa",
]
for path in refined_paths:
if path.exists() and path.stat().st_size > 0:
species_ids = read_fasta_ids(path)
ids.update(species_ids)
counts[species] = len(species_ids)
break
manifest_candidates = [workspace / "final" / f"{prefix}.identify.manifest.json"]
scope_id = str(scope_data.get("scope_id") or "").strip()
if scope_id:
manifest_candidates.append(workspace / "final" / f"{prefix}.{sanitize_token(scope_id)}.identify.manifest.json")
identify_manifest = next((path for path in manifest_candidates if path.exists() and path.stat().st_size > 0), None)
if identify_manifest:
try:
manifest = json.loads(identify_manifest.read_text(encoding="utf-8"))
for result in manifest.get("species_results", []):
species = str(result.get("species") or "").strip()
if species not in target_species:
continue
if species in counts:
continue
final_pep = resolve_manifest_path(workspace, str(result.get("final_pep_file") or ""))
if final_pep:
species_ids = read_fasta_ids(final_pep)
ids.update(species_ids)
counts[species] = len(species_ids)
except Exception as exc:
warnings.append(f"Failed to read identify manifest for scope filtering: {exc}")
for species in target_species:
if species in counts:
continue
fallback_paths = [
workspace / "identify" / species / "final" / f"{species}.final_gfmembers.pep.fa",
workspace / "final" / f"{species}.final_gfmembers.pep.fa",
workspace / "final" / f"{species}.final_members.pep.fa",
]
for path in fallback_paths:
if path.exists() and path.stat().st_size > 0:
species_ids = read_fasta_ids(path)
ids.update(species_ids)
counts[species] = len(species_ids)
break
if species not in counts:
warnings.append(f"No target member FASTA found for scope species {species}")
return ids, counts, warnings
def collect_target_ids_by_species(workspace: Path, prefix: str, scope_data: dict | None) -> tuple[dict[str, set[str]], list[str]]:
"""Return target member IDs grouped by target species for per-species outputs."""
if not scope_data:
return {}, []
target_species = [str(item) for item in scope_data.get("target_species") or [] if str(item).strip()]
comparative_species = [str(item) for item in scope_data.get("comparative_species") or [] if str(item).strip()]
for item in comparative_species:
if item not in target_species:
target_species.append(item)
grouped: dict[str, set[str]] = {species: set() for species in target_species}
warnings: list[str] = []
family = str(scope_data.get("family") or (prefix.split(".")[-1] if "." in prefix else "")).strip()
for species in target_species:
refined_paths = [
workspace / "final" / f"Target.{species}.{family}.final_members.pep.fa",
workspace / "final" / f"{species}.{family}.final_members.pep.fa",
workspace / "final" / f"{species}.final_members.pep.fa",
]
for path in refined_paths:
if path.exists() and path.stat().st_size > 0:
grouped[species].update(read_fasta_ids(path))
break
manifest_candidates = [workspace / "final" / f"{prefix}.identify.manifest.json"]
scope_id = str(scope_data.get("scope_id") or "").strip()
if scope_id:
manifest_candidates.append(workspace / "final" / f"{prefix}.{sanitize_token(scope_id)}.identify.manifest.json")
identify_manifest = next((path for path in manifest_candidates if path.exists() and path.stat().st_size > 0), None)
if identify_manifest:
try:
manifest = json.loads(identify_manifest.read_text(encoding="utf-8"))
for result in manifest.get("species_results", []):
species = str(result.get("species") or "").strip()
if species not in grouped or grouped[species]:
continue
final_pep = resolve_manifest_path(workspace, str(result.get("final_pep_file") or ""))
if final_pep:
grouped[species].update(read_fasta_ids(final_pep))
except Exception as exc:
warnings.append(f"Failed to read identify manifest for per-species subfamily outputs: {exc}")
for species in target_species:
if grouped.get(species):
continue
fallback_paths = [
workspace / "identify" / species / "final" / f"{species}.final_gfmembers.pep.fa",
workspace / "final" / f"Target.{species}.{scope_data.get('family', '')}.final_members.pep.fa",
workspace / "final" / f"{species}.{scope_data.get('family', '')}.final_members.pep.fa",
workspace / "final" / f"{species}.final_members.pep.fa",
]
for path in fallback_paths:
if path.exists() and path.stat().st_size > 0:
grouped[species].update(read_fasta_ids(path))
break
if not grouped.get(species):
warnings.append(f"No target member FASTA found for per-species subfamily output: {species}")
return grouped, warnings
def rewrite_tsv_rows(path: Path, fieldnames: list[str], rows: list[dict[str, str]]) -> None:
with path.open("w", encoding="utf-8", newline="") as handle:
writer = csv.DictWriter(handle, fieldnames=fieldnames, delimiter="\t")
writer.writeheader()
writer.writerows(rows)
def filter_reference_classification_outputs(
assignment_out: Path,
subfamily_out: Path,
reference_subfamily: Path,
reference_prefix: str,
allowed_target_ids: set[str],
) -> dict[str, int | bool]:
if not allowed_target_ids:
return {"applied": False, "assignment_rows_before": 0, "assignment_rows_after": 0, "dropped_assignment_rows": 0}
with assignment_out.open(encoding="utf-8", newline="") as handle:
reader = csv.DictReader(handle, delimiter="\t")
fieldnames = list(reader.fieldnames or [])
rows = list(reader)
allowed_target_aliases = set()
for gene_id in allowed_target_ids:
allowed_target_aliases.update(gene_id_aliases(gene_id))
filtered_rows = [row for row in rows if row.get("target_gene", "") in allowed_target_aliases]
if fieldnames:
rewrite_tsv_rows(assignment_out, fieldnames, filtered_rows)
reference_labels = read_reference_gene_labels(reference_subfamily, reference_prefix)
subfamily_rows: list[tuple[str, str]] = []
before_subfamily = 0
with subfamily_out.open(encoding="utf-8", errors="ignore") as handle:
for raw in handle:
line = raw.strip()
if not line or line.startswith("#"):
continue
before_subfamily += 1
parts = line.split()
if len(parts) < 2:
continue
label, subfamily = parts[0], parts[1]
label_aliases = gene_id_aliases(label)
if label_aliases & allowed_target_aliases or label in reference_labels or label.lower() in reference_labels:
subfamily_rows.append((label, subfamily))
with subfamily_out.open("w", encoding="utf-8") as handle:
for label, subfamily in sorted(subfamily_rows):
handle.write(f"{label}\t{subfamily}\n")
return {
"applied": True,
"assignment_rows_before": len(rows),
"assignment_rows_after": len(filtered_rows),
"dropped_assignment_rows": len(rows) - len(filtered_rows),
"subfamily_rows_before": before_subfamily,
"subfamily_rows_after": len(subfamily_rows),
}
def write_per_species_subfamily_outputs(
*,
workspace: Path,
final_dir: Path,
prefix: str,
output_prefix: str,
assignment_out: Path,
subfamily_out: Path,
confidence_out: Path,
reference_subfamily: Path,
reference_prefix: str,
scope_data: dict | None,
) -> dict:
"""Write one subfamily result set per target species.
The combined output remains the canonical full-scope artifact. These
per-species files make downstream single-species plotting and manual file
inspection unambiguous.
"""
if not scope_data:
return {"written": 0, "files": [], "warnings": []}
family = str(scope_data.get("family") or "").strip() or (prefix.split(".")[-1] if "." in prefix else "")
target_species = [str(item) for item in scope_data.get("target_species") or [] if str(item).strip()]
if not family or not target_species:
return {"written": 0, "files": [], "warnings": ["missing family or target_species for per-species outputs"]}
ids_by_species, warnings = collect_target_ids_by_species(workspace, prefix, scope_data)
reference_labels = read_reference_gene_labels(reference_subfamily, reference_prefix)
with assignment_out.open(encoding="utf-8", newline="") as handle:
reader = csv.DictReader(handle, delimiter="\t")
assignment_fields = list(reader.fieldnames or [])
assignment_rows = list(reader)
with confidence_out.open(encoding="utf-8", newline="") as handle:
reader = csv.DictReader(handle, delimiter="\t")
confidence_fields = list(reader.fieldnames or [])
confidence_rows = list(reader)
subfamily_rows: list[tuple[str, str]] = []
with subfamily_out.open(encoding="utf-8", errors="ignore") as handle:
for raw in handle:
line = raw.strip()
if not line or line.startswith("#"):
continue
parts = line.split()
if len(parts) >= 2:
subfamily_rows.append((parts[0], parts[1]))
written_files: list[str] = []
for species in target_species:
species_ids = ids_by_species.get(species) or set()
if not species_ids:
continue
species_id_aliases = set()
for gene_id in species_ids:
species_id_aliases.update(gene_id_aliases(gene_id))
species_prefix = f"{species}.{family}.target_plus_reference"
species_assignment = final_dir / f"{species_prefix}.reference_classification.assignment.tsv"
species_subfamily = final_dir / f"{species_prefix}.reference_classification.subfamily.tsv"
species_confidence = final_dir / f"{species_prefix}.reference_classification.confidence.tsv"
species_compat = final_dir / f"{species_prefix}.geneID.subfamily"
species_manifest = final_dir / f"{species_prefix}.reference_classification.manifest.json"
species_assignment_rows = [row for row in assignment_rows if row.get("target_gene", "") in species_id_aliases]
species_confidence_rows = [row for row in confidence_rows if row.get("gene_id", "") in species_id_aliases]
species_subfamily_rows = [
(label, subfamily)
for label, subfamily in subfamily_rows
if gene_id_aliases(label) & species_id_aliases or label in reference_labels or label.lower() in reference_labels
]
if assignment_fields:
rewrite_tsv_rows(species_assignment, assignment_fields, species_assignment_rows)
if confidence_fields:
rewrite_tsv_rows(species_confidence, confidence_fields, species_confidence_rows)
with species_subfamily.open("w", encoding="utf-8") as handle:
for label, subfamily in sorted(species_subfamily_rows):
handle.write(f"{label}\t{subfamily}\n")
shutil.copyfile(species_subfamily, species_compat)
manifest_payload = {
"manifest_version": "1.0",
"stage": "reference_classification",
"status": "success",
"prefix": f"{species}.{family}",
"output_prefix": species_prefix,
"parent_output_prefix": output_prefix,
"timestamp": datetime.now().isoformat(timespec="seconds"),
"target_species": [species],
"reference_species": scope_data.get("reference_species") or [],
"comparative_species": [],
"scope_mode": "target_plus_reference",
"scope_id": f"{family}.target-{sanitize_token(species)}.ref-{'-'.join(sanitize_token(x) for x in scope_data.get('reference_species') or [])}.target_plus_reference",
"analysis_route": "subfamily_prediction",
"artifacts": {
"assignment_table": {"path": f"final/{species_assignment.name}", "type": "tsv"},
"subfamily_file": {"path": f"final/{species_subfamily.name}", "type": "tsv"},
"geneID_subfamily_compat": {"path": f"final/{species_compat.name}", "type": "tsv"},
"confidence_table": {"path": f"final/{species_confidence.name}", "type": "tsv"},
},
"metrics": {
"target_gene_count": len(species_assignment_rows),
"subfamily_rows": len(species_subfamily_rows),
"source": "split_from_parent_reference_classification",
},
}
species_manifest.write_text(json.dumps(manifest_payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
written_files.extend([
str(species_assignment),
str(species_subfamily),
str(species_confidence),
str(species_compat),
str(species_manifest),
])
return {"written": len(written_files), "files": written_files, "warnings": warnings}
def tree_tip_counts(tree_file: Path, reference_subfamily: Path, target_prefix: str, reference_prefix: str) -> dict[str, int]:
try:
from Bio import Phylo
ref_genes = read_reference_gene_ids(reference_subfamily)
tree = Phylo.read(tree_file, "newick")
terminals = tree.get_terminals()
except Exception:
return {"total": 0, "target": 0, "reference": 0}
reference = 0
target = 0
for tip in terminals:
name = tip.name or ""
raw = name.split("|", 1)[1] if "|" in name else name
is_ref = name.startswith(f"{reference_prefix}|") or raw in ref_genes or raw.lower() in ref_genes
if is_ref:
reference += 1
elif name.startswith(f"{target_prefix}|") or not name.startswith(f"{reference_prefix}|"):
target += 1
return {"total": len(terminals), "target": target, "reference": reference}
def write_joint_fasta(
target_proteins: Path,
reference_proteins: Path,
reference_subfamily: Path,
reference_prefix: str,
out_fasta: Path,
) -> tuple[int, int]:
ref_genes = read_reference_gene_ids(reference_subfamily)
target_count = 0
ref_count = 0
seen: set[str] = set()
out_fasta.parent.mkdir(parents=True, exist_ok=True)
with out_fasta.open("w", encoding="utf-8") as out:
for header, seq in fasta_records(target_proteins):
gene = fasta_id(header)
clean_seq = clean_protein_sequence(seq)
if not gene or not clean_seq:
continue
if gene in seen:
continue
seen.add(gene)
out.write(f">{gene}\n{clean_seq}\n")
target_count += 1
for header, seq in fasta_records(reference_proteins):
gene = fasta_id(header)
if gene not in ref_genes and gene.lower() not in ref_genes:
continue
clean_seq = clean_protein_sequence(seq)
if not gene or not clean_seq:
continue
label = f"{reference_prefix}|{gene}"
if label in seen:
continue
seen.add(label)
out.write(f">{label}\n{clean_seq}\n")
ref_count += 1
return target_count, ref_count
def run_logged(cmd: list[str], stdout_path: Path | None = None) -> None:
print("[CMD]", " ".join(cmd))
if stdout_path is None:
result = subprocess.run(cmd, text=True, capture_output=True, cwd=str(PROJECT_ROOT), check=False)
else:
with stdout_path.open("w", encoding="utf-8") as out_handle:
result = subprocess.run(cmd, text=True, stdout=out_handle, stderr=subprocess.PIPE, cwd=str(PROJECT_ROOT), check=False)
if result.stdout:
print(result.stdout)
if result.stderr:
print(result.stderr, file=sys.stderr)
if result.returncode != 0:
raise RuntimeError(f"Command failed ({result.returncode}): {' '.join(cmd)}")
def run_subfamily_tree_plot(
tree_file: Path,
subfamily_file: Path,
final_dir: Path,
prefix: str,
layout: str = "both",
target_species: str = "",
reference_species: str = "",
) -> dict[str, object]:
"""Draw the final phylogeny again with reference-guided subfamily labels."""
script = PROJECT_ROOT / "scripts" / "tree_plot.R"
rscript = Path("/home/xzq/miniconda3/envs/test4py311/bin/Rscript")
if not rscript.exists():
rscript = Path("Rscript")
output_pdf = final_dir / f"{prefix}.reference_classification.tree.pdf"
expected = [
output_pdf,
output_pdf.with_suffix(".png"),
final_dir / f"{prefix}.reference_classification.tree.circular_tree.pdf",
final_dir / f"{prefix}.reference_classification.tree.circular_tree.png",
]
if not script.exists():
return {"success": False, "error": f"Tree plot script not found: {script}", "output_files": []}
if not tree_file.exists():
return {"success": False, "error": f"Tree file not found: {tree_file}", "output_files": []}
if not subfamily_file.exists():
return {"success": False, "error": f"Subfamily file not found: {subfamily_file}", "output_files": []}
cmd = [
str(rscript),
str(script),
str(tree_file.resolve()),
str(subfamily_file.resolve()),
str(output_pdf.resolve()),
layout,
]
if target_species or reference_species:
cmd.extend([target_species, reference_species])
print("[CMD]", " ".join(cmd))
result = subprocess.run(cmd, text=True, capture_output=True, cwd=str(final_dir), check=False)
if result.stdout:
print(result.stdout)
if result.stderr:
print(result.stderr, file=sys.stderr)
output_files = [str(path.resolve()) for path in expected if path.exists() and path.stat().st_size > 0]
return {
"success": result.returncode == 0,
"command": " ".join(cmd),
"returncode": result.returncode,
"stdout": result.stdout[-1000:] if result.stdout else "",
"stderr": result.stderr[-1000:] if result.stderr else "",
"output_files": output_files,
"layout": layout,
"tree_file": str(tree_file.resolve()),
"subfamily_file": str(subfamily_file.resolve()),
"error": "" if result.returncode == 0 else (result.stderr or result.stdout or "tree plotting failed")[-1000:],
}
def read_subfamily_labels(subfamily_file: Path) -> set[str]:
labels: set[str] = set()
with subfamily_file.open(encoding="utf-8", errors="ignore") as handle:
for raw in handle:
line = raw.strip()
if not line or line.startswith("#"):
continue
parts = line.split()
if parts:
labels.add(parts[0])
return labels
def find_publication_phylogeny_tree(final_dir: Path, species_prefix: str, reference_species: str) -> Path | None:
species = species_prefix.split(".", 1)[0]
patterns = [
f"{species_prefix}.*target-{species}.ref-{reference_species}.target_plus_reference.tree.nwk",
f"{species_prefix}.*target-{species}.ref-{reference_species}.target_plus_reference.iqtree.treefile",
f"{species_prefix}*target-{species}*ref-{reference_species}*target_plus_reference.tree.nwk",
]
candidates: list[Path] = []
for pattern in patterns:
candidates.extend(final_dir.glob(pattern))
candidates = [p for p in candidates if p.exists() and p.stat().st_size > 0 and "reference_classification" not in p.name]
if not candidates:
return None
candidates.sort(key=lambda item: (".tree.nwk" not in item.name, -item.stat().st_mtime))
return candidates[0]
def normalize_reference_tip_labels_for_plot(tree_file: Path, reference_subfamily: Path, reference_prefix: str, output_tree: Path) -> tuple[Path, dict[str, object]]:
from Bio import Phylo
ref_genes: set[str] = set()
with reference_subfamily.open(encoding="utf-8") as handle:
for raw_line in handle:
line = raw_line.strip()
if not line or line.startswith("#"):
continue
label = line.split()[0]
if label.startswith(f"{reference_prefix}|"):
gene = label.split("|", 1)[1]
ref_genes.add(gene)
ref_genes.add(gene.lower())
tree = Phylo.read(tree_file, "newick")
renamed = 0
for tip in tree.get_terminals():
name = tip.name or ""
if name.startswith(f"{reference_prefix}|"):
continue
raw = name.split("|", 1)[1] if "|" in name else name
if raw in ref_genes or raw.lower() in ref_genes:
tip.name = f"{reference_prefix}|{raw}"
renamed += 1
output_tree.parent.mkdir(parents=True, exist_ok=True)
Phylo.write(tree, output_tree, "newick")
return output_tree, {"source_tree": str(tree_file), "renamed_reference_tips": renamed}
def prune_tree_to_labels(tree_file: Path, keep_labels: set[str], output_tree: Path) -> tuple[bool, int, int, str]:
"""Write a tree containing only tips listed in keep_labels."""
try:
from Bio import Phylo
tree = Phylo.read(tree_file, "newick")
original_tips = list(tree.get_terminals())
keep_names = {tip.name for tip in original_tips if tip.name in keep_labels}
if len(keep_names) < 3:
return False, len(original_tips), len(keep_names), "fewer than 3 matching tree tips"
pruned = copy.deepcopy(tree)
for tip in list(pruned.get_terminals()):
if tip.name not in keep_names:
pruned.prune(tip)
output_tree.parent.mkdir(parents=True, exist_ok=True)
Phylo.write(pruned, output_tree, "newick")
return True, len(original_tips), len(keep_names), ""
except Exception as exc:
return False, 0, 0, str(exc)
def tree_label_match_summary(tree_file: Path, keep_labels: set[str]) -> dict[str, object]:
"""Summarize whether a tree exactly matches the expected subfamily labels."""
try:
from Bio import Phylo
tree = Phylo.read(tree_file, "newick")
tips = [tip.name or "" for tip in tree.get_terminals()]
tip_set = set(tips)
missing = sorted(label for label in keep_labels if label not in tip_set)
extra = sorted(name for name in tip_set if name not in keep_labels)
return {
"success": True,
"total_tips": len(tips),
"matched_tips": sum(1 for name in tips if name in keep_labels),
"expected_tips": len(keep_labels),
"missing_labels": missing,
"extra_labels": extra,
}
except Exception as exc:
return {
"success": False,
"error": str(exc),
"total_tips": 0,
"matched_tips": 0,
"expected_tips": len(keep_labels),
"missing_labels": sorted(keep_labels),
"extra_labels": [],
}
def run_per_species_subfamily_tree_plots(
*,
tree_file: Path,
final_dir: Path,
per_species_outputs: dict,
layout: str = "both",
reference_species: str = "",
output_prefix: str = "",
) -> dict:
"""Create short-name target+reference subfamily confirmation plots per target species."""
files = [Path(str(path)) for path in per_species_outputs.get("files") or []]
subfamily_files = [
path for path in files
if path.name.endswith(".target_plus_reference.geneID.subfamily")
or path.name.endswith(".target_plus_reference.reference_classification.subfamily.tsv")
]
# Prefer the compatibility name; it is what downstream single-species stages
# already search for.
preferred: dict[str, Path] = {}
for path in subfamily_files:
if not path.exists() or path.stat().st_size == 0:
continue
prefix = path.name
if ".target_plus_reference." not in prefix:
continue
species_prefix = prefix.split(".target_plus_reference.", 1)[0]
if species_prefix not in preferred or path.name.endswith(".geneID.subfamily"):
preferred[species_prefix] = path
results: list[dict[str, object]] = []
for species_prefix, subfamily_file in sorted(preferred.items()):
keep_labels = read_subfamily_labels(subfamily_file)
plot_prefix = output_prefix if output_prefix and len(preferred) == 1 else f"{species_prefix}.target_plus_reference"
pruned_tree = final_dir / f"{plot_prefix}.reference_classification.tree.nwk"
bootstrap_source_tree = find_publication_phylogeny_tree(final_dir, species_prefix, reference_species)
tree_source_info: dict[str, object] = {}
if bootstrap_source_tree is not None:
normalized_tree = final_dir / f"{species_prefix}.target_plus_reference.reference_classification.tree.normalized_tmp.nwk"
try:
_, tree_source_info = normalize_reference_tip_labels_for_plot(
bootstrap_source_tree,
subfamily_file,
reference_species,
normalized_tree,
)
ok, original_tips, kept_tips, error = prune_tree_to_labels(normalized_tree, keep_labels, pruned_tree)
tree_source_info.update({
"normalized_tree": str(normalized_tree),
"pruned_tree": str(pruned_tree),
"expected_tips": len(keep_labels),
})
except Exception as exc:
ok, original_tips, kept_tips, error = prune_tree_to_labels(tree_file, keep_labels, pruned_tree)
tree_source_info = {"source_tree": str(tree_file), "fallback_reason": str(exc)}
finally:
try:
normalized_tree.unlink(missing_ok=True)
except Exception:
pass
else:
ok, original_tips, kept_tips, error = prune_tree_to_labels(tree_file, keep_labels, pruned_tree)
tree_source_info = {"source_tree": str(tree_file), "fallback_reason": "publication phylogeny tree not found"}
if ok and kept_tips != len(keep_labels):
ok = False
error = f"pruned tree kept {kept_tips}/{len(keep_labels)} expected subfamily labels"
if not ok:
results.append({
"prefix": species_prefix,
"success": False,
"subfamily_file": str(subfamily_file),
"tree_file": str(tree_file),
"original_tips": original_tips,
"kept_tips": kept_tips,
"error": error,
})
continue
target_species = species_prefix.split(".", 1)[0]
plot_result = run_subfamily_tree_plot(
tree_file=pruned_tree,
subfamily_file=subfamily_file,
final_dir=final_dir,
prefix=plot_prefix,
layout=layout,
target_species=target_species,
reference_species=reference_species,
)
plot_result.update({
"prefix": species_prefix,
"pruned_tree": str(pruned_tree),
"original_tips": original_tips,
"kept_tips": kept_tips,
"tree_source": tree_source_info,
})
results.append(plot_result)
manifest_path = final_dir / f"{plot_prefix}.reference_classification.manifest.json"
if manifest_path.exists():
try:
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
artifacts = manifest.setdefault("artifacts", {})
artifacts["subfamily_tree"] = {
"path": f"final/{pruned_tree.name}",
"type": "newick",
"description": "Per-target target+reference tree used for subfamily confirmation plot",
}
for output_file in plot_result.get("output_files", []) or []:
p = Path(str(output_file))
plot_kind = "circular" if ".circular_tree" in p.name else "rectangular"
artifacts[f"subfamily_tree_plot_{plot_kind}_{p.suffix.lower().lstrip('.')}"] = {
"path": f"final/{p.name}",
"type": "image" if p.suffix.lower() == ".png" else "pdf",
"description": "Per-target reference-guided subfamily confirmation tree",
}
metrics = manifest.setdefault("metrics", {})
metrics["subfamily_tree_plot_success"] = bool(plot_result.get("success"))
metrics["subfamily_tree_kept_tips"] = kept_tips
manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
except Exception as exc:
print(f"[WARNING] Failed to update per-species manifest {manifest_path}: {exc}", file=sys.stderr)
return {
"success": any(bool(item.get("success")) for item in results),
"count": len(results),
"results": results,
}
def latest_stage_manifest(final_dir: Path, output_prefix: str, stage: str) -> Path | None:
exact = final_dir / f"{output_prefix}.{stage}.manifest.json"
if exact.exists() and exact.stat().st_size > 0:
return exact
candidates = sorted(
final_dir.glob(f"*.{stage}.manifest.json"),
key=lambda p: p.stat().st_mtime,
reverse=True,
)
for path in candidates:
if output_prefix in path.name and path.stat().st_size > 0:
return path
for path in candidates:
if path.stat().st_size <= 0:
continue
try:
data = json.loads(path.read_text(encoding="utf-8"))
except Exception:
continue
if data.get("stage") == stage and data.get("status") in {"success", "partial_success"}:
return path
return None
def latest_species_evidence_manifest(final_dir: Path, species_prefix: str, stage: str) -> Path | None:
"""Find motif/domain/structure evidence manifest for subfamily confirmation.
Prefer the target+reference evidence scope so reference genes appear in the
combo evidence panels. Fall back to target-only products when a reference
structure input is unavailable.
"""
species = species_prefix.split(".", 1)[0]
candidates = sorted(
final_dir.glob(f"*.{stage}.manifest.json"),
key=lambda p: p.stat().st_mtime,
reverse=True,
)
fallback: Path | None = None
for path in candidates:
if path.stat().st_size <= 0:
continue
try:
data = json.loads(path.read_text(encoding="utf-8"))
except Exception:
continue
if data.get("stage") != stage or data.get("status") not in {"success", "partial_success"}:
continue
target_species = [str(x) for x in data.get("target_species") or data.get("scope", {}).get("target_species") or []]
reference_species = [str(x) for x in data.get("reference_species") or data.get("scope", {}).get("reference_species") or []]
scope_mode = str(data.get("scope_mode") or data.get("scope", {}).get("scope_mode") or "")
analysis_route = str(data.get("analysis_route") or data.get("parameters", {}).get("analysis_route") or "")
output_prefix = str(data.get("output_prefix") or data.get("parameters", {}).get("output_prefix") or "")
prefix = str(data.get("prefix") or "")
name = path.name
species_match = (
target_species == [species]
or species in target_species
or output_prefix.startswith(f"{species_prefix}.")
or prefix == species_prefix
)
evidence_match = (
analysis_route == "subfamily_evidence"
or (scope_mode == "target_plus_reference" and bool(reference_species))
or ".target_plus_reference" in output_prefix
)
target_only_match = scope_mode == "target_only" or ".target_only" in output_prefix or ".target_only" in name
if species_match and evidence_match:
return path
if fallback is None and species_match:
if target_only_match or not reference_species:
fallback = path
return fallback
def manifest_artifact_path(workspace: Path, manifest: dict, keys: list[str]) -> Path | None:
artifacts = manifest.get("artifacts") or {}
if not isinstance(artifacts, dict):
return None
for key in keys:
item = artifacts.get(key)
if not isinstance(item, dict):
continue
text = str(item.get("full_path") or item.get("path") or "").strip()
if not text:
continue
path = Path(text)
if not path.is_absolute():
path = workspace / path
if path.exists() and path.stat().st_size > 0:
return path.resolve()
return None
def _strip_species_alias(label: str) -> str:
text = str(label or "").strip().strip("'\"")
if "|" in text:
text = text.split("|", 1)[1]
return text
def _canonical_evidence_id(label: str) -> str:
text = _strip_species_alias(label)
text = text.strip().strip("'\"")
# Tree/Newick writers replace ':' with '_' for labels such as mRNA:HanXR...,
# while MEME/hmmscan/GFF keep the original ':' form. Treat both as the same
# typed accession and strip repeated type wrappers like transcript:mRNA:.
while True:
updated = re.sub(r"^(gene|transcript|CDS|mRNA)[:_]", "", text, flags=re.IGNORECASE)
if updated == text:
break
text = updated
text = re.sub(r"_[PT][0-9]+$", "", text, flags=re.IGNORECASE)
text = re.sub(r"\.[0-9]+$", "", text)
return text.lower()
def parse_meme_sequence_ids(meme_xml: Path) -> set[str]:
try:
import xml.etree.ElementTree as ET
root = ET.parse(meme_xml).getroot()
except Exception:
return set()
ids: set[str] = set()
for elem in root.iter():
if not str(elem.tag).endswith("sequence"):
continue
for key in ("name", "id"):
value = elem.attrib.get(key)
if value:
ids.add(_canonical_evidence_id(value))
break
return ids
def parse_domtblout_sequence_ids(domtblout_file: Path) -> set[str]:
ids: set[str] = set()
try:
with domtblout_file.open(encoding="utf-8", errors="ignore") as handle:
for raw in handle:
line = raw.strip()
if not line or line.startswith("#"):
continue
parts = line.split()
if len(parts) >= 4:
ids.add(_canonical_evidence_id(parts[3]))
except OSError:
return set()
return ids
def parse_gff_gene_ids(gff3_file: Path) -> set[str]:
ids: set[str] = set()
try:
with gff3_file.open(encoding="utf-8", errors="ignore") as handle:
for raw in handle:
if not raw.strip() or raw.startswith("#"):
continue
parts = raw.rstrip("\n").split("\t")
if len(parts) < 9:
continue
attrs = {}
for item in parts[8].split(";"):
if "=" in item:
k, v = item.split("=", 1)
attrs[k] = v
for key in ("ID", "Parent", "Name", "gene_id", "transcript_id", "protein_id"):
value = attrs.get(key)
if value:
for token in value.split(","):
ids.add(_canonical_evidence_id(token))
except OSError:
return set()
return ids
def evidence_id_match_summary(
inputs: dict[str, Path],
expected_labels: Iterable[str],
reference_species: str = "",
) -> dict[str, object]:
expected = {_canonical_evidence_id(x) for x in expected_labels if str(x).strip()}
ref_tokens = [x.lower() for x in re.split(r"[,;\s]+", reference_species or "") if x]
def is_reference(label: str) -> bool:
lowered = str(label or "").lower()
return any(lowered == token or lowered.startswith(token + "|") or lowered.startswith(token) for token in ref_tokens)
target_expected = {label for label in expected if not is_reference(label)} if ref_tokens else expected
summaries: dict[str, object] = {
"expected_count": len(expected),
"target_expected_count": len(target_expected),
"stages": {},
}
parsers = {
"meme_xml": parse_meme_sequence_ids,
"domtblout_file": parse_domtblout_sequence_ids,
"gff3_file": parse_gff_gene_ids,
}
ok = True
for key, parser in parsers.items():
path = inputs.get(key)
ids = parser(path) if path else set()
matched = expected & ids
matched_target = target_expected & ids
stage_ok = bool(ids) and len(matched) > 0 and len(matched_target) > 0
ok = ok and stage_ok
summaries["stages"][key] = {
"path": str(path) if path else "",
"id_count": len(ids),
"matched_expected": len(matched),
"matched_target_expected": len(matched_target),
"missing_expected": max(0, len(expected - ids)),
"missing_target_expected": max(0, len(target_expected - ids)),
}
summaries["success"] = ok
return summaries
def resolve_combo_evidence_inputs(workspace: Path, final_dir: Path, species_prefix: str) -> tuple[dict[str, Path], list[str]]:
warnings: list[str] = []
resolved: dict[str, Path] = {}
specs = {
"motif": ("meme_xml", ["meme_xml"]),
"structure": ("gff3_file", ["structure_gff3_filtered"]),
"domain": ("domtblout_file", ["domtblout"]),
}
for stage, (out_key, artifact_keys) in specs.items():
manifest_path = latest_species_evidence_manifest(final_dir, species_prefix, stage)
if not manifest_path:
warnings.append(f"Missing {stage} evidence manifest for {species_prefix} evidence combo")
continue
try:
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
except Exception as exc:
warnings.append(f"Failed to read {stage} manifest {manifest_path}: {exc}")
continue
path = manifest_artifact_path(workspace, manifest, artifact_keys)
if not path:
warnings.append(f"Missing {stage} artifact {artifact_keys}")
continue
resolved[out_key] = path
return resolved, warnings
def run_evidence_combo_plot(
*,
tree_file: Path,
subfamily_file: Path,
meme_xml: Path,
gff3_file: Path,
domtblout_file: Path,
output_pdf: Path,
target_species: str = "",
reference_species: str = "",
) -> dict[str, object]:
script = PROJECT_ROOT / "scripts" / "combo_tree_motif_structure_domain.R"
rscript = Path("/home/xzq/miniconda3/envs/gfpy312/bin/Rscript")
if not rscript.exists():
rscript = Path("/home/xzq/miniconda3/envs/test4py311/bin/Rscript")
if not script.exists():
return {"success": False, "error": f"Combo script not found: {script}", "output_files": []}
cmd = [
str(rscript),
str(script),
str(tree_file.resolve()),
str(subfamily_file.resolve()),
str(meme_xml.resolve()),
str(gff3_file.resolve()),
str(domtblout_file.resolve()),
"1e-5",
"0",
str(output_pdf.resolve()),
"motif,structure,domain",
target_species,
reference_species,
]
print("[CMD]", " ".join(cmd))
result = subprocess.run(cmd, text=True, capture_output=True, cwd=str(PROJECT_ROOT), check=False, timeout=300)
if result.stdout:
print(result.stdout[-2000:])
if result.stderr:
print(result.stderr[-2000:], file=sys.stderr)
output_png = output_pdf.with_suffix(".png")
output_files = [str(path.resolve()) for path in (output_pdf, output_png) if path.exists() and path.stat().st_size > 0]
return {
"success": bool(output_files) and result.returncode == 0,
"command": " ".join(cmd),
"returncode": result.returncode,
"stdout": result.stdout[-1000:] if result.stdout else "",
"stderr": result.stderr[-1000:] if result.stderr else "",
"output_files": output_files,
"error": "" if result.returncode == 0 else (result.stderr or result.stdout or "combo plotting failed")[-1000:],
}
def run_per_species_evidence_combo_plots(
*,
workspace: Path,
final_dir: Path,
per_species_tree_plots: dict,
reference_species: str = "",
) -> dict[str, object]:
warnings: list[str] = []
results: list[dict[str, object]] = []
for item in per_species_tree_plots.get("results") or []:
if not item.get("success"):
continue
species_prefix = str(item.get("prefix") or "").strip()
inputs, species_warnings = resolve_combo_evidence_inputs(workspace, final_dir, species_prefix)
warnings.extend(species_warnings)
required = {"meme_xml", "gff3_file", "domtblout_file"}
if set(inputs) != required:
results.append({
"prefix": species_prefix,
"success": False,
"error": "; ".join(species_warnings) or "missing target-only evidence inputs",
"output_files": [],
})
continue
pruned_tree = Path(str(item.get("pruned_tree") or ""))
subfamily_file = Path(str(item.get("subfamily_file") or ""))
if not species_prefix or not pruned_tree.exists() or not subfamily_file.exists():
continue
keep_labels = read_subfamily_labels(subfamily_file)
tree_match = tree_label_match_summary(pruned_tree, keep_labels)
if (
not tree_match.get("success")
or tree_match.get("total_tips") != tree_match.get("expected_tips")
or tree_match.get("matched_tips") != tree_match.get("expected_tips")
):
results.append({
"prefix": species_prefix,
"success": False,
"error": (
f"combo tree/subfamily scope mismatch: tree tips={tree_match.get('total_tips')} "
f"matched={tree_match.get('matched_tips')} expected={tree_match.get('expected_tips')}"
),
"tree_match": tree_match,
"output_files": [],
})
continue
evidence_match = evidence_id_match_summary(inputs, keep_labels, reference_species=reference_species)
if not evidence_match.get("success"):
results.append({
"prefix": species_prefix,
"success": False,
"error": f"combo evidence scope mismatch: {json.dumps(evidence_match, ensure_ascii=False)}",
"evidence_match": evidence_match,
"output_files": [],
})
continue
output_pdf = final_dir / f"{species_prefix}.target_plus_reference.combo.pdf"
target_species = species_prefix.split(".", 1)[0]
combo_result = run_evidence_combo_plot(
tree_file=pruned_tree,
subfamily_file=subfamily_file,
meme_xml=inputs["meme_xml"],
gff3_file=inputs["gff3_file"],
domtblout_file=inputs["domtblout_file"],
output_pdf=output_pdf,
target_species=target_species,
reference_species=reference_species,
)
combo_result.update({"prefix": species_prefix})
results.append(combo_result)
manifest_path = final_dir / f"{species_prefix}.target_plus_reference.reference_classification.manifest.json"
if manifest_path.exists():
try:
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
artifacts = manifest.setdefault("artifacts", {})
for output_file in combo_result.get("output_files", []) or []:
p = Path(str(output_file))
artifacts[f"evidence_combo_{p.suffix.lower().lstrip('.')}"] = {
"path": f"final/{p.name}",
"type": "image" if p.suffix.lower() == ".png" else "pdf",
"description": "Per-target tree + subfamily + motif/structure/domain combo",
}
metrics = manifest.setdefault("metrics", {})
metrics["evidence_combo_success"] = bool(combo_result.get("success"))
manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
except Exception as exc:
warnings.append(f"Failed to update evidence combo manifest {manifest_path}: {exc}")
return {
"success": any(bool(item.get("success")) for item in results),
"count": len(results),
"results": results,
"warnings": warnings,
}
def build_neighbor_joining_tree(msa_path: Path, tree_path: Path) -> None:
from Bio import AlignIO, Phylo
from Bio.Phylo.TreeConstruction import DistanceCalculator, DistanceTreeConstructor
alignment = AlignIO.read(msa_path, "fasta")
calculator = DistanceCalculator("identity")
distance_matrix = calculator.get_distance(alignment)
constructor = DistanceTreeConstructor()
tree = constructor.nj(distance_matrix)
tree_path.parent.mkdir(parents=True, exist_ok=True)
Phylo.write(tree, tree_path, "newick")
def build_joint_tree_if_needed(
tree: Path,
workspace: Path,
input_dir: Path,
analysis_dir: Path,
prefix: str,
target_prefix: str,
reference_prefix: str,
reference_subfamily: Path,
scope_data: dict | None = None,
output_prefix: str | None = None,
) -> tuple[Path, dict[str, str | int | bool]]:
counts = tree_tip_counts(tree, reference_subfamily, target_prefix, reference_prefix)
info: dict[str, str | int | bool] = {
"auto_built_joint_tree": False,
"input_tree_total_tips": counts["total"],
"input_tree_target_tips": counts["target"],
"input_tree_reference_tips": counts["reference"],
}
if counts["target"] > 0 and counts["reference"] > 0:
return tree, info
family = prefix.split(".", 1)[1] if "." in prefix else ""
final_dir = workspace / "final"
output_prefix = output_prefix or prefix
final_joint_tree = final_dir / f"{output_prefix}.reference_joint.tree.nwk"
if final_joint_tree.exists() and final_joint_tree.stat().st_size > 0:
joint_counts = tree_tip_counts(final_joint_tree, reference_subfamily, target_prefix, reference_prefix)
if joint_counts["target"] > 0 and joint_counts["reference"] > 0:
info.update({
"auto_built_joint_tree": True,
"reused_joint_tree": True,
"joint_tree": str(final_joint_tree),
"joint_tree_target_tips": joint_counts["target"],
"joint_tree_reference_tips": joint_counts["reference"],
})
return final_joint_tree, info
target_proteins = find_target_proteins(final_dir, prefix, workspace=workspace, scope_data=scope_data)
reference_proteins = find_reference_proteins(workspace, input_dir, reference_prefix, family)
if target_proteins is None or reference_proteins is None:
missing = []
if target_proteins is None:
missing.append(f"{prefix}.final_members.pep.clean.fa or scoped phylogeny input FASTA")
if reference_proteins is None:
missing.append(f"{reference_prefix}.{family}.ncbi_fetch.protein.fa or {reference_prefix}.{family}.pep.fa")
raise SystemExit(
"Reference classification requires a joint tree containing target and reference genes. "
f"The current tree has target={counts['target']}, reference={counts['reference']}; "
f"auto-build could not proceed because missing: {', '.join(missing)}."
)
joint_fasta = analysis_dir / f"{output_prefix}.reference_joint.pep.fa"
target_count, ref_count = write_joint_fasta(target_proteins, reference_proteins, reference_subfamily, reference_prefix, joint_fasta)
if target_count == 0 or ref_count == 0:
raise SystemExit(
"Reference classification could not build a joint FASTA: "
f"target sequences={target_count}, reference sequences matching subfamily table={ref_count}."
)
joint_msa = analysis_dir / f"{output_prefix}.reference_joint.msa.fa"
run_logged(["mafft", "--auto", "--thread", "-1", str(joint_fasta)], stdout_path=joint_msa)
built_tree = analysis_dir / f"{output_prefix}.reference_joint.nj.treefile"
# Classification only needs a target+reference topology to transfer labels.
# The publication-grade tree remains the responsibility of the phylogeny stage.
build_neighbor_joining_tree(joint_msa, built_tree)
if not built_tree.exists():
raise SystemExit(f"Auto-built joint tree missing: {built_tree}")
shutil.copyfile(built_tree, final_joint_tree)
info.update({
"auto_built_joint_tree": True,
"joint_fasta": str(joint_fasta),
"joint_tree": str(final_joint_tree),
"joint_target_sequences": target_count,
"joint_reference_sequences": ref_count,
})
return final_joint_tree, info
def write_confidence_table(assignment_tsv: Path, confidence_tsv: Path) -> int:
rows_written = 0
with assignment_tsv.open(encoding="utf-8", newline="") as in_handle, confidence_tsv.open(
"w", encoding="utf-8", newline=""
) as out_handle:
reader = csv.DictReader(in_handle, delimiter="\t")
writer = csv.DictWriter(out_handle, fieldnames=["gene_id", "subfamily", "confidence", "note"], delimiter="\t")
writer.writeheader()
for row in reader:
writer.writerow(
{
"gene_id": row.get("target_gene", ""),
"subfamily": row.get("assigned_subfamily", ""),
"confidence": row.get("confidence", ""),
"note": row.get("classification_note", ""),
}
)
rows_written += 1
return rows_written
def summarize_assignment_table(assignment_tsv: Path) -> dict[str, object]:
assigned = Counter()
confidence = Counter()
notes = Counter()
rows = 0
with assignment_tsv.open(encoding="utf-8", newline="") as handle:
reader = csv.DictReader(handle, delimiter="\t")
for row in reader:
rows += 1
if row.get("assigned_subfamily"):
assigned[row["assigned_subfamily"]] += 1
if row.get("confidence"):
confidence[row["confidence"]] += 1
if row.get("classification_note"):
notes[row["classification_note"]] += 1
return {
"target_gene_count": rows,
"confidence_counts": dict(sorted(confidence.items())),
"assigned_subfamily_counts": dict(sorted(assigned.items())),
"classification_note_counts": dict(sorted(notes.items())),
}
def infer_target_prefix(tree_file: Path, reference_prefix: str, fallback_prefix: str) -> str:
try:
from Bio import Phylo
tree = Phylo.read(tree_file, "newick")
counts = Counter()
for tip in tree.get_terminals():
name = tip.name or ""
if "|" not in name:
continue
species_prefix = name.split("|", 1)[0]
if species_prefix != reference_prefix:
counts[species_prefix] += 1
if len(counts) == 1:
return next(iter(counts))
except Exception:
pass
return fallback_prefix
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--prefix", required=True)
parser.add_argument("--workspace", required=True, type=Path)
parser.add_argument("--input-dir", type=Path)
parser.add_argument("--tree", type=str, default="")
parser.add_argument("--scope-file", type=str, default="", help="Gene-family scope JSON from resolve_scope.py")
parser.add_argument("--analysis-route", default="", help="Explicit analysis route identity for manifest/status")
parser.add_argument("--member-set-id", default="", help="Explicit member-set identity for manifest/status")
parser.add_argument("--reference-subfamily", type=str, default="")
parser.add_argument("--target-prefix", default="")
parser.add_argument("--reference-prefix", default="At")
parser.add_argument("--write-gene-subfamily", action="store_true", default=True)
parser.add_argument("--no-write-gene-subfamily", action="store_false", dest="write_gene_subfamily")
parser.add_argument("--overwrite-gene-subfamily", action="store_true")
parser.add_argument("--log-file", type=Path)
args = parser.parse_args()
workspace = args.workspace.resolve()
input_dir = (args.input_dir or (PROJECT_ROOT / "input" / workspace.name)).resolve()
final_dir = workspace / "final"
analysis_dir = workspace / "reference_classification" / "analysis"
final_dir.mkdir(parents=True, exist_ok=True)
analysis_dir.mkdir(parents=True, exist_ok=True)
base_dirs = [PROJECT_ROOT, workspace, input_dir, final_dir]
scope_file = resolve_path(args.scope_file, base_dirs)
scope_data: dict | None = None
if args.scope_file:
if scope_file is None:
raise SystemExit(f"Scope file not found: {args.scope_file}")
scope_data = read_scope(scope_file)
scope_id = str((scope_data or {}).get("scope_id") or "").strip()
output_prefix = user_output_prefix_for(
"reference_classification",
args.prefix,
scope_data or {},
analysis_route=args.analysis_route or "subfamily_prediction",
member_set_id=args.member_set_id,
) if scope_id else args.prefix
tree = resolve_path(args.tree, base_dirs) if args.tree else final_dir / f"{args.prefix}.tree.nwk"
if not tree or not tree.exists():
raise SystemExit(f"Tree file not found: {tree}")
reference_subfamily = resolve_path(args.reference_subfamily, base_dirs)
if reference_subfamily is None:
reference_subfamily = find_reference_subfamily(workspace, input_dir, args.prefix, args.reference_prefix)
if reference_subfamily is None or not reference_subfamily.exists():
raise SystemExit(
"Reference subfamily file not found. Provide --reference-subfamily, "
"or upload a two-column *.subfamily/TSV file to input/final."
)
target_prefix = args.target_prefix.strip()
if not target_prefix:
target_prefix = infer_target_prefix(tree, args.reference_prefix, args.prefix.split(".")[0])
tree, joint_info = build_joint_tree_if_needed(
tree,
workspace,
input_dir,
analysis_dir,
args.prefix,
target_prefix,
args.reference_prefix,
reference_subfamily,
scope_data=scope_data,
output_prefix=output_prefix,
)
assignment_out = final_dir / f"{output_prefix}.reference_classification.assignment.tsv"
subfamily_out = final_dir / f"{output_prefix}.reference_classification.subfamily.tsv"
confidence_out = final_dir / f"{output_prefix}.reference_classification.confidence.tsv"
manifest_out = final_dir / f"{output_prefix}.reference_classification.manifest.json"
report_out = final_dir / f"{output_prefix}.reference_classification_report.md"
script = PROJECT_ROOT / "scripts" / "reference_classification_from_tree.py"
cmd = [
sys.executable,
str(script),
"--tree",
str(tree),
"--reference-subfamily",
str(reference_subfamily),
"--target-prefix",
target_prefix,
"--reference-prefix",
args.reference_prefix,
"--assignment-out",
str(assignment_out),
"--subfamily-out",
str(subfamily_out),
"--manifest-out",
str(manifest_out),
"--report-out",
str(report_out),
]
print("[CMD]", " ".join(cmd))
result = subprocess.run(cmd, text=True, capture_output=True, cwd=str(PROJECT_ROOT), check=False)
if result.stdout:
print(result.stdout)
if result.stderr:
print(result.stderr, file=sys.stderr)
if result.returncode != 0:
return result.returncode
allowed_target_ids, scope_target_counts, scope_filter_warnings = collect_target_ids_from_scope(workspace, args.prefix, scope_data)
filter_info = filter_reference_classification_outputs(
assignment_out,
subfamily_out,
reference_subfamily,
args.reference_prefix,
allowed_target_ids,
)
confidence_rows = write_confidence_table(assignment_out, confidence_out)
assignment_summary = summarize_assignment_table(assignment_out)
per_species_outputs = write_per_species_subfamily_outputs(
workspace=workspace,
final_dir=final_dir,
prefix=args.prefix,
output_prefix=output_prefix,
assignment_out=assignment_out,
subfamily_out=subfamily_out,
confidence_out=confidence_out,
reference_subfamily=reference_subfamily,
reference_prefix=args.reference_prefix,
scope_data=scope_data,
)
compatibility_out = final_dir / f"{output_prefix}.geneID.subfamily"
compatibility_written = False
if args.write_gene_subfamily:
should_overwrite_compat = args.overwrite_gene_subfamily or args.analysis_route == "subfamily_prediction"
if should_overwrite_compat or not compatibility_out.exists():
shutil.copyfile(subfamily_out, compatibility_out)
compatibility_written = True
else:
print(f"[INFO] Compatibility subfamily exists, not overwritten: {compatibility_out}")
target_species_for_scope = [str(item) for item in (scope_data or {}).get("target_species") or [] if str(item).strip()]
if target_species_for_scope:
tree_plot_result = {
"success": False,
"skipped": True,
"reason": "reference_classification now emits per-classifiable-species confirmation plots only",
"output_files": [],
}
else:
plot_subfamily_file = compatibility_out if compatibility_out.exists() else subfamily_out
tree_plot_result = run_subfamily_tree_plot(
tree_file=tree,
subfamily_file=plot_subfamily_file,
final_dir=final_dir,
prefix=output_prefix,
layout="both",
target_species=target_prefix,
reference_species=args.reference_prefix,
)
if not tree_plot_result.get("success"):
print(f"[WARNING] Reference classification tree plot failed: {tree_plot_result.get('error', '')}", file=sys.stderr)
per_species_tree_plots = run_per_species_subfamily_tree_plots(
tree_file=tree,
final_dir=final_dir,
per_species_outputs=per_species_outputs,
layout="both",
reference_species=args.reference_prefix,
output_prefix=output_prefix,
)
per_species_evidence_combos = run_per_species_evidence_combo_plots(
workspace=workspace,
final_dir=final_dir,
per_species_tree_plots=per_species_tree_plots,
reference_species=args.reference_prefix,
)
manifest = json.loads(manifest_out.read_text(encoding="utf-8"))
manifest["prefix"] = args.prefix
manifest["output_prefix"] = output_prefix
artifacts = manifest.setdefault("artifacts", {})
artifacts["confidence_table"] = {
"path": f"final/{confidence_out.name}",
"type": "tsv",
"description": "Target-gene confidence and classification note table",
}
if compatibility_written or compatibility_out.exists():
artifacts["geneID_subfamily_compat"] = {
"path": f"final/{compatibility_out.name}",
"type": "tsv",
"description": "Backward-compatible two-column subfamily file for existing plotting stages",
}
if tree_plot_result.get("success"):
for output_file in tree_plot_result.get("output_files", []):
plot_path = Path(str(output_file))
key = (
"subfamily_tree_"
+ plot_path.stem.replace(f"{args.prefix}.", "").replace(".", "_")
+ plot_path.suffix.lower().replace(".", "_")
)
artifacts[key] = {
"path": f"final/{plot_path.name}",
"type": "image" if plot_path.suffix.lower() == ".png" else "pdf",
"description": "Reference-guided phylogenetic tree with subfamily labels",
}
for item in per_species_tree_plots.get("results") or []:
for output_file in item.get("output_files", []) or []:
plot_path = Path(str(output_file))
key = "per_species_subfamily_tree_" + plot_path.stem.replace(".", "_") + plot_path.suffix.lower().replace(".", "_")
artifacts[key] = {
"path": f"final/{plot_path.name}",
"type": "image" if plot_path.suffix.lower() == ".png" else "pdf",
"description": "Per-target reference-guided subfamily confirmation tree",
}
if item.get("pruned_tree"):
p = Path(str(item["pruned_tree"]))
artifacts["per_species_subfamily_tree_" + p.stem.replace(".", "_")] = {
"path": f"final/{p.name}",
"type": "newick",
"description": "Per-target target+reference tree used for subfamily confirmation plot",
}
for item in per_species_evidence_combos.get("results") or []:
for output_file in item.get("output_files", []) or []:
plot_path = Path(str(output_file))
key = "per_species_evidence_combo_" + plot_path.stem.replace(".", "_") + plot_path.suffix.lower().replace(".", "_")
artifacts[key] = {
"path": f"final/{plot_path.name}",
"type": "image" if plot_path.suffix.lower() == ".png" else "pdf",
"description": "Per-target tree + subfamily + motif/structure/domain combo",
}
metrics = manifest.setdefault("metrics", {})
metrics.update(assignment_summary)
metrics["confidence_rows"] = confidence_rows
metrics["scope_filter"] = filter_info
metrics["per_species_subfamily_outputs"] = per_species_outputs
metrics["per_species_subfamily_tree_plots"] = per_species_tree_plots
metrics["per_species_evidence_combo_plots"] = per_species_evidence_combos
metrics["joint_subfamily_tree_plot_success"] = bool(tree_plot_result.get("success"))
metrics["per_species_subfamily_tree_plot_success"] = bool(per_species_tree_plots.get("success"))
metrics["subfamily_tree_plot_success"] = bool(tree_plot_result.get("success") or per_species_tree_plots.get("success"))
if scope_target_counts:
metrics["scope_target_member_counts"] = scope_target_counts
if scope_filter_warnings:
manifest.setdefault("warnings", []).extend(scope_filter_warnings)
if per_species_outputs.get("warnings"):
manifest.setdefault("warnings", []).extend(per_species_outputs.get("warnings") or [])
for index, file_path in enumerate(per_species_outputs.get("files") or [], start=1):
p = Path(str(file_path))
artifacts[f"per_species_subfamily_{index}"] = {
"path": f"final/{p.name}",
"type": "tsv" if p.suffix.lower() in {".tsv", ".subfamily"} or "subfamily" in p.name.lower() else "json",
"description": "Per-target-species subfamily output split from reference classification",
}
manifest.setdefault("parameters", {})["target_prefix"] = target_prefix
manifest.setdefault("parameters", {})["output_prefix"] = output_prefix
manifest.setdefault("parameters", {})["reference_prefix"] = args.reference_prefix
manifest.setdefault("parameters", {})["reference_subfamily_file"] = str(reference_subfamily)
manifest.setdefault("parameters", {})["subfamily_tree_plot"] = tree_plot_result
if scope_data:
manifest["scope_id"] = scope_data.get("scope_id", "")
manifest["scope_mode"] = scope_data.get("scope_mode", "")
manifest["scope_file"] = str(scope_file)
manifest["target_species"] = scope_data.get("target_species", [])
manifest["reference_species"] = scope_data.get("reference_species", [])
manifest["comparative_species"] = scope_data.get("comparative_species", [])
manifest["selected_species"] = scope_data.get("selected_species", [])
manifest.setdefault("inputs", {})["target_species"] = scope_data.get("target_species", [])
manifest.setdefault("inputs", {})["reference_species"] = scope_data.get("reference_species", [])
manifest.setdefault("inputs", {})["comparative_species"] = scope_data.get("comparative_species", [])
manifest.setdefault("inputs", {})["selected_species"] = scope_data.get("selected_species", [])
manifest["scope"] = {
"family": scope_data.get("family", ""),
"prefix": scope_data.get("prefix", ""),
"scope_id": scope_data.get("scope_id", ""),
"scope_mode": scope_data.get("scope_mode", ""),
"target_species": scope_data.get("target_species", []),
"reference_species": scope_data.get("reference_species", []),
"comparative_species": scope_data.get("comparative_species", []),
"selected_species": scope_data.get("selected_species", []),
}
manifest.setdefault("parameters", {}).update(joint_info)
context = build_stage_context(
stage="reference_classification",
prefix=args.prefix,
scope_data=scope_data or {},
scope_file=scope_file or "",
output_prefix=output_prefix,
subfamily_file=subfamily_out,
analysis_route=args.analysis_route,
member_set_id=args.member_set_id,
)
apply_stage_context(manifest, context)
if joint_info.get("auto_built_joint_tree"):
manifest.setdefault("artifacts", {})["joint_tree"] = {
"path": f"final/{Path(str(joint_info['joint_tree'])).name}",
"type": "newick",
"description": "Auto-built target + reference tree used for reference classification",
}
manifest_out.write_text(json.dumps(manifest, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
if scope_data and scope_data.get("scope_id"):
scoped_manifest = write_legacy_manifest_copy(manifest_out, args.prefix, str(scope_data.get("scope_id") or ""), "reference_classification")
if scoped_manifest:
print(f"legacy_manifest={scoped_manifest}")
print(f"assignment_out={assignment_out}")
print(f"subfamily_out={subfamily_out}")
print(f"confidence_out={confidence_out}")
print(f"manifest_out={manifest_out}")
return 0
if __name__ == "__main__":
raise SystemExit(main())