Tutorial: Fast self-forced trajectories

An essential component of the FastEMRIWaveforms (FEW) framework is the rapid evaluation of the inspiral trajectory of the secondary compact object. FEW provides the tools required for users to generate and visualise these trajectories. Users can also implement their own trajectory models to incorporate additional physics (such as environmental effects or modifications to gravity). We will demonstrate these capabilities in this tutorial.

ODE classes

Parameter conventions

FEW constructs trajectories by integrating the equations of motion of the compact binary; these are a system of ordinary differential equations (ODEs) that describe the evolution of the following parameters:

Parameter	Definition
\(p\)	(Dimensionless) semi-latus rectum
\(e\)	Eccentricity
\(x_I\)	Cosine of the inclination of the orbital plane
\(\Phi_\phi\)	Azimuthal GW phase
\(\Phi_\theta\)	Polar GW phase
\(\Phi_r\)	Radial GW phase

The ODE system can be integrated efficiently with adaptive Runge-Kutta techniques. These solvers obtain trajectories by integrating the right-hand side (RHS) of the ODE with respect to a time parameter.

Stock trajectory models

In order to handle the necessary information for evaluating the RHS of the ODE, FEW defines ODE classes. The following stock options are available:

Model	Features	Notes
KerrEqEccFlux	Adiabatic; equatorial eccentric inspiral with spinning primary	Our most up-to-date model
SchwarzEccFlux	Adiabatic; inspiral with non-spinning primary	An older model (currently outdated)
PN5	Post-Newtonian; eccentric inclined inspirals with spinning primary	Most complete model, but inaccurate for larger \(e\) and/or lower \(p\)

Obtaining the right-hand side of the trajectory ODE

As an example, let’s examine the SchwarzEccFlux model:

from few.trajectory.ode import SchwarzEccFlux

rhs = SchwarzEccFlux()

Before evaluating the RHS, we must first define the parameters of the system that remain fixed during an inspiral. At adiabatic order, these are the component masses (\(m_1\) and \(m_2\)) and the dimensionless spin parameter of the primary (\(a\)):

m1 = 1e6  # Solar masses
m2 = 1e1  # Solar masses
a = 0.0  # For a Schwarzschild inspiral, the spin parameter is zero

rhs.add_fixed_parameters(m1, m2, a)

We can then access the ODE derivatives:

p = 10.0
e = 0.3
x = 1.0  # Schwarzschild inspiral is equatorial by definition

pdot, edot, xIdot, Omega_phi, Omega_theta, Omega_r = rhs([p, e, x])
pdot, edot, xIdot, Omega_phi, Omega_theta, Omega_r

(-0.017771723696175704,
 -0.000779281500397855,
 0.0,
 0.028647063536752972,
 0.028647063536752972,
 0.018040932375307846)

FEW integrates trajectories on the radiation-reaction timescale \(t_\mathrm{rr} = t \nu\), where \(\nu = \mu / M\) is the mass ratio of the system (for \(\mu\) and \(M\) as the reduced and total mass respectively). The RHS of the ODE is therefore defined such that the \((\dot{p}, \dot{e}, \dot{x}_I)\) are scaled by \(\nu^{-1}\).

Checking the properties of an ODE system

The ODE class also defines a number of properties that describe the physical and systematic assumptions of the model:

from few.trajectory.ode.base import get_ode_properties

get_ode_properties(rhs)

{'convert_Y': False,
 'equatorial': True,
 'circular': False,
 'supports_ELQ': True,
 'background': 'Schwarzschild',
 'separatrix_buffer_dist': 0.1,
 'nparams': 6,
 'flux_output_convention': 'pex'}

Checking the model domain of validity

ODE models that interpolate over pre-computed data can only be evaluated over a finite domain. The bounds of this domain can be verified with class methods. For instance, to check the maximum eccentricity at \(p=10\) for this model, we write

rhs.max_e(10.)

0.755

Some models, such as KerrEccEqFlux, perform eccentricity tapering at low \(p\) as this pre-computing data is prohibitively expensive in this region. Examining it over a wide range of \(p\) for a single value of \(a\), we can observe the taper kicking in at \(\sim p=13\).

from few.trajectory.ode import KerrEccEqFlux
import matplotlib.pyplot as plt
import numpy as np

rhs_kerr = KerrEccEqFlux()

spin = 0.8

p_lim_check = np.linspace(4., 20, 101)
e_max = np.asarray([rhs_kerr.max_e(p_here, a=spin) for p_here in p_lim_check])
plt.plot(p_lim_check, e_max)
plt.xlabel("Semi-latus rectum")
plt.ylabel("Maximum eccentricity")

Text(0, 0.5, 'Maximum eccentricity')

Trajectory evaluation

The EMRIInspiral class

Trajectories are evaluated using the EMRIInspiral class, which interfaces with the underlying integrator classes and methods in order to integrate the trajectory from its initial conditions until either the maximum alloted duration elapses or pre-determined stopping conditions are satisfied.

