RateJAX#

class exo_skryer.rate_jax.RateJAX(thermo: NASA9ThermoJAX, C: float = 0.00025, N: float = 0.0001, O: float = 0.0005, fHe: float = 0.0)[source]#

Bases: object

RATE-style thermochemical equilibrium solver implemented in JAX.

This class computes equilibrium abundances for a reduced H/C/N/O chemistry network over a 1D (T, p) profile. It is designed to be usable inside JIT-compiled forward models:

Uses jax.vmap to solve each layer independently
Avoids SciPy; uses optimistix for root finding where needed
Returns a VMR dictionary keyed by species name

Parameters:

thermoNASA9ThermoJAX: NASA-9 thermo evaluator created by load_nasa9_cache.
C, N, Ofloat: Elemental abundances (number ratios relative to H₂, following the original RATE conventions).
fHefloat: Helium fraction factor used to compute He from H-bearing species.

Methods Summary

`HCNO_poly8_CO`(f, k1, k2, k3, k4, k5, k6)	JAX version of original HCNO_poly8_CO (CO is the root variable).
`HCNO_poly8_H2O`(f, k1, k2, k3, k4, k5, k6)	JAX version of original HCNO_poly8_H2O (H2O is the root variable).
`HCO_poly6_CO`(f, k1, k2, k3, k4)	HCO chemistry, polynomial in CO.
`HCO_poly6_H2O`(f, k1, k2, k3, k4)	HCO chemistry, polynomial in H2O.
`g_over_RT`(spec, T)
`kprime0`(T, p)	Equilibrium constant for hydrogen dissociation: H2 ↔ 2H
`kprime1`(T, p)	Equilibrium constant for methane-water reaction: CH₄ + H₂O ↔ CO + 3H₂
`kprime2`(T)	Equilibrium constant for carbon dioxide reduction: CO₂ + H₂ ↔ CO + H₂O
`kprime3`(T, p)	Equilibrium constant for acetylene formation: 2CH₄ ↔ C₂H₂ + 3H₂
`kprime4`(T, p)	Equilibrium constant for ethylene-acetylene: C₂H₄ ↔ C₂H₂ + H₂
`kprime5`(T, p)	Equilibrium constant for ammonia dissociation: 2NH₃ ↔ N₂ + 3H₂
`kprime6`(T, p)	Equilibrium constant for HCN formation: NH₃ + CH₄ ↔ HCN + 3H₂
`newton_raphson_bounded`(A, guess, vmax[, ...])	Robust polynomial root finding with bounded domain using Optimistix.
`solve_profile`(T, p[, return_diagnostics])	Solve thermochemical equilibrium across a 1D T-p profile.
`solve_rest`(H2O, CO, f, k1, k2, k3, k4, k5, k6)	JAX version of solve_rest for a single layer.
`top`(T, C, N, O)	Turnover pressure: transition between CO-dominated and H2O-dominated chemistry.

Methods Documentation

HCNO_poly8_CO(f, k1, k2, k3, k4, k5, k6)[source]#: JAX version of original HCNO_poly8_CO (CO is the root variable).

HCNO_poly8_H2O(f, k1, k2, k3, k4, k5, k6)[source]#: JAX version of original HCNO_poly8_H2O (H2O is the root variable).

HCO_poly6_CO(f, k1, k2, k3, k4)[source]#: HCO chemistry, polynomial in CO. Now returns 7 coefficients (last one is 0.0) for JAX compatibility.

HCO_poly6_H2O(f, k1, k2, k3, k4)[source]#: HCO chemistry, polynomial in H2O. Now returns 7 coefficients (last one is 0.0) for JAX compatibility.

g_over_RT(spec: str, T: Array) → Array[source]#

kprime0(T: Array, p: Array) → Array[source]#

Equilibrium constant for hydrogen dissociation: H2 ↔ 2H

K’₀ = exp(-ΔG/RT) / p where ΔG = 2·G(H) - G(H₂)

Parameters:

Tarray: Temperature [K]
parray: Pressure [bar]

Returns:

K’₀array: Modified equilibrium constant [bar⁻¹]

kprime1(T: Array, p: Array) → Array[source]#

Equilibrium constant for methane-water reaction: CH₄ + H₂O ↔ CO + 3H₂

K’₁ = exp(-ΔG/RT) / p² where ΔG = G(CO) + 3·G(H₂) - G(CH₄) - G(H₂O)

This is the key reaction controlling the C/O ratio in hot atmospheres.

Parameters:

Tarray: Temperature [K]
parray: Pressure [bar]

Returns:

K’₁array: Modified equilibrium constant [bar⁻²]

kprime2(T: Array) → Array[source]#

Equilibrium constant for carbon dioxide reduction: CO₂ + H₂ ↔ CO + H₂O

