From 6faab61ebe8e3007c441a53ef29e655f1a755d1b Mon Sep 17 00:00:00 2001 From: Vittoria Tommasini Date: Tue, 26 May 2026 18:14:12 -0700 Subject: [PATCH] Add LOO and plotting diagnostics and unit tests Simplified parallelization and tests, and decreased number of test_data points to make tests run faster Shortened tests, fixed multiprocessing hang in loo_crossval that was causing tests to time out --- .github/workflows/Tests.yaml | 2 + pyproject.toml | 1 + src/SimulationSupport/gpr/diagnostics.py | 252 +++++++++++++++++++++++ tests/test_diagnostics.py | 126 ++++++++++++ 4 files changed, 381 insertions(+) create mode 100644 src/SimulationSupport/gpr/diagnostics.py create mode 100644 tests/test_diagnostics.py diff --git a/.github/workflows/Tests.yaml b/.github/workflows/Tests.yaml index 6f27ff3..7e9f94d 100644 --- a/.github/workflows/Tests.yaml +++ b/.github/workflows/Tests.yaml @@ -16,6 +16,8 @@ jobs: - uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} + - name: Install system dependencies + run: sudo apt-get update && sudo apt-get install -y libfftw3-dev - name: Install dependencies run: | python -m pip install --upgrade pip diff --git a/pyproject.toml b/pyproject.toml index 8502ad0..de6be2c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,7 @@ dependencies = [ "torch", "gpytorch", "pandas", + "sxs", ] [project.optional-dependencies] diff --git a/src/SimulationSupport/gpr/diagnostics.py b/src/SimulationSupport/gpr/diagnostics.py new file mode 100644 index 0000000..93bb9ee --- /dev/null +++ b/src/SimulationSupport/gpr/diagnostics.py @@ -0,0 +1,252 @@ +# Distributed under the MIT License. +# See LICENSE.txt for details. + +""" +Gaussian Process Regression machine learning diagnostic functions library. +Contains all the functions necessary to validate and plot the GPR model used +to predict better low-eccentricity orbital parameter initial guesses. +""" + +import logging +import multiprocessing as mp +import os + +import matplotlib.pyplot as plt +import numpy as np +import torch + +from SimulationSupport.gpr import predict_with_gpr_model, train_gpr_model + +logger = logging.getLogger(__name__) + + +# Leave-one-out parallelization set up +def _loo_single(i, X, Y): + """ + Run a single LOO iteration for index i. + + Trains a GPR on all points except index i, then predicts the held-out point. + Called in parallel by loo_crossval. + + Args: + i (int): Index of the held-out point + X (np.ndarray): Input features, with shape (N, D) + Y (np.ndarray): Target variable, with shape (N, ) + + Returns: + pred_mean (float): Predicted mean for the held-out point + pred_std (float): Predicted std dev for the held-out point + """ + N = len(Y) + + # Create train and test split + # Boolean mask: all True except index i (held out point) + train_mask = np.ones(N, dtype=bool) + train_mask[i] = False + + X_train = X[train_mask] + Y_train = Y[train_mask] + # Slice preserves the 2D shape expected by the model + X_test = X[i : i + 1] + + # Train and predict + model_loo, likelihood_loo = train_gpr_model(X_train, Y_train) + pred_mean, pred_std = predict_with_gpr_model( + X_test, model_loo, likelihood_loo + ) + + return pred_mean[0], pred_std[0] + + +# Leave-one-out cross-validation +def loo_crossval( + X: np.ndarray, + Y: np.ndarray, + target_name="Target", + n_jobs=None, +): + """ + Perform Leave-One-Out Cross-Validation for a GPR model. + + Trains N models (each omitting one point), predicts the held-out point, + collect predictions and uncertainties, and then computes and plots summary metrics. + This gives an unbiased estimate of generalization performance. + + Args: + X (np.ndarray): Input features, with shape (N, D) + Y (np.ndarray): Target variable, with shape (N, ) + target_name (str): Label for plots and print outputs + n_jobs (int): Number of parallel workers (None uses all available cores) + + Returns: + predictions_loo (np.ndarray): LOO predicted values, with shape (N, ) + uncertainties_loo (np.ndarray): LOO predicted std devs, with shape (N, ) + """ + N = len(Y) + + # Force single worker when GPU is used as parallel processes can't share a GPU + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + if device.type == "cuda": + n_jobs = 1 + + logger.info( + f"Processing {N} LOO iterations for {target_name} using" + f" {n_jobs or os.cpu_count()} parallel workers..." + ) + + # Build argument tuples for each LOO fold + # because pool.starmap needs one tuple of positional args per call + args = [(i, X, Y) for i in range(N)] + + # Each worker is a new Python process and reimports the modules it needs + # Avoids deadlock + ctx = mp.get_context("spawn") + if n_jobs == 1: + # Run sequentially when there is a single worker + results = [_loo_single(*a) for a in args] + else: + # Run LOO folds in parallel with multiple workers + with ctx.Pool(processes=n_jobs) as pool: + results = pool.starmap(_loo_single, args) + + # Unpack results back into the predictions and uncertainties arrays + predictions_loo, uncertainties_loo = np.array(results).T + + return predictions_loo, uncertainties_loo + + +def plot_loo_crossval(Y, predictions_loo, target_name="Target", plot=False): + """ + Compute summary statistics for Leave-One-Out Cross-Validation results, and optionally + plot the correlation. + + Args: + Y (np.ndarray): True target values, with shape (N, ) + predictions_loo (np.ndarray): LOO predicted values, with shape (N, ) + target_name (str): Label for plots and log outputs + plot (bool): Whether to produce a correlation plot. Default is False. + + Returns: + rmse_loo (float): Root mean squared error of the LOO predictions + mae_loo (float): Mean absolute error of the LOO predictions + r_squared_loo (float): R^2 computed from the Pearson correlation + """ + + Y_loo = Y # Same as the original Y for the multi input case + + # Calculate metrics - always computed, regardless of whether plot is requested + # R^2 is computed from the Pearson correlation coefficient - + # equivalent to the coefficient of determination for a linear fit through the origin + correlation = np.corrcoef(Y_loo, predictions_loo)[0, 1] + r_squared_loo = correlation**2 + + # Metrics with goal values + rmse_loo = np.sqrt(np.mean((Y_loo - predictions_loo) ** 2)) + mae_loo = np.mean(np.abs(Y_loo - predictions_loo)) + y_range = Y_loo.max() - Y_loo.min() + + # Plot correlation + if plot: + plt.figure(figsize=(8, 6)) + plt.scatter(Y_loo, predictions_loo, alpha=0.6, s=20) + + # y = x reference line: perfect predictions would lie exactly on this line + min_val = min(Y_loo.min(), predictions_loo.min()) + max_val = max(Y_loo.max(), predictions_loo.max()) + plt.plot( + [min_val, max_val], + [min_val, max_val], + "r--", + lw=2, + label="Perfect Correlation", + ) + + # Labels and formatting + plt.xlabel(f"True Δ{target_name}", fontsize=12) + plt.ylabel(f"LOO Predicted Δ{target_name}", fontsize=12) + plt.title(f"LOO: GPR Predictions vs True ({target_name})", fontsize=14) + plt.grid(True, alpha=0.3) + plt.legend() + + # Display R^2 + plt.text( + 0.95, + 0.95, + f"R² = {r_squared_loo:.4f}", + transform=plt.gca().transAxes, + fontsize=12, + bbox=dict(boxstyle="round", facecolor="white", alpha=0.8), + horizontalalignment="right", + ) + plt.tight_layout() + plt.show() + + # Log metrics with goal values for quick analysis + logger.info(f"Leave-one-out Cross Validation Results ({target_name})") + logger.info( + f"RMSE: {rmse_loo:.6f} ({100 * rmse_loo / y_range:.6f}% of target" + " range; goal: < 1 % of target range, lower is better)" + ) + logger.info( + f"MAE: {mae_loo:.6f} ({100 * mae_loo / y_range:.6f}% of target range;" + " goal: < 1 % of target range, lower is better)" + ) + logger.info( + f"R²: {r_squared_loo:.4f} (goal: > 0.95 excellent, > 0.90 good, < 0.70" + " poor)" + ) + logger.info(f"\n Dataset size: {len(Y_loo)} points") + logger.info( + f"Each model is trained on {len(Y_loo)-1} points, and tested on 1 point" + ) + logger.info("This provides an unbiased generalization estimate.") + + return rmse_loo, mae_loo, r_squared_loo + + +# Leave-one-out residual computation and plotting +def plot_loo_residuals(Y_loo, predictions_loo, target_name="Target", show=True): + """ + Calculate LOO prediction residuals, plot a histogram, and print summary statistics. + + Args: + Y_loo (np.ndarray): True target values from LOO cross validation + predictions_loo (np.ndarray): Predicted values from LOO cross validation + target_name (str): Name of target variable + show (bool): Whether to produce a residual plot. Default is True. + + Returns: + residuals_loo (np.ndarray): Per-point residuals (true - predicted), with shape (N,) + + """ + + # Compute residuals: LOO prediction error + # Residual = true - predicted + residuals_loo = Y_loo - predictions_loo + + # Make histogram + plt.figure(figsize=(8, 5)) + plt.hist(residuals_loo, bins=20, color="skyblue", edgecolor="k", alpha=0.8) + # Vertical line at zero: residuals centered here indicate no systematic bias; + # ideally the histogram is centered on this line; a shifted distribution suggests + # the model is over- or under-predicting + plt.axvline(0, color="r", linestyle="--", label="Zero Error") + + plt.title(f"LOO Residuals Histogram for {target_name}") + plt.xlabel(" Residuals", fontsize=16) + plt.ylabel("Count", fontsize=16) + plt.tick_params(axis="both", which="major", labelsize=14) + plt.grid(True, alpha=0.3) + plt.legend(fontsize=14) + plt.tight_layout() + if show: + plt.show() + + # Print statistics + logger.info(f"Residual statistics for {target_name}:") + logger.info(f"Mean residual: {np.mean(residuals_loo):.4e}") + logger.info(f"Std of residuals: {np.std (residuals_loo):.4e}") + logger.info(f"Max residual: {np.max (residuals_loo):.4e}") + logger.info(f"Min residual: {np.min (residuals_loo):.4e}") + + return residuals_loo diff --git a/tests/test_diagnostics.py b/tests/test_diagnostics.py new file mode 100644 index 0000000..a99ef64 --- /dev/null +++ b/tests/test_diagnostics.py @@ -0,0 +1,126 @@ +# Distributed under the MIT License. +# See LICENSE.txt for details. + +""" +Unit tests for diagnostics.py + +Run with: + python -m unittest test_diagnostics.py -v +""" + +import os +import unittest + +import matplotlib +import numpy as np +import pandas as pd + +TEST_DATA_PATH = os.path.join(os.path.dirname(__file__), "test_data.csv") + +# Use a non-interactive backend so the tests can run in CLI without plot outputs +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +from SimulationSupport.gpr.diagnostics import ( + loo_crossval, + plot_loo_crossval, + plot_loo_residuals, +) + + +def make_df(): + """ + Load test data from a CSV file containing the first 25 rows of the q87d subset + of the SXS catalog, so the tests use a realistic data input. + """ + return pd.read_csv(TEST_DATA_PATH) + + +# Single class with one test function for each diagnostics function +class TestDiagnostics(unittest.TestCase): + def tearDown(self): + plt.close("all") + + # Test loo_crossval + def test_loo_crossval(self): + df = make_df() + input_columns = [ + "initial_separation", + "mass_ratio", + "S1x", + "S1y", + "S1z", + "S2x", + "S2y", + "S2z", + ] + X = df[input_columns].values + Y = df["initial_orbital_frequency"].values + + preds, uncertainties = loo_crossval(X, Y, target_name="omega") + + # Test that the function returns (predictions, uncertainties) + result = (preds, uncertainties) + self.assertEqual(len(result), 2) + + # Test that the outputs are arrays + self.assertIsInstance(preds, np.ndarray) + self.assertIsInstance(uncertainties, np.ndarray) + + # Test that there is one prediction per data point + self.assertEqual(preds.shape, Y.shape) + + # Test that the GP predictive uncertainties are non-negative + self.assertTrue(np.all(uncertainties >= 0)) + + # Test plot_loo_crossval + def test_plot_loo_crossval(self): + Y = np.array([1.0, 2.0, 3.0, 4.0, 5.0]) + preds = np.array([1.1, 1.9, 3.1, 3.8, 5.2]) + + # Call plot_loo_crossval once and reuse the result for each test + rmse, mae, r2 = plot_loo_crossval(Y, preds) + + # Test that the function returns (rmse, mae, r2) + result = (rmse, mae, r2) + self.assertEqual(len(result), 3) + + # Test that outputs are floats + for scalar in result: + self.assertIsInstance(float(scalar), float) + + # Test that R^2 is a squared correlation and lies in [0,1] + self.assertGreaterEqual(r2, 0.0) + self.assertLessEqual(r2, 1.0) + + # Test that perfect predictions give R^2 = 1, RMSE = 0 = MAE + rmse_perf, mae_perf, r2_perf = plot_loo_crossval(Y, Y.copy()) + self.assertAlmostEqual(r2_perf, 1.0, places=5) + self.assertAlmostEqual(rmse_perf, 0.0, places=10) + self.assertAlmostEqual(mae_perf, 0.0, places=10) + + # Test plot_loo_residuals + # No GPR calls; show = False is passed in all tests to suppress plots + def test_plot_loo_residuals(self): + Y = np.array([1.0, 2.0, 3.0, 4.0, 5.0]) + preds = np.array([1.1, 1.9, 3.1, 3.8, 5.2]) + + # Test that residual = true - predicted for each point + residuals = plot_loo_residuals(Y, preds, show=False) + np.testing.assert_allclose(residuals, Y - preds, rtol=1e-10) + + # Test that the output shape matches the input arrays + Y_shape = np.array([1.0, 2.0, 3.0]) + preds_shape = np.array([1.0, 2.0, 3.0]) + residuals_shape = plot_loo_residuals(Y_shape, preds_shape, show=False) + self.assertEqual(residuals_shape.shape, Y_shape.shape) + + # Test that if the predictions match the truth exactly, all the residuals are 0 + Y_perf = np.array([1.0, 2.0, 3.0, 4.0]) + preds_perf = Y_perf.copy() + residuals_perf = plot_loo_residuals(Y_perf, preds_perf, show=False) + np.testing.assert_allclose(residuals_perf, 0.0, atol=1e-10) + + +if __name__ == "__main__": + unittest.main(verbosity=2)