diff --git a/support/Pipelines/Bbh/EccentricityControl.py b/support/Pipelines/Bbh/EccentricityControl.py index 9e503efb20ef..c8b4294e71ea 100644 --- a/support/Pipelines/Bbh/EccentricityControl.py +++ b/support/Pipelines/Bbh/EccentricityControl.py @@ -3,18 +3,21 @@ import logging import warnings +from pathlib import Path +from typing import Optional, Union import matplotlib.pyplot as plt import pandas as pd import yaml +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 +26,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 +73,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 +142,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..8cd45a8cc14f 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 PipelineStep, 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, @@ -361,6 +376,14 @@ def generate_id( "instead of a run directory (-o)." ), ) +@click.option( + "--eccentricity-control", + is_flag=True, + help=( + "Perform eccentricity reduction script that finds current eccentricity" + "and better guesses for the input orbital parameters." + ), +) @click.option( "--pipeline-dir", "-d", 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..1c8d0899e6d1 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 @@ -63,23 +62,50 @@ def create_h5_file(self): def create_yaml_file(self): # Define YAML data and write it to the file + metadata = { + "Next": { + "With": { + "control_params": { + "mass_A": 1.0, + "mass_B": 1.0, + "spin_A": [0.0, 0.0, 0.0], + "spin_B": [0.0, 0.0, 0.0], + }, + "refinement_level": 1, + "polynomial_order": 10, + } + } + } + data1 = { "Background": { - "Binary": {"AngularVelocity": 0.01, "Expansion": 0.001} + "Binary": { + "AngularVelocity": 0.01, + "Expansion": 0.001, + "XCoords": [10.0, -10.0], # Example values for XCoords + } } } + + # Pass both metadata and data1 to the YAML file with open(self.id_input_file_path, "w") as yaml_file: - # Keep first dictionary in this list empty to match - # the structure of the real file - yaml.dump_all([{}, data1], yaml_file) + yaml.dump_all([metadata, data1], yaml_file) # Test the eccentricity control function with the created files def test_eccentricity_control(self): + output_path = os.path.join(self.test_dir, "output.pdf") + # Call the function with updated parameters 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, + output=output_path, + ) + # Add checks if necessary + self.assertTrue( + os.path.exists(os.path.join(self.test_dir, "output.pdf")) ) diff --git a/tests/support/Pipelines/Bbh/Test_InitialData.py b/tests/support/Pipelines/Bbh/Test_InitialData.py index 3a72505e8129..6e94939f7883 100644 --- a/tests/support/Pipelines/Bbh/Test_InitialData.py +++ b/tests/support/Pipelines/Bbh/Test_InitialData.py @@ -134,6 +134,7 @@ def test_cli(self): "-d", str(self.test_dir / "Pipeline"), "--evolve", + "--eccentricity-control", "--no-submit", ] ) @@ -141,7 +142,7 @@ def test_cli(self): self.assertEqual(e.code, 0) with open( self.test_dir - / "Pipeline/001_InitialData/ControlParams_000/InitialData.yaml", + / "Pipeline/000_InitialData/ControlParams_000/InitialData.yaml", "r", ) as open_input_file: metadata = next(yaml.safe_load_all(open_input_file)) @@ -165,6 +166,7 @@ def test_cli(self): "linear_momentum": [0.0, 0.0, 0.0], }, "evolve": True, + "eccentricity_control": True, "scheduler": "None", "copy_executable": "None", "submit_script_template": "None", diff --git a/tests/support/Pipelines/Bbh/Test_Inspiral.py b/tests/support/Pipelines/Bbh/Test_Inspiral.py index 960db65f97dd..ff791780b1ef 100644 --- a/tests/support/Pipelines/Bbh/Test_Inspiral.py +++ b/tests/support/Pipelines/Bbh/Test_Inspiral.py @@ -159,7 +159,7 @@ def test_cli(self): except SystemExit as e: self.assertEqual(e.code, 0) with open( - self.test_dir / "Pipeline/002_Inspiral/Segment_0000/Inspiral.yaml", + self.test_dir / "Pipeline/000_Inspiral/Segment_0000/Inspiral.yaml", "r", ) as open_input_file: metadata = next(yaml.safe_load_all(open_input_file)) @@ -180,6 +180,50 @@ def test_cli(self): }, }, ) + # Test with eccentricity control + try: + start_inspiral_command( + common_args + + [ + "-d", + str(self.test_dir / "Pipeline"), + "--eccentricity-control", + "--no-submit", + ] + ) + except SystemExit as e: + self.assertEqual(e.code, 0) + with open( + self.test_dir / "Pipeline/001_Inspiral/Segment_0000/Inspiral.yaml", + "r", + ) as open_input_file: + metadata = next(yaml.safe_load_all(open_input_file)) + self.maxDiff = None + modulename = "spectre.Pipelines.Bbh.EccentricityControl" + self.assertEqual( + metadata["Next"], + { + "Run": modulename + ":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", + "id_input_file_path": str( + self.id_dir.resolve() / "InitialData.yaml" + ), + "pipeline_dir": str(self.test_dir.resolve() / "Pipeline"), + "scheduler": "None", + "copy_executable": "None", + "submit_script_template": "None", + "submit": True, + }, + }, + ) if __name__ == "__main__": diff --git a/tests/support/Python/Test_DirectoryStructure.py b/tests/support/Python/Test_DirectoryStructure.py index 21fecfea8f9e..e51f85900b47 100644 --- a/tests/support/Python/Test_DirectoryStructure.py +++ b/tests/support/Python/Test_DirectoryStructure.py @@ -9,8 +9,10 @@ from spectre.Informer import unit_test_build_path from spectre.support.DirectoryStructure import ( Checkpoint, + PipelineStep, Segment, list_checkpoints, + list_pipeline_steps, list_segments, ) from spectre.support.Logging import configure_logging @@ -64,6 +66,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)