orbix.kepler.core ================= .. py:module:: orbix.kepler.core .. autoapi-nested-parse:: 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 (https://github.com/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 ---------- .. autoapisummary:: orbix.kepler.core.if3 orbix.kepler.core.if5 orbix.kepler.core.if7 orbix.kepler.core.if9 orbix.kepler.core.if11 orbix.kepler.core.if13 orbix.kepler.core.if15 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_jit orbix.kepler.core.E_solve_vec orbix.kepler.core.E_solve_trig_jit orbix.kepler.core.E_solve_trig_vec orbix.kepler.core.solve_trig_jit orbix.kepler.core.solve_trig_vec orbix.kepler.core.compute_dE_vectorized Functions --------- .. autoapisummary:: orbix.kepler.core.E_solve orbix.kepler.core.E_solve_trig orbix.kepler.core.solve_trig orbix.kepler.core.diff_solve_trig orbix.kepler.core._diff_solve_trig_fwd orbix.kepler.core._diff_solve_trig_bwd orbix.kepler.core.shortsin orbix.kepler.core.cut_M orbix.kepler.core.getbounds orbix.kepler.core.init_E_poly orbix.kepler.core.init_E_coeffs orbix.kepler.core.dE_num_denom orbix.kepler.core.dE_2nd orbix.kepler.core.dE_3rd orbix.kepler.core.compute_dE_single orbix.kepler.core.le_E orbix.kepler.core.le_E_trig orbix.kepler.core.he_E orbix.kepler.core.he_E_trig orbix.kepler.core.Etrig_1 orbix.kepler.core.Etrig_2 orbix.kepler.core.Etrig_3 orbix.kepler.core.Etrig orbix.kepler.core.fast_sinE_cosE orbix.kepler.core.identity_solver orbix.kepler.core.identity_solver_trig Module Contents --------------- .. py:data:: if3 :value: 0.16666666666666666 .. py:data:: if5 :value: 0.008333333333333333 .. py:data:: if7 :value: 0.0001984126984126984 .. py:data:: if9 :value: 2.7557319223985893e-06 .. py:data:: if11 :value: 2.505210838544172e-08 .. py:data:: if13 :value: 1.6059043836821613e-10 .. py:data:: if15 :value: 7.647163731819816e-13 .. py:data:: pi .. py:data:: pi_d_12 .. py:data:: pi_d_6 .. py:data:: pi_d_4 .. py:data:: pi_d_3 .. py:data:: fivepi_d_12 .. py:data:: pi_d_2 .. py:data:: sevenpi_d_12 .. py:data:: twopi_d_3 .. py:data:: threepi_d_4 .. py:data:: fivepi_d_6 .. py:data:: elevenpi_d_12 .. py:function:: E_solve(M, e) Vectorized orvara solver for eccentric anomaly. :param M: Mean anomaly. Shape: (n,). :type M: jnp.ndarray :param e: Eccentricity. :type e: float :returns: Eccentric anomaly. Shape: (n,). :rtype: 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). .. py:data:: E_solve_jit :value: None .. py:data:: E_solve_vec :value: None .. py:function:: E_solve_trig(M, e) Vectorized orvara solver for eccentric anomaly and trigonometric functions. :param M: Mean anomaly. Shape: (n,). :type M: jnp.ndarray :param e: Eccentricity. :type e: float :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,). :rtype: E (jnp.ndarray) .. py:data:: E_solve_trig_jit :value: None .. py:data:: E_solve_trig_vec :value: None .. py:function:: solve_trig(M, e) Wrapper around E_solve_trig that returns only (sinE, cosE). :param M: Mean anomaly. Shape: (n,). :type M: jnp.ndarray :param e: Eccentricity. :type e: float :returns: Sine of the eccentric anomaly. Shape: (n,). cosE (jnp.ndarray): Cosine of the eccentric anomaly. Shape: (n,). :rtype: 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). .. py:data:: solve_trig_jit :value: None .. py:data:: solve_trig_vec :value: None .. py:function:: diff_solve_trig(M, e) Solve Kepler's equation, returning (sinE, cosE) with exact gradients. Drop-in replacement for :func:`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). :param M: Mean anomaly. Shape: (n,). :type M: jnp.ndarray :param e: Eccentricity. :type e: float :returns: Sine of the eccentric anomaly. Shape: (n,). cosE (jnp.ndarray): Cosine of the eccentric anomaly. Shape: (n,). :rtype: 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). .. py:function:: _diff_solve_trig_fwd(M, e) Forward pass: solve and save residuals for the backward pass. .. py:function:: _diff_solve_trig_bwd(res, g) 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)``. .. py:function:: shortsin(x) Approximates the sine function using a short polynomial. This is only valid between [0, pi]. .. py:function:: cut_M(M) Cut M to be between 0 and pi. Also returns the sign of the eccentric anomaly. :param M: Mean anomalies (rad). Shape: (n,). :type M: jnp.ndarray :returns: Sign of the eccentric anomaly. Shape: (n,). _M (jnp.ndarray): Modified mean anomalies. Shape: (n,). :rtype: Esigns (jnp.ndarray) .. py:function:: getbounds(e) Create bounds and coefficients for the eccentric anomaly polynomial. :param e: Eccentricity :type e: float :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) :rtype: tuple .. py:function:: init_E_poly(M, e) 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. :param M: Mean anomaly in radians. :type M: jnp.ndarray :param e: Eccentricity of the orbit. :type e: float :returns: Initial estimate of the eccentric anomaly in radians. :rtype: jnp.ndarray .. py:function:: init_E_coeffs(M, bounds, coeffs) Create the initial guess for the eccentric anomaly using the polynomials. .. py:function:: dE_num_denom(M, E, e_inv, sinE, cosE) Compute the numerator and denominator for dE. .. py:function:: dE_2nd(M, E, e_inv, sinE, cosE) Compute the second order approximation of dE. .. py:function:: dE_3rd(M, E, e_inv, sinE, cosE) Compute the third order approximation of dE. .. py:function:: compute_dE_single(M, init_E_val, e_inv_val, sinE_val, cosE_val) Computes dE for a single element based on the condition M > 0.4. :param M: Single element from _M. :type M: float :param init_E_val: Corresponding element from init_E. :type init_E_val: float :param e_inv_val: Inverse of eccentricity. :type e_inv_val: float :param sinE_val: Sine of E. :type sinE_val: float :param cosE_val: Cosine of E. :type cosE_val: float :returns: Computed dE for the element. :rtype: float .. py:data:: compute_dE_vectorized .. py:function:: le_E(M, e) Inverts Kepler's time equation for elliptical orbits using Orvara's method. :param M: Mean anomalies (rad). Shape: (n,). :type M: jnp.ndarray :param e: Eccentricity. Must satisfy 0 <= e < 1. :type e: float :returns: Eccentric anomalies (rad). Shape: (n,). :rtype: - E (jnp.ndarray) .. py:function:: le_E_trig(M, e) Inverts Kepler's time equation for elliptical orbits using Orvara's method. Also returns the sine and cosine of the eccentric anomaly. :param M: Mean anomalies (rad). Shape: (n,). :type M: jnp.ndarray :param e: Eccentricity. Must satisfy 0 <= e < 1. :type e: float :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,). :rtype: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray] .. py:function:: he_E(M, e) Inverts Kepler's time equation for elliptical orbits with e > 0.78. :param M: Mean anomalies (rad). Shape: (n,). :type M: jnp.ndarray :param e: Eccentricity. Must satisfy 0 <= e < 1. :type e: float :returns: Eccentric anomalies (rad). Shape: (n,). :rtype: - E (jnp.ndarray) .. py:function:: he_E_trig(M, e) Inverts Kepler's time equation for elliptical orbits with e > 0.78. :param M: Mean anomalies (rad). Shape: (n,). :type M: jnp.ndarray :param e: Eccentricity. Must satisfy 0 <= e < 1. :type e: float :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,). :rtype: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray] .. py:function:: Etrig_1(E) When E <= pi_d_4. .. py:function:: Etrig_2(E) When E > pi_d_4 and E < three_pi_d_4. .. py:function:: Etrig_3(E) When E > pi_d_2 and E > three_pi_d_4. .. py:function:: Etrig(i, E) Apply the correct trigonometric function based on the index. .. py:function:: fast_sinE_cosE(E) Compute the sine and cosine of the eccentric anomaly using shortsin. .. py:function:: identity_solver(M, e) Returns M as E when e is 0. .. py:function:: identity_solver_trig(M, e) Returns M as E when e is 0.