Bind PN BinaryTrajectories in Python

tests/Unit/Helpers/PointwiseFunctions/PostNewtonian/CMakeLists.txt

@@ -25,3 +25,5 @@ target_link_libraries(
   PUBLIC
   DataStructures
   )
+
+add_subdirectory(Python)

tests/Unit/Helpers/PointwiseFunctions/PostNewtonian/Python/Bindings.cpp

@@ -0,0 +1,47 @@
+// Distributed under the MIT License.
+// See LICENSE.txt for details.
+
+#include <array>
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <string>
+
+#include "DataStructures/DataVector.hpp"
+#include "DataStructures/Tensor/Tensor.hpp"
+#include "Helpers/PointwiseFunctions/PostNewtonian/BinaryTrajectories.hpp"
+#include "Utilities/ErrorHandling/SegfaultHandler.hpp"
+
+namespace py = pybind11;
+
+PYBIND11_MODULE(_Pybindings, m) {  // NOLINT
+  enable_segfault_handler();
+  py::module::import("spectre.DataStructures");
+  py::module::import("spectre.DataStructures.Tensor");
+  py::class_<BinaryTrajectories>(m, "BinaryTrajectories")
+      .def(py::init<double, const std::array<double, 3>&, bool>(),
+           py::arg("initial_separation"),
+           py::arg("center_of_mass_velocity") =
+               std::array<double, 3>{{0.0, 0.0, 0.0}},
+           py::arg("newtonian") = false)
+      .def("separation", &BinaryTrajectories::separation<double>,
+           py::arg("time"))
+      .def("separation", &BinaryTrajectories::separation<DataVector>,
+           py::arg("time"))
+      .def("orbital_frequency", &BinaryTrajectories::orbital_frequency<double>,
+           py::arg("time"))
+      .def("orbital_frequency",
+           &BinaryTrajectories::orbital_frequency<DataVector>, py::arg("time"))
+      .def("angular_velocity", &BinaryTrajectories::angular_velocity<double>,
+           py::arg("time"))
+      .def("angular_velocity",
+           &BinaryTrajectories::angular_velocity<DataVector>, py::arg("time"))
+      .def("positions", &BinaryTrajectories::positions<double>, py::arg("time"))
+      .def("positions", &BinaryTrajectories::positions<DataVector>,
+           py::arg("time"))
+      .def("positions_no_expansion",
+           &BinaryTrajectories::positions_no_expansion<double>, py::arg("time"))
+      .def("positions_no_expansion",
+           &BinaryTrajectories::positions_no_expansion<DataVector>,
+           py::arg("time"));
+}

tests/Unit/Helpers/PointwiseFunctions/PostNewtonian/Python/CMakeLists.txt

@@ -0,0 +1,21 @@
+# Distributed under the MIT License.
+# See LICENSE.txt for details.
+
+set(LIBRARY "PyPostNewtonianHelpers")
+
+spectre_python_add_module(
+  PostNewtonian
+  MODULE_PATH "testing"
+  LIBRARY_NAME ${LIBRARY}
+  SOURCES
+  Bindings.cpp
+  PYTHON_FILES
+  __init__.py
+  )
+
+spectre_python_link_libraries(
+  ${LIBRARY}
+  PRIVATE
+  PostNewtonianHelpers
+  pybind11::module
+  )

tests/Unit/Helpers/PointwiseFunctions/PostNewtonian/Python/__init__.py

@@ -0,0 +1,4 @@
+# Distributed under the MIT License.
+# See LICENSE.txt for details.
+
+from ._Pybindings import *

tests/Unit/Helpers/Tests/PointwiseFunctions/PostNewtonian/CMakeLists.txt

@@ -14,3 +14,10 @@ target_link_libraries(
   PRIVATE
   PostNewtonianHelpers
   )
+
+spectre_add_python_bindings_test(
+  "Unit.PostNewtonian.Python.BinaryTrajectories"
+  "Test_BinaryTrajectories.py"
+  "unit;Python"
+  PyPostNewtonianHelpers
+  )

tests/Unit/Helpers/Tests/PointwiseFunctions/PostNewtonian/Test_BinaryTrajectories.py

