| """Train a supervised Partial Least Squares drought-response model.""" |
|
|
| from pathlib import Path |
| import json |
|
|
| 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 |
| from sklearn.cross_decomposition import PLSRegression |
| from sklearn.pipeline import Pipeline |
| from sklearn.preprocessing import StandardScaler |
|
|
|
|
| PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| DATA_DIR = PROJECT_ROOT / "data" / "processed" |
| RESULTS_DIR = PROJECT_ROOT / "results" / "pls" |
| FIGURES_DIR = RESULTS_DIR / "figures" |
| MODEL_DIR = PROJECT_ROOT / "models" |
|
|
| COMPONENT_GRID = [1, 2, 3, 5, 10] |
| N_INPUT_GENES = 2_000 |
| MAX_ITERATIONS = 500 |
|
|
|
|
| def load_data() -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: |
| expression = pd.read_csv( |
| DATA_DIR / "expression_log2_cpm.tsv.gz", sep="\t", index_col=0 |
| ) |
| groups = pd.read_csv(DATA_DIR / "sample_groups.csv").set_index("sample_group_id") |
| groups = groups.loc[expression.columns] |
| joined = expression.T.join(groups[["ecotype", "condition"]]) |
| means = joined.groupby(["ecotype", "condition"]).mean() |
| control = means.xs("control", level="condition").sort_index() |
| drought = means.xs("drought", level="condition").sort_index() |
| delta = drought - control |
| splits = pd.read_csv(DATA_DIR / "ecotype_splits.csv").set_index("ecotype") |
| return control, delta, splits |
|
|
|
|
| def select_input_genes(control: pd.DataFrame, count: int = N_INPUT_GENES) -> pd.Index: |
| return control.var(axis=0).nlargest(count).index |
|
|
|
|
| def make_model(n_components: int) -> Pipeline: |
| return Pipeline( |
| [ |
| ("scale_inputs", StandardScaler()), |
| ( |
| "pls", |
| PLSRegression( |
| n_components=n_components, |
| scale=False, |
| max_iter=MAX_ITERATIONS, |
| tol=1e-6, |
| ), |
| ), |
| ] |
| ) |
|
|
|
|
| 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(actual: pd.DataFrame, predicted: np.ndarray, model: 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, |
| "mse": np.mean(error**2), |
| "mae": np.mean(np.abs(error)), |
| "pearson": pearson, |
| "spearman": spearman, |
| } |
| ) |
| return pd.DataFrame(rows) |
|
|
|
|
| def summarize(metrics: pd.DataFrame) -> pd.DataFrame: |
| summary = metrics.groupby("model")[["mse", "mae", "pearson", "spearman"]].agg( |
| ["mean", "median", "std"] |
| ) |
| summary.columns = ["_".join(column) for column in summary.columns] |
| return summary.reset_index() |
|
|
|
|
| def main() -> None: |
| control, delta, splits = load_data() |
| train_ids = splits.index[splits["split"].eq("train")] |
| validation_ids = splits.index[splits["split"].eq("validation")] |
| test_ids = splits.index[splits["split"].eq("test")] |
|
|
| training_genes = select_input_genes(control.loc[train_ids]) |
| validation_rows = [] |
| for components in COMPONENT_GRID: |
| model = make_model(components) |
| model.fit(control.loc[train_ids, training_genes], delta.loc[train_ids]) |
| prediction = model.predict(control.loc[validation_ids, training_genes]) |
| score = evaluate(delta.loc[validation_ids], prediction, "pls") |
| validation_rows.append( |
| { |
| "n_components": components, |
| "n_input_genes": len(training_genes), |
| "validation_mse": score["mse"].mean(), |
| "validation_mae": score["mae"].mean(), |
| "validation_pearson": score["pearson"].mean(), |
| "validation_spearman": score["spearman"].mean(), |
| "maximum_iterations_used": int( |
| np.max(model.named_steps["pls"].n_iter_) |
| ), |
| } |
| ) |
|
|
| validation = pd.DataFrame(validation_rows).sort_values( |
| ["validation_mse", "validation_mae", "n_components"] |
| ) |
| selected_components = int(validation.iloc[0]["n_components"]) |
|
|
| |
| |
| development_ids = train_ids.append(validation_ids) |
| development_genes = select_input_genes(control.loc[development_ids]) |
| final_model = make_model(selected_components) |
| final_model.fit( |
| control.loc[development_ids, development_genes], |
| delta.loc[development_ids], |
| ) |
| pls_prediction = final_model.predict(control.loc[test_ids, development_genes]) |
|
|
| development_mean = delta.loc[development_ids].mean(axis=0).to_numpy() |
| mean_prediction = np.tile(development_mean, (len(test_ids), 1)) |
| zero_prediction = np.zeros_like(pls_prediction) |
| pca_ridge_metrics = pd.read_csv( |
| PROJECT_ROOT / "results" / "reduced_rank_ridge" / "test_metrics_by_ecotype.csv" |
| ).loc[lambda frame: frame["model"].eq("pca_ridge")] |
|
|
| test_metrics = pd.concat( |
| [ |
| evaluate(delta.loc[test_ids], zero_prediction, "zero_change"), |
| evaluate(delta.loc[test_ids], mean_prediction, "development_mean"), |
| pca_ridge_metrics, |
| evaluate(delta.loc[test_ids], pls_prediction, "pls"), |
| ], |
| ignore_index=True, |
| ) |
| test_summary = summarize(test_metrics) |
|
|
| RESULTS_DIR.mkdir(parents=True, exist_ok=True) |
| FIGURES_DIR.mkdir(parents=True, exist_ok=True) |
| MODEL_DIR.mkdir(parents=True, exist_ok=True) |
| validation.to_csv(RESULTS_DIR / "validation_components.csv", index=False) |
| test_metrics.to_csv(RESULTS_DIR / "test_metrics_by_ecotype.csv", index=False) |
| test_summary.to_csv(RESULTS_DIR / "test_summary.csv", index=False) |
| pd.Series(development_genes, name="gene_id").to_csv( |
| RESULTS_DIR / "selected_input_genes.csv", index=False |
| ) |
| joblib.dump( |
| {"input_genes": development_genes, "pipeline": final_model}, |
| MODEL_DIR / "pls_model.joblib", |
| ) |
| (MODEL_DIR / "pls_model.json").write_text( |
| json.dumps( |
| { |
| "model": "Partial Least Squares regression", |
| "selected_components": selected_components, |
| "selected_input_genes": len(development_genes), |
| "input_gene_selection": "highest variance in development controls", |
| "selection_metric": "mean validation MSE across ecotypes", |
| "development_ecotypes": len(development_ids), |
| "test_ecotypes": len(test_ids), |
| "maximum_iterations_used_final_fit": int( |
| np.max(final_model.named_steps["pls"].n_iter_) |
| ), |
| }, |
| indent=2, |
| ), |
| encoding="utf-8", |
| ) |
|
|
| sns.set_theme(style="whitegrid", context="notebook") |
| fig, axes = plt.subplots(1, 2, figsize=(10.5, 4.3)) |
| ordered_validation = validation.sort_values("n_components") |
| sns.lineplot( |
| data=ordered_validation, |
| x="n_components", |
| y="validation_mse", |
| marker="o", |
| color="#356859", |
| ax=axes[0], |
| ) |
| axes[0].set( |
| title="PLS component selection", |
| xlabel="PLS components", |
| ylabel="Validation MSE", |
| ) |
|
|
| sns.boxplot( |
| data=test_metrics, |
| x="model", |
| y="mae", |
| hue="model", |
| legend=False, |
| palette=["#9AA4A6", "#3D6FB4", "#C58A2B", "#7A5195"], |
| ax=axes[1], |
| ) |
| axes[1].set(title="Supervised-model test error", xlabel="Model", ylabel="MAE") |
| axes[1].tick_params(axis="x", rotation=18) |
| fig.tight_layout() |
| fig.savefig(FIGURES_DIR / "pls_model_comparison.png", dpi=220) |
| plt.close(fig) |
|
|
| print("Selected using validation data:") |
| print(f"Input genes: {N_INPUT_GENES}") |
| print(f"PLS components: {selected_components}") |
| print(f"Validation MSE: {validation.iloc[0]['validation_mse']:.6f}") |
| print("\nFinal test means:") |
| print( |
| test_summary[["model", "mse_mean", "mae_mean", "pearson_mean", "spearman_mean"]] |
| .to_string(index=False) |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
|
|
