import warnings
import numpy as np
from .utils import _Z, _flow_A, _flow_B, _flow_mixed
[docs]
class GeodesicIntegrator:
"""
Geodesic Integrator, based on [1]_.
This module uses Forward Mode Automatic Differentiation
to calculate metric derivatives to machine precision
leading to stable simulations.
References
----------
.. [1] Christian, Pierre and Chan, Chi-Kwan;
"FANTASY: User-Friendly Symplectic Geodesic Integrator
for Arbitrary Metrics with Automatic Differentiation";
`2021 ApJ 909 67 <https://doi.org/10.3847/1538-4357/abdc28>`__
"""
# TODO: Update arXiv attributions to ApJ (See #572)
def __init__(
self,
metric,
metric_params,
q0,
p0,
time_like=True,
steps=100,
delta=0.5,
rtol=1e-2,
atol=1e-2,
order=2,
omega=1.0,
suppress_warnings=False,
):
"""
Constructor
Parameters
----------
metric : callable
Metric Function. Currently, these metrics are supported:
1. Schwarzschild
2. Kerr
3. KerrNewman
metric_params : array_like
Tuple of parameters to pass to the metric
E.g., ``(a,)`` for Kerr
q0 : array_like
Initial 4-Position
p0 : array_like
Initial 4-Momentum
time_like : bool, optional
Determines type of Geodesic
``True`` for Time-like geodesics
``False`` for Null-like geodesics
Defaults to ``True``
steps : int
Number of integration steps
Defaults to ``50``
delta : float
Initial integration step-size
Defaults to ``0.5``
rtol : float
Relative Tolerance
Defaults to ``1e-2``
atol : float
Absolute Tolerance
Defaults to ``1e-2``
order : int
Integration Order
Defaults to ``2``
omega : float
Coupling between Hamiltonian Flows
Smaller values imply smaller integration error, but too
small values can make the equation of motion non-integrable.
For non-capture trajectories, ``omega = 1.0`` is recommended.
For trajectories, that either lead to a capture or a grazing
geodesic, a decreased value of ``0.01`` or less is recommended.
Defaults to ``1.0``
suppress_warnings : bool
Whether to suppress warnings during simulation
Warnings are shown for every step, where numerical errors
exceed specified tolerance (controlled by ``rtol`` and ``atol``)
Defaults to ``False``
Raises
------
NotImplementedError
If ``order`` is not in [2, 4, 6, 8]
"""
ORDERS = {
2: self._ord_2,
4: self._ord_4,
6: self._ord_6,
8: self._ord_8,
}
self.metric = metric
self.metric_params = metric_params
self.q0 = q0
self.p0 = p0
self.time_like = time_like
self.steps = steps
self.delta = delta
self.omega = omega
if order not in ORDERS:
raise NotImplementedError(
f"Order {order} integrator has not been implemented."
)
self.order = order
self.integrator = ORDERS[order]
self.rtol = rtol
self.atol = atol
self.suppress_warnings = suppress_warnings
self.step_num = 0
self.res_list = [q0, p0, q0, p0]
self.results = list()
def __str__(self):
return f"""{self.__class__.__name__}(\n\
metric : {self.metric}\n\
metric_params : {self.metric_params}\n\
q0 : {self.q0},\n\
p0 : {self.p0},\n\
time_like : {self.time_like},\n\
steps : {self.steps},\n\
delta : {self.delta},\n\
omega : {self.omega},\n\
order : {self.order},\n\
rtol : {self.rtol},\n\
atol : {self.atol}\n\
suppress_warnings : {self.suppress_warnings}
)"""
def __repr__(self):
return self.__str__()
def _ord_2(self, q1, p1, q2, p2, delta):
"""
Order 2 Integration Scheme
References
----------
.. [1] Christian, Pierre and Chan, Chi-Kwan;
"FANTASY : User-Friendly Symplectic Geodesic Integrator
for Arbitrary Metrics with Automatic Differentiation";
`2021 ApJ 909 67 <https://doi.org/10.3847/1538-4357/abdc28>`__
"""
dl, omg = delta, self.omega
g = self.metric
g_prms = self.metric_params
HA1 = np.array(
[
q1,
_flow_A(g, g_prms, q1, p1, q2, p2, 0.5 * dl)[1],
_flow_A(g, g_prms, q1, p1, q2, p2, 0.5 * dl)[0],
p2,
]
)
HB1 = np.array(
[
_flow_B(g, g_prms, HA1[0], HA1[1], HA1[2], HA1[3], 0.5 * dl)[0],
HA1[1],
HA1[2],
_flow_B(g, g_prms, HA1[0], HA1[1], HA1[2], HA1[3], 0.5 * dl)[1],
]
)
HC = _flow_mixed(HB1[0], HB1[1], HB1[2], HB1[3], dl, omg)
HB2 = np.array(
[
_flow_B(g, g_prms, HC[0], HC[1], HC[2], HC[3], 0.5 * dl)[0],
HC[1],
HC[2],
_flow_B(g, g_prms, HC[0], HC[1], HC[2], HC[3], 0.5 * dl)[1],
]
)
HA2 = np.array(
[
HB2[0],
_flow_A(g, g_prms, HB2[0], HB2[1], HB2[2], HB2[3], 0.5 * dl)[1],
_flow_A(g, g_prms, HB2[0], HB2[1], HB2[2], HB2[3], 0.5 * dl)[0],
HB2[3],
]
)
return HA2
def _ord_4(self, q1, p1, q2, p2, delta):
"""
Order 4 Integration Scheme
References
----------
.. [1] Yoshida, Haruo,
"Construction of higher order symplectic integrators";
Physics Letters A, vol. 150, no. 5-7, pp. 262-268, 1990.
`DOI: <https://doi.org/10.1016/0375-9601(90)90092-3>`__
"""
dl = delta
Z0, Z1 = _Z(self.order)
step1 = self._ord_2(q1, p1, q2, p2, dl * Z1)
step2 = self._ord_2(step1[0], step1[1], step1[2], step1[3], dl * Z0)
step3 = self._ord_2(step2[0], step2[1], step2[2], step2[3], dl * Z1)
return step3
def _ord_6(self, q1, p1, q2, p2, delta):
"""
Order 6 Integration Scheme
References
----------
.. [1] Yoshida, Haruo,
"Construction of higher order symplectic integrators";
Physics Letters A, vol. 150, no. 5-7, pp. 262-268, 1990.
`DOI: <https://doi.org/10.1016/0375-9601(90)90092-3>`__
"""
dl = delta
Z0, Z1 = _Z(self.order)
step1 = self._ord_4(q1, p1, q2, p2, dl * Z1)
step2 = self._ord_4(step1[0], step1[1], step1[2], step1[3], dl * Z0)
step3 = self._ord_4(step2[0], step2[1], step2[2], step2[3], dl * Z1)
return step3
def _ord_8(self, q1, p1, q2, p2, delta):
"""
Order 8 Integration Scheme
References
----------
.. [1] Yoshida, Haruo,
"Construction of higher order symplectic integrators";
Physics Letters A, vol. 150, no. 5-7, pp. 262-268, 1990.
`DOI: <https://doi.org/10.1016/0375-9601(90)90092-3>`__
"""
dl = delta
Z0, Z1 = _Z(self.order)
step1 = self._ord_6(q1, p1, q2, p2, dl * Z1)
step2 = self._ord_6(step1[0], step1[1], step1[2], step1[3], dl * Z0)
step3 = self._ord_6(step2[0], step2[1], step2[2], step2[3], dl * Z1)
return step3
[docs]
def step(self):
"""
Advances integration by one step
"""
rl = self.res_list
arr = self.integrator(rl[0], rl[1], rl[2], rl[3], self.delta)
self.res_list = arr
self.step_num += 1
# Stability check
if not self.suppress_warnings:
g = self.metric
g_prms = self.metric_params
q1 = arr[0]
p1 = arr[1]
# Ignoring
# q_2 = arr[2]
# p_2 = arr[3]
const = -int(self.time_like)
# g.p.p ~ -1 or 0 (const)
if not np.allclose(
g(q1, *g_prms) @ p1 @ p1, const, rtol=self.rtol, atol=self.atol
):
warnings.warn(
f"Numerical error has exceeded specified tolerance at step = {self.step_num}.",
RuntimeWarning,
)
self.results.append(self.res_list)