The EMRIInspiral class offers a number of features that can be adjusted according to the desired output. We will demonstrate these below, working with the KerrEccEqFlux trajectory model as an example.

from few.trajectory.inspiral import EMRIInspiral

# You can also instantiate this as EMRIInspiral(func="KerrEccEqFlux") to save an import.
traj_model = EMRIInspiral(func=KerrEccEqFlux)

m1 = 1e6
m2 = 1e2
a = 0.9  # This model supports a spinning primary compact object
p0 = 10.0
e0 = 0.8
xI0 = 1.0  # +1 for prograde, -1 for retrograde inspirals

T = 4.0  # duration of trajectory in years (as defined by few.utils.constants.YRSID_SI)

traj_pars = [m1, m2, a, p0, e0, xI0]

import matplotlib.pyplot as plt
import numpy as np

t, p, e, xI, Phi_phi, Phi_theta, Phi_r = traj_model(*traj_pars, T=T)

fig, axes = plt.subplots(2, 3)
plt.subplots_adjust(wspace=0.35)
fig.set_size_inches(14, 8)
axes = axes.ravel()

ylabels = [r"$e$", r"$p$", r"$e$", r"$\Phi_\phi$", r"$\Phi_\theta$", r"$\Phi_r$"]
xlabels = [r"$p$", r"$t$", r"$t$", r"$t$", r"$t$", r"$t$"]
ys = [e, p, e, Phi_phi, Phi_theta, Phi_r]
xs = [p, t, t, t, t, t]

for i, (ax, x, y, xlab, ylab) in enumerate(zip(axes, xs, ys, xlabels, ylabels)):
    ax.plot(x, y)
    ax.set_xlabel(xlab, fontsize=16)
    ax.set_ylabel(ylab, fontsize=16)

Trajectory stopping conditions

In principle, the inspiral trajectory continues until \(p\) intersects with the separatrix, \(p_\mathrm{sep}(a, e, x_I)\). However, FEW truncates the inspiral at the buffer point \(p_\mathrm{stop} = p_\mathrm{sep} + \Delta p\) for numerical stability, performing a root-finding operation to place the last trajectory point within \(\sim 10^{-8}\) of \(p_\mathrm{stop}\).

The size of \(\Delta p\) can vary based on the trajectory model, and can be obtained directly from the ODE object:

from few.utils.geodesic import get_separatrix

Delta_p = KerrEccEqFlux().separatrix_buffer_dist
print("Delta_p:", Delta_p)
p_sep = get_separatrix(a, e[-1], xI[-1])  # the separatrix at the trajectory end-point.
p[-1] - (p_sep + Delta_p)

Delta_p: 0.002

8.388845174067683e-13

Output for a requested time duration

By using the keyword argument T, the trajectory to be computed until it reaches a desired time (in years).

import matplotlib.pyplot as plt
import numpy as np

T = 0.1  # duration of trajectory in years (as defined by few.utils.constants.YRSID_SI)

t_trunc, p_trunc, e_trunc, xI_trunc, Phi_phi_trunc, Phi_theta_trunc, Phi_r_trunc = traj_model(*traj_pars, T=T)

t[-1], t_trunc[-1]

(17366327.489808623, 3155814.97635456)

Output on a requested time grid

In addition to the sparse output of the adaptive solver, the user can also construct trajectories on pre-defined time grids. For uniform grids, the sampling cadence \(\mathrm{d}t\) can be supplied and the grid will be constructed automatically. Non-uniform grids can be supplied via the new_t keyword argument. In either case, the fix_t keyword argument can be used to truncate the returned trajectory at its end-point.

dt = 100.0
new_t = np.arange(0, 2e7, dt)

# a warning will be thrown if the new_t array goes beyond the time array output from the trajectory
t1, p1, e1, x1, Phi_phi1, Phi_theta1, Phi_r1 = traj_model(
    *traj_pars, T=T, new_t=new_t, upsample=True
)

# you can cut the excess on these arrays by setting fix_t to True
t2, p2, e2, x2, Phi_phi2, Phi_theta2, Phi_r2 = traj_model(
    *traj_pars, T=T, new_t=new_t, upsample=True, fix_t=True
)

plt.plot(t1, Phi_phi1)
plt.plot(t2, Phi_phi2, ls="--")

print("t1 max:", t1.max(), "t2 max:", t2.max())

t1 max: 19999900.0 t2 max: 4713600.0

Return trajectory in dimensionless time coordinates

Trajectories are returned in coordinate time by default, but may be obtained in dimensionless time units by passing is_coordinate_time=False:

t, p, e, x, Phi_phi, Phi_theta, Phi_r = traj_model(*traj_pars, in_coordinate_time=False)

plt.plot(t, Phi_phi, label=r"$\Phi_\phi$")
plt.plot(t, Phi_r, label=r"$\Phi_r$")
plt.legend(frameon=False)
plt.show()
print("Dimensionless time step:", t[1] - t[0])

Dimensionless time step: 454.85645146086296

Integrating trajectories with a fixed time-step

