Files
PlantAI-dev/apps/cli/run_subfamily_confirm.py
T
potaserandClaude Opus 5 7097ecdec5 Add memory directory
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-07 09:39:30 +08:00

577 lines
25 KiB
Python
Executable File

#!/usr/bin/env python3
"""Create a confirmed subfamily table from predicted or user-edited labels."""
from __future__ import annotations
import argparse
import csv
import hashlib
import json
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Iterable, Tuple
PROJECT_ROOT = Path(__file__).resolve().parents[2]
import sys
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, user_output_prefix_for, write_legacy_manifest_copy
def _read_table(path: Path) -> list[dict[str, str]]:
text = path.read_text(encoding="utf-8", errors="ignore").splitlines()
rows = [line for line in text if line.strip() and not line.lstrip().startswith("#")]
if not rows:
return []
delimiter = "\t" if "\t" in rows[0] else ","
first = [x.strip() for x in rows[0].split(delimiter)]
header_tokens = {x.lower() for x in first}
has_header = bool(header_tokens & {"gene_id", "gene", "id", "subfamily", "predicted_subfamily", "confirmed_subfamily"})
if has_header:
reader = csv.DictReader(rows, delimiter=delimiter)
return [{str(k or "").strip(): str(v or "").strip() for k, v in row.items()} for row in reader]
parsed = []
for line in rows:
parts = [x.strip() for x in line.split(delimiter)]
if len(parts) >= 2:
parsed.append({"gene_id": parts[0], "confirmed_subfamily": parts[1], "predicted_subfamily": parts[1]})
return parsed
def _row_value(row: dict[str, str], names: Iterable[str]) -> str:
lowered = {k.lower(): v for k, v in row.items()}
for name in names:
value = lowered.get(name.lower())
if value:
return value.strip()
return ""
def _load_labels(path: Path) -> Dict[str, dict[str, str]]:
labels: Dict[str, dict[str, str]] = {}
for row in _read_table(path):
gene_id = _row_value(row, ["gene_id", "gene", "id", "target_gene", "label"])
confirmed = _row_value(row, ["confirmed_subfamily", "subfamily", "assigned_subfamily", "prediction"])
predicted = _row_value(row, ["predicted_subfamily", "assigned_subfamily", "subfamily", "prediction"])
confidence = _row_value(row, ["confidence", "score"])
note = _row_value(row, ["note", "status", "comment"])
if not gene_id or not confirmed:
continue
labels[gene_id] = {
"gene_id": gene_id,
"predicted_subfamily": predicted or confirmed,
"confirmed_subfamily": confirmed,
"confidence": confidence,
"note": note,
}
return labels
def _read_id_list(path: Path) -> list[str]:
ids: list[str] = []
seen: set[str] = set()
for line in path.read_text(encoding="utf-8", errors="ignore").splitlines():
gene_id = line.strip()
if not gene_id or gene_id.startswith("#"):
continue
if gene_id in seen:
raise ValueError(f"Duplicate ID in {path}: {gene_id}")
seen.add(gene_id)
ids.append(gene_id)
return ids
def _read_fasta_ids(path: Path) -> list[str]:
ids: list[str] = []
seen: set[str] = set()
with path.open("r", encoding="utf-8", errors="ignore") as handle:
for line in handle:
if not line.startswith(">"):
continue
gene_id = line[1:].strip().split()[0]
if not gene_id:
continue
if gene_id in seen:
raise ValueError(f"Duplicate FASTA ID in {path}: {gene_id}")
seen.add(gene_id)
ids.append(gene_id)
return ids
def _sha256(path: Path) -> str:
h = hashlib.sha256()
with path.open("rb") as handle:
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
h.update(chunk)
return h.hexdigest()
def _atomic_promote_pair(curated_id: Path, curated_pep: Path, final_id: Path, final_pep: Path) -> None:
tmp_id = final_id.with_name(f".{final_id.name}.tmp")
tmp_pep = final_pep.with_name(f".{final_pep.name}.tmp")
tmp_id.write_text(curated_id.read_text(encoding="utf-8", errors="ignore"), encoding="utf-8")
tmp_pep.write_text(curated_pep.read_text(encoding="utf-8", errors="ignore"), encoding="utf-8")
tmp_ids = _read_id_list(tmp_id)
tmp_fasta_ids = _read_fasta_ids(tmp_pep)
if sorted(tmp_ids) != sorted(tmp_fasta_ids):
tmp_id.unlink(missing_ok=True)
tmp_pep.unlink(missing_ok=True)
raise RuntimeError(f"Temporary final member validation failed for {final_id.stem}")
tmp_id.replace(final_id)
tmp_pep.replace(final_pep)
def _member_prefix_candidates(prefix: str, output_prefix: str = "", compact_prefix: str = "") -> list[str]:
candidates: list[str] = []
def add(value: str) -> None:
value = str(value or "").strip()
if value and value not in candidates:
candidates.append(value)
for value in (prefix, output_prefix, compact_prefix):
add(value)
parts = [part for part in str(value or "").split(".") if part]
if len(parts) == 2:
add(f"Target.{parts[0]}.{parts[1]}")
elif len(parts) >= 3 and parts[0] == "Target":
add(".".join(parts[1:3]))
return candidates
def _promote_curated_members_to_final(final_dir: Path, prefixes: Iterable[str], labels: Dict[str, dict[str, str]]) -> list[dict[str, Any]]:
label_ids = set(labels)
promotions: list[dict[str, Any]] = []
for member_prefix in prefixes:
curated_id = final_dir / f"{member_prefix}.curated_members.id"
curated_pep = final_dir / f"{member_prefix}.curated_members.pep.fa"
if not curated_id.exists() and not curated_pep.exists():
continue
if not curated_id.exists() or not curated_pep.exists():
raise FileNotFoundError(
f"Incomplete curated members for {member_prefix}: expected both {curated_id.name} and {curated_pep.name}"
)
curated_ids = _read_id_list(curated_id)
if not curated_ids:
raise ValueError(f"Curated member ID file is empty: {curated_id}")
fasta_ids = _read_fasta_ids(curated_pep)
if sorted(curated_ids) != sorted(fasta_ids):
raise ValueError(
f"Curated member ID/FASTA mismatch for {member_prefix}: "
f"{len(curated_ids)} IDs vs {len(fasta_ids)} FASTA records"
)
missing_labels = [gene_id for gene_id in curated_ids if gene_id not in label_ids]
if missing_labels:
preview = ", ".join(missing_labels[:5])
raise ValueError(
f"Confirmed subfamily labels do not cover curated members for {member_prefix}: "
f"{len(missing_labels)} missing ({preview})"
)
final_id = final_dir / f"{member_prefix}.final_members.id"
final_pep = final_dir / f"{member_prefix}.final_members.pep.fa"
_atomic_promote_pair(curated_id, curated_pep, final_id, final_pep)
promoted_ids = _read_id_list(final_id)
promoted_fasta_ids = _read_fasta_ids(final_pep)
if sorted(promoted_ids) != sorted(curated_ids) or sorted(promoted_fasta_ids) != sorted(curated_ids):
raise RuntimeError(f"Post-promotion final member validation failed for {member_prefix}")
promotions.append({
"prefix": member_prefix,
"count": len(curated_ids),
"curated_id": str(curated_id),
"curated_pep": str(curated_pep),
"final_id": str(final_id),
"final_pep": str(final_pep),
"id_checksum": _sha256(final_id),
"pep_checksum": _sha256(final_pep),
"source_curated_id_checksum": _sha256(curated_id),
"source_curated_pep_checksum": _sha256(curated_pep),
})
return promotions
def _find_source(final_dir: Path, input_dir: Path, prefix: str, explicit: str = "", output_prefix: str = "", source_prefixes: Iterable[str] = ()) -> Tuple[Path, str]:
candidates = []
if explicit:
explicit_path = Path(explicit).expanduser()
if explicit_path.is_absolute():
candidates.append((explicit_path, "explicit"))
else:
workspace = final_dir.parent
candidates.extend([
(Path(explicit), "explicit"),
(workspace / explicit_path, "explicit"),
(final_dir / explicit_path, "explicit"),
(input_dir / explicit_path, "explicit"),
(final_dir / explicit_path.name, "explicit"),
(input_dir / explicit_path.name, "explicit"),
])
for base in [*source_prefixes, output_prefix, prefix]:
if not base:
continue
candidates.extend([
(final_dir / f"{base}.confirmed_subfamily.edited.tsv", "edited"),
(input_dir / f"{base}.confirmed_subfamily.edited.tsv", "edited"),
(final_dir / f"{base}.reference_classification.subfamily.tsv", "reference_classification"),
(final_dir / f"{base}.geneID.subfamily", "geneID_subfamily"),
(input_dir / f"{base}.geneID.subfamily", "input_geneID_subfamily"),
])
for path, source in candidates:
if path.exists() and path.stat().st_size > 0:
return path, source
raise FileNotFoundError(
"No subfamily source found. Provide --edited-subfamily or run reference_classification first."
)
def _compact_confirmed_prefix(args_prefix: str, output_prefix: str, scope_data: dict[str, Any]) -> str:
family = str(
scope_data.get("family")
or scope_data.get("gene_family")
or scope_data.get("geneFamily")
or ""
).strip()
targets = scope_data.get("target_species") or scope_data.get("targets") or []
if isinstance(targets, str):
targets = [targets]
targets = [str(item or "").strip() for item in targets if str(item or "").strip()]
for candidate in (args_prefix, output_prefix):
parts = [part for part in str(candidate or "").split(".") if part]
if len(parts) >= 2 and parts[0] not in {"Target", "Multi"}:
return f"{parts[0]}.{parts[1]}"
if not family and len(parts) >= 2:
family = parts[1]
if len(targets) == 1 and family:
return f"{targets[0]}.{family}"
return ""
def _load_assignment_confidence(final_dir: Path, prefix: str, output_prefix: str = "") -> Dict[str, dict[str, str]]:
assignment = final_dir / f"{output_prefix or prefix}.reference_classification.assignment.tsv"
if not assignment.exists() and output_prefix:
assignment = final_dir / f"{prefix}.reference_classification.assignment.tsv"
if not assignment.exists():
return {}
out: Dict[str, dict[str, str]] = {}
for row in _read_table(assignment):
gene_id = _row_value(row, ["gene_id", "target_gene", "gene", "id"])
if not gene_id:
continue
out[gene_id] = {
"predicted_subfamily": _row_value(row, ["assigned_subfamily", "predicted_subfamily", "subfamily"]),
"confidence": _row_value(row, ["confidence", "score"]),
"note": _row_value(row, ["note", "status"]),
}
return out
def _read_manifest(path_text: str) -> dict[str, Any]:
path_text = str(path_text or "").strip()
if not path_text:
return {}
path = Path(path_text).expanduser()
if not path.exists():
return {}
try:
data = json.loads(path.read_text(encoding="utf-8"))
except Exception:
return {}
return data if isinstance(data, dict) else {}
def _family_from_scope_or_prefix(scope_id: str, prefix: str, manifest: dict[str, Any]) -> str:
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 {}
for source in (manifest.get("family"), stage_context.get("family"), scope_obj.get("family")):
text = str(source or "").strip()
if text:
return text
scope_id = str(scope_id or "").strip()
if ".target-" in scope_id:
return scope_id.split(".target-", 1)[0].strip()
parts = [part for part in str(prefix or "").split(".") if part]
if len(parts) >= 2:
return parts[1]
return ""
def _scope_data_from_prediction_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 {}
target_species = [str(x) for x in (pick("target_species", []) or []) if str(x).strip()]
reference_species = [str(x) for x in (pick("reference_species", []) or []) if str(x).strip()]
comparative_species = [str(x) for x in (pick("comparative_species", []) or []) if str(x).strip()]
selected_species = [str(x) for x in (pick("selected_species", []) or []) if str(x).strip()]
if not selected_species:
selected_species = [*target_species, *reference_species, *comparative_species]
return {
"family": _family_from_scope_or_prefix(scope_id, prefix, manifest),
"prefix": pick("prefix", prefix),
"scope_id": scope_id,
"scope_mode": str(pick("scope_mode", "") or ""),
"target_species": target_species,
"reference_species": reference_species,
"comparative_species": comparative_species,
"selected_species": selected_species,
}
def main() -> int:
parser = argparse.ArgumentParser(description="Confirm predicted subfamily labels for downstream plotting.")
parser.add_argument("--workspace", required=True)
parser.add_argument("--prefix", required=True)
parser.add_argument("--input-dir", default="")
parser.add_argument("--edited-subfamily", default="", help="User-edited table. Supports gene_id/subfamily or confirmed_subfamily columns.")
parser.add_argument("--scope-file", default="")
parser.add_argument("--analysis-route", default="confirmed_subfamily")
parser.add_argument("--member-set-id", default="")
parser.add_argument("--source-prediction-manifest", default="", help="Reference-classification manifest that produced the labels being confirmed.")
parser.add_argument("--source-prediction-prefix", default="", help="Output prefix of the source subfamily prediction.")
args = parser.parse_args()
workspace = Path(args.workspace).resolve()
final_dir = workspace / "final"
input_dir = Path(args.input_dir).resolve() if args.input_dir else workspace.parent.parent / "input" / workspace.name
final_dir.mkdir(parents=True, exist_ok=True)
scope_data = {}
scope_file = Path(args.scope_file).expanduser() if args.scope_file else None
if scope_file and scope_file.exists():
scope_data = read_scope(scope_file)
source_prediction_manifest = _read_manifest(args.source_prediction_manifest)
source_params = source_prediction_manifest.get("parameters") if isinstance(source_prediction_manifest.get("parameters"), dict) else {}
source_prediction_prefix = str(
args.source_prediction_prefix
or source_prediction_manifest.get("output_prefix")
or source_params.get("output_prefix")
or ""
).strip()
if not scope_data and source_prediction_manifest:
scope_data = _scope_data_from_prediction_manifest(source_prediction_manifest, args.prefix)
output_prefix = user_output_prefix_for(
"subfamily_confirm",
args.prefix,
scope_data,
analysis_route=args.analysis_route or "confirmed_subfamily",
member_set_id=args.member_set_id,
) if scope_data.get("scope_id") and not source_prediction_manifest else args.prefix
prediction_prefix = user_output_prefix_for(
"reference_classification",
args.prefix,
scope_data,
analysis_route="subfamily_prediction",
member_set_id=args.member_set_id,
) if scope_data.get("scope_id") else args.prefix
if source_prediction_prefix:
prediction_prefix = source_prediction_prefix
source_path, source_type = _find_source(
final_dir,
input_dir,
args.prefix,
args.edited_subfamily,
output_prefix=output_prefix,
source_prefixes=[prediction_prefix],
)
labels = _load_labels(source_path)
assignment_prefix = source_prediction_prefix or output_prefix
assignment = _load_assignment_confidence(final_dir, args.prefix, output_prefix=assignment_prefix)
for gene_id, meta in labels.items():
if gene_id in assignment:
meta["predicted_subfamily"] = assignment[gene_id].get("predicted_subfamily") or meta["predicted_subfamily"]
meta["confidence"] = meta.get("confidence") or assignment[gene_id].get("confidence", "")
meta["note"] = meta.get("note") or assignment[gene_id].get("note", "")
if not labels:
raise SystemExit(f"No valid confirmed subfamily labels parsed from {source_path}")
confirmed_stem = output_prefix if output_prefix.endswith(".confirmed_subfamily") else f"{output_prefix}.confirmed_subfamily"
confirmed_out = final_dir / f"{confirmed_stem}.tsv"
review_out = final_dir / f"{confirmed_stem}.review.tsv"
report_out = final_dir / f"{confirmed_stem}_report.md"
manifest_out = final_dir / f"{output_prefix}.subfamily_confirm.manifest.json"
with confirmed_out.open("w", encoding="utf-8") as handle:
for gene_id in sorted(labels):
handle.write(f"{gene_id}\t{labels[gene_id]['confirmed_subfamily']}\n")
with review_out.open("w", encoding="utf-8", newline="") as handle:
writer = csv.DictWriter(
handle,
delimiter="\t",
fieldnames=["gene_id", "predicted_subfamily", "confirmed_subfamily", "confidence", "source", "note"],
)
writer.writeheader()
for gene_id in sorted(labels):
row = labels[gene_id]
writer.writerow({
"gene_id": gene_id,
"predicted_subfamily": row.get("predicted_subfamily", ""),
"confirmed_subfamily": row.get("confirmed_subfamily", ""),
"confidence": row.get("confidence", ""),
"source": source_type,
"note": row.get("note", ""),
})
compact_confirmed_out = None
compact_review_out = None
compact_prefix = _compact_confirmed_prefix(args.prefix, output_prefix, scope_data)
current_prefix = confirmed_stem[: -len(".confirmed_subfamily")] if confirmed_stem.endswith(".confirmed_subfamily") else confirmed_stem
if compact_prefix and compact_prefix != current_prefix:
compact_confirmed_out = final_dir / f"{compact_prefix}.confirmed_subfamily.tsv"
compact_review_out = final_dir / f"{compact_prefix}.confirmed_subfamily.review.tsv"
compact_confirmed_out.write_text(confirmed_out.read_text(encoding="utf-8"), encoding="utf-8")
compact_review_out.write_text(review_out.read_text(encoding="utf-8"), encoding="utf-8")
member_promotions = _promote_curated_members_to_final(
final_dir,
_member_prefix_candidates(args.prefix, output_prefix, compact_prefix),
labels,
)
counts: Dict[str, int] = {}
for row in labels.values():
sub = row["confirmed_subfamily"]
counts[sub] = counts.get(sub, 0) + 1
report_out.write_text(
"# Confirmed Subfamily Report\n\n"
f"- Prefix: {args.prefix}\n"
f"- Source: {source_path}\n"
f"- Confirmed genes: {len(labels)}\n"
f"- Subfamilies: {', '.join(f'{k}={v}' for k, v in sorted(counts.items()))}\n",
encoding="utf-8",
)
manifest = {
"manifest_version": "1.0",
"stage": "subfamily_confirm",
"status": "success",
"prefix": args.prefix,
"output_prefix": output_prefix,
"timestamp": datetime.now().isoformat(),
"parameters": {
"source": str(source_path),
"source_type": source_type,
"scope_file": args.scope_file,
"source_prediction_manifest": str(Path(args.source_prediction_manifest).expanduser()) if args.source_prediction_manifest else "",
"source_prediction_prefix": source_prediction_prefix,
"source_prediction_scope_id": scope_data.get("scope_id", ""),
"analysis_route": args.analysis_route or "confirmed_subfamily",
"member_set_id": args.member_set_id,
},
"metrics": {
"confirmed_genes": len(labels),
"subfamily_counts": dict(sorted(counts.items())),
"final_members_count": member_promotions[0]["count"] if member_promotions else 0,
},
"artifacts": {
"confirmed_subfamily": {"path": f"final/{confirmed_out.name}", "full_path": str(confirmed_out), "type": "tsv"},
"review_table": {"path": f"final/{review_out.name}", "full_path": str(review_out), "type": "tsv"},
"report": {"path": f"final/{report_out.name}", "full_path": str(report_out), "type": "markdown"},
"manifest": {"path": f"final/{manifest_out.name}", "full_path": str(manifest_out), "type": "json"},
},
}
manifest["parameters"]["final_members_source_stage"] = "subfamily_confirm" if member_promotions else ""
manifest["parameters"]["final_members_promoted"] = bool(member_promotions)
if member_promotions:
primary_promotion = member_promotions[0]
manifest["member_promotions"] = member_promotions
manifest["parameters"].update({
"member_set_source_stage": "subfamily_confirm",
"member_set_source_kind": "confirmed_final_members",
"member_set_source_path": primary_promotion["final_id"],
"member_count": primary_promotion["count"],
"member_id_checksum": primary_promotion["id_checksum"],
"member_pep_checksum": primary_promotion["pep_checksum"],
"source_curated_path": primary_promotion["curated_id"],
"source_curated_checksum": primary_promotion["source_curated_id_checksum"],
})
manifest["artifacts"]["final_members_id"] = {
"path": f"final/{Path(primary_promotion['final_id']).name}",
"full_path": primary_promotion["final_id"],
"type": "id_list",
"sha256": primary_promotion["id_checksum"],
"members": primary_promotion["count"],
}
manifest["artifacts"]["final_members_pep"] = {
"path": f"final/{Path(primary_promotion['final_pep']).name}",
"full_path": primary_promotion["final_pep"],
"type": "fasta",
"sha256": primary_promotion["pep_checksum"],
"members": primary_promotion["count"],
}
else:
manifest.setdefault("warnings", []).append(
"No curated_members.* files were found; final_members were not promoted by subfamily_confirm."
)
if compact_confirmed_out and compact_review_out:
manifest["artifacts"]["compact_confirmed_subfamily"] = {
"path": f"final/{compact_confirmed_out.name}",
"full_path": str(compact_confirmed_out),
"type": "tsv",
}
manifest["artifacts"]["compact_review_table"] = {
"path": f"final/{compact_review_out.name}",
"full_path": str(compact_review_out),
"type": "tsv",
}
if args.scope_file:
try:
if not scope_data:
scope_data = read_scope(Path(args.scope_file))
manifest["scope_id"] = scope_data.get("scope_id", "")
manifest["scope_mode"] = scope_data.get("scope_mode", "")
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["scope"] = scope_data
except Exception as exc:
manifest.setdefault("warnings", []).append(f"Could not read scope file: {exc}")
context = build_stage_context(
stage="subfamily_confirm",
prefix=args.prefix,
scope_data=scope_data,
scope_file=args.scope_file,
output_prefix=output_prefix,
subfamily_file=confirmed_out,
analysis_route=args.analysis_route or "confirmed_subfamily",
member_set_id=args.member_set_id,
)
apply_stage_context(manifest, context)
manifest_out.write_text(json.dumps(manifest, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
if scope_data.get("scope_id"):
write_legacy_manifest_copy(manifest_out, args.prefix, str(scope_data.get("scope_id") or ""), "subfamily_confirm")
print(f"confirmed_subfamily={confirmed_out}")
print(f"review_table={review_out}")
if compact_confirmed_out:
print(f"compact_confirmed_subfamily={compact_confirmed_out}")
print(f"manifest={manifest_out}")
return 0
if __name__ == "__main__":
raise SystemExit(main())