Use PostNewtonian.jl for initial orbital parameters

Replaces the dependence on SpEC to compute low-eccentricity initial orbital parameters with the `sxs` package. It provides the PN approximations at higher PN order than SpEC, is much faster, and avoids spurious output from old Fortran code (LSODA) that was used in SpEC. The PN approximations are provided through the `PostNewtonian.jl` Julia package, so Julia will be downloaded on first use, which may take a few minutes.
sxs-collaboration · Aug 29, 2024 · 55e160d · 55e160d
1 parent 189bc8c
commit 55e160d
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 29 deletions.
diff --git a/support/Pipelines/EccentricityControl/InitialOrbitalParameters.py b/support/Pipelines/EccentricityControl/InitialOrbitalParameters.py
@@ -10,6 +10,43 @@
 logger = logging.getLogger(__name__)
 
 
+# The following two functions are modernized versions of those in SpEC's
+# `ZeroEccParamsFromPN.py`. They use higher PN orders (whichever are implemented
+# in the PostNewtonian module), are much faster, and avoid spurious output from
+# old Fortran code (LSODA) that was used in SpEC's `ZeroEccParamsFromPN.py`.
+# They are consistent with SpEC up to 2.5 PN order, as tested by Mike Boyle (see
+# https://github.com/moble/PostNewtonian.jl/issues/41).
+#
+# Since these functions use Julia through Python bindings, they will download
+# Julia and precompile the packages on first use, which may take a few minutes
+# (see https://moble.github.io/PostNewtonian.jl/dev/interface/python/).
+
+
+def omega_and_adot(r, q, chiA, chiB):
+    from sxs.julia import PostNewtonian
+
+    pn = PostNewtonian.BBH(
+        np.array(
+            [1.0 / (1.0 + q), q / (1.0 + q), *chiA, *chiB, 1, 0, 0, 0, 1, 0]
+        )
+    )
+    pn.state[12] = PostNewtonian.separation_inverse(r, pn)
+    return PostNewtonian.Omega(pn), PostNewtonian.separation_dot(pn) / r
+
+
+def num_orbits_and_time_to_merger(q, chiA0, chiB0, omega0):
+    from sxs.julia import PNWaveform
+
+    pn_waveform = PNWaveform(
+        M1=1.0 / (1.0 + q),
+        M2=q / (1.0 + q),
+        chi1=chiA0,
+        chi2=chiB0,
+        Omega_i=omega0,
+    )
+    return 0.5 * pn_waveform.orbital_phase[-1] / np.pi, pn_waveform.time[-1]
+
+
 def initial_orbital_parameters(
     mass_ratio: float,
     dimensionless_spin_a: Sequence[float],
@@ -76,8 +113,8 @@ def initial_orbital_parameters(
         )
         return separation, orbital_angular_velocity, radial_expansion_velocity
 
-    # The functions from SpEC currently work only for zero eccentricity. We will
-    # need to generalize this for eccentric orbits.
+    # The functions from the PostNewtonian module currently work only for zero
+    # eccentricity. We will need to generalize this for eccentric orbits.
     assert eccentricity == 0.0, (
         "Initial orbital parameters can currently only be computed for zero"
         " eccentricity."
@@ -97,25 +134,13 @@ def initial_orbital_parameters(
         " 'time_to_merger'."
     )
 
-    # Import functions from SpEC until we have ported them over. These functions
-    # call old Fortran code (LSODA) through scipy.integrate.odeint, which leads
-    # to lots of noise in stdout. When porting these functions, we should
-    # modernize them to use scipy.integrate.solve_ivp.
-    try:
-        from ZeroEccParamsFromPN import nOrbitsAndTotalTime, omegaAndAdot
-    except ImportError:
-        raise ImportError(
-            "Importing from SpEC failed. Make sure you have pointed "
-            "'-D SPEC_ROOT' to a SpEC installation when configuring the build "
-            "with CMake."
-        )
-
     # Find an omega0 that gives the right number of orbits or time to merger
     if num_orbits is not None or time_to_merger is not None:
+        logger.info("Finding orbital angular velocity...")
         opt_result = minimize(
             lambda x: (
                 abs(
-                    nOrbitsAndTotalTime(
+                    num_orbits_and_time_to_merger(
                         q=mass_ratio,
                         chiA0=dimensionless_spin_a,
                         chiB0=dimensionless_spin_b,
@@ -140,14 +165,14 @@ def initial_orbital_parameters(
 
     # Find the separation that gives the desired orbital angular velocity
     if orbital_angular_velocity is not None:
+        logger.info("Finding separation...")
         opt_result = minimize(
             lambda x: abs(
-                omegaAndAdot(
+                omega_and_adot(
                     r=x[0],
                     q=mass_ratio,
                     chiA=dimensionless_spin_a,
                     chiB=dimensionless_spin_b,
-                    rPrime0=1.0,  # Choice also made in SpEC
                 )[0]
                 - orbital_angular_velocity
             ),
@@ -163,12 +188,11 @@ def initial_orbital_parameters(
         logger.debug(f"Found initial separation: {separation}")
 
     # Find the radial expansion velocity
-    new_orbital_angular_velocity, radial_expansion_velocity = omegaAndAdot(
+    new_orbital_angular_velocity, radial_expansion_velocity = omega_and_adot(
         r=separation,
         q=mass_ratio,
         chiA=dimensionless_spin_a,
         chiB=dimensionless_spin_b,
-        rPrime0=1.0,  # Choice also made in SpEC
     )
     if orbital_angular_velocity is None:
         orbital_angular_velocity = new_orbital_angular_velocity
@@ -181,7 +205,7 @@ def initial_orbital_parameters(
         )
 
     # Estimate number of orbits and time to merger
-    num_orbits, time_to_merger = nOrbitsAndTotalTime(
+    num_orbits, time_to_merger = num_orbits_and_time_to_merger(
         q=mass_ratio,
         chiA0=dimensionless_spin_a,
         chiB0=dimensionless_spin_b,

diff --git a/tests/support/Pipelines/EccentricityControl/CMakeLists.txt b/tests/support/Pipelines/EccentricityControl/CMakeLists.txt
@@ -9,7 +9,9 @@ spectre_add_python_bindings_test(
   None
   TIMEOUT 60)
 
-if (SpEC_FOUND)
+# Disable this test if using jemalloc as a shared library, because Julia
+# doesn't like the 'LD_PRELOAD' trick to load jemalloc for Pybindings.
+if (NOT "${JEMALLOC_LIB_TYPE}" STREQUAL SHARED)
   spectre_add_python_bindings_test(
     "support.Pipelines.EccentricityControl.InitialOrbitalParameters"
     Test_InitialOrbitalParameters.py

diff --git a/tests/support/Pipelines/EccentricityControl/Test_InitialOrbitalParameters.py b/tests/support/Pipelines/EccentricityControl/Test_InitialOrbitalParameters.py
@@ -4,17 +4,18 @@
 import logging
 import unittest
 
+import numpy as np
 import numpy.testing as npt
 
 from spectre.Pipelines.EccentricityControl.InitialOrbitalParameters import (
     initial_orbital_parameters,
 )
-from support.Python.Logging import configure_logging
+from spectre.support.Logging import configure_logging
 
 
 class TestInitialOrbitalParameters(unittest.TestCase):
     def test_initial_orbital_parameters(self):
-        # Expected results are computed from SpEC's ZeroEccParamsFromPN.py
+        np.set_printoptions(precision=14)
         npt.assert_allclose(
             initial_orbital_parameters(
                 mass_ratio=1.0,
@@ -34,7 +35,7 @@ def test_initial_orbital_parameters(self):
                 eccentricity=0.0,
                 separation=16.0,
             ),
-            [16.0, 0.014474280975952748, -4.117670632867514e-05],
+            [16.0, 0.014454484323416913, -4.236562633362394e-05],
         )
         npt.assert_allclose(
             initial_orbital_parameters(
@@ -44,7 +45,7 @@ def test_initial_orbital_parameters(self):
                 eccentricity=0.0,
                 orbital_angular_velocity=0.015,
             ),
-            [15.6060791015625, 0.015, -4.541705362753467e-05],
+            [15.59033203125, 0.015, -4.696365029012517e-05],
         )
         npt.assert_allclose(
             initial_orbital_parameters(
@@ -54,7 +55,7 @@ def test_initial_orbital_parameters(self):
                 eccentricity=0.0,
                 orbital_angular_velocity=0.015,
             ),
-            [15.6060791015625, 0.015, -4.541705362753467e-05],
+            [15.59033203125, 0.015, -4.696365029012517e-05],
         )
         npt.assert_allclose(
             initial_orbital_parameters(
@@ -64,7 +65,7 @@ def test_initial_orbital_parameters(self):
                 eccentricity=0.0,
                 num_orbits=20,
             ),
-            [16.0421142578125, 0.014419921875000002, -4.0753460821644916e-05],
+            [15.71142578125, 0.014835205078125004, -4.554164727449197e-05],
         )
         npt.assert_allclose(
             initial_orbital_parameters(
@@ -74,7 +75,7 @@ def test_initial_orbital_parameters(self):
                 eccentricity=0.0,
                 time_to_merger=6000,
             ),
-            [16.1357421875, 0.01430025219917298, -3.9831982447244026e-05],
+            [16.0909423828125, 0.01433787536621094, -4.14229775202535e-05],
         )