orbix.kepler.core

orbix.kepler.core#

JIT-compatible functions to solve Kepler’s equation by vectorizing the orvara solver.

The main entry points here are E_solve and E_solve_trig. Use jax.jit on them if you are using them in a performance-critical function. If you cannot be bothered to do so, use the jitted versions E_solve_jit or E_solve_trig_jit. Alternatively, to call on arrays of M and e values use the vectorized versions E_solve_vec or E_solve_trig_vec.

E_solve takes an array of mean anomaly values and a single eccentricity and returns the eccentric anomaly.

E_solve_trig takes the same input and returns the eccentric anomaly and its sine and cosine (which are calculated in the course of E_solve anyways).

The system works by defining separate functions for different ranges of eccentricities and then selecting between them with jnp.select: - e = 0 -> identity_solver

  • Returns input mean anomaly as eccentric anomaly

  • 0 < e < 0.78 -> le_E

    • Low Eccentricity

  • e > 0.78 -> he_E

    • High Eccentricity

Then there are equivalent functions for the trigonometric functions: - e = 0 -> identity_solver_trig - 0 < e < 0.78 -> le_E_trig - e > 0.78 -> he_E_trig

Acknowledgements: Portions of this code are adapted from orvara (t-brandt/orvara). orvara is distributed under a BSD 3-clause license and is Copyright (c) 2021, Timothy Brandt, Trent Dupuy, Yiting Li, G. Mirek Brandt, Yunlin Zeng, Daniel Michalik, and Virginia Raposo-Pulido.

Attributes#

if3

if5

if7

if9

if11

if13

if15

pi

pi_d_12

pi_d_6

pi_d_4

pi_d_3

fivepi_d_12

pi_d_2

sevenpi_d_12

twopi_d_3

threepi_d_4

fivepi_d_6

elevenpi_d_12

E_solve_jit

E_solve_vec

E_solve_trig_jit

E_solve_trig_vec

solve_trig_jit

solve_trig_vec

compute_dE_vectorized

Functions#

E_solve(M, e)

Vectorized orvara solver for eccentric anomaly.

E_solve_trig(M, e)

Vectorized orvara solver for eccentric anomaly and trigonometric functions.

solve_trig(M, e)

Wrapper around E_solve_trig that returns only (sinE, cosE).

diff_solve_trig(M, e)

Solve Kepler's equation, returning (sinE, cosE) with exact gradients.

_diff_solve_trig_fwd(M, e)

Forward pass: solve and save residuals for the backward pass.

_diff_solve_trig_bwd(res, g)

Backward pass: exact gradients via the Implicit Function Theorem.

shortsin(x)

Approximates the sine function using a short polynomial.

cut_M(M)

Cut M to be between 0 and pi.

getbounds(e)

Create bounds and coefficients for the eccentric anomaly polynomial.

init_E_poly(M, e)

Initial guess for the eccentric anomaly.

init_E_coeffs(M, bounds, coeffs)

Create the initial guess for the eccentric anomaly using the polynomials.

dE_num_denom(M, E, e_inv, sinE, cosE)

Compute the numerator and denominator for dE.

dE_2nd(M, E, e_inv, sinE, cosE)

Compute the second order approximation of dE.

dE_3rd(M, E, e_inv, sinE, cosE)

Compute the third order approximation of dE.

compute_dE_single(M, init_E_val, e_inv_val, sinE_val, ...)

Computes dE for a single element based on the condition M > 0.4.

le_E(M, e)

Inverts Kepler's time equation for elliptical orbits using Orvara's method.

le_E_trig(M, e)

Inverts Kepler's time equation for elliptical orbits using Orvara's method.

he_E(M, e)

Inverts Kepler's time equation for elliptical orbits with e > 0.78.

he_E_trig(M, e)

Inverts Kepler's time equation for elliptical orbits with e > 0.78.

Etrig_1(E)

When E <= pi_d_4.

Etrig_2(E)

When E > pi_d_4 and E < three_pi_d_4.

Etrig_3(E)

When E > pi_d_2 and E > three_pi_d_4.

Etrig(i, E)

Apply the correct trigonometric function based on the index.

fast_sinE_cosE(E)

Compute the sine and cosine of the eccentric anomaly using shortsin.

identity_solver(M, e)

Returns M as E when e is 0.

identity_solver_trig(M, e)

Returns M as E when e is 0.

Module Contents#

