Performance benchmarks#
This notebook records rough wall-clock timings for squishyplanet’s transit models: a spherical planet, an oblate planet, and a ringed planet, for both forward evaluations and gradients.
A few caveats to keep in mind:
The absolute numbers below were generated on a single Apple Silicon CPU core and depend heavily on the machine. The ratios between the models are the meaningful takeaway.
All of the lightcurve functions are just-in-time compiled by
JAX. The first call to each pays a one-time compilation cost, which we report separately from the steady-state per-call cost.Timings jitter from run to run; we report the median of repeated calls.
[8]:
import time
import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
import numpy as np
from squishyplanet import OblateSystem, RingedSystem
def first_call(fn):
"""Wall-clock time of the first (compiling) call, in seconds."""
start = time.perf_counter()
jax.block_until_ready(fn())
return time.perf_counter() - start
def per_call(fn, n_repeats=100):
"""Median wall-clock time per call after compilation, in ms."""
jax.block_until_ready(fn())
times = []
for _ in range(n_repeats):
start = time.perf_counter()
jax.block_until_ready(fn())
times.append(time.perf_counter() - start)
return np.median(times) * 1e3
Forward models#
We’ll put all three planets on the same orbit and evaluate 1,000 points that all fall within the transit – the worst case for the transit solver, since out-of-transit points are masked and cost almost nothing (demonstrated below).
[9]:
base = {
"t_peri": -2.5,
"period": 10.0,
"a": 5.0,
"i": jnp.pi / 2 - 0.01,
"r": 0.1,
"obliq": 0.4,
"prec": 0.6,
"tidally_locked": False,
"ld_u_coeffs": jnp.array([0.3, 0.1]),
}
times_in_transit = jnp.linspace(-0.3, 0.3, 1000)
systems = {
"spherical": OblateSystem(times=times_in_transit, **base),
"oblate": OblateSystem(times=times_in_transit, f1=0.1, f2=0.05, **base),
"ringed": RingedSystem(
times=times_in_transit, ring_inner_r=0.15, ring_outer_r=0.3, **base
),
}
forward_times = {}
print(f"{'model':<12}{'compile + first call [s]':>26}{'per call [ms]':>16}")
for name, system in systems.items():
compile_s = first_call(system.lightcurve)
steady_ms = per_call(system.lightcurve)
forward_times[name] = steady_ms
print(f"{name:<12}{compile_s:>26.2f}{steady_ms:>16.2f}")
model compile + first call [s] per call [ms]
spherical 0.18 3.02
oblate 0.17 2.95
ringed 0.54 49.58
[10]:
ratio = forward_times["ringed"] / forward_times["spherical"]
print(f"ringed / spherical: {ratio:.1f}x")
ringed / spherical: 16.4x
The spherical and oblate models cost the same: both trace the same code path, which handles an arbitrary projected ellipse. The ringed model is roughly an order of magnitude slower per in-transit point. That is the price of the extra geometry solved at every timestep: intersections between the planet, both ring edges, and the stellar limb, plus up to five boundary-integral terms instead of one (see engine.ringed_transit). Reducing this gap is a known optimization target.
Out-of-transit points are nearly free#
Every timestep is first checked against a bounding circle around the planet + ring; points that cannot be in transit skip the solver entirely. Here is the same ringed system evaluated at 1,000 points spread over a window about six times wider than the transit, so only ~17% of the points are in transit:
[11]:
times_wide = jnp.linspace(-2.5, 2.5, 1000)
ringed_wide = RingedSystem(
times=times_wide, ring_inner_r=0.15, ring_outer_r=0.3, **base
)
first_call(ringed_wide.lightcurve)
wide_ms = per_call(ringed_wide.lightcurve)
print(f"1000 points, all in transit: {forward_times['ringed']:.2f} ms")
print(f"1000 points, ~17% in transit: {wide_ms:.2f} ms")
1000 points, all in transit: 49.58 ms
1000 points, ~17% in transit: 10.36 ms
The cost scales with the number of in-transit points, not the total number of points, so densely-sampled baseline does not slow the model down much.
Gradients#
Both classes wrap lightcurve in a custom_vjp that computes derivatives via forward-mode differentiation (one Jacobian column per input parameter), which is far faster here than naive reverse-mode. The cost therefore grows with the number of parameters you differentiate with respect to. We’ll time the gradient of a scalar summary of the lightcurve with respect to two parameters for each model:
[12]:
spherical = systems["spherical"]
ringed = systems["ringed"]
def spherical_loss(params):
return jnp.sum(spherical.lightcurve(params))
def ringed_loss(params):
return jnp.sum(ringed.lightcurve(params))
spherical_params = {"r": jnp.array([0.1]), "i": jnp.array([jnp.pi / 2 - 0.01])}
ringed_params = {"r": jnp.array([0.1]), "ring_outer_r": jnp.array([0.3])}
grad_cases = {
"spherical": (jax.jit(jax.grad(spherical_loss)), spherical_params),
"ringed": (jax.jit(jax.grad(ringed_loss)), ringed_params),
}
print(f"{'model':<12}{'compile + first call [s]':>26}{'per call [ms]':>16}")
for name, (grad_fn, params) in grad_cases.items():
compile_s = first_call(lambda: grad_fn(params))
steady_ms = per_call(lambda: grad_fn(params))
print(f"{name:<12}{compile_s:>26.2f}{steady_ms:>16.2f}")
model compile + first call [s] per call [ms]
spherical 0.83 5.06
ringed 2.39 80.61
Summary#
The spherical and oblate models cost the same; expect a few ms per 1,000 in-transit points on a CPU.
The ringed model costs roughly 15-20x the planet-only model per in-transit point. Out-of-transit points are masked and nearly free for all models.
Gradients via the built-in
custom_vjpcost a small multiple of a forward evaluation (roughly 1.7x here with two parameters), and that multiple grows with the number of parameters being differentiated since each adds a Jacobian column. Differentiating with respect to a handful of parameters is cheap; differentiating with respect to dozens is not.Compilation is a one-time cost per (system, number of timesteps) combination – under a second for the forward models, a few seconds for gradients.