gp_fromdata.py

#!/usr/bin/env python3
"""
Train a GP surrogate model on a chosen dataset using the BoTorch-based GPSurrogate.

Usage examples:

./gp_fromdata.py --n_train=200 --kernel=rbf --isotropic
./gp_fromdata.py --n_train=200 --kernel=matern
./gp_fromdata.py --n_train=200 --kernel=matern --normalize_y --plot
./gp_fromdata.py --n_train=300 --kernel=matern --log
"""

import argparse
import time
from datetime import datetime
from pathlib import Path

import numpy as np
from sklearn.metrics import mean_absolute_error, root_mean_squared_error as rmse

from surmod import data_processing

from surmod.gaussian_process import GPSurrogate


def parse_arguments():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="Train GP surrogate models on datasets from data/.",
    )

    parser.add_argument(
        "-d",
        "--dataset",
        type=str,
        choices=list(data_processing.DATASET_CONFIG.keys()),
        default="JAG",
        help="Which dataset to use (default: JAG).",
    )

    parser.add_argument(
        "-tr",
        "--n_train",
        type=int,
        default=50,
        help="Number of train samples.",
    )

    parser.add_argument(
        "-te",
        "--n_test",
        type=int,
        default=500,
        help="Number of test samples.",
    )

    parser.add_argument(
        "-ny",
        "--normalize_y",
        action="store_true",
        help="Standardize outputs (maps to GPSurrogate.scale_outputs).",
    )

    parser.add_argument(
        "-k",
        "--kernel",
        type=str,
        choices=["matern", "rbf", "periodic"],
        default="matern",
        help="Kernel type for GPSurrogate.",
    )

    parser.add_argument(
        "-i",
        "--isotropic",
        action="store_true",
        help="Use isotropic kernel (shared lengthscale). Default is ARD.",
    )

    parser.add_argument(
        "--scale_inputs",
        dest="scale_inputs",
        action="store_true",
        default=True,
        help="Normalize inputs to unit cube (GPSurrogate.scale_inputs).",
    )

    parser.add_argument(
        "--no-scale_inputs",
        dest="scale_inputs",
        action="store_false",
        help="Disable input normalization.",
    )

    parser.add_argument(
        "--lengthscale_bounds",
        type=float,
        nargs=2,
        default=(1e-2, 100.0),
        metavar=("LOW", "HIGH"),
        help="Bounds for kernel lengthscale constraint.",
    )

    parser.add_argument(
        "--noise_bounds",
        type=float,
        nargs=2,
        default=(1e-16, 1e-1),
        metavar=("LOW", "HIGH"),
        help="Bounds for likelihood noise constraint.",
    )

    parser.add_argument(
        "-l",
        "--log",
        action="store_true",
        help="Append results to output log file.",
    )

    parser.add_argument(
        "-p",
        "--plot",
        action="store_true",
        help="Create observed vs predicted parity plot with 95 percent intervals.",
    )

    parser.add_argument(
        "--LHD",
        action="store_true",
        help="Use an LHD design (passed into split_data if supported).",
    )

    parser.add_argument(
        "-s",
        "--seed",
        type=int,
        default=42,
        help="Random number generator seed.",
    )

    return parser.parse_args()


def log_results(log_message: str, path_to_log: Path) -> None:
    path_to_log.parent.mkdir(parents=True, exist_ok=True)
    with open(path_to_log, "a", encoding="utf-8") as f:
        f.write(log_message)


