Automate multiple iterations of eccentricity reduction

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,
+    )

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,

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") }}

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 (

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:
 

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,
         )

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)

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,
         )