jnkepler.jaxttv package

Submodules

jnkepler.jaxttv.conversion module

Coordinate transformations and orbital-element utilities.

jnkepler.jaxttv.conversion.a2cm_map(xast, vast, masses)

Vectorized version of astrocentric_to_cm. Takes similar arguments as astrocentric_to_cm but with additional array axes over which astrocentric_to_cm is mapped.

Original documentation:

conversion from astrocentric to CoM

Args:

xast: astrocentric positions (Norbit, xyz) vast: astrocentric velocities (Norbit, xyz) masses: masses of the bodies (Nbody,)

Returns:
tuple:

  • CoM positions (Nbody, xyz); now star (index 0) is added

  • CoM velocities (Nbody, xyz); now star (index 0) is added

jnkepler.jaxttv.conversion.astrocentric_to_cm(xast, vast, masses)[source]

conversion from astrocentric to CoM

Parameters:

  • xast – astrocentric positions (Norbit, xyz)

  • vast – astrocentric velocities (Norbit, xyz)

  • masses – masses of the bodies (Nbody,)

Returns:

  • CoM positions (Nbody, xyz); now star (index 0) is added

  • CoM velocities (Nbody, xyz); now star (index 0) is added

Return type:

tuple

jnkepler.jaxttv.conversion.cm_to_astrocentric(x, v, a, j)[source]

astrocentric x/v/a of the jth orbit (planet) from CoM x/v/a

Parameters:

  • x – CoM positions (Nstep, Nbody, xyz)

  • v – CoM velocities (Nstep, Nbody, xyz)

  • a – CoM accelerations (Nstep, Nbody, xyz)

  • j – orbit (planet) index

Returns:

  • astrocentric position of jth orbit (Nstep, xyz)

  • astrocentric velocity of jth orbit (Nstep, xyz)

  • astrocentric acceleration of jth orbit (Nstep, xyz)

Return type:

tuple

jnkepler.jaxttv.conversion.elements_to_xv(porb, ecc, inc, omega, lnode, u, mass)[source]

convert single set of orbital elements to position and velocity

Parameters:

  • porb – orbital period (day)

  • ecc – eccentricity

  • inc – inclination (radian)

  • omega – argument of periastron (radian)

  • lnode – longitude of ascending node (radian)

  • u – eccentric anomaly (radian)

  • mass – mass in Kepler’s 3rd law

Returns:

  • xout: positions (xyz, )

  • vout: velocities (xyz, )

Return type:

tuple

jnkepler.jaxttv.conversion.j2a_map(xjac, vjac, masses)

Vectorized version of jacobi_to_astrocentric. Takes similar arguments as jacobi_to_astrocentric but with additional array axes over which jacobi_to_astrocentric is mapped.

Original documentation:

conversion from Jacobi to astrocentric

Args:

xjac: jacobi positions (Norbit, xyz) vjac: jacobi velocities (Norbit, xyz) masses: masses of the bodies (Nbody,)

Returns:
tuple:

  • astrocentric positions (Norbit, xyz)

  • astrocentric velocities (Norbit, xyz)

jnkepler.jaxttv.conversion.jacobi_to_astrocentric(xjac, vjac, masses)[source]

conversion from Jacobi to astrocentric

Parameters:

  • xjac – jacobi positions (Norbit, xyz)

  • vjac – jacobi velocities (Norbit, xyz)

  • masses – masses of the bodies (Nbody,)

Returns:

  • astrocentric positions (Norbit, xyz)

  • astrocentric velocities (Norbit, xyz)

Return type:

tuple

jnkepler.jaxttv.conversion.m_to_u(M, ecc)[source]
jnkepler.jaxttv.conversion.reduce_angle(M)[source]

get angles between -pi and pi

Parameters:

M – angle (radian)

Returns:

angle mapped to [-pi, pi)

jnkepler.jaxttv.conversion.tic_to_m(tic, period, ecc, omega, t_epoch)[source]

convert time of inferior conjunction to mean anomaly M_epoch

Parameters:

  • tic – time of inferior conjunction

  • period – orbital period

  • ecc – eccentricity

  • omega – argument of periastron

  • t_epoch – time to which osculating elemetns are referred

Returns:

mean anomaly at t_epoch

jnkepler.jaxttv.conversion.tic_to_u(tic, period, ecc, omega, t_epoch)[source]

convert time of inferior conjunction to eccentric anomaly u

Parameters:

  • tic – time of inferior conjunction

  • period – orbital period

  • ecc – eccentricity

  • omega – argument of periastron

  • t_epoch – time to which osculating elemetns are referred