The adaptive stepping of the solver can be disabled with the DENSE_STEPPING keyword argument, which switches to taking uniform steps of size dt. This is more expensive and memory-intensive than adaptive stepping.

T_dense = 0.005

t_dense, p, e, x, Phi_phi, Phi_theta, Phi_r = traj_model(
    *traj_pars, T=T_dense, DENSE_STEPPING=1
)
t_adaptive, p, e, x, Phi_phi, Phi_theta, Phi_r = traj_model(*traj_pars, T=T_dense)
t_dense.shape, t_adaptive.shape

((15780,), (6,))

We can directly verify the accuracy of the adaptive solution by evaluating its output on the densely-stepped trajectory’s time grid. The phase errors will be larger than the error in the orbital elements by a factor of \(M / \mu\) due to conventions chosen during integration:

t_dense, p, e, x, Phi_phi, Phi_theta, Phi_r = traj_model(
    *traj_pars, T=T_dense, DENSE_STEPPING=1
)
t_adaptive, p_adaptive, e, x, Phi_phi_adaptive, Phi_theta, Phi_r = traj_model(
    *traj_pars, T=T_dense, new_t=t_dense, upsample=True
)

plt.semilogy(t_dense, abs(Phi_phi_adaptive - Phi_phi), label="Phase error")
plt.semilogy(t_dense, abs(p_adaptive - p), label="p error")
plt.legend(frameon=False)
plt.show()

We can also set a maximum time-step during adaptive integration:

t_adaptive = traj_model(*traj_pars, T=1.)[0]
t_adaptive_capped = traj_model(*traj_pars, T=1., max_step_size=86400)[0]
np.max(np.diff(t_adaptive)), np.max(np.diff(t_adaptive_capped))

(1869829.3568359115, 79645.72327024862)

Evolving trajectories backwards in time

Trajectories can also be evolved in reverse via the integrate_backwards keyword argument. Trajectories evolved in either direction will agree up to the numerical tolerance of the integrator:

traj_pars_back = traj_pars.copy()

T = 0.55

# forward
t, p, e, xI, Phi_phi, Phi_theta, Phi_r = traj_model(*traj_pars, T=T)

traj_pars_back[3] = p[-1]
traj_pars_back[4] = e[-1]
traj_pars_back[5] = xI[-1]

# backward
t_back, p_back, e_back, xI_back, Phi_phi_back, Phi_theta_back, Phi_r_back = traj_model(
    *traj_pars_back, T=T, integrate_backwards=True
)

plt.plot(p, e)
plt.plot(p_back, e_back, ls="--")
p_back[-1] - p[0]

-5.704414718366024e-11

Setting the integrator error tolerance

The integrator tolerance (which is an absolute tolerance) can be adjusted by setting the err keyword argument. Higher error tolerances reduce the number of points in the trajectory, but at the cost of numerical accuracy. Note that err is a local error tolerance, not a global one, and must therefore be set conservatively with respect to the desired global error tolerance. The default setting of \(10^{-11}\) is a reasonable value and should not need to be changed for typical scenarios.

p = traj_model(*traj_pars, T=T, err=1e-11)[1]
p_hightol = traj_model(*traj_pars, T=T, err=1e-9)[1]

p.size, p_hightol.size, (p[-1] - p_hightol[-1]).item()

(53, 32, 2.7758617626716386e-09)

Integrating trajectories in \((E, L, Q)\) coordinates

Trajectories can also be solved in the constants of motion parameterisation \((E, L, Q)\) via the integrate_constants_of_motion keyword. This is supported for all stock trajectory models. Given input \((p_0, e_0, x_{I,0})\), the integrator converts these initial conditions to \((E, L, Q)\) and integrates these quantities alongside the inspiral phases. The output trajectory is converted back to \((p, e, x_I)\) coordinates if the keyword argument convert_pex=True (which it is by default).

Integrating in \((E, L, Q)\) can offer some advantages, especially near the separatrix (see Hughes 2024 for more information). Typically the trajectory also consists of fewer points than its \((p, e, x_I)\) counterpart, and is therefore faster to compute. However, it can also be numerically unstable for very small eccentricities, where small errors in the constants of motion can be amplified when transformed back to \((p, e, x_I)\). Usage of this method should be considered somewhat experimental, especially for eccentric and inclined (generic) inspirals.

m1 = 1e6
m2 = 1e2
a = 0.0
p0 = 10.0
e0 = 0.7
xI0 = 1.0
T = 1.0

traj_model_ELQ = EMRIInspiral(func=SchwarzEccFlux, integrate_constants_of_motion=True)
traj_model_pex = EMRIInspiral(func=SchwarzEccFlux, integrate_constants_of_motion=False)

fig, axes = plt.subplots(2, 3, figsize=(14, 8))
plt.subplots_adjust(wspace=0.35)
axes = axes.ravel()
ylabels = [r"$e$", r"$p$", r"$e$", r"$\Phi_\phi$", r"$\Phi_\theta$", r"$\Phi_r$"]
xlabels = [r"$p$", r"$t$", r"$t$", r"$t$", r"$t$", r"$t$"]

