Merge pull request #6295 from vtommasini/singleecccontrol

Single iteration of eccentricity control reduction

support/Pipelines/Bbh/CMakeLists.txt

@@ -7,6 +7,7 @@ spectre_python_add_module(
   PYTHON_FILES
   __init__.py
   ControlId.py
+  EccentricityControl.py
   FindHorizon.py
   InitialData.py
   InitialData.yaml

support/Pipelines/Bbh/EccentricityControl.py

@@ -0,0 +1,119 @@
+# Distributed under the MIT License.
+# See LICENSE.txt for details.
+
+import logging
+import warnings
+
+import matplotlib.pyplot as plt
+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"
+)
+
+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
+):
+    """Eccentricity reduction post inspiral.
+
+    This function can be called after the inspiral has run (see the 'Next'
+    section of the Inspiral.yam 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.
+
+    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
+    """
+    # 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
+
+    fit_result = ecout["fit result"]
+
+    # Prepare data from fit result
+    data = {
+        "Attribute": [
+            "Eccentricity",
+            "Omega Update",
+            "Expansion Update",
+            "Previous Omega",
+            "Previous Expansion",
+            "Updated Omega",
+            "Updated Expansion",
+        ],
+        "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"],
+        ],
+    }
+
+    # 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)

support/Pipelines/Bbh/Inspiral.py

@@ -290,6 +290,7 @@ def start_inspiral(
     ] = INSPIRAL_INPUT_FILE_TEMPLATE,
     id_horizons_path: Optional[Union[str, Path]] = None,
     continue_with_ringdown: 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,
@@ -314,6 +315,10 @@ def start_inspiral(
         "The BBH pipeline is still experimental. Please review the"
         " generated input files."
     )
+    assert not (continue_with_ringdown and eccentricity_control), (
+        "Cannot enable both 'continue_with_ringdown' and"
+        " 'eccentricity_control'. Choose which of the two to perform next."
+    )
 
     # Determine inspiral parameters from initial data
     if id_run_dir is None:
@@ -374,6 +379,8 @@ def start_inspiral(
         **inspiral_params,
         **scheduler_kwargs,
         continue_with_ringdown=continue_with_ringdown,
+        eccentricity_control=eccentricity_control,
+        id_input_file_path=Path(id_input_file_path).resolve(),
         pipeline_dir=pipeline_dir,
         run_dir=run_dir,
         segments_dir=segments_dir,
@@ -462,6 +469,14 @@ def start_inspiral(
         " formed."
     ),
 )
+@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",

support/Pipelines/Bbh/Inspiral.yaml

@@ -15,6 +15,15 @@ Next:
     copy_executable: {{ copy_executable | default("None") }}
     submit_script_template: {{ submit_script_template | default("None") }}
     submit: True
+{% elif eccentricity_control %}
+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
+    id_input_file_path: {{id_input_file_path }}
 {% endif %}
 
 ---

tests/support/Pipelines/Bbh/CMakeLists.txt

@@ -1,6 +1,13 @@
 # Distributed under the MIT License.
 # See LICENSE.txt for details.
 
+spectre_add_python_bindings_test(
+  "support.Pipelines.Bbh.EccentricityControl"
+  Test_EccentricityControl.py
+  "Python"
+  None
+  TIMEOUT 20)
+
 spectre_add_python_bindings_test(
   "support.Pipelines.Bbh.FindHorizon"
   Test_FindHorizon.py

tests/support/Pipelines/Bbh/Test_EccentricityControl.py

@@ -0,0 +1,142 @@
+# Distributed under the MIT License.
+# See LICENSE.txt for details.
+# Unit test for eccentricity control
+
+import logging
+import os
+import shutil
+import unittest
+
+import numpy as np
+import yaml
+
+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.support.Logging import configure_logging
+
+
+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"
+        )
+        self.h5_filename = os.path.join(
+            self.test_dir, "TestEccentricityControlData.h5"
+        )
+        self.id_input_file_path = os.path.join(self.test_dir, "Inspiral.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)
+        # Create HDF5 and YAML files for the test
+        self.create_h5_file()
+        self.create_yaml_file()
+
+    # Clean up and remove test directory after tests are done
+    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}"
+        )
+
+    def create_yaml_file(self):
+        # Define YAML data and write it to the file
+        data1 = {
+            "Background": {
+                "Binary": {"AngularVelocity": 0.01, "Expansion": 0.001}
+            }
+        }
+        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)
+
+    # Test the eccentricity control function with the created files
+    def test_eccentricity_control(self):
+        eccentricity_control(
+            h5_file=self.h5_filename,
+            id_input_file_path=self.id_input_file_path,
+            tmin=0,
+            tmax=10,
+        )
+
+
+if __name__ == "__main__":
+    configure_logging(log_level=logging.DEBUG)
+    unittest.main(verbosity=2)