orbix.kepler.core.if3 = 0.16666666666666666#
orbix.kepler.core.if5 = 0.008333333333333333#
orbix.kepler.core.if7 = 0.0001984126984126984#
orbix.kepler.core.if9 = 2.7557319223985893e-06#
orbix.kepler.core.if11 = 2.505210838544172e-08#
orbix.kepler.core.if13 = 1.6059043836821613e-10#
orbix.kepler.core.if15 = 7.647163731819816e-13#
orbix.kepler.core.pi#
orbix.kepler.core.pi_d_12#
orbix.kepler.core.pi_d_6#
orbix.kepler.core.pi_d_4#
orbix.kepler.core.pi_d_3#
orbix.kepler.core.fivepi_d_12#
orbix.kepler.core.pi_d_2#
orbix.kepler.core.sevenpi_d_12#
orbix.kepler.core.twopi_d_3#
orbix.kepler.core.threepi_d_4#
orbix.kepler.core.fivepi_d_6#
orbix.kepler.core.elevenpi_d_12#
orbix.kepler.core.E_solve(M, e)[source]#

Vectorized orvara solver for eccentric anomaly.

Parameters:

  • M (jnp.ndarray) – Mean anomaly. Shape: (n,).

  • e (float) – Eccentricity.

Returns:

Eccentric anomaly. Shape: (n,).

Return type:

E (jnp.ndarray)

The solver contract is 0 <= e < 1; e >= 1 or e < 0 silently produces NaN or garbage (unchecked to keep the hot path branch-free).

orbix.kepler.core.E_solve_jit = None#
orbix.kepler.core.E_solve_vec = None#
orbix.kepler.core.E_solve_trig(M, e)[source]#

Vectorized orvara solver for eccentric anomaly and trigonometric functions.

Parameters:

  • M (jnp.ndarray) – Mean anomaly. Shape: (n,).

  • e (float) – Eccentricity.

Returns:

Eccentric anomaly. Shape: (n,). sinE (jnp.ndarray): Sine of the eccentric anomaly. Shape: (n,). cosE (jnp.ndarray): Cosine of the eccentric anomaly. Shape: (n,).

Return type:

E (jnp.ndarray)

orbix.kepler.core.E_solve_trig_jit = None#
orbix.kepler.core.E_solve_trig_vec = None#
orbix.kepler.core.solve_trig(M, e)[source]#

Wrapper around E_solve_trig that returns only (sinE, cosE).

Parameters:

  • M (jnp.ndarray) – Mean anomaly. Shape: (n,).

  • e (float) – Eccentricity.

Returns:

Sine of the eccentric anomaly. Shape: (n,). cosE (jnp.ndarray): Cosine of the eccentric anomaly. Shape: (n,).

Return type:

sinE (jnp.ndarray)

The solver contract is 0 <= e < 1; e >= 1 or e < 0 silently produces NaN or garbage (unchecked to keep the hot path branch-free).

orbix.kepler.core.solve_trig_jit = None#
orbix.kepler.core.solve_trig_vec = None#
orbix.kepler.core.diff_solve_trig(M, e)[source]#

Solve Kepler’s equation, returning (sinE, cosE) with exact gradients.

Drop-in replacement for solve_trig() that supports reverse-mode autodiff (jax.grad, jax.vjp). Gradients come from the Implicit Function Theorem on M = E - e*sin(E), computed from (sinE, cosE, e) alone (no extra trig calls, no iterative re-solves).

Parameters:

  • M (jnp.ndarray) – Mean anomaly. Shape: (n,).

  • e (float) – Eccentricity.

Returns:

Sine of the eccentric anomaly. Shape: (n,). cosE (jnp.ndarray): Cosine of the eccentric anomaly. Shape: (n,).

Return type:

sinE (jnp.ndarray)

The solver contract is 0 <= e < 1; e >= 1 or e < 0 silently produces NaN or garbage (unchecked to keep the hot path branch-free).

orbix.kepler.core._diff_solve_trig_fwd(M, e)[source]#

Forward pass: solve and save residuals for the backward pass.

orbix.kepler.core._diff_solve_trig_bwd(res, g)[source]#

Backward pass: exact gradients via the Implicit Function Theorem.

From M = E - e*sinE: dE/dM = 1/(1 - e*cosE) and dE/de = sinE/(1 - e*cosE); chained through (sinE, cosE).

orbix.kepler.core.shortsin(x)[source]#

Approximates the sine function using a short polynomial.

This is only valid between [0, pi].

orbix.kepler.core.cut_M(M)[source]#

Cut M to be between 0 and pi.

Also returns the sign of the eccentric anomaly.

Parameters:

M (jnp.ndarray) – Mean anomalies (rad). Shape: (n,).

Returns:

Sign of the eccentric anomaly. Shape: (n,). _M (jnp.ndarray):

Modified mean anomalies. Shape: (n,).

Return type:

Esigns (jnp.ndarray)

orbix.kepler.core.getbounds(e)[source]#

Create bounds and coefficients for the eccentric anomaly polynomial.

Parameters:

e (float) – Eccentricity

Returns:

bounds (jnp.ndarray):

Array of bounds for the eccentric anomaly intervals. Shape: (13,)

coeffs (jnp.ndarray)

