diff --git a/cmake/SetupSpec.cmake b/cmake/SetupSpec.cmake index b6762c6a5e69e..67efcec09c81e 100644 --- a/cmake/SetupSpec.cmake +++ b/cmake/SetupSpec.cmake @@ -6,7 +6,8 @@ find_package(SpEC) # Make SpEC scripts available in Python. These can be used until we have ported # them to SpECTRE. if (SPEC_ROOT) - set(PYTHONPATH "${SPEC_ROOT}/Support/Python:${PYTHONPATH}") + set(PYTHONPATH "${SPEC_ROOT}/Support/Python:\ +${SPEC_ROOT}/Support/DatDataManip:${PYTHONPATH}") endif() if (NOT SpEC_FOUND) diff --git a/support/Pipelines/Bbh/EccentricityControl.py b/support/Pipelines/Bbh/EccentricityControl.py index 9e503efb20ef6..6f31d8d9aa048 100644 --- a/support/Pipelines/Bbh/EccentricityControl.py +++ b/support/Pipelines/Bbh/EccentricityControl.py @@ -2,118 +2,71 @@ # See LICENSE.txt for details. import logging -import warnings -import matplotlib.pyplot as plt +import click import pandas as pd -import yaml -from spectre.Pipelines.EccentricityControl.EccentricityControl import ( - coordinate_separation_eccentricity_control, -) - -# Suppress specific RuntimeWarnings -warnings.filterwarnings( - "ignore", message="Number of calls to functionhas reached maxfev" +from spectre.Pipelines.EccentricityControl.EccentricityControlParams import ( + eccentricity_control_params, + eccentricity_control_params_options, ) logger = logging.getLogger(__name__) -def eccentricity_control( - h5_file, - id_input_file_path, - subfile_name_aha="ApparentHorizons/ControlSystemAhA_Centers.dat", - subfile_name_ahb="ApparentHorizons/ControlSystemAhB_Centers.dat", - tmin=500, - tmax=None, # use all available data - output=None, # reads from Inspiral.yaml -): +def eccentricity_control(h5_files, id_input_file_path, **kwargs): """Eccentricity reduction post inspiral. This function can be called after the inspiral has run (see the 'Next' - section of the Inspiral.yam file). + section of the Inspiral.yaml file). This function does the following: - - Reads orbital parameters from the 'id_input_file_path'. - - - For now, the time boundaries are set to start at 500, and at the end of - the simulation to use all available data, but we will make tmin and tmax - more dynamic later, so the user can change those variables according to - interest. + - Get the new orbital parameters by calling the function + 'eccentricity_control_params' in + 'spectre.Pipelines.EccentricityControl.EccentricityControl'. + - Print the new orbital parameters in a tabular format. This will be updated + to start a new inspiral with the updated parameters. Arguments: h5_file: file that contains the trajectory data id_input_file_path: path to the input file of the initial data run - subfile_name_aha: subfile for black hole A; optional parameter - subfile_name_ahb: subfile for black hole B; optional parameter - tmin: starting point for eccentricity reduction script; defaults to - 500 if not specified - tmax: stopping point for eccentricity reduction script; set to use all - available data by default - output: outputs to terminal plus makes pdf file, if specified + **kwargs: additional arguments to be forwarded to + 'eccentricity_control_params'. """ - # Read and process the initial data input file - with open(id_input_file_path, "r") as open_input_file: - _, id_input_file = yaml.safe_load_all(open_input_file) - binary_data = id_input_file["Background"]["Binary"] - orbital_angular_velocity = binary_data["AngularVelocity"] - radial_expansion_velocity = binary_data["Expansion"] - - if output: - fig = plt.figure() - else: - fig = None - # Find the current eccentricity and extract new parameters to put into - # generate-id. Call the coordinate_separation_eccentricity_control - # function; extract only the ["H4"] , i.e. nonlinear fit. - ecout = coordinate_separation_eccentricity_control( - h5_file=h5_file, - subfile_name_aha=subfile_name_aha, - subfile_name_ahb=subfile_name_ahb, - tmin=tmin, - tmax=tmax, - angular_velocity_from_xcts=orbital_angular_velocity, - expansion_from_xcts=radial_expansion_velocity, - fig=fig, - )["H4"] - - if output: - fig.savefig(output) - - # Print the fit result - # extract new parameters to put into generate_id + # generate-id + ( + eccentricity, + ecc_std_dev, + new_orbital_params, + ) = eccentricity_control_params(h5_files, id_input_file_path, **kwargs) - fit_result = ecout["fit result"] - - # Prepare data from fit result + # Create DataFrame to display data in tabular format data = { "Attribute": [ "Eccentricity", - "Omega Update", - "Expansion Update", - "Previous Omega", - "Previous Expansion", - "Updated Omega", - "Updated Expansion", + "Eccentricity error", + "Updated Omega0", + "Updated adot0", ], "Value": [ - fit_result["eccentricity"], - fit_result["xcts updates"]["omega update"], - fit_result["xcts updates"]["expansion update"], - fit_result["previous xcts values"]["omega"], - fit_result["previous xcts values"]["expansion"], - fit_result["updated xcts values"]["omega"], - fit_result["updated xcts values"]["expansion"], + eccentricity, + ecc_std_dev, + new_orbital_params["Omega0"], + new_orbital_params["adot0"], ], } - - # Create DataFrame to display data in tabular format df = pd.DataFrame(data) # Print header line print("=" * 40) # Display table print(df.to_string(index=False)) print("=" * 40) + + +@click.command(name="eccentricity-control", help=eccentricity_control.__doc__) +@eccentricity_control_params_options +def eccentricity_control_command(**kwargs): + _rich_traceback_guard = True # Hide traceback until here + eccentricity_control(**kwargs) diff --git a/support/Pipelines/Bbh/Inspiral.yaml b/support/Pipelines/Bbh/Inspiral.yaml index b82e1a6178afa..ee93d60db70c3 100644 --- a/support/Pipelines/Bbh/Inspiral.yaml +++ b/support/Pipelines/Bbh/Inspiral.yaml @@ -19,11 +19,14 @@ Next: Next: Run: spectre.Pipelines.Bbh.EccentricityControl:eccentricity_control With: - h5_file: ./BbhReductions.h5 - subfile_name_aha: ApparentHorizons/ControlSystemAhA_Centers.dat - subfile_name_ahb: ApparentHorizons/ControlSystemAhB_Centers.dat - output: ecc_control.pdf + h5_files: ./BbhReductions.h5 id_input_file_path: {{id_input_file_path }} + subfile_name_aha_trajectories: ApparentHorizons/ControlSystemAhA_Centers.dat + subfile_name_ahb_trajectories: ApparentHorizons/ControlSystemAhB_Centers.dat + subfile_name_aha_quantities: ObservationAhA.dat + subfile_name_ahb_quantities: ObservationAhB.dat + target_eccentricity: 0.0 + plot_output_dir: ./ {% endif %} --- diff --git a/support/Pipelines/Bbh/__init__.py b/support/Pipelines/Bbh/__init__.py index 49e873360b02d..42d32dc3b8c56 100644 --- a/support/Pipelines/Bbh/__init__.py +++ b/support/Pipelines/Bbh/__init__.py @@ -11,6 +11,7 @@ class Bbh(click.MultiCommand): def list_commands(self, ctx): return [ + "eccentricity-control", "find-horizon", "generate-id", "postprocess-id", @@ -19,15 +20,19 @@ def list_commands(self, ctx): ] def get_command(self, ctx, name): - if name == "find-horizon": + if name in ["eccentricity-control", "ecc-control"]: + from .EccentricityControl import eccentricity_control_command + + return eccentricity_control_command + elif name == "find-horizon": from .FindHorizon import find_horizon_command return find_horizon_command - if name == "generate-id": + elif name == "generate-id": from .InitialData import generate_id_command return generate_id_command - if name == "postprocess-id": + elif name == "postprocess-id": from .PostprocessId import postprocess_id_command return postprocess_id_command diff --git a/support/Pipelines/EccentricityControl/CMakeLists.txt b/support/Pipelines/EccentricityControl/CMakeLists.txt index 4d4939f8d9ac7..72df2cca81be7 100644 --- a/support/Pipelines/EccentricityControl/CMakeLists.txt +++ b/support/Pipelines/EccentricityControl/CMakeLists.txt @@ -6,6 +6,6 @@ spectre_python_add_module( MODULE_PATH Pipelines PYTHON_FILES __init__.py - EccentricityControl.py + EccentricityControlParams.py InitialOrbitalParameters.py ) diff --git a/support/Pipelines/EccentricityControl/EccentricityControl.py b/support/Pipelines/EccentricityControl/EccentricityControl.py deleted file mode 100644 index 2569dc8fb3d91..0000000000000 --- a/support/Pipelines/EccentricityControl/EccentricityControl.py +++ /dev/null @@ -1,598 +0,0 @@ -#!/usr/bin/env python - -# Distributed under the MIT License. -# See LICENSE.txt for details. - -import logging -import os -import sys - -import click -import h5py -import matplotlib.pyplot as plt -import numpy as np -import rich -import scipy -from scipy import io, optimize - -from spectre.Visualization.Plot import ( - apply_stylesheet_command, - show_or_save_plot_command, -) -from spectre.Visualization.ReadH5 import available_subfiles - -logger = logging.getLogger(__name__) - - -# Read .dat files from Reductions data: Inertial centers, Mass, Spin -def extract_data_from_file(h5_file, subfile_name, functions, x_axis): - """Extract data from '.dat' datasets in H5 files""" - - with h5py.File(h5_file, "r") as h5file: - # Print available subfiles and exit - if subfile_name is None: - import rich.columns - - rich.print( - rich.columns.Columns( - available_subfiles(h5file, extension=".dat") - ) - ) - return - - # Open subfile - if not subfile_name.endswith(".dat"): - subfile_name += ".dat" - dat_file = h5file.get(subfile_name) - if dat_file is None: - raise click.UsageError( - f"Unable to open dat file '{subfile_name}'. Available " - f"files are:\n {available_subfiles(h5file, extension='.dat')}" - ) - - # Read legend from subfile - legend = list(dat_file.attrs["Legend"]) - - # Select x-axis - if x_axis is None: - x_axis = legend[0] - elif x_axis not in legend: - raise click.UsageError( - f"Unknown x-axis '{x_axis}'. Available columns are: {legend}" - ) - - num_obs = len(dat_file[:, legend.index(x_axis)]) - out_table = np.reshape(dat_file[:, legend.index(x_axis)], (num_obs, 1)) - - # Assemble in table - for function, label in functions: - if function not in legend: - raise click.UsageError( - f"Unknown function '{function}'. " - f"Available functions are: {legend}" - ) - - out_table = np.append( - out_table, - np.reshape(dat_file[:, legend.index(function)], (num_obs, 1)), - axis=1, - ) - - return out_table - - -# Compute separation norm -def compute_separation(h5_file, subfile_name_aha, subfile_name_ahb): - """Compute coordinate separation""" - - functions = [ - ["InertialCenter_x", "x"], - ["InertialCenter_y", "y"], - ["InertialCenter_z", "z"], - ] - x_axis = "Time" - - # Extract data - ObjectA_centers = extract_data_from_file( - h5_file=h5_file, - subfile_name=subfile_name_aha, - functions=functions, - x_axis=x_axis, - ) - - ObjectB_centers = extract_data_from_file( - h5_file=h5_file, - subfile_name=subfile_name_ahb, - functions=functions, - x_axis=x_axis, - ) - - if ( - subfile_name_aha is None - or subfile_name_ahb is None - or subfile_name_aha == subfile_name_ahb - ): - raise click.UsageError( - f"Dat files '{subfile_name_aha}' and '{subfile_name_aha}' are the" - " same or at least one of them is missing. Choose files for" - " different objects. Available files are:\n" - f" {available_subfiles(h5_file, extension='.dat')}" - ) - - # Compute separation - num_obs = min(len(ObjectA_centers[:, 0]), len(ObjectB_centers[:, 0])) - - # Separation vector - separation_vec = ( - ObjectA_centers[:num_obs, 1:] - ObjectB_centers[:num_obs, 1:] - ) - - # Compute separation norm - separation_norm = np.zeros((num_obs, 2)) - separation_norm[:, 0] = ObjectA_centers[:num_obs, 0] - separation_norm[:, 1] = np.linalg.norm(separation_vec, axis=1) - - return separation_norm - - -def compute_time_derivative_of_separation_in_window(data, tmin=None, tmax=None): - """Compute time derivative of separation on time window""" - traw = data[:, 0] - sraw = data[:, 1] - - # Compute separation derivative - dsdtraw = (sraw[2:] - sraw[0:-2]) / (traw[2:] - traw[0:-2]) - - trawcut = traw[1:-1] - - # Apply time window: select values in [tmin, tmax] - if tmin == None and tmax == None: - if traw[-1] < 200: - which_indices = trawcut > 20 - else: - which_indices = trawcut > 60 - elif tmax == None: - which_indices = trawcut > tmin - else: - which_indices = np.logical_and(trawcut > tmin, trawcut < tmax) - - dsdt = dsdtraw[which_indices] - t = trawcut[which_indices] - - return t, dsdt - - -def fit_model(x, y, model): - """Fit coordinate separation""" - F = model["function"] - inparams = model["initial guess"] - - errfunc = lambda p, x, y: F(p, x) - y - p, success = optimize.leastsq(errfunc, inparams[:], args=(x, y)) - - # Compute rms error of fit - e2 = (errfunc(p, x, y)) ** 2 - rms = np.sqrt(sum(e2) / np.size(e2)) - - return dict([("parameters", p), ("rms", rms), ("success", success)]) - - -def compute_coord_sep_updates( - x, y, model, initial_separation, initial_xcts_values=None -): - """Compute updates for eccentricity control""" - fit_results = fit_model(x=x, y=y, model=model) - - amplitude, omega, phase = fit_results["parameters"][:3] - - # Compute updates for Omega and expansion and compute eccentricity - dOmg = amplitude / 2.0 / initial_separation * np.sin(phase) - dadot = -amplitude / initial_separation * np.cos(phase) - - # The amplitude could be negative due to degeneracy with phase shifts - # of pi/2 - # Make sure eccentricity estimate is positive - eccentricity = np.abs(amplitude / initial_separation / omega) - - fit_results["eccentricity"] = eccentricity - fit_results["xcts updates"] = dict( - [("omega update", dOmg), ("expansion update", dadot)] - ) - - # Update xcts parameters if given - if initial_xcts_values is not None: - xcts_omega, xcts_expansion = initial_xcts_values - fit_results["previous xcts values"] = dict( - [("omega", xcts_omega), ("expansion", xcts_expansion)] - ) - updated_xcts_omega = xcts_omega + dOmg - updated_xcts_expansion = xcts_expansion + dadot - fit_results["updated xcts values"] = dict( - [ - ("omega", updated_xcts_omega), - ("expansion", updated_xcts_expansion), - ] - ) - - return fit_results - - -def coordinate_separation_eccentricity_control_digest( - h5_file, x, y, data, functions, fig=None -): - """Print and plot for eccentricity control""" - - if fig is not None: - traw = data[:, 0] - sraw = data[:, 1] - # Plot coordinate separation - ((ax1, ax2), (ax3, ax4)) = fig.subplots(2, 2) - fig.suptitle( - h5_file, - color="b", - size="large", - ) - ax2.plot(traw, sraw, "k", label="s", linewidth=2) - ax2.set_title("coordinate separation " + r"$ D $") - - # Plot derivative of coordinate separation - ax1.plot(x, y, "k", label=r"$ dD/dt $", linewidth=2) - ax1.set_title(r"$ dD/dt $") - - ax4.set_axis_off() - - logger.info("Eccentricity control summary") - - for name, func in functions.items(): - expression = func["label"] - rms = func["fit result"]["rms"] - - logger.info( - f"==== Function fitted to dD/dt: {expression:30s}, rms =" - f" {rms:4.3g} ====" - ) - - F = func["function"] - eccentricity = func["fit result"]["eccentricity"] - - # Print fit parameters - p = func["fit result"]["parameters"] - logger.info("Fit parameters:") - if np.size(p) == 3: - logger.info( - f"Oscillatory part: (B, w)=({p[0]:4.3g}, {p[1]:7.4f}), ", - f"Polynomial part: ({p[2]:4.2g})", - ) - else: - logger.info( - f"Oscillatory part: (B, w, phi)=({p[0]:4.3g}, {p[1]:6.4f}," - f" {p[2]:6.4f}), " - ) - if np.size(p) >= 4: - logger.info(f"Polynomial part: ({p[3]:4.2g}, ") - for q in p[4:]: - logger.info(f"{q:4.2g}") - logger.info(")") - - # Print eccentricity - logger.info(f"Eccentricity based on fit: {eccentricity:9.6f}") - # Print suggested updates based on fit - xcts_updates = func["fit result"]["xcts updates"] - dOmg = xcts_updates["omega update"] - dadot = xcts_updates["expansion update"] - logger.info("Suggested updates based on fit:") - logger.info(f"(dOmega, dadot) = ({dOmg:+13.10f}, {dadot:+8.6g})") - - if "updated xcts values" in func["fit result"]: - xcts_omega = func["fit result"]["updated xcts values"]["omega"] - xcts_expansion = func["fit result"]["updated xcts values"][ - "expansion" - ] - logger.info("Updated Xcts values based on fit:") - logger.info( - f"(Omega, adot) = ({(xcts_omega):13.10f}," - f" {(xcts_expansion):13.10g})" - ) - - # Plot - if fig is not None: - errfunc = lambda p, x, y: F(p, x) - y - # Plot dD/dt - ax1.plot( - x, - F(p, x), - label=( - f"{expression:s} \n rms = {rms:2.1e}, ecc =" - f" {eccentricity:4.5f}" - ), - ) - ax_handles, ax_labels = ax1.get_legend_handles_labels() - - # Plot residual - ax3.plot(x, errfunc(p, x, y), label=expression) - ax3.set_title("Residual") - - ax4.legend(ax_handles, ax_labels) - - -def coordinate_separation_eccentricity_control( - h5_file, - subfile_name_aha: str = "ApparentHorizons/ControlSystemAhA_Centers.dat", - subfile_name_ahb: str = "ApparentHorizons/ControlSystemAhB_Centers.dat", - tmin: float = None, - tmax: float = None, - angular_velocity_from_xcts: float = None, - expansion_from_xcts: float = None, - fig: plt.Figure = None, -): - """Compute updates based on fits to the coordinate separation for manual - eccentricity control - - This routine applies a time window. (To avoid large initial transients - and to allow the user to specify about 2 to 3 orbits of data.) - - Computes the coordinate separations between Objects A and B, as well as - a finite difference approximation to the time derivative. - - It fits different models to the time derivative. (The amplitude of the - oscillations is related to the eccentricity.) - - This function returns a dictionary containing data for the fits to all - the models considered below. For each model, the results of the fit as - well as the suggested updates for omega and the expansion provided. - - The updates are computed using Newtonian estimates. - See ArXiv:gr-qc/0702106 and ArXiv:0710.0158 for more details. - - A summary is printed to screen and if a matplotlib figure is provided, a - plot is generated. The latter is useful to decide between the updates of - different models (look for small residuals at early times). - - See OmegaDoEccRemoval.py in SpEC for improved eccentricity control. - - """ - - data = compute_separation( - h5_file=h5_file, - subfile_name_aha=subfile_name_aha, - subfile_name_ahb=subfile_name_ahb, - ) - - # Get initial separation from data (unwindowed) - initial_separation = data[:, 1][0] - - # Compute derivative in time window - t, dsdt = compute_time_derivative_of_separation_in_window( - data=data, tmin=tmin, tmax=tmax - ) - - # Collect initial xcts values (if given) - if ( - angular_velocity_from_xcts is not None - and expansion_from_xcts is not None - ): - initial_xcts_values = ( - angular_velocity_from_xcts, - expansion_from_xcts, - ) - else: - initial_xcts_values = None - - # Define functions to model the time derivative of the separation - functions = dict([]) - - # ==== Restricted fit ==== - functions["H1"] = dict( - [ - ("label", "B*cos(w*t+np.pi/2)+const"), - ( - "function", - lambda p, t: p[0] * np.cos(p[1] * t + np.pi / 2) + p[3], - ), - ("initial guess", [0, 0.010, np.pi / 2, 0]), - ] - ) - - # ==== const + cos ==== - functions["H2"] = dict( - [ - ("label", "B*cos(w*t+phi)+const"), - ("function", lambda p, t: p[0] * np.cos(p[1] * t + p[2]) + p[3]), - ("initial guess", [0, 0.010, 0, 0]), - ] - ) - - # ==== linear + cos ==== - functions["H3"] = dict( - [ - ("label", "B*cos(w*t+phi)+linear"), - ( - "function", - lambda p, t: p[3] + p[4] * t + p[0] * np.cos(p[1] * t + p[2]), - ), - ( - "initial guess", - [ - 0, - 0.017, - 0, - 0, - 0, - ], - ), - ] - ) - - # ==== quadratic + cos ==== - functions["H4"] = dict( - [ - ("label", "B*cos(w*t+phi)+quadratic"), - ( - "function", - lambda p, t: p[3] - + p[4] * t - + p[5] * t**2 - + p[0] * np.cos(p[1] * t + p[2]), - ), - ( - "initial guess", - [ - 0, - 0.017, - 0, - 0, - 0, - 0, - ], - ), - ] - ) - - # Fit and compute updates - for name, func in functions.items(): - # We will handle H4 separately - if name == "H4": - continue - - func["fit result"] = compute_coord_sep_updates( - x=t, - y=dsdt, - model=func, - initial_separation=initial_separation, - initial_xcts_values=initial_xcts_values, - ) - - # ==== quadratic + cos ==== - # Replace the initial guess with that of the linear solve - iguess_len = len(functions["H3"]["initial guess"]) - functions["H4"]["initial guess"][0:iguess_len] = functions["H3"][ - "fit result" - ]["parameters"][0:iguess_len] - - functions["H4"]["fit result"] = compute_coord_sep_updates( - x=t, - y=dsdt, - model=functions["H4"], - initial_separation=initial_separation, - initial_xcts_values=initial_xcts_values, - ) - - # Print results and plot - coordinate_separation_eccentricity_control_digest( - h5_file=h5_file, - x=t, - y=dsdt, - data=data, - functions=functions, - fig=fig, - ) - - return functions - - -@click.command(name="eccentricity-control") -@click.argument( - "h5_file", - type=click.Path(exists=True, file_okay=True, dir_okay=False, readable=True), -) -@click.option( - "--subfile-name-aha", - "-A", - default="ApparentHorizons/ControlSystemAhA_Centers.dat", - show_default=True, - help=( - "Name of subfile containing the apparent horizon centers for object A." - ), -) -@click.option( - "--subfile-name-ahb", - "-B", - default="ApparentHorizons/ControlSystemAhB_Centers.dat", - show_default=True, - help=( - "Name of subfile containing the apparent horizon centers for object B." - ), -) -@click.option( - "--tmin", - type=float, - help=( - "The lower time bound to start the fit. Used to remove initial junk" - " and transients in the coordinate separations. Default tmin=20 (or 60)" - " for tmax<200 (or >200)." - ), -) -@click.option( - "--tmax", - type=float, - help=( - "The upper time bound to start the fit. A reasonable value would" - " include 2-3 orbits." - ), -) -@click.option( - "--angular-velocity-from-xcts", - type=float, - help="Value of the angular velocity used in the Xcts file.", -) -@click.option( - "--expansion-from-xcts", - type=float, - help="Value of the expansion velocity (adot) used in the Xcts file.", -) -@apply_stylesheet_command() -@show_or_save_plot_command() -def eccentricity_control_command( - h5_file, - subfile_name_aha, - subfile_name_ahb, - tmin, - tmax, - angular_velocity_from_xcts, - expansion_from_xcts, -): - """Compute updates based on fits to the coordinate separation for manual - eccentricity control - - Usage: - - Select an appropriate time window without large initial transients and about - 2 to 3 orbits of data. - - This script uses the coordinate separations between Objects A and B to - compute a finite difference approximation to the time derivative. - It then fits different models to it. - - For each model, the suggested updates dOmega and dadot based on Newtonian - estimates are printed. Note that when all models fit the data adequately, - their updates are similar. When they differ, examine the output plot to - find a model that is good fit and has small residuals (especially at early - times). - - Finally, replace the updated values to the angular velocity and expansion - parameters (respectively) in the Xcts input file, or use the suggested - updates to compute them (if the initial xcts parameters were not provided). - - See ArXiv:gr-qc/0702106 and ArXiv:0710.0158 for more details. - - Limitations: - - 1) These eccentricity updates work only for non-precessing binaries. - 2) The time window is manually specified by the user. - 3) The coordinate separation is used, instead of the proper distance. - - See OmegaDoEccRemoval.py in SpEC for improved eccentricity control. - - """ - fig = plt.figure() - functions = coordinate_separation_eccentricity_control( - h5_file=h5_file, - subfile_name_aha=subfile_name_aha, - subfile_name_ahb=subfile_name_ahb, - tmin=tmin, - tmax=tmax, - angular_velocity_from_xcts=angular_velocity_from_xcts, - expansion_from_xcts=expansion_from_xcts, - fig=fig, - ) - return fig diff --git a/support/Pipelines/EccentricityControl/EccentricityControlParams.py b/support/Pipelines/EccentricityControl/EccentricityControlParams.py new file mode 100644 index 0000000000000..3216f0170b6cf --- /dev/null +++ b/support/Pipelines/EccentricityControl/EccentricityControlParams.py @@ -0,0 +1,288 @@ +#!/usr/bin/env python + +# Distributed under the MIT License. +# See LICENSE.txt for details. + +import functools +import logging +from pathlib import Path +from typing import Dict, Literal, Optional, Sequence, Tuple, Union + +import click +import h5py +import numpy as np +import pandas as pd +import yaml + +from spectre.Visualization.PlotTrajectories import import_A_and_B +from spectre.Visualization.ReadH5 import to_dataframe + +logger = logging.getLogger(__name__) + +# Input orbital parameters that can be controlled +OrbitalParams = Literal[ + "Omega0", + "adot0", + "D0", +] + +DEFAULT_AHA_TRAJECTORIES = "ApparentHorizons/ControlSystemAhA_Centers.dat" +DEFAULT_AHB_TRAJECTORIES = "ApparentHorizons/ControlSystemAhB_Centers.dat" + + +def eccentricity_control_params( + h5_files: Union[Union[str, Path], Sequence[Union[str, Path]]], + id_input_file_path: Union[str, Path], + subfile_name_aha_trajectories: str = DEFAULT_AHA_TRAJECTORIES, + subfile_name_ahb_trajectories: str = DEFAULT_AHB_TRAJECTORIES, + subfile_name_aha_quantities: str = "ObservationAhA.dat", + subfile_name_ahb_quantities: str = "ObservationAhB.dat", + tmin: Optional[float] = None, + tmax: Optional[float] = None, + target_eccentricity: float = 0.0, + target_mean_anomaly_fraction: Optional[float] = None, + plot_output_dir: Optional[Union[str, Path]] = None, +) -> Tuple[float, float, Dict[OrbitalParams, float]]: + """Get new orbital parameters for a binary system to control eccentricity. + + The eccentricity is estimated from the trajectories of the binary objects + and updates to the orbital parameters are suggested to drive the orbit to + the target eccentricity, using SpEC's OmegaDotEccRemoval.py. Currently + supports only circular target orbits (target eccentricity = 0). + + Arguments: + h5_files: Paths to the H5 files containing the trajectory data (e.g. + BbhReductions.h5). + id_input_file_path: Path to the initial data input file from which the + evolution started. This file contains the initial data parameters that + are being controlled. + subfile_name_aha_trajectories: (Optional) Name of the subfile containing + the apparent horizon centers for object A. + subfile_name_ahb_trajectories: (Optional) Name of the subfile containing + the apparent horizon centers for object B. + subfile_name_aha_quantities: (Optional) Name of the subfile containing the + quantities measured on apparent horizon A (masses and spins). + subfile_name_ahb_quantities: (Optional) Name of the subfile containing the + quantities measured on apparent horizon B (masses and spins). + tmin: (Optional) The lower time bound for the eccentricity estimate. + Used to remove initial junk and transients in the data. + tmax: (Optional) The upper time bound for the eccentricity estimate. + A reasonable value would include 2-3 orbits. + target_eccentricity: (Optional) The target eccentricity to drive the + orbit to. Default is 0.0 (circular orbit). + target_mean_anomaly_fraction: (Optional) The target mean anomaly of the + orbit divided by 2 pi, so it is a number between 0 and 1. The value 0 + corresponds to the pericenter of the orbit (closest approach), the value + 0.5 corresponds to the apocenter of the orbit (farthest distance), and + the value 1 corresponds to the pericenter again. Currently this is + unused because only an eccentricity of 0 is supported. + plot_output_dir: (Optional) Output directory for plots. + + Returns: + Tuple of eccentricity estimate, eccentricity error, and dictionary of + new orbital parameters. + """ + # Import from SpEC + from OmegaDotEccRemoval import ComputeOmegaAndDerivsFromFile, performAllFits + + # Make sure h5_files is a sequence + if isinstance(h5_files, (str, Path)): + h5_files = [h5_files] + + # Read initial data parameters from input file + with open(id_input_file_path, "r") as open_input_file: + _, id_input_file = yaml.safe_load_all(open_input_file) + id_binary = id_input_file["Background"]["Binary"] + Omega0 = id_binary["AngularVelocity"] + adot0 = id_binary["Expansion"] + D0 = id_binary["XCoords"][1] - id_binary["XCoords"][0] + + # Load trajectory data + traj_A, traj_B = import_A_and_B( + h5_files, subfile_name_aha_trajectories, subfile_name_ahb_trajectories + ) + if tmin is not None: + traj_A = traj_A[traj_A[:, 0] >= tmin] + traj_B = traj_B[traj_B[:, 0] >= tmin] + if tmax is not None: + traj_A = traj_A[traj_A[:, 0] <= tmax] + traj_B = traj_B[traj_B[:, 0] <= tmax] + + # Load horizon parameters from evolution data at reference time (tmin) + def get_horizons_data(reductions_file): + with h5py.File(reductions_file, "r") as open_h5file: + horizons_data = [] + for ab, subfile_name in zip( + "AB", [subfile_name_aha_quantities, subfile_name_ahb_quantities] + ): + ah_subfile = open_h5file.get(subfile_name) + if ah_subfile is not None: + horizons_data.append( + to_dataframe(ah_subfile) + .set_index("Time") + .add_prefix(f"Ah{ab} ") + ) + if not horizons_data: + return pd.DataFrame() + return pd.concat(horizons_data, axis=1) + + horizon_params = pd.concat(map(get_horizons_data, h5_files)) + if tmin is not None: + horizon_params = horizon_params[horizon_params.index >= tmin] + if horizon_params.empty: + logger.warning( + "No horizon data found in time range. " + "Using initial data masses and ignoring spins." + ) + mA = id_binary["ObjectRight"]["KerrSchild"]["Mass"] + mB = id_binary["ObjectLeft"]["KerrSchild"]["Mass"] + sA = sB = None + else: + mA = horizon_params["AhA ChristodoulouMass"].iloc[0] + mB = horizon_params["AhB ChristodoulouMass"].iloc[0] + if "AhA DimensionfulSpinVector_x" in horizon_params.columns: + sA = [ + horizon_params[f"AhA DimensionfulSpinVector_{xyz}"] + for xyz in "xyz" + ] + sB = [ + horizon_params[f"AhB DimensionfulSpinVector_{xyz}"] + for xyz in "xyz" + ] + else: + logger.warning("No horizon spins found in data, ignoring spins.") + sA = sB = None + + # Call into SpEC's OmegaDotEccRemoval.py + t, Omega, dOmegadt, OmegaVec = ComputeOmegaAndDerivsFromFile(traj_A, traj_B) + eccentricity, delta_Omega0, delta_adot0, delta_D0, ecc_std_dev, _ = ( + performAllFits( + XA=traj_A, + XB=traj_B, + t=t, + Omega=Omega, + dOmegadt=dOmegadt, + OmegaVec=OmegaVec, + mA=mA, + mB=mB, + sA=sA, + sB=sB, + IDparam_omega0=Omega0, + IDparam_adot0=adot0, + IDparam_D0=D0, + tmin=tmin or 0.0, + tmax=tmax or np.inf, + tref=tmin or 0.0, + opt_freq_filter=True, + opt_varpro=True, + opt_type="bbh", + opt_tmin=tmin, + opt_improved_Omega0_update=True, + check_periastron_advance=True, + plot_output_dir=plot_output_dir, + Source="", + ) + ) + logger.info( + f"Eccentricity estimate is {eccentricity:g} +/- {ecc_std_dev:e}." + " Update orbital parameters as follows" + f" for target eccentricity {target_eccentricity:g} (choose two):\n" + f"Omega0 += {delta_Omega0:e} -> {Omega0 + delta_Omega0:.8g}\n" + f"adot0 += {delta_adot0:e} -> {adot0 + delta_adot0:e}\n" + f"D0 += {delta_D0:e} -> {D0 + delta_D0:.8g}" + ) + return ( + eccentricity, + ecc_std_dev, + { + "Omega0": Omega0 + delta_Omega0, + "adot0": adot0 + delta_adot0, + "D0": D0 + delta_D0, + }, + ) + + +def eccentricity_control_params_options(f): + """CLI options for the 'eccentricity_control_params' function. + + These options can be used by CLI commands that call the + 'eccentricity_control_params' function. + """ + + @click.argument( + "h5_files", + nargs=-1, + type=click.Path( + exists=True, file_okay=True, dir_okay=False, readable=True + ), + ) + @click.option( + "--subfile-name-aha-trajectories", + default=DEFAULT_AHA_TRAJECTORIES, + show_default=True, + help=( + "Name of subfile containing the apparent horizon centers for" + " object A." + ), + ) + @click.option( + "--subfile-name-ahb-trajectories", + default=DEFAULT_AHB_TRAJECTORIES, + show_default=True, + help=( + "Name of subfile containing the apparent horizon centers for" + " object B." + ), + ) + @click.option( + "--subfile-name-aha-quantities", + default="ObservationsAhA.dat", + show_default=True, + help=( + "Name of subfile containing the quantities measured on apparent" + " horizon A (masses and spins)." + ), + ) + @click.option( + "--subfile-name-ahb-quantities", + default="ObservationsAhB.dat", + show_default=True, + help=( + "Name of subfile containing the quantities measured on apparent" + " horizon A (masses and spins)." + ), + ) + @click.option( + "--id-input-file", + "-i", + "id_input_file_path", + required=True, + help="Input file with initial data parameters.", + ) + @click.option( + "--tmin", + type=float, + help=( + "The lower time bound for the eccentricity estimate. Used to remove" + " initial junk and transients in the data." + ), + ) + @click.option( + "--tmax", + type=float, + help=( + "The upper time bound for the eccentricity estimate. A reasonable" + " value would include 2-3 orbits." + ), + ) + @click.option( + "-o", + "--plot-output-dir", + type=click.Path(writable=True), + help="Output directory for plots.", + ) + @functools.wraps(f) + def wrapper(*args, **kwargs): + return f(*args, **kwargs) + + return wrapper diff --git a/support/Pipelines/EccentricityControl/__init__.py b/support/Pipelines/EccentricityControl/__init__.py index a658f9724b466..32e61c9134fa5 100644 --- a/support/Pipelines/EccentricityControl/__init__.py +++ b/support/Pipelines/EccentricityControl/__init__.py @@ -1,9 +1,2 @@ # Distributed under the MIT License. # See LICENSE.txt for details. - -import click - -from .EccentricityControl import eccentricity_control_command - -if __name__ == "__main__": - eccentricity_control_command(help_option_names=["-h", "--help"]) diff --git a/tests/support/Pipelines/Bbh/Test_EccentricityControl.py b/tests/support/Pipelines/Bbh/Test_EccentricityControl.py index 9c34bdcca2d3e..e4901dbb19c31 100644 --- a/tests/support/Pipelines/Bbh/Test_EccentricityControl.py +++ b/tests/support/Pipelines/Bbh/Test_EccentricityControl.py @@ -6,26 +6,34 @@ import os import shutil import unittest +from pathlib import Path import numpy as np import yaml +from click.testing import CliRunner import spectre.IO.H5 as spectre_h5 from spectre.Informer import unit_test_build_path, unit_test_src_path -from spectre.Pipelines.Bbh.EccentricityControl import eccentricity_control +from spectre.Pipelines.Bbh.EccentricityControl import ( + eccentricity_control, + eccentricity_control_command, +) from spectre.support.Logging import configure_logging +from spectre.testing.Pipelines.MockBinaryData import write_mock_trajectory_data class TestEccentricityControl(unittest.TestCase): # Set up and prepare test directory and file paths def setUp(self): self.test_dir = os.path.join( - unit_test_build_path(), "Pipelines", "EccentricityControl" + unit_test_build_path(), "Pipelines/Bbh/EccentricityControl" ) self.h5_filename = os.path.join( self.test_dir, "TestEccentricityControlData.h5" ) - self.id_input_file_path = os.path.join(self.test_dir, "Inspiral.yaml") + self.id_input_file_path = os.path.join( + self.test_dir, "InitialData.yaml" + ) # Clean up any existing test directory and create new one shutil.rmtree(self.test_dir, ignore_errors=True) os.makedirs(self.test_dir, exist_ok=True) @@ -38,88 +46,29 @@ def tearDown(self): shutil.rmtree(self.test_dir) def create_h5_file(self): - logging.info(f"Creating HDF5 file: {self.h5_filename}") - # Define parameters for sample data - nsamples = 100 - dt = 0.02 - x0, y0, z0, z1 = 0.35, 0.35, 0, -9.0e-6 - cosAmp, sinAmp, cosFreq, sinFreq = 7.43, 7.43, 0.0173, 0.0172 - - # Define functions to generate data in position vs time format - def SpiralA(t): - return np.array( - [ - x0 + cosAmp * np.cos(-cosFreq * t), - y0 - sinAmp * np.sin(sinFreq * t), - z0 + z1 * (1 - 0.1) * t, - ] - ) - - def SpiralB(t): - return np.array( - [ - -x0 + cosAmp * np.cos(np.pi + cosFreq * t), - -y0 + sinAmp * np.sin(sinFreq * t), - z0 + z1 * (1 + 0.1) * t, - ] - ) - - # Generate time table and sample data - tTable = np.arange(0, (nsamples + 1) * dt, dt) - AhA_data = np.array([[t, *SpiralA(t), *SpiralA(t)] for t in tTable]) - AhB_data = np.array([[t, *SpiralB(t), *SpiralB(t)] for t in tTable]) - - # Create and populate the HDF5 files with data - with spectre_h5.H5File(self.h5_filename, "w") as h5_file: - dataset_AhA = h5_file.insert_dat( - "ApparentHorizons/ControlSystemAhA_Centers", - legend=[ - "Time", - "GridCenter_x", - "GridCenter_y", - "GridCenter_z", - "InertialCenter_x", - "InertialCenter_y", - "InertialCenter_z", - ], - version=0, - ) - for data_point in AhA_data: - dataset_AhA.append(data_point) - logging.debug( - f"Appended {len(AhA_data)} data points to AhA dataset." - ) - h5_file.close_current_object() - - dataset_AhB = h5_file.insert_dat( - "ApparentHorizons/ControlSystemAhB_Centers", - legend=[ - "Time", - "GridCenter_x", - "GridCenter_y", - "GridCenter_z", - "InertialCenter_x", - "InertialCenter_y", - "InertialCenter_z", - ], - version=0, - ) - for data_point in AhB_data: - dataset_AhB.append(data_point) - logging.debug( - f"Appended {len(AhB_data)} data points to AhB dataset." - ) - h5_file.close_current_object() - - logging.info( - f"Successfully created and populated HDF5 file: {self.h5_filename}" + self.initial_separation = 16.0 + self.angular_velocity = 0.015625 + self.eccentricity = 0.0 + write_mock_trajectory_data( + self.h5_filename, + t=np.arange(0, 2000, 2.0), + separation=self.initial_separation, ) def create_yaml_file(self): # Define YAML data and write it to the file data1 = { "Background": { - "Binary": {"AngularVelocity": 0.01, "Expansion": 0.001} + "Binary": { + "AngularVelocity": self.angular_velocity, + "Expansion": -1e-6, + "XCoords": [ + self.initial_separation / 2.0, + -self.initial_separation / 2.0, + ], + "ObjectLeft": {"KerrSchild": {"Mass": 0.5}}, + "ObjectRight": {"KerrSchild": {"Mass": 0.5}}, + }, } } with open(self.id_input_file_path, "w") as yaml_file: @@ -130,12 +79,26 @@ def create_yaml_file(self): # Test the eccentricity control function with the created files def test_eccentricity_control(self): eccentricity_control( - h5_file=self.h5_filename, + h5_files=self.h5_filename, id_input_file_path=self.id_input_file_path, - tmin=0, - tmax=10, ) + def test_cli(self): + runner = CliRunner() + result = runner.invoke( + eccentricity_control_command, + [ + self.h5_filename, + "-i", + self.id_input_file_path, + "-o", + self.test_dir, + ], + catch_exceptions=False, + ) + self.assertEqual(result.exit_code, 0, result.output) + self.assertTrue((Path(self.test_dir) / "FigureEccRemoval.pdf").exists()) + if __name__ == "__main__": configure_logging(log_level=logging.DEBUG) diff --git a/tests/support/Pipelines/CMakeLists.txt b/tests/support/Pipelines/CMakeLists.txt index 06e9fda336aa2..e776e4e8709a7 100644 --- a/tests/support/Pipelines/CMakeLists.txt +++ b/tests/support/Pipelines/CMakeLists.txt @@ -1,5 +1,12 @@ # Distributed under the MIT License. # See LICENSE.txt for details. +spectre_python_add_module( + Pipelines + MODULE_PATH "testing" + PYTHON_FILES + MockBinaryData.py +) + add_subdirectory(Bbh) add_subdirectory(EccentricityControl) diff --git a/tests/support/Pipelines/EccentricityControl/CMakeLists.txt b/tests/support/Pipelines/EccentricityControl/CMakeLists.txt index db8024cfdcc37..3548e47061c2e 100644 --- a/tests/support/Pipelines/EccentricityControl/CMakeLists.txt +++ b/tests/support/Pipelines/EccentricityControl/CMakeLists.txt @@ -1,15 +1,13 @@ # Distributed under the MIT License. # See LICENSE.txt for details. - -spectre_add_python_bindings_test( - "support.Pipelines.EccentricityControl.EccentricityControl" - Test_EccentricityControl.py - "Python" - None - TIMEOUT 60) - if (SpEC_FOUND) + spectre_add_python_bindings_test( + "support.Pipelines.EccentricityControl.EccentricityControlParams" + Test_EccentricityControlParams.py + "Python" + None + TIMEOUT 60) spectre_add_python_bindings_test( "support.Pipelines.EccentricityControl.InitialOrbitalParameters" Test_InitialOrbitalParameters.py diff --git a/tests/support/Pipelines/EccentricityControl/Test_EccentricityControl.py b/tests/support/Pipelines/EccentricityControl/Test_EccentricityControl.py deleted file mode 100644 index dd568e4660fa4..0000000000000 --- a/tests/support/Pipelines/EccentricityControl/Test_EccentricityControl.py +++ /dev/null @@ -1,262 +0,0 @@ -# Distributed under the MIT License. -# See LICENSE.txt for details. - -import logging -import os -import shutil -import unittest - -import numpy as np -import numpy.testing as npt -from click.testing import CliRunner - -import spectre.IO.H5 as spectre_h5 -from spectre.Informer import unit_test_build_path -from spectre.Pipelines.EccentricityControl import eccentricity_control_command -from spectre.Pipelines.EccentricityControl.EccentricityControl import ( - coordinate_separation_eccentricity_control, -) -from spectre.support.Logging import configure_logging - - -class TestEccentricityControl(unittest.TestCase): - def setUp(self): - self.test_dir = os.path.join( - unit_test_build_path(), "Pipelines", "EccentricityControl" - ) - self.h5_filename = os.path.join( - self.test_dir, "TestPlotTrajectoriesReductions.h5" - ) - shutil.rmtree(self.test_dir, ignore_errors=True) - os.makedirs(self.test_dir, exist_ok=True) - self.create_h5_file() - - def tearDown(self): - shutil.rmtree(self.test_dir) - - def create_h5_file(self): - """Create the H5""" - logging.info(f"Creating HDF5 file: {self.h5_filename}") - - # Generate mock-up inspiral data - nsamples = 120 - dt = 2.0 - - # Mock eccentric orbit parameters - ecc = 1e-2 - a = 16.0 - initial_phase = 0.0 - - # Some useful definitions - Omega_bar = np.sqrt(1.0 / (a**3)) - r_bar = a * (1 - ecc**2) - - # Origin offset and linear drift in z - x0 = 0.35 - y0 = 0.35 - z0 = 0 - z1 = -9.0e-6 - - # Save expected values for eccentricity oscillations - self.initial_separation = r_bar - self.initial_phase = initial_phase - self.ecc = ecc - self.frequency = Omega_bar - self.amplitude = r_bar * ecc * Omega_bar - logging.info(f"Expected eccentric parameters") - logging.info(f"Frequency: {self.frequency}") - logging.info(f"Amplitude of oscillations: {self.amplitude}") - - # Linear in eccentricity orbit approximation - def r(t): - return r_bar * (1 - ecc * np.cos(Omega_bar * t)) - - def phase(t): - return ( - Omega_bar * t - + (2.0 * ecc / Omega_bar) * np.sin(Omega_bar * t) - + initial_phase - ) - - # Define the spirals for equal mass orbits - def SpiralA(t): - return np.array( - [ - x0 + 0.5 * r(t) * np.cos(phase(t)), - y0 + 0.5 * r(t) * np.sin(phase(t)), - z0 + z1 * (1 - 0.1) * t, - ] - ) - - def SpiralB(t): - return np.array( - [ - x0 - 0.5 * r(t) * np.cos(phase(t)), - y0 - 0.5 * r(t) * np.sin(phase(t)), - z0 + z1 * (1 - 0.1) * t, - ] - ) - - # Generate time samples - tTable = np.arange(0, (nsamples + 1) * dt, dt) - - # Map time to spiral points - AhA_data = np.array([[t, *SpiralA(t), *SpiralA(t)] for t in tTable]) - AhB_data = np.array([[t, *SpiralB(t), *SpiralB(t)] for t in tTable]) - - with spectre_h5.H5File(self.h5_filename, "w") as h5_file: - # Insert dataset for AhA - dataset_AhA = h5_file.insert_dat( - "ApparentHorizons/ControlSystemAhA_Centers.dat", - legend=[ - "Time", - "GridCenter_x", - "GridCenter_y", - "GridCenter_z", - "InertialCenter_x", - "InertialCenter_y", - "InertialCenter_z", - ], # Legend for the dataset - version=0, # Version number - ) - # Populate dataset with AhA data - for data_point in AhA_data: - dataset_AhA.append(data_point) - # Close dataset for AhA - h5_file.close_current_object() - - # Insert dataset for AhB - dataset_AhB = h5_file.insert_dat( - "ApparentHorizons/ControlSystemAhB_Centers.dat", - legend=[ - "Time", - "GridCenter_x", - "GridCenter_y", - "GridCenter_z", - "InertialCenter_x", - "InertialCenter_y", - "InertialCenter_z", - ], # Legend for the dataset - version=0, # Version number - ) - # Populate dataset with AhB data - for data_point in AhB_data: - dataset_AhB.append(data_point) - # Close dataset for AhB - h5_file.close_current_object() - logging.info( - f"Successfully created and populated HDF5 file: {self.h5_filename}" - ) - - def test_cli(self): - output_filename = os.path.join(self.test_dir, "output.pdf") - runner = CliRunner() - result = runner.invoke( - eccentricity_control_command, - [ - self.h5_filename, - "-A", - "ApparentHorizons/ControlSystemAhA_Centers.dat", - "-B", - "ApparentHorizons/ControlSystemAhB_Centers.dat", - "--tmin", - 0.0, - "--tmax", - 500.0, - "--angular-velocity-from-xcts", - 0.0173, - "--expansion-from-xcts", - -1e-6, - "-o", - output_filename, - ], - catch_exceptions=False, - ) - self.assertEqual(result.exit_code, 0, result.output) - self.assertTrue(os.path.exists(output_filename)) - - def test_output_parameters(self): - mock_angular_velocity_from_xcts = 0.0156 - mock_expansion_from_xcts = -1e-6 - - functions = coordinate_separation_eccentricity_control( - h5_file=self.h5_filename, - subfile_name_aha="ApparentHorizons/ControlSystemAhA_Centers.dat", - subfile_name_ahb="ApparentHorizons/ControlSystemAhB_Centers.dat", - tmin=0.0, - tmax=1200.0, - angular_velocity_from_xcts=mock_angular_velocity_from_xcts, - expansion_from_xcts=mock_expansion_from_xcts, - fig=None, - ) - - for name, func in functions.items(): - logging.info(f"Test fit function: {name}") - # Values from fit - ampl, freq, phase = func["fit result"]["parameters"][:3] - updates_from_fit = func["fit result"]["xcts updates"] - updated_xcts_values_from_fit = func["fit result"][ - "updated xcts values" - ] - eccentricity_from_fit = func["fit result"]["eccentricity"] - - # Expected orbital values - expected_magnitude_of_frequency = self.frequency - expected_magnitude_of_oscillations = self.amplitude - expected_magnitude_of_phase = np.pi / 2 - - tol = 1e-6 - self.assertAlmostEqual( - np.abs(ampl), expected_magnitude_of_oscillations, delta=tol - ) - self.assertAlmostEqual( - np.abs(freq), expected_magnitude_of_frequency, delta=tol - ) - self.assertAlmostEqual( - np.abs(phase), expected_magnitude_of_phase, delta=tol - ) - - # Expected updates - expected_dOmg = ( - self.amplitude - / 2.0 - / self.initial_separation - * np.sin(self.initial_phase) - ) - expected_dadot = ( - -self.amplitude - / self.initial_separation - * np.cos(self.initial_phase) - ) - - # Expected updates - expected_updated_Omg = ( - mock_angular_velocity_from_xcts + expected_dOmg - ) - expected_updated_adot = mock_expansion_from_xcts + expected_dadot - - # Updates from fit - dOmg = updates_from_fit["omega update"] - dadot = updates_from_fit["expansion update"] - updated_Omg = updated_xcts_values_from_fit["omega"] - updated_adot = updated_xcts_values_from_fit["expansion"] - - tol_for_updates = 5e-4 - self.assertAlmostEqual(expected_dOmg, dOmg, delta=tol_for_updates) - self.assertAlmostEqual(expected_dadot, dadot, delta=tol_for_updates) - self.assertAlmostEqual( - expected_updated_Omg, updated_Omg, delta=tol_for_updates - ) - self.assertAlmostEqual( - expected_updated_adot, updated_adot, delta=tol_for_updates - ) - - tol_for_ecc = 5e-4 - self.assertAlmostEqual( - self.ecc, eccentricity_from_fit, delta=tol_for_ecc - ) - - -if __name__ == "__main__": - configure_logging(log_level=logging.DEBUG) - unittest.main(verbosity=2) diff --git a/tests/support/Pipelines/EccentricityControl/Test_EccentricityControlParams.py b/tests/support/Pipelines/EccentricityControl/Test_EccentricityControlParams.py new file mode 100644 index 0000000000000..cd56e40655492 --- /dev/null +++ b/tests/support/Pipelines/EccentricityControl/Test_EccentricityControlParams.py @@ -0,0 +1,85 @@ +# Distributed under the MIT License. +# See LICENSE.txt for details. + +import logging +import os +import shutil +import unittest +from pathlib import Path + +import numpy as np +import numpy.testing as npt +import yaml + +from spectre.Informer import unit_test_build_path +from spectre.Pipelines.EccentricityControl.EccentricityControlParams import ( + eccentricity_control_params, +) +from spectre.support.Logging import configure_logging +from spectre.testing.Pipelines.MockBinaryData import write_mock_trajectory_data + + +class TestEccentricityControlParams(unittest.TestCase): + def setUp(self): + self.test_dir = os.path.join( + unit_test_build_path(), "Pipelines", "EccentricityControl" + ) + self.h5_filename = os.path.join( + self.test_dir, "TestPlotTrajectoriesReductions.h5" + ) + shutil.rmtree(self.test_dir, ignore_errors=True) + os.makedirs(self.test_dir, exist_ok=True) + + # Write mock trajectory data to an H5 file + self.initial_separation = 16.0 + self.angular_velocity = 0.015625 + self.eccentricity = 0.0 + write_mock_trajectory_data( + self.h5_filename, + t=np.arange(0, 2000, 2.0), + separation=self.initial_separation, + ) + + # Write a mock initial data input file + self.id_input_file_path = os.path.join( + self.test_dir, "InitialData.yaml" + ) + with open(self.id_input_file_path, "w") as open_input_file: + yaml.safe_dump_all( + [ + {}, + { + "Background": { + "Binary": { + "AngularVelocity": self.angular_velocity, + "Expansion": -1e-6, + "XCoords": [ + self.initial_separation / 2.0, + -self.initial_separation / 2.0, + ], + "ObjectLeft": {"KerrSchild": {"Mass": 0.5}}, + "ObjectRight": {"KerrSchild": {"Mass": 0.5}}, + }, + } + }, + ], + open_input_file, + ) + + def tearDown(self): + shutil.rmtree(self.test_dir) + + def test_ecc_control_params(self): + ecc, ecc_std_dev, param_updates = eccentricity_control_params( + h5_files=[self.h5_filename], + id_input_file_path=self.id_input_file_path, + tmin=0.0, + tmax=1200.0, + plot_output_dir=self.test_dir, + ) + self.assertAlmostEqual(ecc, self.eccentricity) + + +if __name__ == "__main__": + configure_logging(log_level=logging.DEBUG) + unittest.main(verbosity=2) diff --git a/tests/support/Pipelines/MockBinaryData.py b/tests/support/Pipelines/MockBinaryData.py new file mode 100644 index 0000000000000..627481cd18c0a --- /dev/null +++ b/tests/support/Pipelines/MockBinaryData.py @@ -0,0 +1,54 @@ +# Distributed under the MIT License. +# See LICENSE.txt for details. + +import numpy as np + +import spectre.IO.H5 as spectre_h5 + + +def write_mock_trajectory_data( + reductions_filename: str, + t: np.ndarray, + separation: float, +): + """Generate mock-up inspiral data + + Simple circular Newtonian orbits for an equal-mass binary. Can be expanded + to more general orbits if needed. + """ + omega = np.sqrt(1.0 / (separation**3)) + phase = omega * t + centers_a = np.array( + [ + 0.5 * separation * np.cos(phase), + 0.5 * separation * np.sin(phase), + 0.0 * phase, + ] + ) + centers_b = np.array( + [ + -0.5 * separation * np.cos(phase), + -0.5 * separation * np.sin(phase), + 0.0 * phase, + ] + ) + + # Write to H5 file + with spectre_h5.H5File(reductions_filename, "a") as h5_file: + for ab, centers in zip("AB", [centers_a, centers_b]): + datfile = h5_file.insert_dat( + f"ApparentHorizons/ControlSystemAh{ab}_Centers.dat", + legend=[ + "Time", + "GridCenter_x", + "GridCenter_y", + "GridCenter_z", + "InertialCenter_x", + "InertialCenter_y", + "InertialCenter_z", + ], + version=0, + ) + for t_i, x_i in zip(t, centers.T): + datfile.append([t_i, *x_i, *x_i]) + h5_file.close_current_object()