diff --git a/src/Visualization/Python/CMakeLists.txt b/src/Visualization/Python/CMakeLists.txt index 828a17ffdc9d..80a2a44f92cf 100644 --- a/src/Visualization/Python/CMakeLists.txt +++ b/src/Visualization/Python/CMakeLists.txt @@ -16,6 +16,7 @@ spectre_python_add_module( PlotCce.py PlotControlSystem.py PlotDatFile.py + PlotEccentricityControl.py PlotEllipticConvergence.py PlotMemoryMonitors.py PlotPowerMonitors.py diff --git a/src/Visualization/Python/PlotEccentricityControl.py b/src/Visualization/Python/PlotEccentricityControl.py new file mode 100644 index 000000000000..aef72a5bffaa --- /dev/null +++ b/src/Visualization/Python/PlotEccentricityControl.py @@ -0,0 +1,159 @@ +#!/usr/bin/env python + +# Distributed under the MIT License. +# See LICENSE.txt for details. + +import logging +from pathlib import Path +from typing import Sequence, Union + +import click +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import yaml +from matplotlib.colors import LogNorm +from scipy.interpolate import CloughTocher2DInterpolator + +from spectre.Visualization.Plot import ( + apply_stylesheet_command, + show_or_save_plot_command, +) + +logger = logging.getLogger(__name__) + + +def plot_eccentricity_control( + ecc_params_files: Sequence[Union[str, Path]], +): + """Plot eccentricity control iterations. + + \f + Arguments: + ecc_params_files: YAML files containing the output of + 'eccentricity_control_params()'. + + Returns: + The figure containing the plot. + """ + # Load the eccentricity parameters + ecc_params = [] + for ecc_params_file in ecc_params_files: + with open(ecc_params_file, "r") as open_file: + ecc_params.append(yaml.safe_load(open_file)) + ecc_params = pd.DataFrame(ecc_params) + + # Set up the figure + fig = plt.figure(figsize=(4.5, 4)) + plt.xlabel(r"Angular velocity $\Omega_0$") + plt.ylabel(r"Expansion $\dot{a}_0$") + plt.grid(zorder=0) + plt.ticklabel_format(axis="both", style="sci", scilimits=(0, 0)) + + # Plot measured eccentricity contours + ecc_norm = LogNorm(vmin=1e-4, vmax=1) + if len(ecc_params) > 2: + ecc_intrp = CloughTocher2DInterpolator( + ecc_params[["Omega0", "Adot0"]], + np.log10(ecc_params["Eccentricity"]), + rescale=True, + ) + x, y = np.meshgrid( + np.linspace(*ecc_params["Omega0"].agg(["min", "max"]), 200), + np.linspace(*ecc_params["Adot0"].agg(["min", "max"]), 200), + ) + z = ecc_intrp(x, y) + ecc_contours = plt.contourf( + x, + y, + 10**z, + levels=10.0 ** np.arange(-4, 1), + norm=ecc_norm, + cmap="Blues", + zorder=10, + ) + else: + ecc_contours = plt.cm.ScalarMappable( + norm=ecc_norm, + cmap="Blues", + ) + plt.colorbar(ecc_contours, label="Eccentricity", ax=plt.gca()) + cbar_ax = fig.axes[1] + plt.gca().use_sticky_edges = False + + # Plot the path through parameter space + plt.plot( + ecc_params["Omega0"], + ecc_params["Adot0"], + color="black", + zorder=20, + ) + + # Place a marker at each iteration + scatter_kwargs = dict( + color="white", + marker="o", + linewidth=1, + edgecolor="black", + s=100, + zorder=25, + ) + annotate_kwargs = dict( + textcoords="offset points", + xytext=(0, -0.5), + ha="center", + va="center", + fontsize=8, + zorder=30, + ) + plt.scatter(ecc_params["Omega0"], ecc_params["Adot0"], **scatter_kwargs) + cbar_ax.scatter( + np.repeat(0.5, len(ecc_params)), + ecc_params["Eccentricity"], + **scatter_kwargs, + ) + for i in range(len(ecc_params)): + plt.annotate( + f"{i}", + ecc_params.iloc[i][["Omega0", "Adot0"]], + **annotate_kwargs, + ) + cbar_ax.annotate( + f"{i}", (0.5, ecc_params["Eccentricity"][i]), **annotate_kwargs + ) + + # Plot location of the next iteration + plt.plot( + ecc_params.iloc[-1][["Omega0", "NewOmega0"]], + ecc_params.iloc[-1][["Adot0", "NewAdot0"]], + color="black", + linestyle="dotted", + zorder=20, + ) + plt.scatter( + ecc_params.iloc[-1]["NewOmega0"], + ecc_params.iloc[-1]["NewAdot0"], + color="black", + marker="o", + s=30, + zorder=40, + ) + return fig + + +@click.command( + name="eccentricity-control", help=plot_eccentricity_control.__doc__ +) +@click.argument( + "ecc_params_files", + nargs=-1, + type=click.Path(exists=True, file_okay=True, dir_okay=False, readable=True), +) +@apply_stylesheet_command() +@show_or_save_plot_command() +def plot_eccentricity_control_command(**kwargs): + return plot_eccentricity_control(**kwargs) + + +if __name__ == "__main__": + plot_eccentricity_control_command(help_option_names=["-h", "--help"]) diff --git a/src/Visualization/Python/__init__.py b/src/Visualization/Python/__init__.py index 86614abc9e31..9024a685420d 100644 --- a/src/Visualization/Python/__init__.py +++ b/src/Visualization/Python/__init__.py @@ -13,6 +13,7 @@ def list_commands(self, ctx): "cce", "control-system", "dat", + "eccentricity-control", "elliptic-convergence", "memory-monitors", "power-monitors", @@ -42,6 +43,12 @@ def get_command(self, ctx, name): from spectre.Visualization.PlotDatFile import plot_dat_command return plot_dat_command + elif name in ["eccentricity-control", "ecc-control"]: + from spectre.Visualization.PlotEccentricityControl import ( + plot_eccentricity_control_command, + ) + + return plot_eccentricity_control_command elif name == "elliptic-convergence": from spectre.Visualization.PlotEllipticConvergence import ( plot_elliptic_convergence_command, diff --git a/support/Pipelines/Bbh/EccentricityControl.py b/support/Pipelines/Bbh/EccentricityControl.py index b403865b5e36..f17d55de73d2 100644 --- a/support/Pipelines/Bbh/EccentricityControl.py +++ b/support/Pipelines/Bbh/EccentricityControl.py @@ -27,6 +27,7 @@ def eccentricity_control( tmin: Optional[float] = 500, tmax: Optional[float] = None, plot_output_dir: Optional[Union[str, Path]] = None, + ecc_params_output_file: Optional[Union[str, Path]] = None, # Scheduler options evolve: bool = True, **scheduler_kwargs, @@ -48,8 +49,6 @@ def eccentricity_control( 'eccentricity_control_params' in 'spectre.Pipelines.EccentricityControl.EccentricityControl'. - - Displays the fit results in a tabular format using a pandas DataFrame. - - If the eccentricity is below a threshold, it prints "Success" and indicates that the simulation can continue. @@ -77,42 +76,17 @@ def eccentricity_control( # Find the current eccentricity and determine new parameters to put into # generate-id - ( - eccentricity, - ecc_std_dev, - new_orbital_params, - ) = eccentricity_control_params( + ecc_params = eccentricity_control_params( h5_files, id_input_file_path, tmin=tmin, tmax=tmax, plot_output_dir=plot_output_dir, + ecc_params_output_file=ecc_params_output_file, ) - # Create DataFrame to display data in tabular format - data = { - "Attribute": [ - "Eccentricity", - "Eccentricity error", - "Updated Omega0", - "Updated adot0", - ], - "Value": [ - eccentricity, - ecc_std_dev, - new_orbital_params["Omega0"], - new_orbital_params["adot0"], - ], - } - df = pd.DataFrame(data) - # Print header line - print("=" * 40) - # Display table - print(df.to_string(index=False)) - print("=" * 40) - # Stop eccentricity control if eccentricity is below threshold - if eccentricity < 0.001: + if ecc_params["Eccentricity"] < 0.001: print("Success") # Should continue the simulation either by restarting from a # checkpoint, or from the volume data - will do later @@ -129,8 +103,8 @@ def eccentricity_control( target_params, # New orbital parameters separation=separation, - orbital_angular_velocity=new_orbital_params["Omega0"], - radial_expansion_velocity=new_orbital_params["adot0"], + orbital_angular_velocity=ecc_params["NewOmega0"], + radial_expansion_velocity=ecc_params["NewAdot0"], # Initial guesses for ID control conformal_mass_a=binary_data["ObjectRight"]["KerrSchild"]["Mass"], conformal_mass_b=binary_data["ObjectLeft"]["KerrSchild"]["Mass"], @@ -163,6 +137,7 @@ def eccentricity_control( @click.option( "--evolve/--no-evolve", default=True, + show_default=True, help=( "Evolve the initial data after generation to continue eccentricity " "control. You can disable this to generate only the new initial data " diff --git a/support/Pipelines/Bbh/Inspiral.yaml b/support/Pipelines/Bbh/Inspiral.yaml index b4bf4e9f315c..747c23e7c264 100644 --- a/support/Pipelines/Bbh/Inspiral.yaml +++ b/support/Pipelines/Bbh/Inspiral.yaml @@ -21,8 +21,9 @@ Next: Run: spectre.Pipelines.Bbh.EccentricityControl:eccentricity_control With: h5_files: "../Segment_*/BbhReductions.h5" - plot_output_dir: ./ id_input_file_path: {{id_input_file_path }} + plot_output_dir: ./ + ecc_params_output_file: "../EccentricityParams.yaml" pipeline_dir: {{ pipeline_dir }} scheduler: {{ scheduler | default("None") }} copy_executable: {{ copy_executable | default("None") }} diff --git a/support/Pipelines/EccentricityControl/EccentricityControlParams.py b/support/Pipelines/EccentricityControl/EccentricityControlParams.py index aa0481ec99d5..a4d30fa88193 100644 --- a/support/Pipelines/EccentricityControl/EccentricityControlParams.py +++ b/support/Pipelines/EccentricityControl/EccentricityControlParams.py @@ -15,16 +15,27 @@ import pandas as pd import yaml +from spectre.support.Yaml import SafeDumper 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[ +# Keys of the dictionary returned by eccentricity_control_params +EccentricityParams = Literal[ + "Eccentricity", + "EccentricityError", "Omega0", - "adot0", + "Adot0", "D0", + "DeltaOmega0", + "DeltaAdot0", + "DeltaD0", + "NewOmega0", + "NewAdot0", + "NewD0", + "Tmin", + "Tmax", ] DEFAULT_AHA_TRAJECTORIES = "ApparentHorizons/ControlSystemAhA_Centers.dat" @@ -41,7 +52,8 @@ def eccentricity_control_params( tmin: Optional[float] = None, tmax: Optional[float] = None, plot_output_dir: Optional[Union[str, Path]] = None, -) -> Tuple[float, float, Dict[OrbitalParams, float]]: + ecc_params_output_file: Optional[Union[str, Path]] = None, +) -> Dict[EccentricityParams, float]: """Get new orbital parameters for a binary system to control eccentricity. The eccentricity is estimated from the trajectories of the binary objects @@ -68,10 +80,10 @@ def eccentricity_control_params( tmax: (Optional) The upper time bound for the eccentricity estimate. A reasonable value would include 2-3 orbits. plot_output_dir: (Optional) Output directory for plots. + ecc_params_output_file: (Optional) Output file for the results. Returns: - Tuple of eccentricity estimate, eccentricity error, and dictionary of - new orbital parameters with the keys listed in 'OrbitalParams'. + Dictionary with the keys listed in 'EccentricityParams'. """ # Import functions from SpEC until we have ported them over try: @@ -206,15 +218,26 @@ def get_horizons_data(reductions_file): 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, - }, - ) + # These keys must correspond to 'EccentricityParams' + ecc_params = { + "Eccentricity": eccentricity, + "EccentricityError": ecc_std_dev, + "Omega0": Omega0, + "Adot0": adot0, + "D0": D0, + "DeltaOmega0": delta_Omega0, + "DeltaAdot0": delta_adot0, + "DeltaD0": delta_D0, + "NewOmega0": Omega0 + delta_Omega0, + "NewAdot0": adot0 + delta_adot0, + "NewD0": D0 + delta_D0, + "Tmin": tmin, + "Tmax": tmax, + } + if ecc_params_output_file: + with open(ecc_params_output_file, "w") as open_file: + yaml.dump(ecc_params, open_file, Dumper=SafeDumper) + return ecc_params def eccentricity_control_params_options(f): @@ -251,7 +274,7 @@ def eccentricity_control_params_options(f): ) @click.option( "--subfile-name-aha-quantities", - default="ObservationsAhA.dat", + default="ObservationAhA.dat", show_default=True, help=( "Name of subfile containing the quantities measured on apparent" @@ -260,7 +283,7 @@ def eccentricity_control_params_options(f): ) @click.option( "--subfile-name-ahb-quantities", - default="ObservationsAhB.dat", + default="ObservationAhB.dat", show_default=True, help=( "Name of subfile containing the quantities measured on apparent" @@ -295,8 +318,22 @@ def eccentricity_control_params_options(f): type=click.Path(writable=True), help="Output directory for plots.", ) + @click.option( + "--ecc-params-output-file", + type=click.Path(writable=True), + help="Output file for the new orbital parameters.", + ) @functools.wraps(f) def wrapper(*args, **kwargs): return f(*args, **kwargs) return wrapper + + +@click.command( + name="eccentricity-control-params", help=eccentricity_control_params.__doc__ +) +@eccentricity_control_params_options +def eccentricity_control_params_command(**kwargs): + _rich_traceback_guard = True # Hide traceback until here + eccentricity_control_params(**kwargs) diff --git a/support/Python/CMakeLists.txt b/support/Python/CMakeLists.txt index 442df6ee11b1..b7c4b51b8fd5 100644 --- a/support/Python/CMakeLists.txt +++ b/support/Python/CMakeLists.txt @@ -14,6 +14,7 @@ spectre_python_add_module( Resubmit.py RunNext.py Schedule.py + Yaml.py ) if(MACHINE) diff --git a/support/Python/Schedule.py b/support/Python/Schedule.py index f4b19765c6a4..b23988b97ad7 100644 --- a/support/Python/Schedule.py +++ b/support/Python/Schedule.py @@ -16,7 +16,6 @@ import numpy as np import yaml from rich.pretty import pretty_repr -from yaml.representer import SafeRepresenter from spectre.support.DirectoryStructure import ( Checkpoint, @@ -26,6 +25,7 @@ ) from spectre.support.Machines import this_machine from spectre.support.RunNext import run_next +from spectre.support.Yaml import SafeDumper from spectre.tools.ValidateInputFile import validate_input_file from spectre.Visualization.ReadInputFile import find_phase_change @@ -119,18 +119,6 @@ def _copy_submit_script_template( return dest -# Write `pathlib.Path` objects to YAML as plain strings -def _path_representer(dumper: yaml.Dumper, path: Path) -> yaml.nodes.ScalarNode: - return dumper.represent_scalar("tag:yaml.org,2002:str", str(path)) - - -# Write `numpy.float64` as regular floats -def _numpy_representer( - dumper: yaml.Dumper, value: np.float64 -) -> yaml.nodes.ScalarNode: - return dumper.represent_scalar("tag:yaml.org,2002:float", str(value)) - - def schedule( input_file_template: Union[str, Path], scheduler: Optional[Union[str, Sequence]], @@ -665,10 +653,7 @@ def schedule( # Write context to file to support resubmissions if segments_dir: with open(run_dir / context_file_name, "w") as open_context_file: - yaml_dumper = yaml.SafeDumper - yaml_dumper.add_multi_representer(Path, _path_representer) - yaml_dumper.add_multi_representer(np.float64, _numpy_representer) - yaml.dump(context, open_context_file, Dumper=yaml_dumper) + yaml.dump(context, open_context_file, Dumper=SafeDumper) # Submit if submit or ( diff --git a/support/Python/Yaml.py b/support/Python/Yaml.py new file mode 100644 index 000000000000..368c852bb756 --- /dev/null +++ b/support/Python/Yaml.py @@ -0,0 +1,24 @@ +# Distributed under the MIT License. +# See LICENSE.txt for details. + +from pathlib import Path + +import numpy as np +import yaml + + +# Write `pathlib.Path` objects to YAML as plain strings +def _path_representer(dumper: yaml.Dumper, path: Path) -> yaml.nodes.ScalarNode: + return dumper.represent_scalar("tag:yaml.org,2002:str", str(path)) + + +# Write `numpy.float64` as regular floats +def _numpy_representer( + dumper: yaml.Dumper, value: np.float64 +) -> yaml.nodes.ScalarNode: + return dumper.represent_scalar("tag:yaml.org,2002:float", str(value)) + + +SafeDumper = yaml.SafeDumper +SafeDumper.add_multi_representer(Path, _path_representer) +SafeDumper.add_multi_representer(np.float64, _numpy_representer) diff --git a/support/Python/__main__.py b/support/Python/__main__.py index a221e8a76311..668a4c173246 100644 --- a/support/Python/__main__.py +++ b/support/Python/__main__.py @@ -25,6 +25,7 @@ def list_commands(self, ctx): "clean-output", "combine-h5", "delete-subfiles", + "eccentricity-control-params", "extend-connectivity", "extract-dat", "extract-input", @@ -66,6 +67,12 @@ def get_command(self, ctx, name): from spectre.IO.H5.DeleteSubfiles import delete_subfiles_command return delete_subfiles_command + elif name in ["eccentricity-control-params", "ecc-control-params"]: + from spectre.Pipelines.EccentricityControl.EccentricityControlParams import ( + eccentricity_control_params_command, + ) + + return eccentricity_control_params_command elif name == "extend-connectivity": from spectre.IO.H5.ExtendConnectivityData import ( extend_connectivity_data_command, diff --git a/tests/Unit/Visualization/Python/CMakeLists.txt b/tests/Unit/Visualization/Python/CMakeLists.txt index 8a268d48b0dc..0b5134f0b5b0 100644 --- a/tests/Unit/Visualization/Python/CMakeLists.txt +++ b/tests/Unit/Visualization/Python/CMakeLists.txt @@ -41,6 +41,13 @@ spectre_add_python_bindings_test( None TIMEOUT 10) +spectre_add_python_bindings_test( + "Unit.Visualization.Python.PlotEccentricityControl" + Test_PlotEccentricityControl.py + "unit;visualization;python" + None + TIMEOUT 10) + spectre_add_python_bindings_test( "Unit.Visualization.Python.PlotEllipticConvergence" Test_PlotEllipticConvergence.py diff --git a/tests/Unit/Visualization/Python/Test_PlotEccentricityControl.py b/tests/Unit/Visualization/Python/Test_PlotEccentricityControl.py new file mode 100644 index 000000000000..df3f75e82d91 --- /dev/null +++ b/tests/Unit/Visualization/Python/Test_PlotEccentricityControl.py @@ -0,0 +1,59 @@ +# Distributed under the MIT License. +# See LICENSE.txt for details. + +import os +import shutil +import unittest + +import numpy as np +import yaml +from click.testing import CliRunner + +from spectre.Informer import unit_test_build_path +from spectre.Visualization.PlotEccentricityControl import ( + plot_eccentricity_control_command, +) + + +class TestPlotEccentricityControl(unittest.TestCase): + def setUp(self): + self.test_dir = os.path.join( + unit_test_build_path(), "Visualization", "PlotEccentricityControl" + ) + os.makedirs(self.test_dir, exist_ok=True) + self.ecc_filenames = [] + for i in range(3): + ecc_filename = os.path.join( + self.test_dir, + f"EccentricityParams{i}.yaml", + ) + with open(ecc_filename, "w") as open_file: + yaml.safe_dump( + { + "Eccentricity": 10 ** (-i), + "Omega0": 0.015 + np.random.rand() * 1e-4, + "Adot0": 1e-4 + np.random.rand() * 1e-4, + "NewOmega0": 0.015 + np.random.rand() * 1e-4, + "NewAdot0": 1e-4 + np.random.rand() * 1e-4, + }, + open_file, + ) + self.ecc_filenames.append(ecc_filename) + + def tearDown(self): + shutil.rmtree(self.test_dir) + + def test_cli(self): + runner = CliRunner() + output_filename = os.path.join(self.test_dir, "output.pdf") + result = runner.invoke( + plot_eccentricity_control_command, + self.ecc_filenames + ["-o", output_filename], + catch_exceptions=False, + ) + self.assertEqual(result.exit_code, 0, result.output) + self.assertTrue(os.path.exists(output_filename)) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/support/Pipelines/Bbh/Test_Inspiral.py b/tests/support/Pipelines/Bbh/Test_Inspiral.py index 484907657a71..4800950de2ca 100644 --- a/tests/support/Pipelines/Bbh/Test_Inspiral.py +++ b/tests/support/Pipelines/Bbh/Test_Inspiral.py @@ -210,10 +210,11 @@ def test_cli(self): "Run": modulename + ":eccentricity_control", "With": { "h5_files": "../Segment_*/BbhReductions.h5", - "plot_output_dir": "./", "id_input_file_path": str( self.id_dir.resolve() / "InitialData.yaml" ), + "plot_output_dir": "./", + "ecc_params_output_file": "../EccentricityParams.yaml", "pipeline_dir": str(self.test_dir.resolve() / "Pipeline"), "scheduler": "None", "copy_executable": "None", diff --git a/tests/support/Pipelines/EccentricityControl/Test_EccentricityControlParams.py b/tests/support/Pipelines/EccentricityControl/Test_EccentricityControlParams.py index 372bbded3b41..c47387b6fb42 100644 --- a/tests/support/Pipelines/EccentricityControl/Test_EccentricityControlParams.py +++ b/tests/support/Pipelines/EccentricityControl/Test_EccentricityControlParams.py @@ -28,6 +28,9 @@ def setUp(self): self.h5_filename = os.path.join( self.test_dir, "TestEccentricityControlData.h5" ) + self.out_filename = os.path.join( + self.test_dir, "TestEccentricityParams.yaml" + ) shutil.rmtree(self.test_dir, ignore_errors=True) os.makedirs(self.test_dir, exist_ok=True) @@ -95,14 +98,17 @@ def tearDown(self): shutil.rmtree(self.test_dir) def test_ecc_control_params(self): - ecc, ecc_std_dev, param_updates = eccentricity_control_params( + ecc_params = eccentricity_control_params( h5_files=self.test_dir + "/TestEccentricity*.h5", id_input_file_path=self.id_input_file_path, tmin=0.0, tmax=1200.0, plot_output_dir=self.test_dir, + ecc_params_output_file=self.out_filename, ) - self.assertAlmostEqual(ecc, 0.0, delta=1e-5) + self.assertAlmostEqual(ecc_params["Eccentricity"], 0.0, delta=1e-5) + with open(self.out_filename, "r") as open_output_file: + self.assertEqual(yaml.safe_load(open_output_file), ecc_params) if __name__ == "__main__":