@@ -0,0 +1,35 @@
+# Distributed under the MIT License.
+# See LICENSE.txt for details.
+
+import unittest
+
+import numpy as np
+import numpy.testing as npt
+
+from spectre.testing.PostNewtonian import BinaryTrajectories
+
+
+class TestBinaryTrajectories(unittest.TestCase):
+    def test_binary_trajectories(self):
+        # This class is tested in C++. Only check here that Pybindings work.
+        binary_trajectories = BinaryTrajectories(initial_separation=16)
+        self.assertEqual(binary_trajectories.separation(0.0), 16.0)
+        self.assertEqual(binary_trajectories.orbital_frequency(0.0), 0.015625)
+        self.assertEqual(binary_trajectories.angular_velocity(0.0), 0.015625)
+        npt.assert_allclose(
+            binary_trajectories.positions(0.0),
+            ([8.0, 0.0, 0.0], [-8.0, 0.0, 0.0]),
+        )
+        # Test vectorized interface
+        times = np.linspace(0.0, 10.0, 10)
+        npt.assert_allclose(
+            binary_trajectories.positions(times),
+            np.transpose(
+                [binary_trajectories.positions(t) for t in times],
+                axes=(1, 2, 0),
+            ),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

tests/Unit/Visualization/Python/Test_PlotTrajectories.py

@@ -13,6 +13,7 @@
 import spectre.IO.H5 as spectre_h5
 from spectre.Informer import unit_test_build_path, unit_test_src_path
 from spectre.support.Logging import configure_logging
+from spectre.testing.PostNewtonian import BinaryTrajectories
 from spectre.Visualization.PlotTrajectories import plot_trajectories_command
 
 
@@ -32,90 +33,27 @@ def tearDown(self):
         shutil.rmtree(self.test_dir)
 
     def create_h5_file(self):
-        """Create the H5"""
-        logging.info(f"Creating HDF5 file: {self.h5_filename}")
-
-        # Generate mock-up inspiral data
-        nsamples = 100
-        dt = 0.02
-        x0 = 0.35
-        y0 = 0.35
-        z0 = 0
-        z1 = -9.0e-6
-        cosAmp = 7.43
-        sinAmp = 7.43
-        cosFreq = 0.0173
-        sinFreq = 0.0172
-
-        # Define the spirals as functions
-        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 samples
-        tTable = np.arange(0, (nsamples + 1) * dt, dt)
-
-        # Map time to spiral points
-        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])
-
+        binary_trajectories = BinaryTrajectories(initial_separation=16)
+        times = np.linspace(0, 200, 100)
+        positions = np.array(binary_trajectories.positions(times))
         with spectre_h5.H5File(self.h5_filename, "w") as h5_file:
-            # Insert dataset for AhA
-            dataset_AhA = h5_file.insert_dat(
-                "ApparentHorizons/ControlSystemAhA_Centers.dat",
-                legend=[
-                    "Time",
-                    "GridCenter_x",
-                    "GridCenter_y",
-                    "GridCenter_z",
-                    "InertialCenter_x",
-                    "InertialCenter_y",
-                    "InertialCenter_z",
-                ],  # Legend for the dataset
-                version=0,  # Version number
-            )
-            # Populate dataset with AhA data
-            for data_point in AhA_data:
-                dataset_AhA.append(data_point)
-            # Close dataset for AhA
-            h5_file.close_current_object()
-
-            # Insert dataset for AhB
-            dataset_AhB = h5_file.insert_dat(
-                "ApparentHorizons/ControlSystemAhB_Centers.dat",
-                legend=[
-                    "Time",
-                    "GridCenter_x",
-                    "GridCenter_y",
-                    "GridCenter_z",
-                    "InertialCenter_x",
-                    "InertialCenter_y",
-                    "InertialCenter_z",
-                ],  # Legend for the dataset
-                version=0,  # Version number
-            )
-            # Populate dataset with AhB data
-            for data_point in AhB_data:
-                dataset_AhB.append(data_point)
-            # Close dataset for AhB
-            h5_file.close_current_object()
-        logging.info(
-            f"Successfully created and populated HDF5 file: {self.h5_filename}"
-        )
+            for i, ab in enumerate("AB"):
+                dataset = h5_file.insert_dat(
+                    f"ApparentHorizons/ControlSystemAh{ab}_Centers.dat",
+                    legend=[
+                        "Time",
+                        "GridCenter_x",
+                        "GridCenter_y",
+                        "GridCenter_z",
+                        "InertialCenter_x",
+                        "InertialCenter_y",
+                        "InertialCenter_z",
+                    ],
+                    version=0,
+                )
+                for t, coords in zip(times, positions[i].T):
+                    dataset.append([t, *coords, *coords])
+                h5_file.close_current_object()
 
     def test_cli(self):
         output_filename = os.path.join(self.test_dir, "output.pdf")

tests/support/Pipelines/Bbh/Test_EccentricityControl.py

@@ -14,6 +14,7 @@
 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
+from spectre.testing.PostNewtonian import BinaryTrajectories
 
 
 class TestEccentricityControl(unittest.TestCase):
@@ -38,82 +39,27 @@ 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
+        binary_trajectories = BinaryTrajectories(initial_separation=16)
+        times = np.arange(0, 1500, 1.0)
+        positions = np.array(binary_trajectories.positions(times))
         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}"
-        )
+            for i, ab in enumerate("AB"):
+                dataset = h5_file.insert_dat(
+                    f"ApparentHorizons/ControlSystemAh{ab}_Centers.dat",
+                    legend=[
+                        "Time",
+                        "GridCenter_x",
+                        "GridCenter_y",
+                        "GridCenter_z",
+                        "InertialCenter_x",
+                        "InertialCenter_y",
+                        "InertialCenter_z",
+                    ],
+                    version=0,
+                )
+                for t, coords in zip(times, positions[i].T):
+                    dataset.append([t, *coords, *coords])
+                h5_file.close_current_object()
 
     def create_yaml_file(self):
         # Define YAML data and write it to the file