"""Grid-based functions to solve Kepler's equation.
These functions are essentially vectorized versions of the functions in
`fixed_e.py`. They work by precomputing a 2D grid of E, sinE, cosE values and
their derivatives (for linear interpolation). Then, for a given (M, e) pair, a
scalar function is created to do the interpolation using the precomputed grid.
This scalar function is then jit-compiled and vectorized over time and planets.
Generally these are 3-5x faster than a vectorized/compiled form of `E_solve`
when computing thousands of epochs but the difference is less pronounced at the
few per-epoch level. The bilinear interpolation is ~1.5x slower than the linear
interpolation, but is accurate to 32 bit precision.
The "linear" kind interpolates in M only; eccentricity snaps to the
nearest grid row, so its error grows like (0.5/n_e) * dE/de and is
worst near e = 1. Use "bilinear" when e accuracy matters.
"""
import sys
from functools import lru_cache, partial
import jax
import jax.numpy as jnp
from hwoutils.constants import two_pi
from jax import lax
from orbix.kepler.core import E_solve_jit
# lru_cache here lets us give the same solver to multiple planets
[docs]
@lru_cache(maxsize=None)
def get_grid_solver(
level="scalar", jit=False, kind="bilinear", E=True, trig=True, n_e=512, n_M=2048
):
"""Helper function to get a grid-based solver and cache it.
Args:
level:
How the solver should be batching things.
- "scalar" means the inputs will be a single (M, e) pair and
the output will be a single E, sinE, cosE value.
- "planet" will vectorize over times first and then over
orbits.
- M: (n_orbits, n_times)
- e: (n_orbits,)
jit:
Whether to jit the solver.
kind:
The kind of solver to use, either "linear" or "bilinear".
E:
Whether to compute the eccentric anomaly.
trig:
Whether to compute the sine and cosine of the eccentric anomaly.
n_e:
The number of eccentricity steps in the grid.
n_M:
The number of mean anomaly steps in the grid.
"""
if not E and not trig:
raise ValueError("get_grid_solver: at least one of E or trig must be True")
if E and trig:
name = "E_trig"
elif E:
name = "E"
elif trig:
name = "trig"
if kind == "bilinear":
name += "_bilin"
elif kind == "linear":
name += "_lin"
else:
raise ValueError(f"Unknown solver kind: {kind}")
# Get the correct function
func = getattr(sys.modules[__name__], name)
func = func(n_e=n_e, n_M=n_M)
if level == "scalar":
pass
elif level == "planet":
# Mean anomaly/time axis map M->(n_times,), e -> scalar
func = jax.vmap(func, in_axes=(0, None))
# orbit axis map M->(n_orbits, n_times), e -> (n_orbits,)
func = jax.vmap(func, in_axes=(0, 0))
else:
raise ValueError(f"Unknown solver level: {level}")
if jit:
func = jax.jit(func)
return func
[docs]
def _setup(n_e: int = 1000, n_M: int = 3600):
"""Setup for grid methods."""
e_grid = jnp.linspace(0.0, 1.0, n_e, endpoint=False)
M_grid = jnp.linspace(0.0, two_pi, n_M, endpoint=False)
dM = M_grid[1] - M_grid[0]
inv_dM = 1.0 / dM
n_M_int = jnp.int32(n_M)
de = e_grid[1] - e_grid[0]
inv_de = 1.0 / de
return e_grid, M_grid, dM, inv_dM, n_M_int, de, inv_de
[docs]
def _E_grids(e_grid, M_grid):
"""Compute E, sinE, cosE grids."""
E_grid = jax.vmap(lambda e_val: E_solve_jit(M_grid, e_val))(e_grid)
return E_grid, jnp.sin(E_grid), jnp.cos(E_grid)
[docs]
def _d_dind_grids(e_grid, E_grid, sinE_grid, cosE_grid, dM):
"""Compute dE/dind, dsinE/dind, dcosE/dind grids."""
dE_dM_grid = 1 / (1 - e_grid.reshape(-1, 1) * jnp.cos(E_grid))
dE_dind_grid = dE_dM_grid * dM
# dsinE/dind = dsinE/dE * dE/dind
dsinE_dind = cosE_grid * dE_dind_grid
# dcosE/dind = dcosE/dE * dE/dind
dcosE_dind = -sinE_grid * dE_dind_grid
return dE_dind_grid, dsinE_dind, dcosE_dind
################################################################################
# Linear interpolation functions
################################################################################
[docs]
def _ind(scalar, inv_d):
"""Returns the indices and fractional difference for linear interpolation."""
ind = scalar * inv_d
i0 = ind.astype(jnp.int32)
di = ind - i0
return i0, di
[docs]
def _lin(tab, dtab, e_ind, M_ind, dM):
"""Linear interpolation for a single (M, e) pair."""
return tab[e_ind, M_ind] + dtab[e_ind, M_ind] * dM
[docs]
def _grid_lin_params(M_scalar, e_scalar, inv_dM, inv_de, n_M_int, n_e):
"""Linear lookup for a single (M, e) pair."""
M_scalar = jnp.mod(M_scalar, two_pi)
M0, dM = _ind(M_scalar, inv_dM)
M0 = M0 % n_M_int
e0 = jnp.clip((e_scalar * inv_de + 0.5).astype(jnp.int32), 0, n_e - 1)
return e0, M0, dM
[docs]
def E_trig_lin(n_e=1024, n_M=4096):
"""Creates a scalar lookup closure for E, sinE, cosE via linear interp 2D grid.
This function precomputes the E, sin(E), cos(E) grids and their derivatives
and closes over them. The returned function is a plain (un-jitted) scalar
closure, `_lookup_scalar(M_scalar, e_scalar) -> (E, sinE, cosE)`; it is not
vectorized or JIT-compiled itself. Callers wrap it with `jax.vmap`/`jax.jit`
as needed (see `get_grid_solver`, which does exactly this).
Args:
n_e:
Number of eccentricity steps in the grid (0 <= e < 1). Default is 1024.
n_M:
Number of mean anomaly steps in the grid (0 <= M < 2pi). Default is 4096.
Returns:
A scalar closure `_lookup_scalar(M_scalar, e_scalar)` that returns a
tuple `(E, sinE, cosE)` of interpolated eccentric anomaly and its sine
and cosine for a single (M, e) pair.
"""
# Setup
e_grid, M_grid, dM, inv_dM, n_M_int, de, inv_de = _setup(n_e, n_M)
E_grid, sinE_grid, cosE_grid = _E_grids(e_grid, M_grid)
dE_dind_grid, dsinE_dind, dcosE_dind = _d_dind_grids(
e_grid, E_grid, sinE_grid, cosE_grid, dM
)
def _lookup_scalar(M_scalar, e_scalar):
"""Performs lookup and interpolation for a single M and e."""
p = _grid_lin_params(M_scalar, e_scalar, inv_dM, inv_de, n_M_int, n_e)
return (
_lin(E_grid, dE_dind_grid, *p),
_lin(sinE_grid, dsinE_dind, *p),
_lin(cosE_grid, dcosE_dind, *p),
)
return _lookup_scalar
[docs]
def E_lin(n_e=1024, n_M=4096):
"""Creates vectorized JIT func for E via linear interp 2D grid."""
# Setup
e_grid, M_grid, dM, inv_dM, n_M_int, de, inv_de = _setup(n_e, n_M)
E_grid, sinE_grid, cosE_grid = _E_grids(e_grid, M_grid)
dE_dind_grid, *_ = _d_dind_grids(e_grid, E_grid, sinE_grid, cosE_grid, dM)
def _lookup_scalar(M_scalar, e_scalar):
"""Performs lookup and interpolation for a single M and e."""
p = _grid_lin_params(M_scalar, e_scalar, inv_dM, inv_de, n_M_int, n_e)
return _lin(E_grid, dE_dind_grid, *p)
return _lookup_scalar
[docs]
def trig_lin(n_e=1024, n_M=4096):
"""Creates vectorized JIT func for sinE, cosE via linear interp 2D grid."""
# Setup
e_grid, M_grid, dM, inv_dM, n_M_int, de, inv_de = _setup(n_e, n_M)
E_grid, sinE_grid, cosE_grid = _E_grids(e_grid, M_grid)
_, dsinE_dind, dcosE_dind = _d_dind_grids(e_grid, E_grid, sinE_grid, cosE_grid, dM)
def _lookup_scalar(M_scalar, e_scalar):
"""Performs lookup and interpolation for a single M and e."""
p = _grid_lin_params(M_scalar, e_scalar, inv_dM, inv_de, n_M_int, n_e)
return _lin(sinE_grid, dsinE_dind, *p), _lin(cosE_grid, dcosE_dind, *p)
return _lookup_scalar
################################################################################
# Bilinear interpolation with stacked grids
################################################################################
[docs]
def _indices_frac(M, e, inv_dM, inv_de, n_M, n_e):
M = jnp.mod(M, two_pi)
# Convert to index space
M_ind = M * inv_dM
e_ind = e * inv_de
# Get integer indices
M_int = M_ind.astype(jnp.int32)
e0 = e_ind.astype(jnp.int32)
# Get fractional differences
dM = M_ind - M_int
de = e_ind - e0
# M is periodic so wrap it
M0 = M_int % n_M
return e0, M0, de, dM
[docs]
def _E_grid_base(n_e: int, n_M: int, *, dtype=jnp.float32):
e_grid = jnp.linspace(0.0, 1.0, n_e, dtype=dtype, endpoint=False)
M_grid = jnp.linspace(0.0, two_pi, n_M, dtype=dtype, endpoint=False)
E_grid = jax.vmap(lambda e: E_solve_jit(M_grid, e))(e_grid)
# Append one more e row, solved just below e=1, so the 2x2 dynamic_slice
# patch is valid in the top e cell (e in [1 - 1/n_e, 1)). Without it the
# slice start row silently clamps and the weights use the wrong cell.
_last_E_row = E_solve_jit(M_grid, jnp.asarray(1.0 - 1e-6, dtype=dtype))
E_grid = jnp.concatenate([E_grid, _last_E_row[jnp.newaxis, :]], axis=0)
# Append last column of 2pi-eps to make the grid complete with one more
# column because we always pull in a square patch of 2x2
_last_E_col = jnp.repeat(two_pi - jnp.finfo(dtype).eps, n_e + 1, axis=0)
E_grid = jnp.concatenate([E_grid, _last_E_col[:, jnp.newaxis]], axis=1)
inv_dM = 1.0 / (M_grid[1] - M_grid[0])
inv_de = 1.0 / (e_grid[1] - e_grid[0])
return E_grid, inv_dM, inv_de
[docs]
def _build_E_grid(n_e: int, n_M: int, *, dtype=jnp.float32):
E_grid, inv_dM, inv_de = _E_grid_base(n_e, n_M, dtype=dtype)
return jax.device_put(E_grid), inv_dM, inv_de
[docs]
def _build_E_trig_grid(n_e: int, n_M: int, *, dtype=jnp.float32):
E_grid, inv_dM, inv_de = _E_grid_base(n_e, n_M, dtype=dtype)
triple = jnp.stack([E_grid, jnp.sin(E_grid), jnp.cos(E_grid)], axis=0)
return jax.device_put(triple), inv_dM, inv_de
[docs]
def _build_trig_grid(n_e: int, n_M: int, *, dtype=jnp.float32):
"""Return sinE-cosE tensor of shape (2, n_e, n_M+1) and inverse steps."""
E_grid, inv_dM, inv_de = _E_grid_base(n_e, n_M, dtype=dtype)
trig = jnp.stack([jnp.sin(E_grid), jnp.cos(E_grid)], axis=0)
return jax.device_put(trig), inv_dM, inv_de
[docs]
def _weights(de, dM, dtype):
return jnp.array(
[[(1 - de) * (1 - dM), (1 - de) * dM], [de * (1 - dM), de * dM]],
dtype=dtype,
)
[docs]
def _scalar_E_trig_bilin(triple, inv_dM, inv_de, n_M, n_e, M_scalar, e_scalar):
# indices and fractions
e0, M0, de, dM = _indices_frac(M_scalar, e_scalar, inv_dM, inv_de, n_M, n_e)
# dynamic slice to avoid multiple table lookups
w = _weights(de, dM, triple.dtype)
# This should get batched to the right shape by vmap
patch = lax.dynamic_slice(triple, (jnp.int32(0), e0, M0), (3, 2, 2))
result = jnp.sum(patch * w, axis=(-2, -1))
return result[0], result[1], result[2]
[docs]
def _scalar_E_bilin(E_grid, inv_dM, inv_de, n_M, n_e, M_scalar, e_scalar):
e0, M0, de, dM = _indices_frac(M_scalar, e_scalar, inv_dM, inv_de, n_M, n_e)
patch = lax.dynamic_slice(E_grid, (e0, M0), (2, 2))
w = _weights(de, dM, E_grid.dtype)
result = jnp.sum(patch * w, axis=(-2, -1))
return result
[docs]
def _scalar_trig_bilin(trig, inv_dM, inv_de, n_M, n_e, M_scalar, e_scalar):
# indices and fractions
e0, M0, de, dM = _indices_frac(M_scalar, e_scalar, inv_dM, inv_de, n_M, n_e)
# dynamic slice to avoid multiple table lookups
w = _weights(de, dM, trig.dtype)
# This should get batched to the right shape by vmap
patch = lax.dynamic_slice(trig, (jnp.int32(0), e0, M0), (2, 2, 2))
result = jnp.sum(patch * w, axis=(-2, -1))
return result[0], result[1]
[docs]
def E_trig_bilin(n_e=1024, n_M=4096):
"""E, sinE, cosE via bilinear interp of packed grid."""
triple, inv_dM, inv_de = _build_E_trig_grid(n_e, n_M)
scalar_fun = partial(_scalar_E_trig_bilin, triple, inv_dM, inv_de, n_M, n_e)
return scalar_fun
[docs]
def E_bilin(n_e=1024, n_M=4096):
"""E via bilinear interp of packed grid."""
E_grid, inv_dM, inv_de = _build_E_grid(n_e, n_M)
scalar_fun = partial(_scalar_E_bilin, E_grid, inv_dM, inv_de, n_M, n_e)
return scalar_fun
[docs]
def trig_bilin(n_e=1024, n_M=4096):
"""sinE, cosE via bilinear interp of packed grid."""
trig, inv_dM, inv_de = _build_trig_grid(n_e, n_M)
scalar_fun = partial(_scalar_trig_bilin, trig, inv_dM, inv_de, n_M, n_e)
return scalar_fun