286 lines
9.5 KiB
Python
Executable File
286 lines
9.5 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Standalone test script: WoLF PSORT subcellular localization heatmap.
|
|
|
|
Default test data:
|
|
workspace/job_1780459186454_21f4cf/pep/analysis/Ls.SPL_wolfpsort.txt
|
|
|
|
Outputs (under --output-dir):
|
|
- subcellular_scores.tsv
|
|
- subcellular_primary.tsv
|
|
- subcellular_heatmap.pdf/png
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Dict, List, Tuple
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
DEFAULT_WOLFPSORT = (
|
|
PROJECT_ROOT
|
|
/ "workspace/job_1780459186454_21f4cf/pep/analysis/Ls.SPL_wolfpsort.txt"
|
|
)
|
|
DEFAULT_OUTPUT = (
|
|
PROJECT_ROOT
|
|
/ "workspace/job_1780471268646_fc6c66/pep/analysis/test_subcellular_heatmap"
|
|
)
|
|
RSCRIPT_BIN = "/home/xia-lab/miniconda3/envs/test4py311/bin/Rscript"
|
|
|
|
LOCATION_LABELS = {
|
|
"nucl": "Nucleus",
|
|
"nucleus": "Nucleus",
|
|
"cyto": "Cytoplasm",
|
|
"cytoplasm": "Cytoplasm",
|
|
"chlo": "Chloroplast",
|
|
"chloroplast": "Chloroplast",
|
|
"mito": "Mitochondrion",
|
|
"mitochondrion": "Mitochondrion",
|
|
"plas": "Plasma membrane",
|
|
"plasma membrane": "Plasma membrane",
|
|
"vacu": "Vacuole",
|
|
"vacuole": "Vacuole",
|
|
"golg": "Golgi apparatus",
|
|
"golgi": "Golgi apparatus",
|
|
"extr": "Extracellular",
|
|
"extracellular": "Extracellular",
|
|
"pero": "Peroxisome",
|
|
"peroxisome": "Peroxisome",
|
|
"e.r.": "Endoplasmic reticulum",
|
|
"er": "Endoplasmic reticulum",
|
|
"cysk": "Cytoskeleton",
|
|
"cytoskeleton": "Cytoskeleton",
|
|
"cyto_nucl": "Cytoplasm / Nucleus",
|
|
}
|
|
|
|
LOCATION_ORDER = [
|
|
"Nucleus",
|
|
"Cytoplasm",
|
|
"Cytoplasm / Nucleus",
|
|
"Chloroplast",
|
|
"Mitochondrion",
|
|
"Plasma membrane",
|
|
"Endoplasmic reticulum",
|
|
"Golgi apparatus",
|
|
"Vacuole",
|
|
"Extracellular",
|
|
"Peroxisome",
|
|
"Cytoskeleton",
|
|
]
|
|
|
|
|
|
def normalize_location_label(label: str) -> str:
|
|
clean = re.sub(r"\s+", " ", str(label or "").strip())
|
|
if not clean:
|
|
return ""
|
|
key = clean.lower().replace("_", " ")
|
|
if key in LOCATION_LABELS:
|
|
return LOCATION_LABELS[key]
|
|
if clean.lower() in LOCATION_LABELS:
|
|
return LOCATION_LABELS[clean.lower()]
|
|
if " " not in clean and "_" in clean:
|
|
parts = [normalize_location_label(part) for part in clean.split("_")]
|
|
parts = [part for part in parts if part]
|
|
if parts:
|
|
return " / ".join(parts)
|
|
return clean[:1].upper() + clean[1:]
|
|
|
|
|
|
def parse_wolfpsort_scores(path: Path) -> Tuple[Dict[str, Dict[str, float]], Dict[str, str]]:
|
|
scores: Dict[str, Dict[str, float]] = {}
|
|
primary: Dict[str, str] = {}
|
|
pattern = re.compile(r"([^:,]+):\s*([\d.]+)")
|
|
|
|
with open(path, encoding="utf-8", errors="ignore") as handle:
|
|
for raw in handle:
|
|
line = raw.strip()
|
|
if not line or line.lower().startswith("protein_id"):
|
|
continue
|
|
parts = line.split("\t")
|
|
if len(parts) < 2:
|
|
continue
|
|
protein_id = parts[0].strip()
|
|
prediction = parts[1].strip()
|
|
loc_scores: Dict[str, float] = {}
|
|
for site, value in pattern.findall(prediction):
|
|
label = normalize_location_label(site.strip())
|
|
if not label:
|
|
continue
|
|
loc_scores[label] = loc_scores.get(label, 0.0) + float(value)
|
|
if not loc_scores:
|
|
continue
|
|
scores[protein_id] = loc_scores
|
|
primary[protein_id] = max(loc_scores.items(), key=lambda item: item[1])[0]
|
|
return scores, primary
|
|
|
|
|
|
def short_label(protein_id: str, max_len: int = 28) -> str:
|
|
if len(protein_id) <= max_len:
|
|
return protein_id
|
|
return protein_id[: max_len - 3] + "..."
|
|
|
|
|
|
def write_scores_tsv(scores: Dict[str, Dict[str, float]], primary: Dict[str, str], path: Path) -> List[str]:
|
|
locations = sorted(
|
|
{loc for loc_scores in scores.values() for loc in loc_scores},
|
|
key=lambda loc: (
|
|
LOCATION_ORDER.index(loc) if loc in LOCATION_ORDER else len(LOCATION_ORDER),
|
|
loc,
|
|
),
|
|
)
|
|
protein_order = sorted(
|
|
scores.keys(),
|
|
key=lambda pid: (
|
|
LOCATION_ORDER.index(primary[pid]) if primary.get(pid) in LOCATION_ORDER else len(LOCATION_ORDER),
|
|
primary.get(pid, ""),
|
|
pid,
|
|
),
|
|
)
|
|
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with open(path, "w", encoding="utf-8") as handle:
|
|
handle.write("protein_id\tprotein_label\tprimary_location\t" + "\t".join(locations) + "\n")
|
|
for protein_id in protein_order:
|
|
loc_scores = scores[protein_id]
|
|
row = [
|
|
protein_id,
|
|
short_label(protein_id),
|
|
primary.get(protein_id, ""),
|
|
]
|
|
row.extend(str(loc_scores.get(loc, 0.0)) for loc in locations)
|
|
handle.write("\t".join(row) + "\n")
|
|
return locations
|
|
|
|
|
|
def write_primary_tsv(primary: Dict[str, str], path: Path) -> None:
|
|
with open(path, "w", encoding="utf-8") as handle:
|
|
handle.write("protein_id\tprimary_location\n")
|
|
for protein_id in sorted(primary):
|
|
handle.write(f"{protein_id}\t{primary[protein_id]}\n")
|
|
|
|
|
|
def build_r_script(scores_tsv: Path, heatmap_pdf: Path, heatmap_png: Path) -> str:
|
|
return f"""
|
|
library(ggplot2)
|
|
library(tidyr)
|
|
library(dplyr)
|
|
|
|
raw <- read.delim("{scores_tsv}", sep="\\t", header=TRUE, check.names=FALSE, stringsAsFactors=FALSE)
|
|
meta_cols <- c("protein_id", "protein_label", "primary_location")
|
|
locations <- setdiff(names(raw), meta_cols)
|
|
|
|
long <- raw %>%
|
|
pivot_longer(cols = all_of(locations), names_to = "location", values_to = "score") %>%
|
|
mutate(
|
|
protein_label = factor(protein_label, levels = rev(unique(raw$protein_label))),
|
|
location = factor(location, levels = locations),
|
|
score = as.numeric(score),
|
|
score_label = ifelse(score > 0, sub("\\\\.0$", "", sprintf("%.1f", score)), ""),
|
|
text_color = ifelse(score >= 7, "white", "#222222")
|
|
)
|
|
|
|
p <- ggplot(long, aes(x = location, y = protein_label, fill = score)) +
|
|
geom_tile(color = "white", linewidth = 0.3) +
|
|
geom_text(
|
|
aes(label = score_label, color = text_color),
|
|
size = 2.0,
|
|
show.legend = FALSE
|
|
) +
|
|
scale_color_identity() +
|
|
scale_fill_gradient(low = "#f7fbff", high = "#08519c", name = "WoLF score") +
|
|
labs(
|
|
title = "Subcellular Localization Heatmap (WoLF PSORT)",
|
|
subtitle = "格内数字 = WoLF PSORT 得分；颜色越深表示该区室得分越高",
|
|
caption = "注：格内数字为 WoLF PSORT 得分；得分 ≥10 通常较可靠；3–5 且多列接近 → 不确定。",
|
|
x = "Predicted compartment",
|
|
y = NULL
|
|
) +
|
|
theme_minimal(base_size = 11) +
|
|
theme(
|
|
axis.text.x = element_text(angle = 45, hjust = 1),
|
|
panel.grid = element_blank(),
|
|
plot.caption = element_text(hjust = 0, size = 9, color = "#444444", margin = margin(t = 8))
|
|
)
|
|
|
|
plot_h <- max(6, 0.28 * length(unique(long$protein_label)) + 2.5)
|
|
plot_w <- max(11, 0.55 * length(locations) + 4)
|
|
ggsave("{heatmap_pdf}", p, width = plot_w, height = plot_h)
|
|
png("{heatmap_png}", width = round(plot_w * 100), height = round(plot_h * 100), res = 150)
|
|
print(p)
|
|
dev.off()
|
|
"""
|
|
|
|
|
|
def run_r_plot(r_script_path: Path) -> None:
|
|
rscript = Path(RSCRIPT_BIN)
|
|
if not rscript.exists():
|
|
raise FileNotFoundError(f"Rscript not found: {rscript}")
|
|
result = subprocess.run([str(rscript), str(r_script_path)], capture_output=True, text=True, timeout=180)
|
|
if result.returncode != 0:
|
|
raise RuntimeError(result.stderr or result.stdout or "R plotting failed")
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="Test WoLF PSORT subcellular heatmap plotting")
|
|
parser.add_argument("--input", type=Path, default=DEFAULT_WOLFPSORT, help="WoLF PSORT result TSV")
|
|
parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT, help="Output directory")
|
|
args = parser.parse_args()
|
|
|
|
input_file = args.input.resolve()
|
|
output_dir = args.output_dir.resolve()
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
if not input_file.exists():
|
|
print(f"[ERROR] Input not found: {input_file}", file=sys.stderr)
|
|
return 1
|
|
|
|
print("=" * 60)
|
|
print("Subcellular localization heatmap test")
|
|
print("=" * 60)
|
|
print(f"Input: {input_file}")
|
|
print(f"Output: {output_dir}")
|
|
|
|
scores, primary = parse_wolfpsort_scores(input_file)
|
|
if not scores:
|
|
print("[ERROR] No WoLF PSORT scores parsed", file=sys.stderr)
|
|
return 1
|
|
|
|
scores_tsv = output_dir / "subcellular_scores.tsv"
|
|
primary_tsv = output_dir / "subcellular_primary.tsv"
|
|
locations = write_scores_tsv(scores, primary, scores_tsv)
|
|
write_primary_tsv(primary, primary_tsv)
|
|
|
|
heatmap_pdf = output_dir / "subcellular_heatmap.pdf"
|
|
heatmap_png = output_dir / "subcellular_heatmap.png"
|
|
r_script_path = output_dir / "plot_subcellular_heatmap.R"
|
|
r_script_path.write_text(build_r_script(scores_tsv, heatmap_pdf, heatmap_png), encoding="utf-8")
|
|
|
|
try:
|
|
run_r_plot(r_script_path)
|
|
except Exception as exc:
|
|
print(f"[ERROR] Plotting failed: {exc}", file=sys.stderr)
|
|
return 1
|
|
|
|
from collections import Counter
|
|
|
|
counts = Counter(primary.values())
|
|
print("\nSummary")
|
|
print(f" Proteins: {len(scores)}")
|
|
print(f" Compartments in matrix: {len(locations)}")
|
|
for loc, count in counts.most_common():
|
|
print(f" Primary {loc}: {count}")
|
|
|
|
print("\nArtifacts")
|
|
for path in [scores_tsv, primary_tsv, heatmap_pdf, heatmap_png]:
|
|
print(f" - {path}")
|
|
print("\nDone.")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|