Returns:

eccentric anomaly at t_epoch

jnkepler.jaxttv.conversion.xv_to_elements(x, v, ki)[source]

convert position/velocity to elements

Parameters:

  • x – position and velocity (Norbit, xyz)

  • v – position and velocity (Norbit, xyz)

  • ki – ‘GM’ in Kepler’s 3rd law (Norbit); depends on what x/v mean (Jacobi, astrocentric, …)

Returns:

  • semi-major axis (au)

  • orbital period (day)

  • eccentricity

  • inclination (radian)

  • argument of periastron (radian)

  • longitude of ascending node (radian)

  • mean anomaly (radian)

Return type:

array

jnkepler.jaxttv.conversion.xvjac_to_xvacm(x, v, masses)[source]

Conversion from Jacobi to center-of-mass

Parameters:

  • xv – positions and velocities in Jacobi coordinates (Nstep, x or v, Norbit, xyz)

  • masses – masses of the bodies (Nbody,), solar unit

Returns:

  • xcm: positions in the CoM frame (Nstep, Norbit)

  • vcm: velocities in the CoM frame

  • acm: accelerations in the CoM frame

Return type:

tuple

jnkepler.jaxttv.conversion.xvjac_to_xvcm(x, v, masses)[source]

Conversion from Jacobi to center-of-mass

Parameters:

  • xv – positions and velocities in Jacobi coordinates (Nstep, x or v, Norbit, xyz)

  • masses – masses of the bodies (Nbody,), solar unit

Returns:

  • xcm: positions in the CoM frame (Nstep, Norbit)

  • vcm: velocities in the CoM frame

Return type:

tuple

jnkepler.jaxttv.findtransit module

Routines for finding transit times.

jnkepler.jaxttv.findtransit.find_transit_params_all(pidxarr, tcobsarr, t, xvjac, masses, nitr=5)[source]

Find transit times and phase-space states using Newton refinement.

Parameters:

  • pidxarr – Orbit indices starting from 0, with shape (Ntransit,).

  • tcobsarr – Flattened array of observed transit times of shape (Ntransit,).

  • t – Time array of shape (Nstep,).

  • xvjac – Jacobi positions and velocities of shape (Nstep, 2, Norbit, 3).

  • masses – Mass array of shape (Nbody,).

  • nitr – Number of Newton-Raphson iterations.

Returns:

Tuple containing

  • transit times as a one-dimensional array, and

  • positions, velocities, and Newton steps from the final iteration.

jnkepler.jaxttv.findtransit.find_transit_params_fast(pidxarr, tcobsarr, t, xvjac, masses)[source]

Find transit times and phase-space states using the fast algorithm.

Parameters:

  • pidxarr – Orbit indices starting from 0, with shape (Ntransit,).

  • tcobsarr – Flattened array of observed transit times of shape (Ntransit,).

  • t – Time array of shape (Nstep,).

  • xvjac – Jacobi positions and velocities of shape (Nstep, 2, Norbit, 3).

  • masses – Mass array of shape (Nbody,).

Returns:

Tuple containing

  • transit times as a one-dimensional array, and

  • positions, velocities, and the time corrections evaluated at tcobsarr.

jnkepler.jaxttv.findtransit.find_transit_times_all(pidxarr, tcobsarr, t, xvjac, masses, nitr=5)[source]

Find transit times for all requested transits using Newton refinement.

Parameters:

  • pidxarr – Orbit indices starting from 0, with shape (Ntransit,).

  • tcobsarr – Flattened array of observed transit times of shape (Ntransit,).

  • t – Time array of shape (Nstep,).

  • xvjac – Jacobi positions and velocities of shape (Nstep, 2, Norbit, 3).

  • masses – Mass array of shape (Nbody,).

  • nitr – Number of Newton-Raphson iterations.

Returns:

Transit times as a one-dimensional array.

jnkepler.jaxttv.findtransit.find_transit_times_fast(pidxarr, tcobsarr, t, xvjac, masses)[source]

Find transit times for all requested transits using the fast algorithm.

Parameters:

  • pidxarr – Orbit indices starting from 0, with shape (Ntransit,).

  • tcobsarr – Flattened array of observed transit times of shape (Ntransit,).

  • t – Time array of shape (Nstep,).

  • xvjac – Jacobi positions and velocities of shape (Nstep, 2, Norbit, 3).

  • masses – Mass array of shape (Nbody,).

Returns:

Transit times as a one-dimensional array.

jnkepler.jaxttv.findtransit.find_transit_times_kepler_all(pidxarr, tcobsarr, t, xvjac, masses, nitr=3)[source]