lens = []
phase_ends = []
for traj_model, label, ls in zip(
    [traj_model_ELQ, traj_model_pex], ["ELQ", "pex"], ["-", "--"]
):
    t, p, e, xI, Phi_phi, Phi_theta, Phi_r = traj_model(m1, m2, a, p0, e0, xI0, T=T)

    # Plot the results for comparison
    ys = [e, p, e, Phi_phi, Phi_theta, Phi_r]
    xs = [p, t, t, t, t, t]

    for i, (ax, x, y, xlab, ylab) in enumerate(zip(axes, xs, ys, xlabels, ylabels)):
        ax.plot(x, y, label=label, ls=ls)
        ax.set_xlabel(xlab, fontsize=16)
        ax.set_ylabel(ylab, fontsize=16)

    lens.append(len(t))
    phase_ends.append(Phi_phi[-1])
axes[0].legend(frameon=False)

plt.show()

print("(ELQ, pex) trajectory lengths:", lens)
print("(ELQ, pex) phase difference:", phase_ends[1] - phase_ends[0])

(ELQ, pex) trajectory lengths: [31, 44]
(ELQ, pex) phase difference: 3.871961234835908e-05

Trajectory-related utilities

Get \(p_0\) based on desired duration of trajectory

If you have a desired length of trajectory to analyze, this function will give you the value of \(p_0\) that corresponds to the proper evolution time.

from few.utils.utility import get_p_at_t

traj = EMRIInspiral(func=SchwarzEccFlux)

# Set initial parameters
m1 = 1e6
m2 = 5e1
a=0.0
e0 = 0.7
xI0=1.0

traj_args = [m1, m2, a, e0, xI0]
traj_kwargs = {}

t_out = 1.5
# Run trajectory
p0_new = get_p_at_t(
    traj,
    t_out,
    traj_args,
    #index_of_a=2, # specifies the position of the index as it appears in the argument of traj
    #index_of_p=3,
    #index_of_e=4,
    #index_of_x=5,
    traj_kwargs=traj_kwargs,
    xtol=1e-15,
    rtol=1e-15,
    bounds=None,
)

print(
    "p0_new = {} will create a waveform that is {} years long, given the other input parameters.".format(
        p0_new, t_out
    )
)

p0_new = 13.05941726854163 will create a waveform that is 1.5 years long, given the other input parameters.

Get \(m_2\) based on desired duration of trajectory

Similarly, if you have a desired length of trajectory to analyze, this function will give you the value of \(m_2\) that corresponds to the proper evolution time.

from few.utils.utility import get_m2_at_t

traj = EMRIInspiral(func=SchwarzEccFlux)

# Set initial parameters
m1 = 1e6
m2 = 5e1
a=0.0
e0 = 0.7
xI0=1.0

traj_args = [m1, 0.0, p0, e0, 1.0]
traj_kwargs = {}
index_of_m2 = 1

t_out = 1.5
# Run trajectory
m2_new = get_m2_at_t(
    traj,
    t_out,
    traj_args,
    index_of_m2 = index_of_m2,
    traj_kwargs=traj_kwargs,
    xtol=1e-15,
    rtol=1e-15,
    bounds=None,
)

print(
    "m2_new = {} will create a waveform that is {} years long, given the other input parameters.".format(
        m2_new, t_out
    )
)

m2_new = 9.810595038283548 will create a waveform that is 1.5 years long, given the other input parameters.

User-defined trajectory models

In addition to FEW’s built-in trajectory models, users can readily implement their own models in the same framework. They can achieve this by:

• Subclassing the ODEBase base class,

• Defining the evaluate_rhs method, and

• Updating any relevant properties of the model (by default, they are set for generic Kerr inspirals).

You do not need to handle any boundary checking in your modified class (e.g., ensuring quantities remain physical) as this is performed automatically by the base class.

Defining a new trajectory model

Let’s implement a Post-Newtonian trajectory in Schwarzschild eccentric as an example:

# Some elliptic functions for evaluating geodesic frequencies
from few.utils.elliptic import EllipK, EllipE, EllipPi

# base classes
from few.trajectory.ode.base import ODEBase


# for common interface with C/mathematica
def Power(x, n):
    return x**n


def Sqrt(x):
    return np.sqrt(x)


