From 8c2ceba3cde140b58d77fa63e83612a0ec5e54ab Mon Sep 17 00:00:00 2001 From: Vittoria Tommasini Date: Mon, 16 Dec 2024 12:53:51 -0800 Subject: [PATCH] Automate multiple iterations of eccentricity reduction --- support/Pipelines/Bbh/EccentricityControl.py | 68 ++++++++++++++-- support/Pipelines/Bbh/InitialData.py | 17 +++- support/Pipelines/Bbh/InitialData.yaml | 1 + support/Pipelines/Bbh/Inspiral.py | 9 ++- support/Pipelines/Bbh/Inspiral.yaml | 14 ++++ support/Pipelines/Bbh/PostprocessId.py | 4 +- support/Python/DirectoryStructure.py | 77 ++++++++++++++++++- .../Pipelines/Bbh/Test_EccentricityControl.py | 2 +- .../support/Python/Test_DirectoryStructure.py | 27 +++++++ 9 files changed, 206 insertions(+), 13 deletions(-) diff --git a/support/Pipelines/Bbh/EccentricityControl.py b/support/Pipelines/Bbh/EccentricityControl.py index 9e503efb20ef..0a6068ffc3a1 100644 --- a/support/Pipelines/Bbh/EccentricityControl.py +++ b/support/Pipelines/Bbh/EccentricityControl.py @@ -3,18 +3,22 @@ import logging import warnings +from pathlib import Path +from typing import Optional, Union import matplotlib.pyplot as plt import pandas as pd import yaml +import spectre.IO.H5 as spectre_h5 +from spectre.Pipelines.Bbh.InitialData import generate_id from spectre.Pipelines.EccentricityControl.EccentricityControl import ( coordinate_separation_eccentricity_control, ) # Suppress specific RuntimeWarnings warnings.filterwarnings( - "ignore", message="Number of calls to functionhas reached maxfev" + "ignore", message="Number of calls to function has reached maxfev" ) logger = logging.getLogger(__name__) @@ -23,29 +27,43 @@ def eccentricity_control( h5_file, id_input_file_path, + pipeline_dir: Union[str, 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 + **scheduler_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. + - Sets the time boundaries for the eccentricity reduction process, starting + at 500 and using all available data by default, with the option to adjust + 'tmin' and 'tmax' dynamically. + + - Calls the 'coordinate_separation_eccentricity_control' function to + calculate the current eccentricity and extract more accurate orbital + parameters. + + - 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. + + - Generates new initial data based on updated orbital parameters using the + 'generate_id' function. Arguments: h5_file: file that contains the trajectory data id_input_file_path: path to the input file of the initial data run + pipeline_dir : directory where the pipeline outputs are stored. 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 @@ -56,10 +74,18 @@ def eccentricity_control( """ # 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) + id_metadata, 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"] + id_params = id_metadata["Next"]["With"] + control_params = id_params["control_params"] + mass_A = control_params["mass_A"] + mass_B = control_params["mass_B"] + spin_A = control_params["spin_A"] + spin_B = control_params["spin_B"] + x_B, x_A = binary_data["XCoords"] + separation = x_A - x_B if output: fig = plt.figure() @@ -117,3 +143,31 @@ def eccentricity_control( # Display table print(df.to_string(index=False)) print("=" * 40) + + if fit_result["eccentricity"] < 0.001: + print("Success") + # Should continue the simulation either by restarting from a + # checkpoint, or from the volume data - will do later + return + + # Generate new initial data based on updated orbital parameters + generate_id( + mass_a=mass_A, + mass_b=mass_B, + dimensionless_spin_a=spin_A, + dimensionless_spin_b=spin_B, + # Orbital parameters + separation=separation, + orbital_angular_velocity=fit_result["updated xcts values"]["omega"], + radial_expansion_velocity=fit_result["updated xcts values"][ + "expansion" + ], + # Scheduling options + control=id_params["control"], + refinement_level=id_params["control_refinement_level"], + polynomial_order=id_params["control_polynomial_order"], + evolve=True, + eccentricity_control=True, + pipeline_dir=pipeline_dir, + **scheduler_kwargs, + ) diff --git a/support/Pipelines/Bbh/InitialData.py b/support/Pipelines/Bbh/InitialData.py index fd218eabae7d..51bc784634e9 100644 --- a/support/Pipelines/Bbh/InitialData.py +++ b/support/Pipelines/Bbh/InitialData.py @@ -12,6 +12,7 @@ from spectre.Pipelines.EccentricityControl.InitialOrbitalParameters import ( initial_orbital_parameters, ) +from spectre.support.DirectoryStructure import list_pipeline_steps from spectre.support.Schedule import schedule, scheduler_options logger = logging.getLogger(__name__) @@ -125,6 +126,7 @@ def generate_id( id_input_file_template: Union[str, Path] = ID_INPUT_FILE_TEMPLATE, control: bool = False, evolve: bool = False, + eccentricity_control: bool = False, pipeline_dir: Optional[Union[str, Path]] = None, run_dir: Optional[Union[str, Path]] = None, segments_dir: Optional[Union[str, Path]] = None, @@ -165,11 +167,15 @@ def generate_id( values. If set to False, the horizon masses and spins in the generated data will differ from the input parameters. (default: False) evolve: Set to True to evolve the initial data after generation. + eccentricity_control: If set to True, an eccentricity reduction script is + run on the initial data to correct the initial orbital parameters. pipeline_dir: Directory where steps in the pipeline are created. Required when 'evolve' is set to True. The initial data will be created in a subdirectory '001_InitialData'. run_dir: Directory where the initial data is generated. Mutually exclusive with 'pipeline_dir'. + segments_dir: Directory where the evolution data is generated. Mutually + exclusive with 'pipeline_dir' and 'run_dir'. out_file_name: Optional. Name of the log file. (Default: "spectre.out") """ logger.warning( @@ -196,8 +202,16 @@ def generate_id( " evolution, etc will be created in the 'pipeline_dir'" " automatically." ) + # If there is a pipeline directory, set run directory as well if pipeline_dir and not run_dir: - run_dir = pipeline_dir / "001_InitialData" + pipeline_steps = list_pipeline_steps(pipeline_dir) + if pipeline_steps: # Check if the list is not empty + run_dir = pipeline_steps[-1].next(label="InitialData").path + else: + run_dir = PipelineStep.first( + directory=pipeline_dir, label="InitialData" + ).path + # If we run a control loop, then run initial data in a subdirectory if control: run_dir = f"{run_dir}/ControlParams_000" out_file_name = f"../{out_file_name}" @@ -225,6 +239,7 @@ def generate_id( **scheduler_kwargs, control=control, evolve=evolve, + eccentricity_control=eccentricity_control, pipeline_dir=pipeline_dir, run_dir=run_dir, segments_dir=segments_dir, diff --git a/support/Pipelines/Bbh/InitialData.yaml b/support/Pipelines/Bbh/InitialData.yaml index 2adfdf724d27..1bffaa4eaada 100644 --- a/support/Pipelines/Bbh/InitialData.yaml +++ b/support/Pipelines/Bbh/InitialData.yaml @@ -28,6 +28,7 @@ Next: center_of_mass: [0., 0., 0.] linear_momentum: [0., 0., 0.] evolve: {{ evolve }} + eccentricity_control: {{ eccentricity_control }} scheduler: {{ scheduler | default("None") }} copy_executable: {{ copy_executable | default("None") }} submit_script_template: {{ submit_script_template | default("None") }} diff --git a/support/Pipelines/Bbh/Inspiral.py b/support/Pipelines/Bbh/Inspiral.py index 12252c6da866..1cdf97797a4d 100644 --- a/support/Pipelines/Bbh/Inspiral.py +++ b/support/Pipelines/Bbh/Inspiral.py @@ -10,6 +10,7 @@ from rich.pretty import pretty_repr import spectre.IO.H5 as spectre_h5 +from spectre.support.DirectoryStructure import PipelineStep, list_pipeline_steps from spectre.support.Schedule import schedule, scheduler_options from spectre.Visualization.ReadH5 import to_dataframe @@ -372,7 +373,13 @@ def start_inspiral( " 'pipeline_dir' automatically." ) if pipeline_dir and not run_dir and not segments_dir: - segments_dir = pipeline_dir / "002_Inspiral" + pipeline_steps = list_pipeline_steps(pipeline_dir) + if pipeline_steps: # Check if the list is not empty + segments_dir = pipeline_steps[-1].next(label="Inspiral").path + else: + segments_dir = PipelineStep.first( + directory=pipeline_dir, label="Inspiral" + ).path # Determine resource allocation if ( diff --git a/support/Pipelines/Bbh/Inspiral.yaml b/support/Pipelines/Bbh/Inspiral.yaml index 55209a14adac..56d71b049c96 100644 --- a/support/Pipelines/Bbh/Inspiral.yaml +++ b/support/Pipelines/Bbh/Inspiral.yaml @@ -24,6 +24,11 @@ Next: subfile_name_ahb: ApparentHorizons/ControlSystemAhB_Centers.dat output: ecc_control.pdf id_input_file_path: {{id_input_file_path }} + pipeline_dir: {{ pipeline_dir }} + scheduler: {{ scheduler | default("None") }} + copy_executable: {{ copy_executable | default("None") }} + submit_script_template: {{ submit_script_template | default("None") }} + submit: True {% endif %} --- @@ -371,6 +376,15 @@ EventsAndTriggersAtSlabs: Value: 2.138 Events: - Completion + # If running eccentricity control, only run short inspiral +{% if eccentricity_control %} + - Trigger: + TimeCompares: + Comparison: GreaterThan + Value: 2000 + Events: + - Completion +{% endif %} EventsAndTriggersAtSteps: diff --git a/support/Pipelines/Bbh/PostprocessId.py b/support/Pipelines/Bbh/PostprocessId.py index 3f83ad8f77e1..94a45e2b2284 100644 --- a/support/Pipelines/Bbh/PostprocessId.py +++ b/support/Pipelines/Bbh/PostprocessId.py @@ -37,6 +37,7 @@ def postprocess_id( control_polynomial_order: int = 6, control_params: Dict[SupportedParams, Union[float, Sequence[float]]] = {}, evolve: bool = False, + eccentricity_control: bool = False, pipeline_dir: Optional[Union[str, Path]] = None, **scheduler_kwargs, ): @@ -153,7 +154,8 @@ def postprocess_id( start_inspiral( id_input_file_path, id_run_dir=id_run_dir, - continue_with_ringdown=True, + continue_with_ringdown=not eccentricity_control, + eccentricity_control=eccentricity_control, pipeline_dir=pipeline_dir, **scheduler_kwargs, ) diff --git a/support/Python/DirectoryStructure.py b/support/Python/DirectoryStructure.py index 446114ee6e9b..7f46c1ba4a5d 100644 --- a/support/Python/DirectoryStructure.py +++ b/support/Python/DirectoryStructure.py @@ -91,9 +91,9 @@ class Segment: NUM_DIGITS = 4 @classmethod - def first(cls, directory: Path) -> "Segment": + def first(cls, directory: Union[str, Path]) -> "Segment": name = "Segment_" + "0" * cls.NUM_DIGITS - return Segment(path=directory / name, id=0) + return Segment(path=Path(directory) / name, id=0) @property def next(self) -> "Segment": @@ -126,3 +126,76 @@ def list_segments(segments_dir: Union[str, Path]) -> List[Segment]: return [] matches = map(Segment.match, segments_dir.iterdir()) return sorted(match for match in matches if match) + + +@dataclass(frozen=True, order=True) +class PipelineStep: + """A step in a pipeline + + Each pipeline step is a numbered directory with a label, e.g. + "000_InitialData". + It can contain a simulation or a pre- or post-processing step + such as archiving. + Here is an example for the directory structure: + + ``` + BASE_DIR/ + 000_InitialData/ + InputFile.yaml + Submit.sh + ... + 001_Inspiral/ + Segment_0000/ + InputFile.yaml + Submit.sh + ... + Segment_0001/ + ... + 002_InitialData/ + ... + ... + ``` + + Note: "InitialData" and "Inspiral" are examples, any name can be used. + """ + + path: Path + id: int + label: str + + NAME_PATTERN = re.compile(r"(\d+)_(.+)") + NUM_DIGITS = 3 + + @classmethod + def first(cls, directory: Union[str, Path], label: str) -> "PipelineStep": + """Create the first directory in a sequence""" + name = "0" * cls.NUM_DIGITS + "_" + label + return PipelineStep(path=Path(directory) / name, id=0, label=label) + + def next(self, label: str) -> "PipelineStep": + """Get the next directory in the sequence""" + next_id = self.id + 1 + next_name = f"{str(next_id).zfill(self.NUM_DIGITS)}_{label}" + return PipelineStep( + path=self.path.resolve().parent / next_name, + id=next_id, + label=label, + ) + + @classmethod + def match(cls, path: Union[str, Path]) -> Optional["PipelineStep"]: + """Checks if the 'path' is a step in the pipeline""" + path = Path(path) + match = re.match(cls.NAME_PATTERN, path.resolve().name) + if not match: + return None + return cls(path=path, id=int(match.group(1)), label=match.group(2)) + + +def list_pipeline_steps(base_dir: Union[str, Path]) -> List[PipelineStep]: + """List all subdirectories in the base directory""" + base_dir = Path(base_dir) + if not base_dir.exists(): + return [] + matches = map(PipelineStep.match, base_dir.iterdir()) + return sorted(match for match in matches if match) diff --git a/tests/support/Pipelines/Bbh/Test_EccentricityControl.py b/tests/support/Pipelines/Bbh/Test_EccentricityControl.py index 52625f3a936d..a636f6f7c3c5 100644 --- a/tests/support/Pipelines/Bbh/Test_EccentricityControl.py +++ b/tests/support/Pipelines/Bbh/Test_EccentricityControl.py @@ -1,6 +1,5 @@ # Distributed under the MIT License. # See LICENSE.txt for details. -# Unit test for eccentricity control import logging import os @@ -78,6 +77,7 @@ def test_eccentricity_control(self): eccentricity_control( h5_file=self.h5_filename, id_input_file_path=self.id_input_file_path, + pipeline_dir=self.test_dir, tmin=0, tmax=10, ) diff --git a/tests/support/Python/Test_DirectoryStructure.py b/tests/support/Python/Test_DirectoryStructure.py index 21fecfea8f9e..47e509c11779 100644 --- a/tests/support/Python/Test_DirectoryStructure.py +++ b/tests/support/Python/Test_DirectoryStructure.py @@ -11,6 +11,7 @@ Checkpoint, Segment, list_checkpoints, + list_pipeline_steps, list_segments, ) from spectre.support.Logging import configure_logging @@ -64,6 +65,32 @@ def test_segments(self): list_segments(self.test_dir), [first_segment, segment, next_segment] ) + def test_pipeline_steps(self): + first_step = PipelineStep.first(self.test_dir, "InitialData") + self.assertEqual( + first_step, + PipelineStep( + path=self.test_dir / "000_InitialData", + id=0, + label="InitialData", + ), + ) + self.assertEqual(PipelineStep.match(first_step.path), first_step) + + next_step = first_step.next("Processing") + self.assertEqual(next_step.id, 1) + self.assertEqual(next_step.path.name, "001_Processing") + self.assertEqual( + next_step.path.resolve().parent, first_step.path.resolve().parent + ) + + self.assertEqual(list_pipeline_steps(self.test_dir), []) + first_step.path.mkdir() + next_step.path.mkdir() + self.assertEqual( + list_pipeline_steps(self.test_dir), [first_step, next_step] + ) + if __name__ == "__main__": configure_logging(log_level=logging.DEBUG)