| """Evaluate the fitted PCA-ridge model on the independent E-MTAB-3279 study.""" |
|
|
| from pathlib import Path |
|
|
| import joblib |
| 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] |
| EXTERNAL_DIR = PROJECT_ROOT / "data" / "raw" / "E-MTAB-3279" |
| PROCESSED_DIR = PROJECT_ROOT / "data" / "processed" |
| RESULTS_DIR = PROJECT_ROOT / "results" / "external_validation" |
| FIGURES_DIR = RESULTS_DIR / "figures" |
|
|
| COUNTS_PATH = EXTERNAL_DIR / "Accessions.snp.htseq.txt" |
| METADATA_PATH = EXTERNAL_DIR / "E-MTAB-3279.sdrf.txt" |
| MODEL_PATH = PROJECT_ROOT / "models" / "pca_ridge_model.joblib" |
|
|
|
|
| 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 metrics(actual: pd.DataFrame, predicted: np.ndarray, model_name: str) -> pd.DataFrame: |
| rows = [] |
| for position, ecotype in enumerate(actual.index): |
| observed = actual.loc[ecotype].to_numpy() |
| estimate = predicted[position] |
| error = estimate - observed |
| pearson, spearman = safe_correlations(observed, estimate) |
| rows.append( |
| { |
| "ecotype": ecotype, |
| "model": model_name, |
| "mse": np.mean(error**2), |
| "mae": np.mean(np.abs(error)), |
| "pearson": pearson, |
| "spearman": spearman, |
| } |
| ) |
| return pd.DataFrame(rows) |
|
|
|
|
| def load_external_expression(training_genes: pd.Index) -> tuple[pd.DataFrame, pd.DataFrame]: |
| counts = pd.read_csv(COUNTS_PATH, sep="\t", index_col=0) |
| counts.index.name = "gene_id" |
| metadata = pd.read_csv(METADATA_PATH, sep="\t", keep_default_na=False) |
| metadata = metadata.drop_duplicates("Assay Name").set_index("Assay Name") |
|
|
| assay_names = counts.columns.str.replace( |
| ".snp.concordant_uniq.htseq", "", regex=False |
| ) |
| |
| |
| assay_names = assay_names.str.replace(r"([CS])2x$", r"\g<1>4", regex=True) |
| if not assay_names.isin(metadata.index).all(): |
| missing = assay_names[~assay_names.isin(metadata.index)].tolist() |
| raise ValueError(f"External count columns missing from metadata: {missing}") |
|
|
| missing_genes = training_genes.difference(counts.index) |
| if len(missing_genes): |
| raise ValueError(f"External counts lack {len(missing_genes)} training genes.") |
|
|
| library_sizes = counts.sum(axis=0) |
| log2_cpm = np.log2(counts.div(library_sizes, axis="columns") * 1_000_000 + 1) |
| log2_cpm = log2_cpm.loc[training_genes] |
|
|
| sample_metadata = metadata.loc[assay_names].reset_index() |
| sample_metadata["count_column"] = counts.columns.to_numpy() |
| sample_metadata["condition"] = sample_metadata[ |
| "Factor Value[growth condition]" |
| ].str.strip() |
| sample_metadata["ecotype"] = sample_metadata["Factor Value[ecotype]"].str.strip() |
|
|
| sample_expression = log2_cpm.T.reset_index(names="count_column").merge( |
| sample_metadata[["count_column", "ecotype", "condition"]], |
| on="count_column", |
| validate="one_to_one", |
| ) |
| ecotype_condition = sample_expression.drop(columns="count_column").groupby( |
| ["ecotype", "condition"] |
| ).mean() |
| control = ecotype_condition.xs("control", level="condition").sort_index() |
| drought = ecotype_condition.xs("mild drought", level="condition").sort_index() |
| if not control.index.equals(drought.index): |
| raise ValueError("External control and drought ecotypes do not align.") |
| return control, drought - control |
|
|
|
|
| def internal_development_mean(training_genes: pd.Index) -> tuple[np.ndarray, pd.DataFrame]: |
| expression = pd.read_csv( |
| PROCESSED_DIR / "expression_log2_cpm.tsv.gz", sep="\t", index_col=0 |
| ) |
| groups = pd.read_csv(PROCESSED_DIR / "sample_groups.csv").set_index("sample_group_id") |
| groups = groups.loc[expression.columns] |
| splits = pd.read_csv(PROCESSED_DIR / "ecotype_splits.csv").set_index("ecotype") |
|
|
| joined = expression.T.join(groups[["ecotype", "condition"]]) |
| means = joined.groupby(["ecotype", "condition"]).mean() |
| internal_delta = means.xs("drought", level="condition") - means.xs( |
| "control", level="condition" |
| ) |
| development_ids = splits.index[splits["split"].isin(["train", "validation"])] |
| return internal_delta.loc[development_ids, training_genes].mean(axis=0).to_numpy(), splits |
|
|
|
|
| def main() -> None: |
| model = joblib.load(MODEL_PATH) |
| training_genes = pd.Index(model.feature_names_in_) |
| control, external_delta = load_external_expression(training_genes) |
| development_mean, internal_splits = internal_development_mean(training_genes) |
|
|
| model_prediction = model.predict(control) |
| mean_prediction = np.tile(development_mean, (len(control), 1)) |
| zero_prediction = np.zeros_like(model_prediction) |
|
|
| result = pd.concat( |
| [ |
| metrics(external_delta, zero_prediction, "zero_change"), |
| metrics(external_delta, mean_prediction, "development_mean"), |
| metrics(external_delta, model_prediction, "pca_ridge"), |
| ], |
| ignore_index=True, |
| ) |
| internal_name = result["ecotype"].replace({"Shahdara": "Sha"}) |
| result["relationship_to_model_development"] = internal_name.map( |
| internal_splits["split"] |
| ).fillna("external_only") |
|
|
| summary = result.groupby("model")[["mse", "mae", "pearson", "spearman"]].agg( |
| ["mean", "median", "std"] |
| ) |
| summary.columns = ["_".join(column) for column in summary.columns] |
| summary = summary.reset_index() |
|
|
| RESULTS_DIR.mkdir(parents=True, exist_ok=True) |
| FIGURES_DIR.mkdir(parents=True, exist_ok=True) |
| result.to_csv(RESULTS_DIR / "external_metrics_by_ecotype.csv", index=False) |
| summary.to_csv(RESULTS_DIR / "external_summary.csv", index=False) |
|
|
| sns.set_theme(style="whitegrid", context="notebook") |
| plot_data = result.loc[result["model"].ne("zero_change")] |
| fig, axes = plt.subplots(1, 2, figsize=(10, 4.2)) |
| sns.boxplot(data=result, x="model", y="mae", hue="model", legend=False, |
| palette=["#9AA4A6", "#3D6FB4", "#C58A2B"], ax=axes[0]) |
| sns.stripplot(data=result, x="model", y="mae", color="#222222", size=4, |
| alpha=0.7, ax=axes[0]) |
| axes[0].set(title="External-study error", xlabel="Model", ylabel="MAE") |
| axes[0].tick_params(axis="x", rotation=15) |
|
|
| sns.barplot(data=plot_data, x="ecotype", y="pearson", hue="model", |
| palette={"development_mean": "#3D6FB4", "pca_ridge": "#C58A2B"}, |
| ax=axes[1]) |
| axes[1].axhline(0, color="#333333", linewidth=0.8) |
| axes[1].set(title="External correlation by ecotype", xlabel="Ecotype", |
| ylabel="Pearson correlation") |
| axes[1].tick_params(axis="x", rotation=30) |
| axes[1].legend(title="Model", frameon=False) |
| fig.tight_layout() |
| fig.savefig(FIGURES_DIR / "external_validation_comparison.png", dpi=220) |
| plt.close(fig) |
|
|
| print(f"External ecotypes: {', '.join(control.index)}") |
| print("Relationship to original split:") |
| print(result.drop_duplicates("ecotype")[["ecotype", "relationship_to_model_development"]].to_string(index=False)) |
| print("\nExternal summary (means across six ecotypes):") |
| print(summary[["model", "mse_mean", "mae_mean", "pearson_mean", "spearman_mean"]].to_string(index=False)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