Find transit times via the legacy TTVFast-style interpolation scheme.

Note

This function is kept for backward compatibility and will be deprecated. It may fail for large dt.

Parameters:

  • pidxarr – Orbit indices starting from 0, with shape (Ntransit,).

  • tcobsarr – Flattened array of observed transit times of shape (Ntransit,).

  • t – Time array of shape (Nstep,).

  • xvjac – Jacobi positions and velocities of shape (Nstep, 2, Norbit, 3).

  • masses – Mass array of shape (Nbody,).

  • nitr – Number of iterations used in the Kepler interpolation step.

Returns:

Transit times as a one-dimensional array.

jnkepler.jaxttv.hermite4 module

4th-order Hermite integrator based on Kokubo, E., & Makino, J. 2004, PASJ, 56, 861 used for transit time computation

jnkepler.jaxttv.hermite4.hermite4_step_map(x, v, masses, dt)

Vectorized version of hermite4_step. Takes similar arguments as hermite4_step but with additional array axes over which hermite4_step is mapped.

Original documentation:

advance the system by a single predictor-corrector step

Args:

x: positions in CoM frame (Norbit, xyz) v: velocities in CoM frame (Norbit, xyz) masses: masses of the bodies (Nbody) dt: timestep

Returns:

new positions, new velocities, ‘new’ accelerations

jnkepler.jaxttv.hermite4.integrate_xv(x, v, masses, times)[source]

Hermite integration of the orbits

Parameters:

  • x – initial CoM positions (Norbit, xyz)

  • v – initial CoM velocities (Norbit, xyz)

  • masses – masses of the bodies (Nbody,), in units of solar mass

  • times – cumulative sum of time steps

Returns:

  • times (initial time omitted)

  • CoM position/velocity array (Nstep, x or v, Norbit, xyz)

Return type:

tuple

jnkepler.jaxttv.infer module

jnkepler.jaxttv.infer.get_flat_param_index(keys, npl, key, planet_idx)[source]

Return the flat index for a given parameter key and planet index.

Parameters:

  • keys – ordered list of parameter keys used in the flattened vector

  • npl – total number of planets in the model

  • key – parameter name, e.g. ‘period’, ‘tic’

  • planet_idx – zero-based planet index

Returns:

flat index into the concatenated parameter vector

Return type:

int

jnkepler.jaxttv.infer.make_phase_to_tic_transform(keys, npl, planet_idx, t_start)[source]

Create a transform that converts the ‘tic’ entry for one planet from orbital phase to time of inferior conjunction.

Assumes the optimizer parameter vector stores:

tic_param = phase

and converts it to:

tic = phase * period + t_start

Parameters:

  • keys – ordered list of parameter keys in the flattened vector

  • npl – total number of planets

  • planet_idx – zero-based index of the target planet

  • t_start – reference start time

Returns:

transform(p_flat) -> transformed p_flat

Return type:

callable

jnkepler.jaxttv.infer.scale_pdic(pdic, param_bounds)[source]

scale parameters using bounds

Parameters:

  • pdic – dict of physical parameters

  • param_bounds – dictionary of (lower bound array, upper bound array)

Returns:

dictionary of scaled parameters

Return type:

dict

jnkepler.jaxttv.infer.ttv_default_parameter_bounds(jttv, npl=None, t0_guess=None, p_guess=None, dtic=0.2, dp_frac=0.01, emax=0.2, mmin=1e-07, mmax=0.001)[source]

Get parameter bounds for TTV optimization.

Parameters:

  • jttv – JaxTTV object.

  • npl (int, optional) – Number of planets. Defaults to jttv.nplanet if None.

  • t0_guess (array-like, optional) – Initial guess for transit times, length must be npl.

  • p_guess (array-like, optional) – Initial guess for orbital periods, length must be npl.

  • dtic (float, optional) – Half-width of bounds around t0_guess for transit time.

  • dp_frac (float, optional) – Fractional width of bounds around p_guess for period.

  • emax (float, optional) – Maximum ecosw/esinw bound.

  • mmin (float, optional) – Minimum mass bound.

  • mmax (float, optional) – Maximum mass bound.

Returns:

Dictionary of parameter bounds with keys as parameter names

and values as [lower_bound_array, upper_bound_array].

Return type:

dict

jnkepler.jaxttv.infer.ttv_optim_curve_fit(*args, **kwargs)[source]

Deprecated wrapper for ttv_optim_least_squares.

