orbix.kepler.shortcuts.grid =========================== .. py:module:: orbix.kepler.shortcuts.grid .. autoapi-nested-parse:: 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 --------- .. autoapisummary:: orbix.kepler.shortcuts.grid.get_grid_solver orbix.kepler.shortcuts.grid._setup orbix.kepler.shortcuts.grid._E_grids orbix.kepler.shortcuts.grid._d_dind_grids orbix.kepler.shortcuts.grid._ind orbix.kepler.shortcuts.grid._lin orbix.kepler.shortcuts.grid._grid_lin_params orbix.kepler.shortcuts.grid.E_trig_lin orbix.kepler.shortcuts.grid.E_lin orbix.kepler.shortcuts.grid.trig_lin orbix.kepler.shortcuts.grid._indices_frac orbix.kepler.shortcuts.grid._E_grid_base orbix.kepler.shortcuts.grid._build_E_grid orbix.kepler.shortcuts.grid._build_E_trig_grid orbix.kepler.shortcuts.grid._build_trig_grid orbix.kepler.shortcuts.grid._weights orbix.kepler.shortcuts.grid._scalar_E_trig_bilin orbix.kepler.shortcuts.grid._scalar_E_bilin orbix.kepler.shortcuts.grid._scalar_trig_bilin orbix.kepler.shortcuts.grid.E_trig_bilin orbix.kepler.shortcuts.grid.E_bilin orbix.kepler.shortcuts.grid.trig_bilin Module Contents --------------- .. py:function:: 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. :param 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,) :param jit: Whether to jit the solver. :param kind: The kind of solver to use, either "linear" or "bilinear". :param E: Whether to compute the eccentric anomaly. :param trig: Whether to compute the sine and cosine of the eccentric anomaly. :param n_e: The number of eccentricity steps in the grid. :param n_M: The number of mean anomaly steps in the grid. .. py:function:: _setup(n_e = 1000, n_M = 3600) Setup for grid methods. .. py:function:: _E_grids(e_grid, M_grid) Compute E, sinE, cosE grids. .. py:function:: _d_dind_grids(e_grid, E_grid, sinE_grid, cosE_grid, dM) Compute dE/dind, dsinE/dind, dcosE/dind grids. .. py:function:: _ind(scalar, inv_d) Returns the indices and fractional difference for linear interpolation. .. py:function:: _lin(tab, dtab, e_ind, M_ind, dM) Linear interpolation for a single (M, e) pair. .. py:function:: _grid_lin_params(M_scalar, e_scalar, inv_dM, inv_de, n_M_int, n_e) Linear lookup for a single (M, e) pair. .. py:function:: 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). :param n_e: Number of eccentricity steps in the grid (0 <= e < 1). Default is 1024. :param 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. .. py:function:: E_lin(n_e=1024, n_M=4096) Creates vectorized JIT func for E via linear interp 2D grid. .. py:function:: trig_lin(n_e=1024, n_M=4096) Creates vectorized JIT func for sinE, cosE via linear interp 2D grid. .. py:function:: _indices_frac(M, e, inv_dM, inv_de, n_M, n_e) .. py:function:: _E_grid_base(n_e, n_M, *, dtype=jnp.float32) .. py:function:: _build_E_grid(n_e, n_M, *, dtype=jnp.float32) .. py:function:: _build_E_trig_grid(n_e, n_M, *, dtype=jnp.float32) .. py:function:: _build_trig_grid(n_e, n_M, *, dtype=jnp.float32) Return sinE-cosE tensor of shape (2, n_e, n_M+1) and inverse steps. .. py:function:: _weights(de, dM, dtype) .. py:function:: _scalar_E_trig_bilin(triple, inv_dM, inv_de, n_M, n_e, M_scalar, e_scalar) .. py:function:: _scalar_E_bilin(E_grid, inv_dM, inv_de, n_M, n_e, M_scalar, e_scalar) .. py:function:: _scalar_trig_bilin(trig, inv_dM, inv_de, n_M, n_e, M_scalar, e_scalar) .. py:function:: E_trig_bilin(n_e=1024, n_M=4096) E, sinE, cosE via bilinear interp of packed grid. .. py:function:: E_bilin(n_e=1024, n_M=4096) E via bilinear interp of packed grid. .. py:function:: trig_bilin(n_e=1024, n_M=4096) sinE, cosE via bilinear interp of packed grid.