Lookup table containing coefficients for the Taylor series expansion. Shape: (13, 6)

Return type:

tuple

orbix.kepler.core.init_E_poly(M, e)[source]#

Initial guess for the eccentric anomaly.

Calculates the initial guess for the eccentric anomaly based on the mean anomaly and eccentricity. Translated from the C implementation into JAX.

Parameters:

  • M (jnp.ndarray) – Mean anomaly in radians.

  • e (float) – Eccentricity of the orbit.

Returns:

Initial estimate of the eccentric anomaly in radians.

Return type:

jnp.ndarray

orbix.kepler.core.init_E_coeffs(M, bounds, coeffs)[source]#

Create the initial guess for the eccentric anomaly using the polynomials.

Parameters:

  • M (jax.numpy.ndarray)

  • bounds (jax.numpy.ndarray)

  • coeffs (jax.numpy.ndarray)

orbix.kepler.core.dE_num_denom(M, E, e_inv, sinE, cosE)[source]#

Compute the numerator and denominator for dE.

orbix.kepler.core.dE_2nd(M, E, e_inv, sinE, cosE)[source]#

Compute the second order approximation of dE.

orbix.kepler.core.dE_3rd(M, E, e_inv, sinE, cosE)[source]#

Compute the third order approximation of dE.

orbix.kepler.core.compute_dE_single(M, init_E_val, e_inv_val, sinE_val, cosE_val)[source]#

Computes dE for a single element based on the condition M > 0.4.

Parameters:

  • M (float) – Single element from _M.

  • init_E_val (float) – Corresponding element from init_E.

  • e_inv_val (float) – Inverse of eccentricity.

  • sinE_val (float) – Sine of E.

  • cosE_val (float) – Cosine of E.

Returns:

Computed dE for the element.

Return type:

float

orbix.kepler.core.compute_dE_vectorized#
orbix.kepler.core.le_E(M, e)[source]#

Inverts Kepler’s time equation for elliptical orbits using Orvara’s method.

Parameters:

  • M (jnp.ndarray) – Mean anomalies (rad). Shape: (n,).

  • e (float) – Eccentricity. Must satisfy 0 <= e < 1.

Returns:

Eccentric anomalies (rad). Shape: (n,).

Return type:

  • E (jnp.ndarray)

orbix.kepler.core.le_E_trig(M, e)[source]#

Inverts Kepler’s time equation for elliptical orbits using Orvara’s method.

Also returns the sine and cosine of the eccentric anomaly.

Parameters:

  • M (jnp.ndarray) – Mean anomalies (rad). Shape: (n,).

  • e (float) – Eccentricity. Must satisfy 0 <= e < 1.

Returns:

  • E (jnp.ndarray): Eccentric anomalies (rad). Shape: (n,).

  • sinE (jnp.ndarray): Sine of eccentric anomalies (rad). Shape: (n,).

  • cosE (jnp.ndarray): Cosine of eccentric anomalies (rad). Shape: (n,).

Return type:

Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]

orbix.kepler.core.he_E(M, e)[source]#

Inverts Kepler’s time equation for elliptical orbits with e > 0.78.

Parameters:

  • M (jnp.ndarray) – Mean anomalies (rad). Shape: (n,).

  • e (float) – Eccentricity. Must satisfy 0 <= e < 1.

Returns:

Eccentric anomalies (rad). Shape: (n,).

Return type:

  • E (jnp.ndarray)

orbix.kepler.core.he_E_trig(M, e)[source]#

Inverts Kepler’s time equation for elliptical orbits with e > 0.78.

Parameters:

  • M (jnp.ndarray) – Mean anomalies (rad). Shape: (n,).

  • e (float) – Eccentricity. Must satisfy 0 <= e < 1.

Returns:

  • E (jnp.ndarray): Eccentric anomalies (rad). Shape: (n,).

  • sinE (jnp.ndarray): Sine of eccentric anomalies (rad). Shape: (n,).

  • cosE (jnp.ndarray): Cosine of eccentric anomalies (rad). Shape: (n,).

Return type:

Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]

orbix.kepler.core.Etrig_1(E)[source]#

When E <= pi_d_4.

orbix.kepler.core.Etrig_2(E)[source]#

When E > pi_d_4 and E < three_pi_d_4.

orbix.kepler.core.Etrig_3(E)[source]#

When E > pi_d_2 and E > three_pi_d_4.

orbix.kepler.core.Etrig(i, E)[source]#

Apply the correct trigonometric function based on the index.

orbix.kepler.core.fast_sinE_cosE(E)[source]#

Compute the sine and cosine of the eccentric anomaly using shortsin.

orbix.kepler.core.identity_solver(M, e)[source]#

Returns M as E when e is 0.

orbix.kepler.core.identity_solver_trig(M, e)[source]#

Returns M as E when e is 0.