# this class defines the right-hand side of the ODE
# we define the method "evaluate_rhs" according to our derivatives
# we set the "equatorial" and "background" properties accordingly
class Schwarzschild_PN(ODEBase):
    @property
    def equatorial(self):
        return True

    @property
    def background(self):
        return "Schwarzschild"

    def evaluate_rhs(self, y):
        # guard against bad integration steps
        p, e, xI = y[:3]

        # perform elliptic calculations
        ellipE = EllipE(4 * e / (p - 6.0 + 2 * e))
        ellipK = EllipK(4 * e / (p - 6.0 + 2 * e))
        ellipPi1 = EllipPi(
            16 * e / (12.0 + 8 * e - 4 * e * e - 8 * p + p * p),
            4 * e / (p - 6.0 + 2 * e),
        )
        ellipPi2 = EllipPi(
            2 * e * (p - 4) / ((1.0 + e) * (p - 6.0 + 2 * e)), 4 * e / (p - 6.0 + 2 * e)
        )

        # Azimuthal frequency
        Omega_phi = (2 * Power(p, 1.5)) / (
            Sqrt(-4 * Power(e, 2) + Power(-2 + p, 2))
            * (
                8
                + (
                    (
                        -2
                        * ellipPi2
                        * (6 + 2 * e - p)
                        * (3 + Power(e, 2) - p)
                        * Power(p, 2)
                    )
                    / ((-1 + e) * Power(1 + e, 2))
                    - (ellipE * (-4 + p) * Power(p, 2) * (-6 + 2 * e + p))
                    / (-1 + Power(e, 2))
                    + (
                        ellipK
                        * Power(p, 2)
                        * (28 + 4 * Power(e, 2) - 12 * p + Power(p, 2))
                    )
                    / (-1 + Power(e, 2))
                    + (
                        4
                        * (-4 + p)
                        * p
                        * (2 * (1 + e) * ellipK + ellipPi2 * (-6 - 2 * e + p))
                    )
                    / (1 + e)
                    + 2
                    * Power(-4 + p, 2)
                    * (
                        ellipK * (-4 + p)
                        + (ellipPi1 * p * (-6 - 2 * e + p)) / (2 + 2 * e - p)
                    )
                )
                / (ellipK * Power(-4 + p, 2))
            )
        )

        # Post-Newtonian calculations
        yPN = pow(Omega_phi, 2.0 / 3.0)

        EdotPN = (
            (96 + 292 * Power(e, 2) + 37 * Power(e, 4))
            / (15.0 * Power(1 - Power(e, 2), 3.5))
            * pow(yPN, 5)
        )
        LdotPN = (
            (4 * (8 + 7 * Power(e, 2)))
            / (5.0 * Power(-1 + Power(e, 2), 2))
            * pow(yPN, 7.0 / 2.0)
        )

        # flux
        Edot = -EdotPN
        Ldot = -LdotPN

        # time derivatives
        pdot = (
            -2
            * (
                Edot
                * Sqrt((4 * Power(e, 2) - Power(-2 + p, 2)) / (3 + Power(e, 2) - p))
                * (3 + Power(e, 2) - p)
                * Power(p, 1.5)
                + Ldot * Power(-4 + p, 2) * Sqrt(-3 - Power(e, 2) + p)
            )
        ) / (4 * Power(e, 2) - Power(-6 + p, 2))

        edot = -(
            (
                Edot
                * Sqrt((4 * Power(e, 2) - Power(-2 + p, 2)) / (3 + Power(e, 2) - p))
                * Power(p, 1.5)
                * (
                    18
                    + 2 * Power(e, 4)
                    - 3 * Power(e, 2) * (-4 + p)
                    - 9 * p
                    + Power(p, 2)
                )
                + (-1 + Power(e, 2))
                * Ldot
                * Sqrt(-3 - Power(e, 2) + p)
                * (12 + 4 * Power(e, 2) - 8 * p + Power(p, 2))
            )
            / (e * (4 * Power(e, 2) - Power(-6 + p, 2)) * p)
        )

        xIdot = 0.0

        Phi_phi_dot = Omega_phi
        Phi_theta_dot = Omega_phi

        Phi_r_dot = (
            p * Sqrt((-6 + 2 * e + p) / (-4 * Power(e, 2) + Power(-2 + p, 2))) * np.pi
        ) / (
            8 * ellipK
            + (
                (-2 * ellipPi2 * (6 + 2 * e - p) * (3 + Power(e, 2) - p) * Power(p, 2))
                / ((-1 + e) * Power(1 + e, 2))
                - (ellipE * (-4 + p) * Power(p, 2) * (-6 + 2 * e + p))
                / (-1 + Power(e, 2))
                + (ellipK * Power(p, 2) * (28 + 4 * Power(e, 2) - 12 * p + Power(p, 2)))
                / (-1 + Power(e, 2))
                + (
                    4
                    * (-4 + p)
                    * p
                    * (2 * (1 + e) * ellipK + ellipPi2 * (-6 - 2 * e + p))
                )
                / (1 + e)
                + 2
                * Power(-4 + p, 2)
                * (
                    ellipK * (-4 + p)
                    + (ellipPi1 * p * (-6 - 2 * e + p)) / (2 + 2 * e - p)
                )
            )
            / Power(-4 + p, 2)
        )

        dydt = [pdot, edot, xIdot, Phi_phi_dot, Phi_theta_dot, Phi_r_dot]

        return dydt

m1 = 1e6
m2 = 1e2
p0 = 10.0
e0 = 0.7
T = 1.0

traj = EMRIInspiral(func=SchwarzEccFlux)
flux = traj(m1, m2, 0.0, p0, e0, 1.0, T=T, dt=10.0)
p,e = flux[1:3]

