373 lines
12 KiB
Python
Executable File
373 lines
12 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""Part A: TF expression time-series clustering across 3 stresses"""
|
||
|
||
import os, sys, re, csv
|
||
import numpy as np
|
||
import pandas as pd
|
||
from collections import defaultdict, OrderedDict
|
||
from pathlib import Path
|
||
|
||
import matplotlib
|
||
|
||
matplotlib.use("Agg")
|
||
import matplotlib.pyplot as plt
|
||
import matplotlib.patches as mpatches
|
||
from matplotlib.colors import LinearSegmentedColormap
|
||
import seaborn as sns
|
||
from scipy.cluster.hierarchy import linkage, leaves_list
|
||
from scipy.stats import zscore
|
||
|
||
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)
|
||
|
||
JOBS = {
|
||
"Salt": "job_1780143735966_0e7710",
|
||
"Drought": "job_1780308105970_346d4e",
|
||
"Heat": "job_1780313949486_3f3732",
|
||
}
|
||
PREFIX = {"Salt": "S", "Drought": "D", "Heat": "H"}
|
||
TIME_POINTS = ["Con", "1", "3", "6", "12", "24", "36", "48"]
|
||
TIME_LABELS = ["Con", "1h", "3h", "6h", "12h", "24h", "36h", "48h"]
|
||
|
||
TF_FAMILIES = {
|
||
"WRKY": "PF03106",
|
||
"NAC": "PF02365",
|
||
"AP2_ERF": "PF00847",
|
||
"MYB": "PF00249",
|
||
"bZIP": "PF00170",
|
||
"bHLH": "PF00010",
|
||
"HSF": "PF00447",
|
||
}
|
||
|
||
|
||
def parse_featurecounts(filepath):
|
||
df = pd.read_csv(filepath, sep="\t", comment="#", skip_blank_lines=False)
|
||
if "Geneid" not in df.columns:
|
||
print(f" ERROR: unexpected format in {filepath}")
|
||
return None
|
||
# Keep only Geneid and count columns (drop Chr, Start, End, Strand, Length)
|
||
count_cols = [
|
||
c
|
||
for c in df.columns
|
||
if c != "Geneid" and c not in ("Chr", "Start", "End", "Strand", "Length")
|
||
]
|
||
counts = df[["Geneid"] + count_cols].copy()
|
||
counts.set_index("Geneid", inplace=True)
|
||
rename_map = {}
|
||
for c in counts.columns:
|
||
base = os.path.basename(c)
|
||
short = base.replace(".sorted.bam", "").replace(".bam", "")
|
||
rename_map[c] = short
|
||
counts.rename(columns=rename_map, inplace=True)
|
||
# Clean gene IDs (remove .1, .2 version suffixes)
|
||
counts.index = counts.index.str.replace(r"\.\d+$", "", regex=True)
|
||
# Aggregate duplicate gene IDs (sum)
|
||
counts = counts.groupby(counts.index).sum()
|
||
return counts
|
||
|
||
|
||
def get_tf_genes(ipr_file, pfam_id):
|
||
genes = set()
|
||
with open(ipr_file) as f:
|
||
for line in f:
|
||
if line.startswith("#"):
|
||
continue
|
||
if f"Name={pfam_id}" in line:
|
||
gene = line.split("\t")[0]
|
||
gene_clean = re.sub(r"\.\d+$", "", gene)
|
||
genes.add(gene_clean)
|
||
return genes
|
||
|
||
|
||
def build_sample_to_condition(job_dir, prefix):
|
||
st_file = os.path.join(job_dir, "rnaseq_sample_qc", "sample_table.tsv")
|
||
mapping = {}
|
||
with open(st_file) as f:
|
||
next(f)
|
||
for line in f:
|
||
parts = line.strip().split("\t")
|
||
if len(parts) < 2:
|
||
continue
|
||
sample, cond = parts[0], parts[1]
|
||
# Map to time point
|
||
if cond == "Con":
|
||
mapping[sample] = "Con"
|
||
else:
|
||
# Extract time point number
|
||
tp = sample.replace(prefix, "").split("-")[0]
|
||
mapping[sample] = tp
|
||
return mapping
|
||
|
||
|
||
def compute_mean_expression(counts_df, sample_map, target_genes):
|
||
avail = [g for g in target_genes if g in counts_df.index]
|
||
if not avail:
|
||
return None
|
||
sub = counts_df.loc[avail]
|
||
# Build time-point groups
|
||
tp_groups = defaultdict(list)
|
||
for sample, tp in sample_map.items():
|
||
if sample in sub.columns:
|
||
tp_groups[tp].append(sample)
|
||
# Compute mean per time point
|
||
result = []
|
||
for tp in TIME_POINTS:
|
||
if tp in tp_groups:
|
||
cols = [c for c in tp_groups[tp] if c in sub.columns]
|
||
if cols:
|
||
mean = sub[cols].mean(axis=1)
|
||
result.append(mean)
|
||
else:
|
||
result.append(pd.Series(0, index=sub.index))
|
||
else:
|
||
result.append(pd.Series(0, index=sub.index))
|
||
return pd.concat(result, axis=1)
|
||
|
||
|
||
def plot_tf_heatmap(expr_dict, tf_name, genes_list):
|
||
"""Create heatmap of z-score expression across all 3 stresses"""
|
||
n_genes = len(genes_list)
|
||
if n_genes == 0:
|
||
print(f" No {tf_name} genes to plot")
|
||
return
|
||
|
||
# Build matrix: genes × (stress_timepoint)
|
||
all_data = []
|
||
for stress in ["Salt", "Drought", "Heat"]:
|
||
if stress not in expr_dict or expr_dict[stress] is None:
|
||
continue
|
||
mat = expr_dict[stress].copy()
|
||
if mat.shape[1] != len(TIME_LABELS):
|
||
print(
|
||
f" WARNING: {stress} has {mat.shape[1]} cols, expected {len(TIME_LABELS)}"
|
||
)
|
||
continue
|
||
mat.columns = [f"{stress}_{t}" for t in TIME_LABELS]
|
||
mat.index = mat.index.astype(str)
|
||
all_data.append(mat)
|
||
|
||
if not all_data:
|
||
return
|
||
|
||
combined = pd.concat(all_data, axis=1)
|
||
# Ensure string column names
|
||
combined.columns = [str(c) for c in combined.columns]
|
||
|
||
# Z-score normalize per gene (row-wise)
|
||
mean_val = combined.mean(axis=1)
|
||
std_val = combined.std(axis=1)
|
||
z = combined.subtract(mean_val, axis=0).div(std_val.replace(0, 1), axis=0)
|
||
z = z.fillna(0)
|
||
|
||
# Filter to our TF genes
|
||
common_genes = [g for g in genes_list if g in z.index]
|
||
if not common_genes or len(common_genes) == 0:
|
||
return
|
||
z = z.loc[common_genes]
|
||
# Ensure it's a DataFrame even with 1 gene
|
||
if isinstance(z, pd.Series):
|
||
z = z.to_frame().T
|
||
|
||
# Hierarchical clustering
|
||
if len(z) > 2:
|
||
try:
|
||
Z = linkage(z.values, method="ward", metric="euclidean")
|
||
order = leaves_list(Z)
|
||
z = z.iloc[order]
|
||
except:
|
||
pass
|
||
|
||
# Plot
|
||
fig_height = max(4, min(20, n_genes * 0.3))
|
||
fig, ax = plt.subplots(figsize=(14, fig_height))
|
||
|
||
# Color bar for stresses
|
||
stress_colors = {"Salt": "#3498DB", "Drought": "#E67E22", "Heat": "#E74C3C"}
|
||
col_colors = []
|
||
for col in z.columns:
|
||
for s in ["Salt", "Drought", "Heat"]:
|
||
if col.startswith(s):
|
||
col_colors.append(stress_colors[s])
|
||
break
|
||
|
||
sns.heatmap(
|
||
z,
|
||
cmap="RdBu_r",
|
||
center=0,
|
||
ax=ax,
|
||
cbar_kws={"label": "Z-score"},
|
||
xticklabels=True,
|
||
yticklabels=len(common_genes) <= 50,
|
||
)
|
||
ax.set_title(
|
||
f"{tf_name} Expression Across Stresses (z-score)",
|
||
fontsize=14,
|
||
fontweight="bold",
|
||
)
|
||
ax.set_ylabel("Genes")
|
||
ax.set_xlabel("Stress × Time Point")
|
||
|
||
# Legend for stress colors
|
||
patches = [mpatches.Patch(color=c, label=s) for s, c in stress_colors.items()]
|
||
ax.legend(handles=patches, loc="upper right", fontsize=10)
|
||
|
||
plt.tight_layout()
|
||
outfile = os.path.join(OUTPUT, f"heatmap_{tf_name.lower()}.png")
|
||
fig.savefig(outfile, dpi=150, bbox_inches="tight")
|
||
plt.close()
|
||
print(f" Saved: {os.path.basename(outfile)}")
|
||
|
||
|
||
def plot_tf_family_trends(expr_dict, tf_name):
|
||
"""Line plot of mean expression trend per stress for a TF family"""
|
||
fig, ax = plt.subplots(figsize=(10, 6))
|
||
colors = {"Salt": "#3498DB", "Drought": "#E67E22", "Heat": "#E74C3C"}
|
||
markers = {"Salt": "o", "Drought": "s", "Heat": "^"}
|
||
|
||
for stress in ["Salt", "Drought", "Heat"]:
|
||
if stress not in expr_dict or expr_dict[stress] is None:
|
||
continue
|
||
mat = expr_dict[stress]
|
||
mean_vals = mat.mean(axis=0).values
|
||
sem_vals = mat.sem(axis=0).values
|
||
|
||
ax.errorbar(
|
||
range(len(TIME_LABELS)),
|
||
mean_vals,
|
||
yerr=sem_vals,
|
||
color=colors[stress],
|
||
marker=markers[stress],
|
||
label=stress,
|
||
linewidth=2,
|
||
capsize=4,
|
||
markersize=6,
|
||
)
|
||
|
||
# Annotate
|
||
for i, v in enumerate(mean_vals):
|
||
ax.annotate(
|
||
f"{v:.0f}",
|
||
(i, v),
|
||
textcoords="offset points",
|
||
xytext=(0, 8),
|
||
ha="center",
|
||
fontsize=7,
|
||
color=colors[stress],
|
||
)
|
||
|
||
ax.set_xticks(range(len(TIME_LABELS)))
|
||
ax.set_xticklabels(TIME_LABELS, fontsize=11)
|
||
ax.set_xlabel("Time Point", fontsize=12)
|
||
ax.set_ylabel("Mean CPM", fontsize=12)
|
||
ax.set_title(f"{tf_name} Family Expression Trend", fontsize=13, fontweight="bold")
|
||
ax.legend(fontsize=11)
|
||
ax.grid(axis="y", alpha=0.3)
|
||
plt.tight_layout()
|
||
|
||
outfile = os.path.join(OUTPUT, f"trend_{tf_name.lower()}.png")
|
||
fig.savefig(outfile, dpi=150, bbox_inches="tight")
|
||
plt.close()
|
||
print(f" Saved: {os.path.basename(outfile)}")
|
||
|
||
|
||
def main():
|
||
ipr_file = os.path.join(
|
||
WORKSPACE, "job_1780143735966_0e7710/annotation/analysis/Os.interproscan.gff3"
|
||
)
|
||
|
||
# Step 1: Read featurecounts for all stresses
|
||
print("Reading featurecounts...")
|
||
count_data = {}
|
||
sample_maps = {}
|
||
for stress, job_id in JOBS.items():
|
||
fc_file = os.path.join(
|
||
WORKSPACE, job_id, "rnaseq_quant/analysis/featurecounts.txt"
|
||
)
|
||
job_dir = os.path.join(WORKSPACE, job_id)
|
||
if os.path.exists(fc_file):
|
||
counts = parse_featurecounts(fc_file)
|
||
if counts is not None:
|
||
count_data[stress] = counts
|
||
# CPM normalization
|
||
lib_sizes = counts.sum(axis=0)
|
||
cpm = counts.div(lib_sizes, axis=1) * 1e6
|
||
count_data[stress] = cpm
|
||
sample_maps[stress] = build_sample_to_condition(job_dir, PREFIX[stress])
|
||
print(f" {stress}: {counts.shape[0]} genes, {counts.shape[1]} samples")
|
||
|
||
# Step 2: Get TF genes
|
||
print("\nGetting TF family members from InterProScan...")
|
||
tf_genes = {}
|
||
for tf_name, pfam_id in TF_FAMILIES.items():
|
||
genes = get_tf_genes(ipr_file, pfam_id)
|
||
print(f" {tf_name}: {len(genes)} genes (PFAM: {pfam_id})")
|
||
tf_genes[tf_name] = sorted(genes)
|
||
|
||
# Step 3: Compute mean expression per TF per stress
|
||
print("\nComputing time-series expression...")
|
||
tf_expr = defaultdict(dict)
|
||
for tf_name, genes in tf_genes.items():
|
||
if not genes:
|
||
continue
|
||
for stress in ["Salt", "Drought", "Heat"]:
|
||
if stress not in count_data:
|
||
continue
|
||
expr = compute_mean_expression(
|
||
count_data[stress], sample_maps[stress], genes
|
||
)
|
||
if expr is not None:
|
||
expr.columns = TIME_LABELS
|
||
tf_expr[tf_name][stress] = expr
|
||
|
||
# Step 4: Generate plots for each TF family
|
||
print("\nGenerating plots...")
|
||
for tf_name in TF_FAMILIES.keys():
|
||
genes = tf_genes.get(tf_name, [])
|
||
if not genes:
|
||
continue
|
||
print(f"\n--- {tf_name} ({len(genes)} genes) ---")
|
||
plot_tf_heatmap(dict(tf_expr.get(tf_name, {})), tf_name, genes)
|
||
plot_tf_family_trends(dict(tf_expr.get(tf_name, {})), tf_name)
|
||
|
||
# Step 5: Combined ALL TF comparison heatmap
|
||
print("\n--- Combined TF comparison ---")
|
||
n_cols = len(TIME_LABELS)
|
||
fig, axes = plt.subplots(len(TF_FAMILIES), 3, figsize=(15, len(TF_FAMILIES) * 3.5))
|
||
|
||
for i, tf_name in enumerate(TF_FAMILIES.keys()):
|
||
for j, stress in enumerate(["Salt", "Drought", "Heat"]):
|
||
ax = axes[i][j] if len(TF_FAMILIES) > 1 else axes[j]
|
||
if stress in tf_expr.get(tf_name, {}):
|
||
mat = tf_expr[tf_name][stress]
|
||
mean_vals = mat.mean(axis=0).values
|
||
sem_vals = mat.sem(axis=0).values
|
||
ax.errorbar(
|
||
range(n_cols),
|
||
mean_vals,
|
||
yerr=sem_vals,
|
||
color=["#3498DB", "#E67E22", "#E74C3C"][j],
|
||
marker="o",
|
||
linewidth=2,
|
||
capsize=3,
|
||
)
|
||
ax.set_xticks(range(n_cols))
|
||
ax.set_xticklabels(TIME_LABELS, fontsize=7, rotation=45)
|
||
if j == 0:
|
||
ax.set_ylabel(f"{tf_name}\n(CPM)", fontsize=9)
|
||
if i == 0:
|
||
ax.set_title(stress, fontsize=11, fontweight="bold")
|
||
|
||
plt.tight_layout()
|
||
outfile = os.path.join(OUTPUT, "tf_family_comparison_grid.png")
|
||
fig.savefig(outfile, dpi=150, bbox_inches="tight")
|
||
plt.close()
|
||
print(f" Saved: {os.path.basename(outfile)}")
|
||
|
||
print(f"\nAll plots saved to: {OUTPUT}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|