def main():
    """
    Trains and evaluates a Gaussian Process (GP) surrogate model on a dataset
    contained in a csv file.
    """
    # Parse command line arguments
    args = parse_arguments()
    dataset = args.dataset
    n_train = args.n_train
    n_test = args.n_test
    normalize_y = args.normalize_y
    kernel = args.kernel
    isotropic = args.isotropic
    scale_inputs = args.scale_inputs
    lengthscale_bounds = tuple(args.lengthscale_bounds)
    noise_bounds = tuple(args.noise_bounds)
    do_log = args.log
    do_plot = args.plot
    seed = args.seed
    use_lhd = args.LHD

    # Set output directories relative to this script
    script_dir = Path(__file__).parent
    results_dir = script_dir / "results"
    plots_dir = script_dir / "plots"

    # Check data availability
    n_samples = n_test + n_train
    if n_samples > 10000:
        raise ValueError(
            f"Requested samples ({n_samples}) exceed existing dataset(s) size limit (10000)."
        )

    # Load and split data
    df = data_processing.load_data(dataset=dataset, n_samples=n_samples, random=False)
    x_train, x_test, y_train, y_test = data_processing.split_data(
        df=df, LHD=use_lhd, n_train=n_train, seed=seed
    )

    # Build and fit BoTorch GP surrogate
    gp = GPSurrogate(
        x_train=x_train,
        y_train=y_train,
        x_test=x_test,
        y_test=y_test,
        kernel=kernel,
        isotropic=isotropic,
        scale_inputs=scale_inputs,
        scale_outputs=normalize_y,
        lengthscale_bounds=lengthscale_bounds,
        noise_bounds=noise_bounds,
    )

    start_time = time.perf_counter()
    gp.fit()
    elapsed_time = time.perf_counter() - start_time

    # Predict on train/test
    pred_train_mean, _pred_train_std = gp.predict(x_train)
    pred_test_mean, pred_test_std = gp.predict(x_test)

    # Metrics (match your previous ones, plus coverage from GPSurrogate.evaluate)
    train_mae = mean_absolute_error(y_train, pred_train_mean)
    test_mae = mean_absolute_error(y_test, pred_test_mean)

    train_rmse = rmse(y_train, pred_train_mean)
    test_rmse = rmse(y_test, pred_test_mean)

    # Max absolute error locations
    train_max_abserr, train_max_input = gp.compute_max_error(
        pred_train_mean, y_train, x_train
    )
    test_max_abserr, test_max_input = gp.compute_max_error(
        pred_test_mean, y_test, x_test
    )

    # 95% confidence interval coverage on test data
    lower = pred_test_mean - 1.96 * pred_test_std
    upper = pred_test_mean + 1.96 * pred_test_std
    coverage = np.mean((y_test >= lower) & (y_test <= upper))

    timestamp = datetime.now().strftime("%m%d_%H%M%S")
    log_lines = [
        f"Run timestamp (%m%d_%H%M%S): {timestamp}",
        f"Test Function: {dataset}",
        f"Number of training points: {n_train}",
        f"Number of testing points: {n_test}",
        f"Kernel: {kernel}",
        f"Isotropic kernel: {isotropic}",
        f"Scale inputs: {scale_inputs}",
        f"Normalize y: {normalize_y}",
        f"Lengthscale bounds: {lengthscale_bounds}",
        f"Noise bounds: {noise_bounds}",
        f"Train RMSE: {train_rmse:.5e}",
        f"Test RMSE: {test_rmse:.5e}",
        f"Test 95% interval coverage: {coverage:.2%}",
        f"Train Max abs err:  {train_max_abserr:.5e} | Location: {train_max_input}",
        f"Test Max abs err:   {test_max_abserr:.5e} | Location: {test_max_input}",
        f"Train Mean abs err: {train_mae:.5e}",
        f"Test Mean abs err:  {test_mae:.5e}",
        f"Training time: {elapsed_time:.3f} seconds",
    ]
    log_message = "\n".join(log_lines) + "\n"

    print(log_message)

    if do_log:
        log_results(
            log_message,
            path_to_log=results_dir / f"{dataset}.txt",
        )

    if do_plot:
        # Uses your class method that calls evaluate() internally
        gp.plot_test_predictions(dataset=dataset, plots_dir=plots_dir)


if __name__ == "__main__":
    main()