traj_PN = EMRIInspiral(func=Schwarzschild_PN)
flux_PN = traj_PN(m1, m2, a, p0, e0, xI0, T=T, dt=10.0)
p_PN, e_PN = flux_PN[1:3]

plt.plot(p, e, label="flux")
plt.plot(p_PN, e_PN, label="pn")
plt.ylabel("e")
plt.xlabel("p")


plt.legend()
plt.show()

Modifying an existing trajectory model

In cases where the desired trajectory model can be expressed as a function of one that is supported out-of-the-box, the inbuilt model can be subclassed and supplemented as required by defining the modify_rhs method. This method takes as input the current system state and the corresponding derivatives produced by the model, which can then be modified as required.

This method does not support a return argument for efficiency; the content of the first argument derivs should be modified in-place.

from few.trajectory.ode.flux import KerrEccEqFlux


class ModifiedKerrEccEqFlux(KerrEccEqFlux):
    def modify_rhs(self, ydot, y):
        # in-place modification of the derivatives
        ydot[0] *= 1 + self.additional_args[0]


m1 = 1e6
m2 = 1e2
a = 0.7
p = 10.0
e = 0.3
xI = 1.0

additional_args = [
    0.01,
]

# Example usage
# Example usage
original_ode = KerrEccEqFlux()
original_ode.add_fixed_parameters(m1, m2, a)

modified_ode = ModifiedKerrEccEqFlux()
modified_ode.add_fixed_parameters(m1, m2, a, additional_args)

original_ode([p, e, xI])[0], modified_ode([p, e, xI])[0]

(-0.011514135697943613, -0.011629277054923049)

Passing additional arguments to a modified trajectory model

When calling a trajectory, the additional_args attribute of the ODE class will be populated with any positional arguments following \(x_I\):

m1 = 1e6
m2 = 1e2
a = 0.7
p0 = 10.0
e0 = 0.3
xI0 = 1.0

traj_modified = EMRIInspiral(func=ModifiedKerrEccEqFlux)
flux_modified = traj_modified(m1, m2, a, p0, e0, xI0, 0.01, T=T, dt=10.0)

traj = EMRIInspiral(func=KerrEccEqFlux)
flux = traj(m1, m2, a, p0, e0, xI0, T=T, dt=10.0)

plt.plot(flux[1], flux[2], label="flux + mod")
plt.plot(flux_modified[1], flux_modified[2], label="flux", ls="--")
plt.ylabel("e")
plt.xlabel("p")

plt.legend()
plt.show()

The flux-based stock models perform interpolation to obtain the integral fluxes \(\dot{E}\), \(\dot{L}_{z}\), and then apply a Jacobian transformation in order to obtain \(\dot{p}\), \(\dot{e}\) if these are the variables of integration (and they are by default). In some cases, the user may wish to adjust \(\dot{E}\), \(\dot{L}_z\) directly if the modifications desired are more easily computed in this parameterisation. This can be achieved similarly to what was shown above, but with the modify_rhs_before_Jacobian method:

# A physically-motivated example: dissipation of energy and angular momentum due to dynamical friction with an accretion disk
# Modify the angular momentum derivative according to Eq(2.2) of https://dx.doi.org/10.1103/PhysRevX.13.021035
class KerrEccEqFluxAccretionDisk(KerrEccEqFlux):
    def modify_rhs(self, ydot, y):
        # in-place modification of the derivatives
        LdotAcc = (
            self.additional_args[0]
            * pow(y[0] / 10.0, self.additional_args[1])
            * 32.0
            / 5.0
            * pow(y[0], -7.0 / 2.0)
        )
        dL_dp = (
            -3 * pow(a, 3)
            + pow(a, 2) * (8 - 3 * y[0]) * np.sqrt(y[0])
            + (-6 + y[0]) * pow(y[0], 2.5)
            + 3 * a * y[0] * (-2 + 3 * y[0])
        ) / (2.0 * pow(2 * a + (-3 + y[0]) * np.sqrt(y[0]), 1.5) * pow(y[0], 1.75))
        # transform back to pdot from Ldot and add GW contribution
        ydot[0] = ydot[0] + LdotAcc / dL_dp


m1 = 1e6
m2 = 5e1
a = 0.9
p = 15.0
# assume circular orbits, for extension to eccentricity see https://arxiv.org/pdf/2411.03436
e = 0.0
x = 1.0
T = 4.0

traj = EMRIInspiral(func=KerrEccEqFlux, integrate_constants_of_motion=False)
flux = traj(m1, m2, a, p0, e0, xI0, T=T, dt=2.0)

traj_accretionDisk = EMRIInspiral(func=KerrEccEqFluxAccretionDisk, integrate_constants_of_motion=False)
flux_accretionDisk = traj_accretionDisk(m1, m2, a, p0, e0, xI0, 0.0, 8.0, T=T, dt=2.0)

plt.plot(flux[0] / 86400, flux[4], label="flux")
plt.plot(flux_accretionDisk[0] / 86400, flux_accretionDisk[4], label="flux + mod", ls="--")
plt.ylabel(r"$\Phi$")
plt.xlabel("t [days]")