jnkepler.jaxttv.infer.ttv_optim_least_squares(jttv, param_bounds_, pinit=None, n_start=1, loss='linear', loss_kwargs=None, jac=False, diff_mode='auto', plot=True, save=None, transit_orbit_idx=None, random_state=None, max_nfev=None, param_transform=None, return_optspace=False)[source]

Simple TTV fit using scipy.optimize.least_squares.

Parameters:

  • jttv – JaxTTV object

  • param_bounds – bounds for parameters, dict of {key: (lower, upper)}

  • pinit – initial guess of parameters (dict)

  • n_start – number of initial guesses when pinit is not provided. If pinit is given, a single optimization is run. Otherwise, multi-start randomizes only the initial planet masses. The first start uses a deterministic midpoint-like initial guess, and later starts randomize only lnpmass.

  • loss –

    loss specification for scipy.optimize.least_squares. This can be:

    • a built-in SciPy loss string such as ‘linear’, ‘soft_l1’, ‘huber’, ‘cauchy’, ‘arctan’

    • ’student_t’ to use Student-t negative log likelihood

    • a custom callable compatible with scipy.optimize.least_squares

  • loss_kwargs –

    optional dict for loss-specific options. Default is None. For loss=’student_t’:

    {‘nu’: …, ‘scale’: …} Default is {‘nu’: 4.0, ‘scale’: 1.0}.

    For built-in or custom non-linear losses:

    {‘f_scale’: …} can be passed through to least_squares.

    For loss=’student_t’, the custom loss returns 2*NLL per data point, so least_squares.cost corresponds to the total NLL. For loss=’linear’, cost = 0.5 * chi2.

  • jac – if True, use a JAX-based analytic Jacobian for the residual function.

  • diff_mode –

    differentiation mode for the analytic Jacobian when jac=True. Must be one of:

    • ’auto’: use ‘fwd’ for transit_time_method=’fast’ and ‘rev’ for transit_time_method=’newton’

    • ’rev’

    • ’fwd’

    Note: for transit_time_method=’newton’, ‘fwd’ is overridden to ‘rev’ because forward-mode differentiation may fail there.

  • plot – if True, TTV models are plotted with data.

  • save – path to save TTV plots.

  • transit_orbit_idx – list of indices to specify which planets are transiting (needed when non-transiting planets are included)

  • random_state – int or np.random.RandomState, for reproducibility. If None, multi-start initializations are not reproducible across runs.

  • max_nfev – maximum number of function evaluations

  • param_transform – optional callable mapping optimizer-space parameters to model-space parameters before residual evaluation. This is useful, for example, when optimizing orbital phase instead of time of inferior conjunction for a non-transiting planet. If jac=True, this should be JAX-compatible.

  • return_optspace – if True, also return the best-fit parameter dictionary in optimizer space as a second output. This is useful for warm starts when param_transform is used.

Returns:

  • If return_optspace=False (default), returns the best-fit JaxTTV parameter dictionary in model space.

  • If return_optspace=True, returns (pdic_opt_modelspace, pdic_opt_optspace).

Return type:

dict or tuple

jnkepler.jaxttv.infer.unscale_pdic(pdic_scaled, param_bounds)[source]

unscale parameters using bounds

Parameters:

  • pdic – dict of scaled parameters

  • param_bounds – dictionary of (lower bound array, upper bound array)

Returns:

dictionary of physical parameters in original scales

Return type:

dict

jnkepler.jaxttv.information module

jnkepler.jaxttv.information.information(jttv, pdic, keys, param_bounds=None)[source]

compute Fisher information matrix for iid gaussian likelihood

Parameters:

  • jttv – JaxTTV object

  • pdic – dict containing parameters; keys must contain {ecosw, esinw, period, tic, lnode, cosi} and {pmass or lnpmass}

  • keys – parameter keys for computing fisher matrix

  • param_bounds – if not None, compute information matrix for parameters scaled by prior widths

Returns:

information matrix computed as grad.T Sigma_inv grad

jnkepler.jaxttv.information.scale_information(matrix, param_bounds, keys)[source]

get information matrix for scaled parameters

Note

This will be deprecated; function ‘information’ above with param_bounds argument does the same.

Parameters:

  • matrix – information matrix

  • param_bounds – dict containing bounds for parameters, 0: lower, 1: upper

  • keys – parameter keys (normally [‘ecosw’, ‘esinw’, ‘mass’, ‘period’, ‘tic’])

Returns:

information matrix for scaled parameters scaled by 1/(upper - lower)

jnkepler.jaxttv.jaxttv module

Main class for modeling TTVs.

