orbix.kepler.shortcuts.grid

orbix.kepler.shortcuts.grid#

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.

Functions#

`get_grid_solver`([level, jit, kind, E, trig, n_e, n_M])	Helper function to get a grid-based solver and cache it.
`_setup`([n_e, n_M])	Setup for grid methods.
`_E_grids`(e_grid, M_grid)	Compute E, sinE, cosE grids.
`_d_dind_grids`(e_grid, E_grid, sinE_grid, cosE_grid, dM)	Compute dE/dind, dsinE/dind, dcosE/dind grids.
`_ind`(scalar, inv_d)	Returns the indices and fractional difference for linear interpolation.
`_lin`(tab, dtab, e_ind, M_ind, dM)	Linear interpolation for a single (M, e) pair.
`_grid_lin_params`(M_scalar, e_scalar, inv_dM, inv_de, ...)	Linear lookup for a single (M, e) pair.
`E_trig_lin`([n_e, n_M])	Creates a scalar lookup closure for E, sinE, cosE via linear interp 2D grid.
`E_lin`([n_e, n_M])	Creates vectorized JIT func for E via linear interp 2D grid.
`trig_lin`([n_e, n_M])	Creates vectorized JIT func for sinE, cosE via linear interp 2D grid.
`_indices_frac`(M, e, inv_dM, inv_de, n_M, n_e)
`_E_grid_base`(n_e, n_M, *[, dtype])
`_build_E_grid`(n_e, n_M, *[, dtype])
`_build_E_trig_grid`(n_e, n_M, *[, dtype])
`_build_trig_grid`(n_e, n_M, *[, dtype])	Return sinE-cosE tensor of shape (2, n_e, n_M+1) and inverse steps.
`_weights`(de, dM, dtype)
`_scalar_E_trig_bilin`(triple, inv_dM, inv_de, n_M, n_e, ...)
`_scalar_E_bilin`(E_grid, inv_dM, inv_de, n_M, n_e, ...)
`_scalar_trig_bilin`(trig, inv_dM, inv_de, n_M, n_e, ...)
`E_trig_bilin`([n_e, n_M])	E, sinE, cosE via bilinear interp of packed grid.
`E_bilin`([n_e, n_M])	E via bilinear interp of packed grid.
`trig_bilin`([n_e, n_M])	sinE, cosE via bilinear interp of packed grid.

Module Contents#

orbix.kepler.shortcuts.grid.get_grid_solver(level='scalar', jit=False, kind='bilinear', E=True, trig=True, n_e=512, n_M=2048)[source]#

Helper function to get a grid-based solver and cache it.

Parameters:

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.

orbix.kepler.shortcuts.grid._setup(n_e=1000, n_M=3600)[source]#

Setup for grid methods.

Parameters:

n_e (int)
n_M (int)

orbix.kepler.shortcuts.grid._E_grids(e_grid, M_grid)[source]#: Compute E, sinE, cosE grids.

orbix.kepler.shortcuts.grid._d_dind_grids(e_grid, E_grid, sinE_grid, cosE_grid, dM)[source]#: Compute dE/dind, dsinE/dind, dcosE/dind grids.

orbix.kepler.shortcuts.grid._ind(scalar, inv_d)[source]#: Returns the indices and fractional difference for linear interpolation.

orbix.kepler.shortcuts.grid._lin(tab, dtab, e_ind, M_ind, dM)[source]#: Linear interpolation for a single (M, e) pair.

orbix.kepler.shortcuts.grid._grid_lin_params(M_scalar, e_scalar, inv_dM, inv_de, n_M_int, n_e)[source]#: Linear lookup for a single (M, e) pair.

orbix.kepler.shortcuts.grid.E_trig_lin(n_e=1024, n_M=4096)[source]#

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).

Parameters:

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.

orbix.kepler.shortcuts.grid.E_lin(n_e=1024, n_M=4096)[source]#: Creates vectorized JIT func for E via linear interp 2D grid.

orbix.kepler.shortcuts.grid.trig_lin(n_e=1024, n_M=4096)[source]#: Creates vectorized JIT func for sinE, cosE via linear interp 2D grid.

orbix.kepler.shortcuts.grid._indices_frac(M, e, inv_dM, inv_de, n_M, n_e)[source]#

orbix.kepler.shortcuts.grid._E_grid_base(n_e, n_M, *, dtype=jnp.float32)[source]#

Parameters:

n_e (int)
n_M (int)

orbix.kepler.shortcuts.grid._build_E_grid(n_e, n_M, *, dtype=jnp.float32)[source]#

Parameters:

n_e (int)
n_M (int)

orbix.kepler.shortcuts.grid._build_E_trig_grid(n_e, n_M, *, dtype=jnp.float32)[source]#

Parameters:

n_e (int)
n_M (int)

orbix.kepler.shortcuts.grid._build_trig_grid(n_e, n_M, *, dtype=jnp.float32)[source]#

Return sinE-cosE tensor of shape (2, n_e, n_M+1) and inverse steps.

Parameters:

n_e (int)
n_M (int)

orbix.kepler.shortcuts.grid._weights(de, dM, dtype)[source]#

orbix.kepler.shortcuts.grid._scalar_E_trig_bilin(triple, inv_dM, inv_de, n_M, n_e, M_scalar, e_scalar)[source]#

orbix.kepler.shortcuts.grid._scalar_E_bilin(E_grid, inv_dM, inv_de, n_M, n_e, M_scalar, e_scalar)[source]#

orbix.kepler.shortcuts.grid._scalar_trig_bilin(trig, inv_dM, inv_de, n_M, n_e, M_scalar, e_scalar)[source]#

orbix.kepler.shortcuts.grid.E_trig_bilin(n_e=1024, n_M=4096)[source]#: E, sinE, cosE via bilinear interp of packed grid.

orbix.kepler.shortcuts.grid.E_bilin(n_e=1024, n_M=4096)[source]#: E via bilinear interp of packed grid.

orbix.kepler.shortcuts.grid.trig_bilin(n_e=1024, n_M=4096)[source]#: sinE, cosE via bilinear interp of packed grid.