Files
PlantAI-rag-new/bioflow_kernel/scripts/cross_stress_analysis.py

435 lines
14 KiB
Python
Executable File

#!/usr/bin/env python3
"""
Cross-stress DEG comparison: PRJNA530826 Salt vs Drought vs Heat
Outputs: Venn intersections, DEG time trends, upSet-like comparison
"""
import os, sys, glob, re
from collections import defaultdict
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib_venn import venn3, venn2
import numpy as np
PROJECT_ROOT = Path(__file__).resolve().parents[1]
WORKSPACE = str(PROJECT_ROOT / "workspace")
OUTPUT = os.path.join(WORKSPACE, "cross_stress_analysis")
os.makedirs(OUTPUT, exist_ok=True)
# Map job to stress
JOBS = {
"Salt": "job_1780143735966_0e7710",
"Drought": "job_1780308105970_346d4e",
"Heat": "job_1780313949486_3f3732",
}
# Map stress prefix for time point comparisons
PREFIX = {"Salt": "S", "Drought": "D", "Heat": "H"}
def read_gene_ids(filepath):
"""Read gene IDs from DE_genes.txt (first column, skip header)"""
genes = set()
with open(filepath) as f:
header = f.readline()
for line in f:
gene = line.strip().split("\t")[0]
if gene:
genes.add(gene)
return genes
def parse_time_comparison(filename):
"""Parse filename to extract comparison type and time point"""
# Patterns: Con_vs_T1, T1_vs_Con, Tx_vs_Ty, time_lrt
f = filename.replace(".DE_genes.txt", "")
return f
def collect_stress_data(job_dir, stress):
"""Collect all DEG info for one stress"""
de_dir = os.path.join(job_dir, "rnaseq_de", "analysis")
if not os.path.isdir(de_dir):
return None
all_deg_files = sorted(glob.glob(os.path.join(de_dir, "*DE_genes.txt")))
result = {
"stress": stress,
"job": job_dir,
"files": {},
"time_lrt_genes": set(),
"con_vs_t1_genes": set(),
"con_vs_t48_genes": set(),
"all_con_vs_tx_genes": set(),
}
for fpath in all_deg_files:
fname = os.path.basename(fpath)
genes = read_gene_ids(fpath)
result["files"][fname] = genes
if fname == "time_lrt.DE_genes.txt":
result["time_lrt_genes"] = genes
# Detect Con_vs_T1 (1h vs Con)
pref = PREFIX[stress]
if fname in [f"Con_vs_{pref}1.DE_genes.txt", f"{pref}1_vs_Con.DE_genes.txt"]:
result["con_vs_t1_genes"] = genes
result["all_con_vs_tx_genes"] |= genes
# Con_vs_T48
if fname == f"Con_vs_{pref}48.DE_genes.txt":
result["con_vs_t48_genes"] = genes
result["all_con_vs_tx_genes"] |= genes
return result
def deduplicate_gene_ids(genes_set):
"""Some genes might appear as LOC_Os and Os (alternate IDs) - normalize"""
# For now just return as-is since all use LOC_Os prefix consistently
return genes_set
def plot_venn(set1, set2, set3, labels, title, filename):
"""Plot 3-way Venn diagram"""
if not set1 or not set2 or not set3:
print(f" WARNING: Empty set for {title}")
return
plt.figure(figsize=(10, 8))
v = venn3([set1, set2, set3], set_labels=labels)
# Add counts to each region
for region in ["100", "010", "001", "110", "101", "011", "111"]:
area = v.get_patch_by_id(region)
label = v.get_label_by_id(region)
if label:
count = int(label.get_text())
# Color the central intersection differently
if region == "111" and area:
area.set_color("#E74C3C")
area.set_alpha(0.6)
plt.title(title, fontsize=14, fontweight="bold")
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT, filename), dpi=150)
plt.close()
print(f" Saved: {filename}")
def plot_deg_time_trends(all_data):
"""Plot DEG counts across time points for each stress"""
plt.figure(figsize=(12, 6))
colors = {"Salt": "#3498DB", "Drought": "#E67E22", "Heat": "#E74C3C"}
markers = {"Salt": "o", "Drought": "s", "Heat": "^"}
time_points = ["T1", "T3", "T6", "T12", "T24", "T36", "T48"]
for stress in ["Salt", "Drought", "Heat"]:
pref = PREFIX[stress]
data = all_data[stress]
# Counts for Tx_vs_T1 files (show cumulative change from 1h)
counts = []
for tp in time_points:
if tp == "T1":
cnt = len(data["con_vs_t1_genes"])
counts.append(cnt)
else:
fname = f"{pref}{tp.replace('T', '')}_vs_{pref}1.DE_genes.txt"
if fname in data["files"]:
counts.append(len(data["files"][fname]))
else:
counts.append(0)
plt.plot(
time_points,
counts,
color=colors[stress],
marker=markers[stress],
label=f"{stress}",
linewidth=2.5,
markersize=8,
)
# Annotate each point
for i, cnt in enumerate(counts):
plt.annotate(
f"{cnt}",
(time_points[i], cnt),
textcoords="offset points",
xytext=(0, 10),
ha="center",
fontsize=8,
color=colors[stress],
)
plt.xlabel("Time Point", fontsize=13)
plt.ylabel("Number of DEGs", fontsize=13)
plt.title(
"DEG Counts Across Time Points (vs 1h baseline)", fontsize=14, fontweight="bold"
)
plt.legend(fontsize=12)
plt.grid(axis="y", alpha=0.3)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT, "DEG_time_trends.png"), dpi=150)
plt.close()
print(" Saved: DEG_time_trends.png")
# Also plot Con_vs_T1 and Con_vs_T48 comparison
fig, ax = plt.subplots(figsize=(8, 6))
x = np.arange(3)
width = 0.25
early = [len(all_data[s]["con_vs_t1_genes"]) for s in ["Salt", "Drought", "Heat"]]
late = [len(all_data[s]["con_vs_t48_genes"]) for s in ["Salt", "Drought", "Heat"]]
bars1 = ax.bar(
x - width / 2, early, width, label="Early (1h vs Con)", color="#3498DB"
)
bars2 = ax.bar(
x + width / 2, late, width, label="Late (48h vs Con)", color="#E74C3C"
)
for bar in bars1:
ax.text(
bar.get_x() + bar.get_width() / 2.0,
bar.get_height() + 100,
f"{int(bar.get_height())}",
ha="center",
va="bottom",
fontsize=9,
)
for bar in bars2:
ax.text(
bar.get_x() + bar.get_width() / 2.0,
bar.get_height() + 100,
f"{int(bar.get_height())}",
ha="center",
va="bottom",
fontsize=9,
)
ax.set_xticks(x)
ax.set_xticklabels(["Salt", "Drought", "Heat"], fontsize=12)
ax.set_ylabel("Number of DEGs", fontsize=12)
ax.set_title("Early vs Late Response Comparison", fontsize=13, fontweight="bold")
ax.legend(fontsize=11)
ax.grid(axis="y", alpha=0.3)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT, "early_vs_late_response.png"), dpi=150)
plt.close()
print(" Saved: early_vs_late_response.png")
def print_venn_summary(name, salt_set, drought_set, heat_set):
"""Print detailed Venn breakdown"""
only_salt = salt_set - drought_set - heat_set
only_drought = drought_set - salt_set - heat_set
only_heat = heat_set - salt_set - drought_set
salt_drought = (salt_set & drought_set) - heat_set
salt_heat = (salt_set & heat_set) - drought_set
drought_heat = (drought_set & heat_set) - salt_set
all_three = salt_set & drought_set & heat_set
print(f"\n{'=' * 60}")
print(f" {name}")
print(f"{'=' * 60}")
print(f" Salt only: {len(only_salt):>6}")
print(f" Drought only: {len(only_drought):>6}")
print(f" Heat only: {len(only_heat):>6}")
print(f" Salt ∩ Drought: {len(salt_drought):>6}")
print(f" Salt ∩ Heat: {len(salt_heat):>6}")
print(f" Drought ∩ Heat: {len(drought_heat):>6}")
print(f" All Three (S∩D∩H): {len(all_three):>6}")
print(f" {'─' * 40}")
print(f" Salt total: {len(salt_set):>6}")
print(f" Drought total: {len(drought_set):>6}")
print(f" Heat total: {len(heat_set):>6}")
return {
"only_salt": only_salt,
"only_drought": only_drought,
"only_heat": only_heat,
"salt_drought": salt_drought,
"salt_heat": salt_heat,
"drought_heat": drought_heat,
"all_three": all_three,
}
def save_gene_lists(sets_dict, name):
"""Save gene lists to files"""
outdir = os.path.join(OUTPUT, f"{name}_gene_lists")
os.makedirs(outdir, exist_ok=True)
for key, genes in sets_dict.items():
if genes:
fpath = os.path.join(outdir, f"{key}.txt")
with open(fpath, "w") as f:
for g in sorted(genes):
f.write(f"{g}\n")
print(f" Saved {key}: {len(genes)} genes → {fpath}")
def main():
# Collect all data
all_data = {}
for stress, job_id in JOBS.items():
job_dir = os.path.join(WORKSPACE, job_id)
print(f"Reading {stress} data from {job_id}...")
data = collect_stress_data(job_dir, stress)
if data:
all_data[stress] = data
print(f" time_lrt: {len(data['time_lrt_genes'])} genes")
print(f" Con_vs_T1: {len(data['con_vs_t1_genes'])} genes")
print(f" Con_vs_T48: {len(data['con_vs_t48_genes'])} genes")
if len(all_data) < 3:
print("ERROR: Not enough stress data collected")
sys.exit(1)
# === ANALYSIS 1: Time LRT Venn ===
print("\n" + "█" * 60)
print(" ANALYSIS 1: Full Temporal Response (time_lrt)")
print("█" * 60)
salt_lrt = all_data["Salt"]["time_lrt_genes"]
drought_lrt = all_data["Drought"]["time_lrt_genes"]
heat_lrt = all_data["Heat"]["time_lrt_genes"]
sets_lrt = print_venn_summary("time_lrt DEGs", salt_lrt, drought_lrt, heat_lrt)
save_gene_lists(sets_lrt, "time_lrt_venn")
plot_venn(
salt_lrt,
drought_lrt,
heat_lrt,
["Salt", "Drought", "Heat"],
"Core Stress Response Genes (timeLRT)",
"venn_timelrt.png",
)
# === ANALYSIS 2: Early Response (Con_vs_T1) Venn ===
print("\n" + "█" * 60)
print(" ANALYSIS 2: Early Response (1h vs Con)")
print("█" * 60)
salt_early = all_data["Salt"]["con_vs_t1_genes"]
drought_early = all_data["Drought"]["con_vs_t1_genes"]
heat_early = all_data["Heat"]["con_vs_t1_genes"]
sets_early = print_venn_summary(
"Early Response (1h vs Con)", salt_early, drought_early, heat_early
)
save_gene_lists(sets_early, "early_venn")
plot_venn(
salt_early,
drought_early,
heat_early,
["Salt", "Drought", "Heat"],
"Early Response DEGs (1h vs Con)",
"venn_early.png",
)
# === ANALYSIS 3: Late Response (Con_vs_T48) Venn ===
if all(s["con_vs_t48_genes"] for s in all_data.values()):
print("\n" + "█" * 60)
print(" ANALYSIS 3: Late Response (48h vs Con)")
print("█" * 60)
salt_late = all_data["Salt"]["con_vs_t48_genes"]
drought_late = all_data["Drought"]["con_vs_t48_genes"]
heat_late = all_data["Heat"]["con_vs_t48_genes"]
sets_late = print_venn_summary(
"Late Response (48h vs Con)", salt_late, drought_late, heat_late
)
save_gene_lists(sets_late, "late_venn")
plot_venn(
salt_late,
drought_late,
heat_late,
["Salt", "Drought", "Heat"],
"Late Response DEGs (48h vs Con)",
"venn_late.png",
)
# === ANALYSIS 4: Time Trend Plot ===
print("\n" + "█" * 60)
print(" ANALYSIS 4: DEG Time Trends")
print("█" * 60)
plot_deg_time_trends(all_data)
# === SUMMARY FILE ===
print("\n" + "█" * 60)
print(" GENERATING SUMMARY REPORT")
print("█" * 60)
with open(os.path.join(OUTPUT, "cross_stress_summary.txt"), "w") as f:
f.write("=" * 70 + "\n")
f.write("PRJNA530826 Cross-Stress DEG Comparison Summary\n")
f.write("Salt (S) vs Drought (D) vs Heat (H)\n")
f.write("=" * 70 + "\n\n")
f.write("DATASET OVERVIEW\n")
f.write("-" * 40 + "\n")
f.write(f"{'Stress':<12} {'Job ID':<30} {'Samples':<10}\n")
f.write(f"{'─' * 52}\n")
for stress in ["Salt", "Drought", "Heat"]:
f.write(f"{stress:<12} {JOBS[stress]:<30} {'24':<10}\n")
f.write("\n\nDEG COUNT SUMMARY\n")
f.write("-" * 40 + "\n")
f.write(f"{'Comparison':<25} {'Salt':<10} {'Drought':<10} {'Heat':<10}\n")
f.write(f"{'─' * 55}\n")
f.write(
f"{'time_lrt':<25} {len(salt_lrt):<10} {len(drought_lrt):<10} {len(heat_lrt):<10}\n"
)
f.write(
f"{'1h vs Con':<25} {len(salt_early):<10} {len(drought_early):<10} {len(heat_early):<10}\n"
)
if all(s["con_vs_t48_genes"] for s in all_data.values()):
f.write(
f"{'48h vs Con':<25} {len(all_data['Salt']['con_vs_t48_genes']):<10} {len(all_data['Drought']['con_vs_t48_genes']):<10} {len(all_data['Heat']['con_vs_t48_genes']):<10}\n"
)
f.write("\n\nVENN INTERSECTIONS (time_lrt)\n")
f.write("-" * 40 + "\n")
for key, label in [
("all_three", "All Three (S∩D∩H)"),
("salt_drought", "Salt ∩ Drought"),
("salt_heat", "Salt ∩ Heat"),
("drought_heat", "Drought ∩ Heat"),
("only_salt", "Salt only"),
("only_drought", "Drought only"),
("only_heat", "Heat only"),
]:
f.write(f"{label:<25} {len(sets_lrt[key]):>7}\n")
f.write("\n\nVENN INTERSECTIONS (early - 1h vs Con)\n")
f.write("-" * 40 + "\n")
for key, label in [
("all_three", "All Three (S∩D∩H)"),
("salt_drought", "Salt ∩ Drought"),
("salt_heat", "Salt ∩ Heat"),
("drought_heat", "Drought ∩ Heat"),
("only_salt", "Salt only"),
("only_drought", "Drought only"),
("only_heat", "Heat only"),
]:
f.write(f"{label:<25} {len(sets_early[key]):>7}\n")
print(f"\n ✅ All results saved to: {OUTPUT}")
print(f" 📄 Summary: cross_stress_summary.txt")
print(f" 📊 Venn diagrams: venn_timelrt.png, venn_early.png, venn_late.png")
print(f" 📈 Time trends: DEG_time_trends.png, early_vs_late_response.png")
if __name__ == "__main__":
main()