class jnkepler.jaxttv.jaxttv.JaxTTV(t_start, t_end, dt, tcobs, p_init, errorobs=None, print_info=True, nitr_kepler=10, transit_time_method='fast', nitr_transit=5)[source]

Bases: Nbody

main class for TTV analysis

check_residuals(par_dict=None, tcmodel=None, jitters=None, student=True, normalize_residuals=True, plot=True, fit_mean=False, save=None, xrange=5)[source]

check the distribution of residuals, fit them with Student’s t distritbution

Parameters:

  • par_dict – dict containing parameters

  • tcmodel – transit time model (1D array)

Returns:

  • dictionary: mean of reisudals, SD of residuals

  • dictionary: parameters of Student’s t dist (lndf, lnvar, mean)

Return type:

tuple

check_timing_precision(par_dict, dtfrac=0.001, nitr_transit=10, nitr_kepler=10)[source]

Compare transit times against a smaller-step Newton calculation.

Note

The baseline calculation uses the current self.transit_time_method at the nominal timestep, while the smaller-step reference always uses the Newton transit finder.

Parameters:

  • params – JaxTTV parameter array

  • dtfrac – (innermost period) * dtfrac is used for the comparison integration

  • nitr_transit – number of Newton iterations for the smaller-step reference calculation

  • nitr_kepler – number of Kepler iterations for the smaller-step reference calculation

Returns:

  • model transit times from the baseline calculation at the nominal timestep

  • model transit times from the smaller-step Newton calculation

Return type:

tuple

get_transit_times_all(par_dict, t_start=None, t_end=None, dt=None, transit_orbit_idx=None)[source]

compute all model transit times between t_start and t_end

Note

This function always uses the Newton transit finder, regardless of self.transit_time_method.

Parameters:

par_dict – dict containing parameters

Returns:

  • 1D flattened array of transit times

  • fractional energy change

Return type:

tuple

get_transit_times_all_list(par_dict, truncate=True, transit_orbit_idx=None)[source]

compute all transit times and retunrs a list

Parameters:

  • par_dict – dict containing parameters

  • truncate – if True, only compute transit times up to the last observed time instead of t_end

Returns:

list of model transit times of length Nplanet; each element is an array of model transit times (length varies for each planet)

Return type:

list

get_transit_times_and_rvs_obs(par_dict, times_rv, transit_orbit_idx=None)[source]

compute model transit times and stellar RVs

Note

The transit times are computed using self.transit_time_method.

Parameters:

  • par_dict – dict containing parameters

  • times_rv – times at which stellar RVs are evaluated

  • transit_orbit_idx – array of idx of the planet (orbit) for each transit times should be evaulated, starting from 0. This must be specified when non-transiting planets exist. If None, all the planets are assumed to be transiting; this is equivalent to setting transit_orbit_idx = jnp.arange(nplanet).

Returns:

  • 1D array: flattened array of transit times

  • 1D array: stellar RVs (m/s)

  • float: fractional energy change

Return type:

tuple

get_transit_times_obs(par_dict, transit_orbit_idx=None)[source]

compute model transit times

Note

This function returns only transit times that are closest to the observed ones. The algorithm is controlled by self.transit_time_method. To get all the transit times, use get_transit_times_all instead.

Parameters:

  • par_dict – dict containing parameters

  • transit_orbit_idx – array of idx of the planet (orbit) for each transit times should be evaulated, starting from 0. This must be specified when non-transiting planets exist. If None, all the planets are assumed to be transiting; this is equivalent to setting transit_orbit_idx = jnp.arange(nplanet).

Returns:

  • 1D array: flattened array of transit times

  • float: fractional energy change

Return type:

tuple

linear_ephemeris()[source]

(Re)derive linear ephemeris when necessary

Returns:

  • array of t0 from linear fitting

  • array of P from linear fitting

Return type:

tuple

plot_model(tcmodellist, tcobslist=None, errorobslist=None, t0_lin=None, p_lin=None, tcmodelunclist=None, tmargin=None, save=None, marker=None, ylims=None, ylims_residual=None, tcmodellist2=None, unit=1440.0, lw=1, ylabel='TTV (min)', xlabel='transit time (day)')[source]

plot transit time model

Parameters:

  • tcmodellist – list of the arrays of model transit times for each planet

  • tcobslist – list of the arrays of observed transit times for each planet

  • errorobslist – list of the arrays of observed transit time errors for each planet

  • t0_lin – linear ephemeris used to show TTVs (n_planet,)

  • p_lin – linear ephemeris used to show TTVs (n_planet,)

  • tcmodelunclist – model uncertainty (same format as tcmodellist)

  • tcmodellist2 (list of arrays, optional) – an optional second set of model transit times to overplot

  • tmargin – margin in x axis

  • save – if not None, plot is saved as “save_planet#.png”

  • marker – marker for model

  • unit – TTV unit (defaults to minutes)

  • ylabel – axis labels in the plots

  • xlabel – axis labels in the plots

  • ylims – y ranges in the plots

  • ylims_residual – y ranges in the plots