K’₂ = exp(-ΔG/RT) where ΔG = G(CO) + G(H₂O) - G(CO₂) - G(H₂)

Parameters:

Tarray: Temperature [K]

Returns:

K’₂array: Equilibrium constant [dimensionless]

kprime3(T: Array, p: Array) → Array[source]#

Equilibrium constant for acetylene formation: 2CH₄ ↔ C₂H₂ + 3H₂

K’₃ = exp(-ΔG/RT) / p² where ΔG = G(C₂H₂) + 3·G(H₂) - 2·G(CH₄)

Important for high-C/O and high-temperature atmospheres.

Parameters:

Tarray: Temperature [K]
parray: Pressure [bar]

Returns:

K’₃array: Modified equilibrium constant [bar⁻²]

kprime4(T: Array, p: Array) → Array[source]#

Equilibrium constant for ethylene-acetylene: C₂H₄ ↔ C₂H₂ + H₂

K’₄ = exp(-ΔG/RT) / p where ΔG = G(C₂H₂) + G(H₂) - G(C₂H₄)

Parameters:

Tarray: Temperature [K]
parray: Pressure [bar]

Returns:

K’₄array: Modified equilibrium constant [bar⁻¹]

kprime5(T: Array, p: Array) → Array[source]#

Equilibrium constant for ammonia dissociation: 2NH₃ ↔ N₂ + 3H₂

K’₅ = exp(-ΔG/RT) / p² where ΔG = G(N₂) + 3·G(H₂) - 2·G(NH₃)

Dominant nitrogen chemistry reaction in hot atmospheres.

Parameters:

Tarray: Temperature [K]
parray: Pressure [bar]

Returns:

K’₅array: Modified equilibrium constant [bar⁻²]

kprime6(T: Array, p: Array) → Array[source]#

Equilibrium constant for HCN formation: NH₃ + CH₄ ↔ HCN + 3H₂

K’₆ = exp(-ΔG/RT) / p² where ΔG = G(HCN) + 3·G(H₂) - G(NH₃) - G(CH₄)

Important when both N and C are abundant at high temperatures.

Parameters:

Tarray: Temperature [K]
parray: Pressure [bar]

Returns:

K’₆array: Modified equilibrium constant [bar⁻²]

classmethod newton_raphson_bounded(A: Array, guess: float, vmax: float, xtol: float = 1e-10, imax: int = 80, kmax: int = 10) → Array[source]#

Robust polynomial root finding with bounded domain using Optimistix.

Uses Newton’s method with automatic differentiation. Tries multiple initial guesses with decreasing scales if needed, then clamps result to [0, vmax].

Parameters:

Aarray, shape (n,): Polynomial coefficients (constant to highest degree)
guessfloat: Initial guess for root
vmaxfloat: Maximum valid value for root
xtolfloat: Relative/absolute tolerance for convergence
imaxint: Maximum iterations per attempt
kmaxint: Maximum number of retry attempts with scaled guesses

Returns:

rootscalar array: Root of polynomial, clamped to [0, vmax]

solve_profile(T: Array, p: Array, return_diagnostics: bool = False) → Dict[str, Array] | Tuple[Dict[str, Array], Dict][source]#

Solve thermochemical equilibrium across a 1D T-p profile.

Parameters:

T1D array [K]: Temperature profile
p1D array [bar]: Pressure profile
return_diagnosticsbool, optional: If True, return (vmr_dict, diagnostics) with convergence info

Returns:

vmrdict[str, jnp.ndarray]: Keys: self.species, each value shape = (nlayers,)
diagnosticsdict, optional: Only returned if return_diagnostics=True. Contains: - ‘n_layers’: number of layers - ‘T_range’: (min, max) temperature - ‘p_range’: (min, max) pressure

Raises:

ValueError: If inputs have incompatible shapes or invalid values

solve_rest(H2O: float, CO: float, f: float, k1: float, k2: float, k3: float, k4: float, k5: float, k6: float) → Array[source]#: JAX version of solve_rest for a single layer. Returns [H2O, CH4, CO, CO2, NH3, C2H2, C2H4, HCN, N2]

static top(T: Array, C: float, N: float, O: float) → Array[source]#

Turnover pressure: transition between CO-dominated and H2O-dominated chemistry.

Computes the pressure where CO and H2O abundances become comparable, based on a polynomial fit to thermochemical equilibrium calculations (Lodders & Fegley 2002).

Parameters:

Tarray: Temperature [K]
Cfloat: Carbon elemental abundance (number ratio relative to H2)
Nfloat: Nitrogen elemental abundance (number ratio relative to H2)
Ofloat: Oxygen elemental abundance (number ratio relative to H2)

Returns:

P_turnoverarray: Turnover pressure [bar], where CO/H2O ~ 1