| """Create ecotype-level splits and evaluate simple drought-response baselines.""" |
|
|
| from pathlib import Path |
|
|
| import matplotlib.pyplot as plt |
| import numpy as np |
| import pandas as pd |
| import seaborn as sns |
| from scipy.stats import pearsonr, spearmanr |
|
|
|
|
| PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| DATA_DIR = PROJECT_ROOT / "data" / "processed" |
| RESULTS_DIR = PROJECT_ROOT / "results" / "baselines" |
| FIGURES_DIR = RESULTS_DIR / "figures" |
|
|
| EXPRESSION_PATH = DATA_DIR / "expression_log2_cpm.tsv.gz" |
| GROUPS_PATH = DATA_DIR / "sample_groups.csv" |
| SPLITS_PATH = DATA_DIR / "ecotype_splits.csv" |
|
|
| RANDOM_SEED = 42 |
| TRAIN_FRACTION = 0.70 |
| VALIDATION_FRACTION = 0.15 |
|
|
| CONTROL = "#356859" |
| DROUGHT = "#C58A2B" |
| BLUE = "#3D6FB4" |
|
|
|
|
| def build_ecotype_matrices() -> tuple[pd.DataFrame, pd.DataFrame]: |
| expression = pd.read_csv(EXPRESSION_PATH, sep="\t", index_col=0) |
| groups = pd.read_csv(GROUPS_PATH).set_index("sample_group_id") |
| groups = groups.loc[expression.columns] |
|
|
| sample_expression = expression.T.join(groups[["ecotype", "condition"]]) |
|
|
| |
| |
| ecotype_condition = sample_expression.groupby(["ecotype", "condition"]).mean() |
| control = ecotype_condition.xs("control", level="condition").sort_index() |
| drought = ecotype_condition.xs("drought", level="condition").sort_index() |
|
|
| if not control.index.equals(drought.index): |
| raise ValueError("Control and drought ecotype sets are different.") |
| if not control.columns.equals(drought.columns): |
| raise ValueError("Control and drought gene sets are different.") |
|
|
| delta = drought - control |
| return control, delta |
|
|
|
|
| def make_splits(ecotypes: pd.Index) -> pd.DataFrame: |
| rng = np.random.default_rng(RANDOM_SEED) |
| shuffled = rng.permutation(np.asarray(sorted(ecotypes))) |
|
|
| n_total = len(shuffled) |
| n_train = round(TRAIN_FRACTION * n_total) |
| n_validation = round(VALIDATION_FRACTION * n_total) |
|
|
| split = np.empty(n_total, dtype=object) |
| split[:n_train] = "train" |
| split[n_train : n_train + n_validation] = "validation" |
| split[n_train + n_validation :] = "test" |
|
|
| splits = pd.DataFrame({"ecotype": shuffled, "split": split}) |
| splits["random_seed"] = RANDOM_SEED |
| return splits.sort_values("ecotype").reset_index(drop=True) |
|
|
|
|
| def safe_correlations(actual: np.ndarray, predicted: np.ndarray) -> tuple[float, float]: |
| if np.allclose(predicted, predicted[0]) or np.allclose(actual, actual[0]): |
| return np.nan, np.nan |
| return pearsonr(actual, predicted).statistic, spearmanr(actual, predicted).statistic |
|
|
|
|
| def evaluate_prediction( |
| actual_delta: pd.DataFrame, |
| predicted_delta: np.ndarray, |
| baseline: str, |
| split: str, |
| ) -> pd.DataFrame: |
| records = [] |
| for ecotype, row in actual_delta.iterrows(): |
| actual = row.to_numpy() |
| error = predicted_delta - actual |
| pearson, spearman = safe_correlations(actual, predicted_delta) |
| records.append( |
| { |
| "ecotype": ecotype, |
| "split": split, |
| "baseline": baseline, |
| "mse": np.mean(error**2), |
| "mae": np.mean(np.abs(error)), |
| "pearson": pearson, |
| "spearman": spearman, |
| } |
| ) |
| return pd.DataFrame(records) |
|
|
|
|
| def summarize_metrics(metrics: pd.DataFrame) -> pd.DataFrame: |
| value_columns = ["mse", "mae", "pearson", "spearman"] |
| summary = ( |
| metrics.groupby(["split", "baseline"])[value_columns] |
| .agg(["mean", "median", "std"]) |
| .reset_index() |
| ) |
| summary.columns = [ |
| "_".join(part for part in column if part) |
| if isinstance(column, tuple) |
| else column |
| for column in summary.columns |
| ] |
| return summary |
|
|
|
|
| def plot_metric_distributions(metrics: pd.DataFrame) -> None: |
| evaluation = metrics.loc[metrics["split"].isin(["validation", "test"])].copy() |
| sns.set_theme(style="whitegrid", context="notebook") |
|
|
| fig, axes = plt.subplots(1, 2, figsize=(10, 4.2)) |
| sns.boxplot( |
| data=evaluation, |
| x="split", |
| y="mae", |
| hue="baseline", |
| palette={"zero_change": "#9AA4A6", "training_mean": BLUE}, |
| ax=axes[0], |
| ) |
| axes[0].set(title="Absolute error by held-out ecotype", xlabel="Split", ylabel="MAE") |
| axes[0].legend(title="Baseline", frameon=False) |
|
|
| correlation = evaluation.loc[evaluation["baseline"].eq("training_mean")] |
| correlation_long = correlation.melt( |
| id_vars=["ecotype", "split"], |
| value_vars=["pearson", "spearman"], |
| var_name="metric", |
| value_name="correlation", |
| ) |
| sns.boxplot( |
| data=correlation_long, |
| x="split", |
| y="correlation", |
| hue="metric", |
| palette={"pearson": CONTROL, "spearman": DROUGHT}, |
| ax=axes[1], |
| ) |
| axes[1].axhline(0, color="#333333", linewidth=0.8) |
| axes[1].set( |
| title="Training-mean response correlation", |
| xlabel="Split", |
| ylabel="Correlation across genes", |
| ) |
| axes[1].legend(title="Metric", frameon=False) |
| fig.tight_layout() |
| fig.savefig(FIGURES_DIR / "baseline_metric_distributions.png", dpi=220) |
| plt.close(fig) |
|
|
|
|
| def plot_representative_test_example( |
| delta: pd.DataFrame, metrics: pd.DataFrame, mean_response: pd.Series |
| ) -> None: |
| test_metrics = metrics.loc[ |
| metrics["split"].eq("test") & metrics["baseline"].eq("training_mean") |
| ].sort_values("pearson") |
| representative = test_metrics.iloc[len(test_metrics) // 2]["ecotype"] |
|
|
| actual = delta.loc[representative] |
| predicted = mean_response |
|
|
| sns.set_theme(style="whitegrid", context="notebook") |
| fig, ax = plt.subplots(figsize=(6.3, 5.3)) |
| ax.hexbin( |
| actual, |
| predicted, |
| gridsize=55, |
| mincnt=1, |
| cmap="YlGnBu", |
| linewidths=0, |
| ) |
| limit = max(abs(actual).max(), abs(predicted).max()) |
| ax.plot([-limit, limit], [-limit, limit], linestyle="--", color="#555555") |
| pearson_value = test_metrics.loc[ |
| test_metrics["ecotype"].eq(representative), "pearson" |
| ].iloc[0] |
| ax.set( |
| title=f"Representative test ecotype: {representative}", |
| xlabel="Actual drought - control expression", |
| ylabel="Training-mean predicted response", |
| ) |
| ax.text( |
| 0.03, |
| 0.97, |
| f"Pearson r = {pearson_value:.3f}", |
| transform=ax.transAxes, |
| va="top", |
| bbox={"facecolor": "white", "alpha": 0.85, "edgecolor": "none"}, |
| ) |
| fig.tight_layout() |
| fig.savefig(FIGURES_DIR / "representative_test_prediction.png", dpi=220) |
| plt.close(fig) |
|
|
|
|
| def main() -> None: |
| _, delta = build_ecotype_matrices() |
| splits = make_splits(delta.index) |
|
|
| split_lookup = splits.set_index("ecotype")["split"] |
| train_ecotypes = split_lookup.index[split_lookup.eq("train")] |
| mean_response = delta.loc[train_ecotypes].mean(axis=0) |
| zero_response = np.zeros(delta.shape[1], dtype=float) |
|
|
| metric_frames = [] |
| for split_name in ["train", "validation", "test"]: |
| ecotypes = split_lookup.index[split_lookup.eq(split_name)] |
| actual = delta.loc[ecotypes] |
| metric_frames.append( |
| evaluate_prediction(actual, zero_response, "zero_change", split_name) |
| ) |
| metric_frames.append( |
| evaluate_prediction( |
| actual, mean_response.to_numpy(), "training_mean", split_name |
| ) |
| ) |
|
|
| metrics = pd.concat(metric_frames, ignore_index=True) |
| summary = summarize_metrics(metrics) |
|
|
| RESULTS_DIR.mkdir(parents=True, exist_ok=True) |
| FIGURES_DIR.mkdir(parents=True, exist_ok=True) |
| splits.to_csv(SPLITS_PATH, index=False) |
| metrics.to_csv(RESULTS_DIR / "baseline_metrics_by_ecotype.csv", index=False) |
| summary.to_csv(RESULTS_DIR / "baseline_summary.csv", index=False) |
| mean_response.rename("mean_delta").to_csv( |
| RESULTS_DIR / "training_mean_response.tsv.gz", |
| sep="\t", |
| compression="gzip", |
| ) |
|
|
| plot_metric_distributions(metrics) |
| plot_representative_test_example(delta, metrics, mean_response) |
|
|
| print("Split sizes:") |
| print(splits["split"].value_counts().reindex(["train", "validation", "test"])) |
| print("\nValidation and test summary (means across ecotypes):") |
| columns = ["split", "baseline", "mse_mean", "mae_mean", "pearson_mean", "spearman_mean"] |
| print(summary.loc[summary["split"].ne("train"), columns].to_string(index=False)) |
| print(f"\nWrote outputs to {RESULTS_DIR}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
|
|