sample_means_and_stds(samples, N=50, truncate=True, original_models=False)[source]

compute mean and standard deviation of transit time models from HMC samples

Parameters:

  • samples – dictionary containing parameter samples (output of mcmc.get_samples())

  • N – number of samples to be used for calculation

  • truncate – if True, only compute transit times up to the last observed time instead of t_end

  • models (original) – if True, just returns a list of transit-time models

Returns:

means and standard deviations of transit time models, list of length(nplanet)

set_tcobs(tcobs, p_init, errorobs=None, print_info=True)[source]

set observed transit times

Note

JaxTTV returns transit times that are closest to the observed times set here, rather than all the transit times between t_start and t_end.

Parameters:

  • tcobs – list of the arrays of transit times for each planet

  • p_init – initial guess for the mean orbital period of each planet

  • errorobs – transit time error (currently assumed to be Gaussian), same format as tcobs

Returns:

attributes of JaxTTV class

tcall_linear(t_start, t_end, truncate=False)[source]

information on all linear transit times between t_start and t_end

Parameters:

  • t_start – start time of integration

  • t_end – end time of integration

Returns:

  • 1D array of orbit (planet) index starting from 1

  • list of linear transit times between t_start and t_end

  • 1D flattend version of tcall_linear

Return type:

tuple

property transit_time_method

Transit-time algorithm used for observed transits.

class jnkepler.jaxttv.jaxttv.Nbody(t_start, t_end, dt, nitr_kepler=10)[source]

Bases: object

superclass for nbody analysis

jnkepler.jaxttv.markley module

Kepler equation solver based on Markley (1995)

jnkepler.jaxttv.markley.get_E(M, e)[source]

compute eccentric anomaly given mean anomaly and eccentricity

Parameters:

  • M – mean anomaly (should be between -pi and pi)

  • e – eccentricity

Returns:

eccentric anomaly

jnkepler.jaxttv.rv module

jnkepler.jaxttv.rv.rv_from_xvjac(times_rv, times, xvjac, masses)[source]

compute stellar RV .. note:: This function will be more efficient if interpolation is performed before conversion to CM frame

Parameters:

  • times_rv – times at which RVs are evaluated

  • times – times for n-body integration (Ntime,)

  • xvjac – jacobi positions and velocities (Ntime, x or v, Norbit, xyz)

  • masses – masses of the bodies (Nbody,), solar unit

Returns:

stellar RVs at times_rvs (m/s), positive when the star is moving away

Return type:

array

jnkepler.jaxttv.symplectic module

