Use PostNewtonian.jl for initial orbital parameters

.github/workflows/Tests.yaml

@@ -574,6 +574,13 @@ jobs:
           python3 -m pip install -r support/Python/requirements.txt \
             -r support/Python/dev_requirements.txt
           python3 -m pip list -v
+      - name: Precompile Julia packages
+        # This happens the first time the `sxs` package is imported. Do it here
+        # to avoid the delay when running the tests. The CI container should
+        # already have the Julia packages precompiled, so this step should only
+        # take a few seconds.
+        run: |
+          python3 -c "from sxs import julia"
       - name: Install ParaView
         if: matrix.test_3d_rendering == 'ON'
         working-directory: /work
@@ -786,13 +793,14 @@ ${{ matrix.build_type }}-pch-${{ matrix.use_pch || 'ON' }}"
           rm -r ./*
       - name: Test formaline tar can be built
         # - We only run the formaline tests in debug mode to reduce total build
-        #   time in CI.
+        #   time in CI. We don't run them with ASAN because then we run out of
+        #   disk space.
         # - We do run for all compilers, though, because formaline injects data
         #   at the linking stage, which means we are somewhat tied to the
         #   compiler version.
         # - We make sure to use the same compiler flags as the full build above
         #   so ccache is able to speed up the build.
-        if: matrix.build_type == 'Debug'
+        if: matrix.build_type == 'Debug' && matrix.ASAN != 'ON'
         working-directory: build
         run: >
           make EvolveBurgers -j${NUMBER_OF_CORES}
@@ -891,7 +899,8 @@ ${{ matrix.build_type }}-pch-${{ matrix.use_pch || 'ON' }}"
       # We install some low-level dependencies with Homebrew. They get picked up
       # by `spack external find`.
       SPECTRE_BREW_DEPS: >-  # Line breaks are spaces, no trailing newline
-        autoconf automake boost catch2 ccache cmake gsl hdf5 openblas yaml-cpp
+        autoconf automake boost catch2 ccache cmake fftw gsl hdf5 openblas
+        yaml-cpp
       # We install these packages with Spack and cache them. The full specs are
       # listed below. This list is only needed to create the cache.
       SPECTRE_SPACK_DEPS: blaze charmpp libxsmm
@@ -988,6 +997,11 @@ ${{ matrix.build_type }}-pch-${{ matrix.use_pch || 'ON' }}"
           source $HOME/spack/share/spack/setup-env.sh
           spack env activate spectre
           pip install -r support/Python/requirements.txt
+      - name: Precompile Julia packages
+        # This happens the first time the `sxs` package is imported. Do it here
+        # to avoid the delay when running the tests.
+        run: |
+          python -c "from sxs import julia"
       # Replace the ccache directory that building the dependencies may have
       # generated with the cached ccache directory.
       - name: Clear ccache from dependencies

containers/Dockerfile.buildenv

@@ -124,7 +124,7 @@ ARG TARGETARCH
 # Install add-apt-repository and basic tools
 RUN if [ ${UBUNTU_VERSION} = 18.04 ] && [ "$TARGETARCH" = "arm64" ]; then \
     echo "Cannot use Ubuntu 18.04 with ARM" && exit 1; fi && apt-get update -y \
-    && apt-get install -y software-properties-common wget git file \
+    && apt-get install -y software-properties-common curl wget git file \
     && if [ ${UBUNTU_VERSION} = 18.04 ]; then \
     add-apt-repository ppa:ubuntu-toolchain-r/test; fi
 
@@ -173,6 +173,7 @@ RUN apt-get update -y \
         libboost-thread-dev libboost-tools-dev libssl-dev \
         libhdf5-dev hdf5-tools \
         libarpack2-dev \
+        libfftw3-dev \
         libbenchmark-dev \
     && if [ ${UBUNTU_VERSION} = 18.04 ]; then \
     wget https://github.com/jbeder/yaml-cpp/archive/refs/tags/0.8.0.tar.gz \
@@ -196,17 +197,10 @@ RUN apt-get update -y \
     apt-get install -y libjemalloc2 libjemalloc-dev libyaml-cpp-dev; \
     fi
 
-# Install Python packages
-# We only install packages that are needed by the build system (e.g. to compile
-# Python bindings or build documentation) or used by Python code that is
-# unit-tested. Any other packages can be installed on-demand.
+# Install Python
 # - We use python-is-python3 because on Ubuntu 20.04 /usr/bin/python was removed
 #   to aid in tracking down anything that depends on python 2. However, many
 #   scripts use `/usr/bin/env python` to find python so restore it.
-# - We install h5py explicitly from binary so that cross compilation is quicker.
-COPY support/Python/requirements.txt requirements.txt
-COPY support/Python/dev_requirements.txt dev_requirements.txt
-ENV DEBIAN_FRONTEND noninteractive
 RUN apt-get update -y \
     && if [ ${UBUNTU_VERSION} = 18.04 ]; then \
     apt-get install -y zlib1g-dev libncurses5-dev libgdbm-dev libnss3-dev \
@@ -215,18 +209,43 @@ RUN apt-get update -y \
     && wget https://www.python.org/ftp/python/3.10.1/Python-3.10.1.tgz \
     && tar -xf Python-3.10.1.tgz && cd ./Python-3.10.1 \
     && ./configure --enable-optimizations && make ${PARALLEL_MAKE_ARG} \
-    && make altinstall && cd ../ \
-    && rm -rf ./Python-3.10.1.tgz ./Python-3.10.1 \
-    && python3.10 -m pip install --upgrade pip \
-    && pip3.10 --no-cache-dir install --only-binary=h5py -r requirements.txt \
-    -r dev_requirements.txt; \
+    && make install && cd ../ \
+    && rm -rf ./Python-3.10.1.tgz ./Python-3.10.1; \
     else \
-    apt-get install -y python3-pip python-is-python3 pkg-config \
-    && pip3 --no-cache-dir install --only-binary=h5py -r requirements.txt \
-    -r dev_requirements.txt; \
-    fi \
+    apt-get install -y python3-pip python-is-python3 pkg-config; \
+    fi
+
+# Install Python packages
+# We only install packages that are needed by the build system (e.g. to compile
+# Python bindings or build documentation) or used by Python code that is
+# unit-tested. Any other packages can be installed on-demand.
+# - Install h5py explicitly from binary so that cross compilation is quicker.
+# - Constrain numpy version to make sure it is binary compatible with h5py.
+COPY support/Python/requirements.txt requirements.txt
+COPY support/Python/dev_requirements.txt dev_requirements.txt
+ENV DEBIAN_FRONTEND noninteractive
+RUN python3 -m pip install --upgrade pip \
+    && pip3 --no-cache-dir install --only-binary=h5py \
+      -r requirements.txt -r dev_requirements.txt "numpy<2.0" \
     && rm requirements.txt dev_requirements.txt
 
+# Install Julia for the SXS package
+# This is optional, as the SXS package will download Julia on first use.
+# However, installing Julia explicitly gives more control over the installation.
+# - Set a consistent path for precompiled Julia packages so they are found when
+#   CI runs as a different user
+ENV JULIA_DEPOT_PATH "/usr/local/julia"
+RUN if [ ${UBUNTU_VERSION} = 22.04 ]; then \
+    curl -fsSL https://install.julialang.org | sh -s -- \
+        -y --add-to-path=false \
+    ; fi
+ENV PATH="$PATH:$JULIA_DEPOT_PATH/bin"
+# Call the SXS package so it precompiles Julia packages on first use, see:
+# https://moble.github.io/PostNewtonian.jl/dev/interface/python/
+RUN if [ ${UBUNTU_VERSION} = 22.04 ]; then \
+    python3 -c "from sxs import julia" \
+    ; fi
+
 # Enable bash-completion by installing it and then adding it to the .bashrc file
 RUN apt-get update -y \
     && apt-get install -y bash-completion \
@@ -497,9 +516,8 @@ ARG PARALLEL_MAKE_ARG=-j4
 
 # vim and emacs for editing files
 # Also ffmpeg for making movies with paraview output pngs
-# paraview needs curl
 RUN apt-get update -y \
-    && apt-get install -y vim emacs-nox ffmpeg curl
+    && apt-get install -y vim emacs-nox ffmpeg
 
 # Install headless paraview so we can run pvserver in the container
 # Note: there is no arm64 linux binary of paraview available, so don't

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,

support/Python/Logging.py

@@ -6,21 +6,21 @@
 import rich.logging
 
 
-def configure_logging(log_level=logging.INFO):
+def configure_logging(log_level: int):
     """
     Configure logging for our python scripts using the 'logging' module
 
     This is factored out into a free function so that any time we need to add
     module-specific logging configuration, we only have to add it to one place.
 
-    Module specific logging info:
-     - Disable 'matplotlib.font_manager' logging for 'logging.DEBUG' or higher
+    Logging verbosity of the spectre module is set to the 'log_level'.
+    For other modules we set the log level to INFO or above, so we don't get
+    debug output from all the modules we import.
     """
     logging.basicConfig(
-        level=logging.INFO if log_level is None else log_level,
+        level=max(log_level, logging.INFO),
         format="%(message)s",
         datefmt="[%X]",
         handlers=[rich.logging.RichHandler()],
     )
-    if log_level is not None and log_level >= logging.DEBUG:
-        logging.getLogger("matplotlib.font_manager").disabled = True
+    logging.getLogger("spectre").setLevel(log_level)

support/Python/__main__.py

@@ -267,7 +267,7 @@ def read_config_file(ctx, param, config_file):
     ),
 )
 def cli(log_level, build_dir, profile, output_profile):
-    configure_logging(log_level=log_level)
+    configure_logging(log_level=log_level or logging.INFO)
     # Format tracebacks with rich
     # - Suppress traceback entries from modules that we don't care about
     rich.traceback.install(

support/Python/requirements.txt

@@ -32,3 +32,7 @@ pyyaml
 # Rich: to format CLI output and tracebacks
 rich >= 12.0.0
 scipy
+# SXS package: to work with SXS data, waveforms, etc. Also to evaluate some
+# post-Newtonian expressions, e.g. for low-eccentricity initial orbital
+# parameters.
+sxs >= 2024.0.3

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

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