plt.legend()
plt.show()

# some nice italian-style plotting
!pip install pastamarkers

Requirement already satisfied: pastamarkers in /Users/c.chapman-bird@bham.ac.uk/miniconda3/envs/few2.0rc1/lib/python3.12/site-packages (0.1.0)
Requirement already satisfied: matplotlib<4.0.0,>=3.4.3 in /Users/c.chapman-bird@bham.ac.uk/miniconda3/envs/few2.0rc1/lib/python3.12/site-packages (from pastamarkers) (3.10.1)
Requirement already satisfied: numpy<2.0.0,>=1.21.2 in /Users/c.chapman-bird@bham.ac.uk/miniconda3/envs/few2.0rc1/lib/python3.12/site-packages (from pastamarkers) (1.26.4)
Requirement already satisfied: contourpy>=1.0.1 in /Users/c.chapman-bird@bham.ac.uk/miniconda3/envs/few2.0rc1/lib/python3.12/site-packages (from matplotlib<4.0.0,>=3.4.3->pastamarkers) (1.3.1)
Requirement already satisfied: cycler>=0.10 in /Users/c.chapman-bird@bham.ac.uk/miniconda3/envs/few2.0rc1/lib/python3.12/site-packages (from matplotlib<4.0.0,>=3.4.3->pastamarkers) (0.12.1)
Requirement already satisfied: fonttools>=4.22.0 in /Users/c.chapman-bird@bham.ac.uk/miniconda3/envs/few2.0rc1/lib/python3.12/site-packages (from matplotlib<4.0.0,>=3.4.3->pastamarkers) (4.56.0)
Requirement already satisfied: kiwisolver>=1.3.1 in /Users/c.chapman-bird@bham.ac.uk/miniconda3/envs/few2.0rc1/lib/python3.12/site-packages (from matplotlib<4.0.0,>=3.4.3->pastamarkers) (1.4.8)
Requirement already satisfied: packaging>=20.0 in /Users/c.chapman-bird@bham.ac.uk/miniconda3/envs/few2.0rc1/lib/python3.12/site-packages (from matplotlib<4.0.0,>=3.4.3->pastamarkers) (24.2)
Requirement already satisfied: pillow>=8 in /Users/c.chapman-bird@bham.ac.uk/miniconda3/envs/few2.0rc1/lib/python3.12/site-packages (from matplotlib<4.0.0,>=3.4.3->pastamarkers) (11.1.0)
Requirement already satisfied: pyparsing>=2.3.1 in /Users/c.chapman-bird@bham.ac.uk/miniconda3/envs/few2.0rc1/lib/python3.12/site-packages (from matplotlib<4.0.0,>=3.4.3->pastamarkers) (3.2.1)
Requirement already satisfied: python-dateutil>=2.7 in /Users/c.chapman-bird@bham.ac.uk/miniconda3/envs/few2.0rc1/lib/python3.12/site-packages (from matplotlib<4.0.0,>=3.4.3->pastamarkers) (2.9.0.post0)
Requirement already satisfied: six>=1.5 in /Users/c.chapman-bird@bham.ac.uk/miniconda3/envs/few2.0rc1/lib/python3.12/site-packages (from python-dateutil>=2.7->matplotlib<4.0.0,>=3.4.3->pastamarkers) (1.17.0)

# Interpolate and compare the two trajectories
from scipy.interpolate import interp1d
from few.utils.utility import get_p_at_t
from few.utils.constants import YRSID_SI
from pastamarkers import markers
import time

traj_args = [m1, m2, a, e0, xI0]
traj_kwargs = {}
t_out = 1.5
results = {"orange": [], "blue": []}
for A, n, marker, col in zip(
    [0.4e-5, 0.1e-5],
    [59 / 10, 8.0],
    [markers.tortellini, markers.ravioli],
    ["orange", "blue"],
):
    for t_out in [2.0, 3.0, 4.0, 5.0, 6.0]:
        # Obtain the initial p value for the desired trajectory duration, see below for futher explanations
        p0_new = get_p_at_t(
            traj,
            t_out,
            traj_args,
            index_of_a=2,
            index_of_p=3,
            index_of_e=4,
            index_of_x=5,
            traj_kwargs=traj_kwargs,
            xtol=1e-12,
            rtol=1e-12,
        )
        print(
            "p0_new = {} will create a waveform that is {} years long, given the other input parameters.".format(
                p0_new, t_out
            )
        )

        print(A)

        T = 10.0
        # T = t_out
        tic = time.time()
        flux = traj(m1, m2, a, p0_new, e0, xI0, T=T, dt=10.0)
        toc = time.time()
        print("Time taken for trajectory:", toc - tic)
        tic = time.time()
        flux_accretionDisk = traj_accretionDisk(m1, m2, a, p0_new, e0, xI0, A, 8.0, T=T, dt=10.0)
        toc = time.time()
        print("Time taken for flux_accretionDisk:", toc - tic)

        interp_flux_accretionDisk = interp1d(flux_accretionDisk[0], flux_accretionDisk[4])
        interp_flux = interp1d(flux[0], flux[4])

        # t = np.arange(len(flux[0]))*dt
        if t_out == 4.0:
            label = f"A={A}, n={n}"
        else:
            label = None

        # use the shortest trajectory as the reference
        t_f = flux[0][-1]
        delta_phi = interp_flux(t_f) - interp_flux_accretionDisk(t_f)
        results[col].append((t_f / YRSID_SI, delta_phi, marker, label, col, toc - tic))