Symplectic integrator. JAX-native implementation following the symplectic approach used in TTVFast (https://github.com/kdeck/TTVFast).

jnkepler.jaxttv.symplectic.integrate_xv(x, v, masses, times, nitr=10)[source]

Integrate Jacobi positions and velocities using the Wisdom-Holman map.

This function currently assumes that times are uniformly spaced. Internally, the map is initialized using the first time interval and then advanced with a fixed step size, so non-uniform times are not supported.

Parameters:

  • x – initial Jacobi positions in Cartesian coordinates (Norbit, xyz)

  • v – initial Jacobi velocities in Cartesian coordinates (Norbit, xyz)

  • masses – masses of the bodies (Nbody,), in units of solar mass

  • times – 1D array of uniformly spaced output times

  • nitr – number of iterations in Kepler’s equation solver

Returns:

  • times at the middle of each map step

  • integrated Jacobi positions and velocities

Return type:

tuple

jnkepler.jaxttv.symplectic.kepler_kick_kepler_map(xjac, vjac, masses, dt, nitr=10)[source]

vmap version of kepler_step + nbody_kicks + kepler_step.

This applies a full KDK-like step written in the order kepler(dt/2) -> kick(dt) -> kepler(dt/2) to each element along the first axis.

Parameters:

  • xjac – Jacobi positions (Ntime, Norbit, xyz)

  • vjac – Jacobi velocities (Ntime, Norbit, xyz)

  • masses – masses of the bodies (Nbody,), in units of solar mass

  • dt – common full time step

Returns:

new Jacobi positions and velocities (Ntime, x or v, Norbit, xyz)

jnkepler.jaxttv.symplectic.kepler_step_map(xjac, vjac, masses, dt, nitr=10)[source]

vmap version of kepler_step; map along the first axis (Ntime)

Parameters:

  • xjac – Jacobi positions (Ntime, Norbit, xyz)

  • vjac – Jacobi velocities (Ntime, Norbit, xyz)

  • masses – masses of the bodies (Nbody,), in units of solar mass

  • dt – common time step

Returns:

new Jacobi positions and velocities (Ntime, x or v, Norbit, xyz)

jnkepler.jaxttv.symplectic.kick_kepler_map(xjac, vjac, masses, dt, nitr=10)[source]

vmap version of nbody_kicks + kepler_step; map along the first axis (Ntime)

Parameters:

  • xjac – jacobi positions (Ntime, Norbit, xyz)

  • vjac – jacobi velocities (Ntime, Norbit, xyz)

  • masses – masses of the bodies (Nbody,), in units of solar mass

  • dt – common time step

Returns:

new jacobi positions and velocities (Ntime, x or v, Norbit, xyz)

jnkepler.jaxttv.ttvfastutils module

Utilities for interfacing with TTVFast models.

This module provides helper functions for constructing TTVFast-compatible parameter sets and model objects. These functions are intended solely for comparison, validation, and conversion purposes, and are not used in the core JAX-based modeling or inference pipelines.

jnkepler.jaxttv.ttvfastutils.get_ttvfast_model(pdic, num_planets, t_start, dt, t_end, skip_planet_idx=[])[source]

compute transit times using ttvfast-python

Parameters:

  • pdic – parameter dataframe from params_for_ttvfast

  • num_planets – number of planets

  • t_start – start time of integration

  • dt – integration time step

  • t_end – end time of integration

Returns:

  • list of transit epochs

  • list of transit times

Return type:

tuple

jnkepler.jaxttv.ttvfastutils.get_ttvfast_model_all(pdic, num_planets, t_start, dt, t_end, skip_planet_idx=[])[source]

compute transit times using ttvfast-python

Parameters:

  • pdic – parameter dataframe from params_for_ttvfast

  • num_planets – number of planets

  • t_start – start time of integration

  • dt – integration time step

  • t_end – end time of integration

  • skip_planet_idx – list of planet idx to be skipped from output (starting from 0)

Returns:

  • list of transit epochs

  • list of transit times

  • list of sky-plane distances (au)

  • list of sky-plane velocities (au/day)

Return type:

tuple

jnkepler.jaxttv.ttvfastutils.get_ttvfast_model_rv(pdic, num_planets, t_start, dt, t_end, times_rv, skip_planet_idx=[])[source]

compute transit times using ttvfast-python

Parameters:

  • pdic – parameter dataframe from params_for_ttvfast

  • num_planets – number of planets

  • t_start – start time of integration

  • dt – integration time step

  • t_end – end time of integration

  • times_rv – times at which RVs are evaluated

Returns:

  • list of transit epochs

  • list of transit times

  • array of RVs

Return type:

tuple

jnkepler.jaxttv.ttvfastutils.params_for_ttvfast(samples, t_epoch, num_planets, WHsplit=True, angles_in_degrees=True, names=['period', 'eccentricity', 'inclination', 'argument', 'longnode', 'mean_anomaly'])[source]

convert JaxTTV samples into TTVFast (or other) format

Parameters:

  • samples – mcmc.get_samples()

  • t_epoch – time at which osculating elements are defined

  • num_planets – number of planets

  • WHsplit – True for TTVFast; False for e.g. REBOUND (cf. Section 2.2 of Rein & Tamayo 2015, MNRAS 452, 376)

  • angles_in_degrees – If True, angles are returned in degrees

Returns:

dataframe containing parameters

jnkepler.jaxttv.utils module

Internal utilities for parameter conversion, initialization, and diagnostics.

jnkepler.jaxttv.utils.convert_elements(par_dict, t_epoch, WHsplit=False)[source]

convert JaxTTV elements into more normal sets of parameters

Parameters:

  • par_dict – parameter dict

  • t_epoch – epoch at which elements are defined

  • WHsplit – elements are converted to coordinates assuming Wisdom-Holman splitting. This should be True when the output is used for TTVFast.

Returns:

  • array: (semi-major axis, period, eccentricity, inclination, argument of periastron, longitude of ascending node, mean anomaly) x (orbits), angles are in radians

  • mass array

Return type:

tuple

jnkepler.jaxttv.utils.dict_to_params(pdic, npl, keys)[source]

Inverse of params_to_dict when each value has length npl.

Parameters:

  • pdic (Mapping[str, array_like]) – Parameter dict, e.g. {‘a’: arr(len=npl), ‘b’: arr(len=npl), …}.

  • keys (Sequence[str]) – Order of parameters; concatenation follows this order.

Returns:

1D parameter array of length len(keys) * npl.

Return type:

ndarray or jax.numpy.ndarray

Raises:

  • KeyError – if a key in keys is missing from pdic.

  • ValueError – if value lengths are inconsistent across keys.

jnkepler.jaxttv.utils.elements_to_pdic(elements, masses, outkeys=None, force_coplanar=True)[source]

convert JaxTTV elements/masses into dictionary

Note

This function is for v<0.2.

Parameters:

  • elements – Jacobi orbital elements (period, ecosw, esinw, cosi, Omega, T_inf_conjunction)

  • masses – masses of the bodies (Nbody,)

  • outkeys – if specified only include these keys in the output

  • force_coplanar – if True, set incl=pi/2 and lnode=0

Returns:

dicionary of the parameters

jnkepler.jaxttv.utils.em_to_dict(elements, masses)[source]

convert arrays of elements and masses in v0.1.0 to parameter dict

Note

This function is mainly for running tests; no longer needed for v>=0.2.

Parameters:

  • elements – elements (JaxTTV format)

  • masses – masses of star + planets (solar units)

Returns:

parameter dict

jnkepler.jaxttv.utils.find_nearest_idx(arr1, arr2)[source]

pick up elements of arr1 nearest to each element in arr2

Parameters:

  • arr1 – array from which elements are picked up

  • arr2 – array of the values for which nearest matches are searched

Returns:

indices of arr1 nearest to each element in arr2

jnkepler.jaxttv.utils.find_nearest_idx_sorted(arr1, arr2)[source]

Return indices in sorted arr1 nearest to each value in arr2.

This implementation assumes arr1 is sorted in ascending order and uses binary search via jnp.searchsorted (O(len(arr2) * log(len(arr1)))).

Ties are resolved by choosing the left candidate.

Parameters:

  • arr1 – 1D array sorted in ascending order.

  • arr2 – 1D array of query values.

Returns:

1D integer array of indices i such that arr1[i] is the nearest value to each element of arr2.

jnkepler.jaxttv.utils.get_energy_diff(xva, masses)[source]

compute fractional energy change given integration result

Parameters:

  • xva – posisions, velocities, accelerations in CoM frame (Nstep, x or v or a, Norbit, xyz)

  • masses – masses of the bodies (Nbody,)

Returns:

fractional change in total energy

jnkepler.jaxttv.utils.get_energy_diff_jac(xvjac, masses, dt)[source]

compute fractional energy change given integration result

Parameters:

  • xvjac – Jacobi posisions and velocities (Nstep, x or v, Norbit, xyz)

  • masses – masses of the bodies (Nbody,)

Returns:

fractional change in total energy

jnkepler.jaxttv.utils.initialize_jacobi_xv(par_dict, t_epoch)[source]

compute initial position/velocity from parameter dict

Note

Here the elements are interpreted as Jacobi elements using the total interior mass (see Section 2.2 of Rein & Tamayo 2015).

Parameters:

  • par_dict – parameter dictionary that needs to contain - either (ecosw, esinw) or (e, omega) - cosi (set to be 0 if not specified) - lnode (set to be 0 if not specified) - either (time of inferior conjunction) or (mean anomaly) - stellar mass (set to be 1 if not specified), solar unit - either (planetary mass) or (ln planetary mass), solar unit, former is used when both are provided

  • t_epoch – epoch at which elements are defined

Returns:

  • Jacobi positions at t_epoch (Norbit, xyz)

  • Jacobi velocities at t_epoch (Norbit, xyz)

  • masses: 1D array of stellar and planetary masses (Nbody,)

Return type:

tuple

jnkepler.jaxttv.utils.params_to_dict(params, npl, keys)[source]

convert 1D parameter array into parameter dict

Parameters:

  • array (parameter) – [arr for key1, arr for key2, …] where len(arr) is the number of planets

  • npl – number of planets

  • keys – parameter keys [key1, key2, …]

Returns:

parameter dict

jnkepler.jaxttv.utils.params_to_elements(params, npl)[source]

convert JaxTTV parameter array into element and mass arrays

Parameters:

  • params – JaxTTV parameter array

  • npl – number of orbits (planets)

Returns:

  • Jacobi orbital elements (period, ecosw, esinw, cosi, Omega, T_inf_conjunction)

  • ln(masses) of the bodies (Nbody,)

Return type:

tuple

Module contents