Diagnostics

.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

pyproject.toml

@@ -17,6 +17,7 @@ dependencies = [
     "torch",
     "gpytorch",
     "pandas",
+    "sxs",
 ]
 
 [project.optional-dependencies]

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

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)