for A, n, marker, col in zip(
    [0.4e-5, 0.1e-5],
    [59 / 10, 8.0],
    [markers.tortellini, markers.ravioli],
    ["orange", "blue"],
):
    results[col] = np.array(results[col])

p0_new = 12.876302295341677 will create a waveform that is 2.0 years long, given the other input parameters.
4e-06
Time taken for trajectory: 0.05723714828491211
Time taken for flux_accretionDisk: 0.09399986267089844
p0_new = 14.288498630733283 will create a waveform that is 3.0 years long, given the other input parameters.
4e-06
Time taken for trajectory: 0.06962704658508301
Time taken for flux_accretionDisk: 0.08910298347473145
p0_new = 15.382138510647557 will create a waveform that is 4.0 years long, given the other input parameters.
4e-06
Time taken for trajectory: 0.07811617851257324
Time taken for flux_accretionDisk: 0.06440186500549316
p0_new = 16.286887612296823 will create a waveform that is 5.0 years long, given the other input parameters.
4e-06
Time taken for trajectory: 0.0824270248413086
Time taken for flux_accretionDisk: 0.09103512763977051
p0_new = 17.064873248338337 will create a waveform that is 6.0 years long, given the other input parameters.
4e-06
Time taken for trajectory: 0.08822298049926758
Time taken for flux_accretionDisk: 0.06623005867004395
p0_new = 12.876302295341677 will create a waveform that is 2.0 years long, given the other input parameters.
1e-06
Time taken for trajectory: 0.0756220817565918
Time taken for flux_accretionDisk: 0.06400394439697266
p0_new = 14.288498630733283 will create a waveform that is 3.0 years long, given the other input parameters.
1e-06
Time taken for trajectory: 0.05479598045349121
Time taken for flux_accretionDisk: 0.0631108283996582
p0_new = 15.382138510647557 will create a waveform that is 4.0 years long, given the other input parameters.
1e-06
Time taken for trajectory: 0.05370283126831055
Time taken for flux_accretionDisk: 0.06365203857421875
p0_new = 16.286887612296823 will create a waveform that is 5.0 years long, given the other input parameters.
1e-06
Time taken for trajectory: 0.05735492706298828
Time taken for flux_accretionDisk: 0.08278703689575195
p0_new = 17.064873248338337 will create a waveform that is 6.0 years long, given the other input parameters.
1e-06
Time taken for trajectory: 0.07090592384338379
Time taken for flux_accretionDisk: 0.11389684677124023

plt.title("Dephasing vs. time to plunge $t_p$ for different accretion disk parameters")
plt.scatter(
    results["orange"][:, 0],
    results["orange"][:, 1],
    color="orange",
    marker=markers.tortellini,
    label=results["orange"][:, 3][2],
    s=200,
    linewidth=0.2,
)
plt.scatter(
    results["blue"][:, 0],
    results["blue"][:, 1],
    color="blue",
    marker=markers.ravioli,
    label=results["blue"][:, 3][2],
    s=200,
    linewidth=0.2,
)
# plt.ylim(1e-3, 1e6)
plt.ylabel(r"$\Delta \Phi_{\Phi}$")
plt.xlabel("$t_p$ [years]")
plt.grid()
# import values from https://dx.doi.org/10.1103/PhysRevX.13.021035 fig 1
plt.scatter(
    4,
    1e2,
    marker=markers.farfalle,
    label="PhysRevX.13.021035",
    s=400,
    linewidth=0.2,
)
plt.legend()
plt.show()

The dephasing \(\Delta \Phi_{\Phi}\) obtained from the difference between the final orbital phase of an evolution with and without migration torques in \(\alpha-\) or \(\beta-\) accretion disks. For \(\alpha-\)disks, the slope is expected to be \(n=8\), for \(\beta-\) disks, \(n=59/10\). The obtained value in previous work for \(t_p=4 yr\) is shown in the Figure as a tortellino (see fig 1 of PhysRevX.13.021035)

plt.scatter(
    results["orange"][:, 0],
    results["orange"][:, -1],
    color="orange",
    marker=markers.tortellini,
    label=results["orange"][:, 3][2],
    s=200,
    linewidth=0.2,
)
plt.scatter(
    results["blue"][:, 0],
    results["blue"][:, -1],
    color="blue",
    marker=markers.ravioli,
    label=results["blue"][:, 3][2],
    s=200,
    linewidth=0.2,
)
# plt.ylim(1e-3, 1e6)
plt.ylabel("Timing [s]")
plt.xlabel("$t_p$ [years]")
plt.grid()
plt.legend()
plt.show()