diff --git a/pyproject.toml b/pyproject.toml index a78c042..ad28e17 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,8 @@ dependencies = [ "pandas", "sxs", "varpro @ git+https://github.com/sxs-collaboration/varpro.git@978106eaf3d7a6a7f0c5f167726d8e0fc59fc95d", + "click", + "rich", ] [project.optional-dependencies] diff --git a/src/SimulationSupport/EccentricityControl/Examples/InitialOrbitalParameters.ipynb b/src/SimulationSupport/EccentricityControl/Examples/InitialOrbitalParameters.ipynb new file mode 100644 index 0000000..19a13a2 --- /dev/null +++ b/src/SimulationSupport/EccentricityControl/Examples/InitialOrbitalParameters.ipynb @@ -0,0 +1,74 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0c6702c4", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "\n", + "from SimulationSupport.EccentricityControl.InitialOrbitalParameters import (\n", + " initial_orbital_parameters,\n", + ")\n", + "\n", + "CHECKPOINT_DIR = Path.cwd() / \"Examples\"\n", + "\n", + "target_params = {\n", + " \"MassRatio\": 8.0,\n", + " \"DimensionlessSpinA\": [3.61e-14, -4.5e-15, 0.-0.8],\n", + " \"DimensionlessSpinB\": [-1.1218e-12, -3.82e-14, -0.8],\n", + " \"Eccentricity\": 0.0,\n", + "}\n", + "\n", + "D0, Omega0, Adot0 = initial_orbital_parameters(\n", + " target_params,\n", + " separation = 13.378235,\n", + " method= \"PN\",\n", + ")\n", + "print(\"PN:\", D0, Omega0, Adot0)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "443ad8b9", + "metadata": {}, + "outputs": [], + "source": [ + "D0, Omega0, Adot0 = initial_orbital_parameters(\n", + " target_params,\n", + " separation = 13.378235,\n", + " method = \"GPR\",\n", + " gpr_checkpoints={\n", + " \"Omega0\": str(CHECKPOINT_DIR / \"gpr_model_omega.pth\"),\n", + " \"Adot0\": str(CHECKPOINT_DIR / \"gpr_model_adot.pth\"),\n", + " },\n", + ")\n", + "print(\"GPR:\", D0, Omega0, Adot0)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "env311 (3.11.6)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.6" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/src/SimulationSupport/EccentricityControl/Examples/gpr_model_adot.pth b/src/SimulationSupport/EccentricityControl/Examples/gpr_model_adot.pth new file mode 100644 index 0000000..24115aa Binary files /dev/null and b/src/SimulationSupport/EccentricityControl/Examples/gpr_model_adot.pth differ diff --git a/src/SimulationSupport/EccentricityControl/Examples/gpr_model_omega.pth b/src/SimulationSupport/EccentricityControl/Examples/gpr_model_omega.pth new file mode 100644 index 0000000..47d69a3 Binary files /dev/null and b/src/SimulationSupport/EccentricityControl/Examples/gpr_model_omega.pth differ diff --git a/src/SimulationSupport/EccentricityControl/InitialOrbitalParameters.py b/src/SimulationSupport/EccentricityControl/InitialOrbitalParameters.py old mode 100644 new mode 100755 index bd2430f..f8dc500 --- a/src/SimulationSupport/EccentricityControl/InitialOrbitalParameters.py +++ b/src/SimulationSupport/EccentricityControl/InitialOrbitalParameters.py @@ -1,11 +1,16 @@ +#!/usr/bin/env python + # Distributed under the MIT License. # See LICENSE.txt for details. """Estimate initial orbital parameters.""" +import json import logging from typing import Optional, Tuple +import click import numpy as np +import rich from scipy.optimize import minimize logger = logging.getLogger(__name__) @@ -16,11 +21,16 @@ def initial_orbital_parameters( separation: Optional[float] = None, orbital_angular_velocity: Optional[float] = None, radial_expansion_velocity: Optional[float] = None, + method: str = "PN", + gpr_checkpoints: Optional[dict] = None, ) -> Tuple[float, float, float]: - r"""Estimate initial orbital parameters from a Post-Newtonian approximation. + r"""Estimate initial orbital parameters from PN or GPR. + + Estimates initial orbital parameters from a Post-Newtonian (PN) approximation, + or from a Gaussian Process Regression (GPR) correction to the PN approximation. Given the target eccentricity and one other orbital parameter, this - routine estimates the initial separation ``D``, orbital angular velocity + routine estimates the initial separation ``D_0``, orbital angular velocity ``Omega_0``, and radial expansion velocity ``adot_0`` for a binary system. The resulting parameters can be fed into an eccentricity control loop to refine the starting parameters. @@ -46,11 +56,25 @@ def initial_orbital_parameters( * ``"NumOrbits"``: Desired number of inspiral orbits until merger. * ``"TimeToMerger"``: Desired time to merger. separation : float, optional - Coordinate separation ``D`` of the black holes. + Coordinate separation ``D_0`` of the black holes. orbital_angular_velocity : float, optional Orbital angular velocity ``Omega_0``. radial_expansion_velocity : float, optional Radial expansion velocity ``adot_0``. + method: str, optional + Either ``PN`` to compute parameters from the PN approximation, + or ``GPR`` to apply a learned correction from a trained GPR model + to the PN approximation + gpr_checkpoints: dict, optional + Required when ``method = GPR``. Maps the quantities to correct to the + path of their trained GPR checkpoint file produced by ``save_gpr_checkpoint``. + Recognized keys are ``Omega0`` and ``Adot0``. Any quantity without + an entry is left at its PN value and no correction is applied. ``D_0`` is + currently never corrected and is always returned at its PN value. Each checkpoint + is checked against the quantity it is supplied for, so passing a checkpoint trained + for a different quantity raises a ``ValueError`` instead of silently producing + a wrong correction. + Returns ------- @@ -66,6 +90,9 @@ def initial_orbital_parameters( ): return separation, orbital_angular_velocity, radial_expansion_velocity + # Unpack the pieces of target_params we need. Everything is derived from this dict + # instead of passed as individual arguments, so callers only have to build one + # dict per system mass_ratio = target_params["MassRatio"] dimensionless_spin_a = np.asarray(target_params["DimensionlessSpinA"]) dimensionless_spin_b = np.asarray(target_params["DimensionlessSpinB"]) @@ -99,19 +126,76 @@ def initial_orbital_parameters( " parameters: 'separation', 'orbital_angular_velocity', 'num_orbits'," " 'time_to_merger'." ) + assert method in ( + "PN", + "GPR", + ), f"Unknown method '{method}'. Choose either 'PN' or 'GPR'." + + # GPR method. This will be modified later to accept + # both non-eccentric and eccentric GPR models. + if method == "GPR": + assert gpr_checkpoints, ( + "The GPR method requires a 'gpr_checkpoints' dict mapping the" + " quantities to correct to their trained checkpoint file paths." + ) + return _initial_orbital_parameters_gpr( + mass_ratio=mass_ratio, + dimensionless_spin_a=dimensionless_spin_a, + dimensionless_spin_b=dimensionless_spin_b, + eccentricity=eccentricity, + separation=separation, + orbital_angular_velocity=orbital_angular_velocity, + num_orbits=num_orbits, + time_to_merger=time_to_merger, + gpr_checkpoints=gpr_checkpoints, + ) + + # Compute the initial orbital parameters from the Post-Newtonian approximation. + return _initial_orbital_parameters_pn( + mass_ratio=mass_ratio, + dimensionless_spin_a=dimensionless_spin_a, + dimensionless_spin_b=dimensionless_spin_b, + eccentricity=eccentricity, + separation=separation, + orbital_angular_velocity=orbital_angular_velocity, + num_orbits=num_orbits, + time_to_merger=time_to_merger, + ) + + +def _initial_orbital_parameters_pn( + mass_ratio, + dimensionless_spin_a, + dimensionless_spin_b, + eccentricity, + separation, + orbital_angular_velocity, + num_orbits, + time_to_merger, +) -> Tuple[float, float, float]: + """Zero-eccentricity initial orbital parameters from the PN approximation. + + This is the PN-only implementation, which also serves as the baseline guess + the GPR models learn to correct.""" - # Import functions from SpEC. These functions currently work only for zero - # eccentricity. We will need to generalize this for eccentric orbits. assert eccentricity == 0.0, ( - "Initial orbital parameters can currently only be computed for zero" - " eccentricity." + "Initial orbital parameters from PN can currently only be computed for" + " zero eccentricity." ) + + # Import functions from SpEC. These functions currently work only for zero + # eccentricity. We will need to generalize this for eccentric orbits. # These functions call old Fortran code (LSODA) through # scipy.integrate.odeint, which leads to lots of noise in stdout. We should # modernize them to use scipy.integrate.solve_ivp. - from .ZeroEccParamsFromPN import nOrbitsAndTotalTime, omegaAndAdot + from SimulationSupport.EccentricityControl.ZeroEccParamsFromPN import ( + nOrbitsAndTotalTime, + omegaAndAdot, + ) - # Find an omega0 that gives the right number of orbits or time to merger + # If the caller specifies a desired number of orbits or time to merger + # instead of an orbital angular velocity, root-find for the omega0 that + # produces it. if num_orbits is not None or time_to_merger is not None: opt_result = minimize( lambda x: ( @@ -139,7 +223,8 @@ def initial_orbital_parameters( f"Found orbital angular velocity: {orbital_angular_velocity}" ) - # Find the separation that gives the desired orbital angular velocity + # Given an orbital angular velocity, either passed in directly, or solved for above, + # root-find for the coordinate separation that produces it under the PN approximation if orbital_angular_velocity is not None: opt_result = minimize( lambda x: abs( @@ -163,13 +248,14 @@ def initial_orbital_parameters( separation = opt_result.x[0] logger.debug(f"Found initial separation: {separation}") - # Find the radial expansion velocity + # Now that we have a separation, either passed in directly, or solved for above, + # find the radial expansion velocity at that separation new_orbital_angular_velocity, radial_expansion_velocity = omegaAndAdot( r=separation, q=mass_ratio, chiA=dimensionless_spin_a, chiB=dimensionless_spin_b, - rPrime0=1.0, # Choice also made in SpEC + rPrime0=1.0, ) if orbital_angular_velocity is None: orbital_angular_velocity = new_orbital_angular_velocity @@ -193,3 +279,297 @@ def initial_orbital_parameters( f" {num_orbits:g}. Time to merger: {time_to_merger:g} M." ) return separation, orbital_angular_velocity, radial_expansion_velocity + + +# Map the user-facing quantity names to the 'output_name' stored inside the +# trained GPR checkpoints. The checkpoint names are decided within the training +# pipelines, so we map them here instead of renaming them. This can be changed in the future. +_CHECKPOINT_OUTPUT_NAMES = {"Omega0": "omega", "Adot0": "adot"} + + +def _apply_gpr_correction( + quantity_name, baseline_value, available_values, checkpoint_path +): + """Load a GPR checkpoint and add its predicted correction to a PN baseline. + + Note: GPR checkpoints predict a correction to the PN baseline, not the direct + quantity itself. + """ + from SimulationSupport.gpr import ( + load_gpr_checkpoint, + predict_with_gpr_model, + ) + + model, likelihood, meta = load_gpr_checkpoint(checkpoint_path) + # Guard against a checkpoint being passed for the wrong quantity, which + # would add the wrong delta and silently produce an incorrect number. + expected_output = _CHECKPOINT_OUTPUT_NAMES[quantity_name] + if meta["output_name"] != expected_output: + raise ValueError( + f"Checkpoint '{checkpoint_path}' was trained to predict" + f" '{meta['output_name']}', but it is being applied to" + f" '{quantity_name}', which expects a checkpoint predicting" + f" '{expected_output}'. Check that the checkpoint matches the" + " quantity." + ) + + # Assemble a raw feature array, in the order the GPR checkpoint expects + try: + raw_x = [available_values[name] for name in meta["input_features"]] + except KeyError as missing_feature: + raise KeyError( + f"GPR checkpoint expects input feature {missing_feature}," + " which is not available. Available features:" + f" {sorted(available_values.keys())}. Update the" + " 'available_values' mapping in this module to match the" + " checkpoint's input_features metadata." + ) from missing_feature + + raw_x = np.asarray([raw_x], dtype=float) + delta_mean, delta_std = predict_with_gpr_model(raw_x, model, likelihood) + # The noneccentric GPR predicts a correction (delta), to the PN approximation. + # This value is then added to the PN baseline to get the final corrected quantity. + corrected_value = baseline_value + float(delta_mean[0]) + logger.debug( + f"GPR correction for {quantity_name}: baseline={baseline_value:g}," + f" delta={float(delta_mean[0]):g} +/- {float(delta_std[0]):g}," + f" corrected={corrected_value:g}" + ) + return corrected_value + + +def _initial_orbital_parameters_gpr( + mass_ratio, + dimensionless_spin_a, + dimensionless_spin_b, + eccentricity, + separation, + orbital_angular_velocity, + num_orbits, + time_to_merger, + gpr_checkpoints, +) -> Tuple[float, float, float]: + """Zero-eccentricity initial orbital parameters, PN baseline, and + GPR correction. + """ + assert eccentricity == 0.0, ( + "Initial orbital parameters from GPR can currently only be computed for" + " zero eccentricity." + ) + # Start from the baseline PN guess, which the GPR models are trained to correct + pn_separation, pn_omega, pn_adot = _initial_orbital_parameters_pn( + mass_ratio=mass_ratio, + dimensionless_spin_a=dimensionless_spin_a, + dimensionless_spin_b=dimensionless_spin_b, + eccentricity=eccentricity, + separation=separation, + orbital_angular_velocity=orbital_angular_velocity, + num_orbits=num_orbits, + time_to_merger=time_to_merger, + ) + # Feature names match SimulationSupport.gpr (see the GPR tutorial notebook + # for a detailed explanation of how to train, save, and load the GPR model + # with real data). Each checkpoint selects the subset of these values it + # was trained on via its 'input_features' metadata. The aligned-spin example + # checkpoints in Examples use 'initial_separation', 'initial_mass_ratio', + # 'initial_dimensionless_spin1_z', and 'initial_dimensionless_spin2_z'. + available_values = { + "initial_separation": pn_separation, + "initial_mass_ratio": mass_ratio, + "initial_dimensionless_spin1_x": dimensionless_spin_a[0], + "initial_dimensionless_spin1_y": dimensionless_spin_a[1], + "initial_dimensionless_spin1_z": dimensionless_spin_a[2], + "initial_dimensionless_spin2_x": dimensionless_spin_b[0], + "initial_dimensionless_spin2_y": dimensionless_spin_b[1], + "initial_dimensionless_spin2_z": dimensionless_spin_b[2], + "pn_guess_omega": pn_omega, + "pn_guess_adot": pn_adot, + } + + corrected = {"Omega0": pn_omega, "Adot0": pn_adot} + for quantity_name, checkpoint_path in gpr_checkpoints.items(): + if quantity_name not in corrected: + raise ValueError( + f"Unknown quantity `{quantity_name}` in `gpr_checkpoints`." + f" Expected one of {','.join(sorted(corrected))}." + ) + corrected[quantity_name] = _apply_gpr_correction( + quantity_name, + corrected[quantity_name], + available_values, + checkpoint_path, + ) + orbital_angular_velocity = corrected["Omega0"] + radial_expansion_velocity = corrected["Adot0"] + logger.info( + "Selected approximately circular orbit using GPR corrected PN guess." + f" D0={pn_separation:g}, Omega0={orbital_angular_velocity:g}," + f" Adot0={radial_expansion_velocity:g}." + ) + return pn_separation, orbital_angular_velocity, radial_expansion_velocity + + +# CLI +# The function can be imported and called from Python directly, or it can be called with the CLI. +@click.command( + name="initial-orbital-parameters", +) +@click.option( + "--mass-ratio", + "-q", + type=float, + required=True, + help=r"Mass ratio, q = M_A / M_B \ge 1, of the two black holes.", +) +@click.option( + "--dimensionless-spin-a", + nargs=3, + type=float, + required=True, + help=( + "Dimensionless spin vector of the larger black hole, for example," + "written as '--dimensionless-spin-a 0.0 0.1 0.1'." + ), +) +@click.option( + "--dimensionless-spin-b", + nargs=3, + type=float, + required=True, + help=( + "Dimensionless spin vector of the smaller black hole, for example," + " written as '--dimensionless-spin-b 0.0 0.0 0.1'." + ), +) +@click.option( + "--eccentricity", + "-e", + type=float, + required=True, + help="Desired orbital eccentricity.", +) +@click.option( + "--mean-anomaly-fraction", + type=float, + help=( + "Mean anomaly divided by 2pi (between 0 and 1). Required if" + " eccentricity is nonzero." + ), +) +@click.option( + "--separation", + "-D", + type=float, + help="Coordinate separation, D_0, between the black holes.", +) +@click.option( + "--orbital-angular-velocity", + "-w", + type=float, + help="Orbital angular velocity, Omega_0.", +) +@click.option( + "--num-orbits", + type=float, + help="Desired number of orbits until merger.", +) +@click.option("--time-to-merger", type=float, help="Desired time until merger.") +@click.option( + "--method", + type=click.Choice(["PN", "GPR"]), + default="PN", + show_default=True, + help=( + "Compute from PN or from GPR, which applies a learned" + " correction to the PN baseline." + ), +) +@click.option( + "--gpr-omega-checkpoint", + type=click.Path(exists=True, dir_okay=False, readable=True), + help="Path to a trained GPR checkpoint providing Omega0 corrections.", +) +@click.option( + "--gpr-adot-checkpoint", + type=click.Path(exists=True, dir_okay=False, readable=True), + help="Path to a trained GPR checkpoint providing Adot0 corrections.", +) +@click.option( + "--output-json", + is_flag=True, + help="Print the result as a JSON file instead of text.", +) +def initial_orbital_parameters_command( + mass_ratio, + dimensionless_spin_a, + dimensionless_spin_b, + eccentricity, + mean_anomaly_fraction, + separation, + orbital_angular_velocity, + num_orbits, + time_to_merger, + method, + gpr_omega_checkpoint, + gpr_adot_checkpoint, + output_json, +): + """Estimate the initial orbital parameters for a BBH evolution. + + Estimates the initial coordinate separation, D_0, orbital angular velocity, Omega_0, and + radial expansion velocity, adot_0, from a Post-Newtonian approximation, optionally + corrected by a trained Gaussian Process Regression (GPR) model. + + Specify the target eccentricity and either '--separation', '--orbital-angular-velocity', + '--num-orbits', or '--time-to-merger'. + """ + _rich_traceback_guard = True + + target_params = { + "MassRatio": mass_ratio, + "DimensionlessSpinA": list(dimensionless_spin_a), + "DimensionlessSpinB": list(dimensionless_spin_b), + "Eccentricity": eccentricity, + } + if mean_anomaly_fraction is not None: + target_params["MeanAnomalyFraction"] = mean_anomaly_fraction + if num_orbits is not None: + target_params["NumOrbits"] = num_orbits + if time_to_merger is not None: + target_params["TimeToMerger"] = time_to_merger + + gpr_checkpoints = None + if method == "GPR": + gpr_checkpoints = {} + if gpr_omega_checkpoint: + gpr_checkpoints["Omega0"] = gpr_omega_checkpoint + if gpr_adot_checkpoint: + gpr_checkpoints["Adot0"] = gpr_adot_checkpoint + if not gpr_checkpoints: + raise click.UsageError( + "'--method GPR' requires either '--gpr-omega-checkpoint'," + " '--gpr-adot-checkpoint', or both." + ) + + D0, Omega0, Adot0 = initial_orbital_parameters( + target_params, + separation=separation, + orbital_angular_velocity=orbital_angular_velocity, + method=method, + gpr_checkpoints=gpr_checkpoints, + ) + + if output_json: + print( + json.dumps({"D0": D0, "Omega0": Omega0, "Adot0": Adot0}, indent=2) + ) + else: + rich.print(f"D0 = {D0}") + rich.print(f"Omega0 = {Omega0}") + rich.print(f"Adot0 = {Adot0}") + + return D0, Omega0, Adot0 + + +if __name__ == "__main__": + initial_orbital_parameters_command(help_option_names=["-h", "--help"]) diff --git a/tests/EccentricityControl/Test_InitialOrbitalParameters.py b/tests/EccentricityControl/Test_InitialOrbitalParameters.py index 4f0ca26..66ff9f3 100644 --- a/tests/EccentricityControl/Test_InitialOrbitalParameters.py +++ b/tests/EccentricityControl/Test_InitialOrbitalParameters.py @@ -1,10 +1,15 @@ # Distributed under the MIT License. # See LICENSE.txt for details. +import json + import numpy.testing as npt +import pytest +from click.testing import CliRunner from SimulationSupport.EccentricityControl.InitialOrbitalParameters import ( initial_orbital_parameters, + initial_orbital_parameters_command, ) @@ -41,13 +46,6 @@ def test_initial_orbital_parameters(): ), [15.6060791015625, 0.015, -4.541705362753467e-05], ) - npt.assert_allclose( - initial_orbital_parameters( - target_params, - orbital_angular_velocity=0.015, - ), - [15.6060791015625, 0.015, -4.541705362753467e-05], - ) npt.assert_allclose( initial_orbital_parameters( {**target_params, "NumOrbits": 20}, @@ -60,3 +58,198 @@ def test_initial_orbital_parameters(): ), [16.1357421875, 0.01430025219917298, -3.9831982447244026e-05], ) + + +def test_initial_orbital_parameters_pn_requires_zero_eccentricity(): + # PN method only supports zero eccentricity + target_params = { + "MassRatio": 1.0, + "MassA": 0.5, + "MassB": 0.5, + "DimensionlessSpinA": [0.0, 0.0, 0.0], + "DimensionlessSpinB": [0.0, 0.0, 0.0], + "Eccentricity": 0.1, + "MeanAnomalyFraction": 0.5, + } + with pytest.raises(AssertionError, match="zero eccentricity"): + initial_orbital_parameters( + target_params, + separation=16.0, + method="PN", + ) + + +def test_initial_orbital_parameters_gpr_requires_zero_eccentricity( + gpr_checkpoint_dir, +): + # GPR method currently only supports zero eccentricity + target_params = { + "MassRatio": 1.0, + "MassA": 0.5, + "MassB": 0.5, + "DimensionlessSpinA": [0.0, 0.0, 0.0], + "DimensionlessSpinB": [0.0, 0.0, 0.0], + "Eccentricity": 0.1, + "MeanAnomalyFraction": 0.5, + } + with pytest.raises(AssertionError, match="zero eccentricity"): + initial_orbital_parameters( + target_params, + separation=16.0, + method="GPR", + gpr_checkpoints={ + "Omega0": str(gpr_checkpoint_dir / "gpr_model_omega.pth") + }, + ) + + +def test_initial_orbital_parameters_gpr_requires_checkpoints(): + # method = "GPR" requires a non-empty gpr_checkpoints dict + target_params = { + "MassRatio": 1.0, + "MassA": 0.5, + "MassB": 0.5, + "DimensionlessSpinA": [0.0, 0.0, 0.0], + "DimensionlessSpinB": [0.0, 0.0, 0.0], + "Eccentricity": 0.0, + } + with pytest.raises(AssertionError, match="gpr_checkpoints"): + initial_orbital_parameters( + target_params, + separation=16.0, + method="GPR", + ) + + +def test_initial_orbital_parameters_gpr_rejects_unknown_quantities( + gpr_checkpoint_dir, +): + """ + Test that the keys of the checkpoint files match the parameter names used. + """ + target_params = { + "MassRatio": 1.0, + "MassA": 0.5, + "MassB": 0.5, + "DimensionlessSpinA": [0.0, 0.0, 0.0], + "DimensionlessSpinB": [0.0, 0.0, 0.0], + "Eccentricity": 0.0, + } + with pytest.raises(ValueError, match="Unknown quantity") as excinfo: + initial_orbital_parameters( + target_params, + separation=16.0, + method="GPR", + gpr_checkpoints={ + "omega": str(gpr_checkpoint_dir / "gpr_model_omega.pth") + }, + ) + assert "Omega0" in str(excinfo.value) + + +def test_initial_orbital_parameters_gpr_rejects_mismatched_checkpoint( + gpr_checkpoint_dir, +): + """ + Test that supplying a checkpoint trained for a different quantity is + prevented, rather than adding to the wrong correction. + """ + target_params = { + "MassRatio": 1.0, + "MassA": 0.5, + "MassB": 0.5, + "DimensionlessSpinA": [0.0, 0.0, 0.0], + "DimensionlessSpinB": [0.0, 0.0, 0.0], + "Eccentricity": 0.0, + } + with pytest.raises(ValueError, match="trained to predict"): + initial_orbital_parameters( + target_params, + separation=16.0, + method="GPR", + gpr_checkpoints={ + "Omega0": str(gpr_checkpoint_dir / "gpr_model_adot.pth") + }, + ) + + +def test_initial_orbital_parameters_gpr_omega_and_adot_correction( + gpr_checkpoint_dir, +): + """ + Test the GPR method with the real, trained checkpoints. The + expected deltas are computed directly from gpr_model_omega.pth + and gpr_model_adot.pth. + """ + target_params = { + "MassRatio": 1.0, + "MassA": 0.5, + "MassB": 0.5, + "DimensionlessSpinA": [0.0, 0.0, 0.0], + "DimensionlessSpinB": [0.0, 0.0, 0.0], + "Eccentricity": 0.0, + } + pn_separation, pn_omega, pn_adot = ( + 16.0, + 0.014474280975952748, + -4.117670632867514e-05, + ) + omega_delta = -3.0704395612701774e-05 + adot_delta = 8.766858081799e-05 + + separation, omega, adot = initial_orbital_parameters( + target_params, + separation=16.0, + method="GPR", + gpr_checkpoints={ + "Omega0": str(gpr_checkpoint_dir / "gpr_model_omega.pth"), + "Adot0": str(gpr_checkpoint_dir / "gpr_model_adot.pth"), + }, + ) + + npt.assert_allclose(separation, pn_separation) + npt.assert_allclose(omega, pn_omega + omega_delta, rtol=1e-4) + npt.assert_allclose(adot, pn_adot + adot_delta, rtol=1e-4) + + +def test_cli_gpr(gpr_checkpoint_dir): + runner = CliRunner() + result = runner.invoke( + initial_orbital_parameters_command, + [ + "--mass-ratio", + "1.0", + "--dimensionless-spin-a", + "0.0", + "0.0", + "0.0", + "--dimensionless-spin-b", + "0.0", + "0.0", + "0.0", + "--eccentricity", + "0.0", + "--separation", + "16.0", + "--method", + "GPR", + "--gpr-omega-checkpoint", + str(gpr_checkpoint_dir / "gpr_model_omega.pth"), + "--gpr-adot-checkpoint", + str(gpr_checkpoint_dir / "gpr_model_adot.pth"), + "--output-json", + ], + catch_exceptions=False, + ) + assert result.exit_code == 0, result.output + output = json.loads(result.output) + npt.assert_allclose( + output["Omega0"], + 0.014474280975952748 - 3.0704395612701774e-05, + rtol=1e-4, + ) + npt.assert_allclose( + output["Adot0"], + -4.117670632867514e-05 + 8.766858081799e-05, + rtol=1e-4, + ) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..1d5fc1f --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,21 @@ +# Distributed under the MIT License. +# See LICENSE.txt for details. + +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parent.parent + + +@pytest.fixture +def gpr_checkpoint_dir(): + """Directory containing example trained GPR checkpoints, + used for in Test_InitialOrbitalParameters.py.""" + return ( + REPO_ROOT + / "src" + / "SimulationSupport" + / "EccentricityControl" + / "Examples" + )