Example of model comparison with Harmonic: TOI-544
In this notebook, we demonstrate how to use the harmonic package to perform Bayesian model comparison for exoplanet radial velocity models. We compare a 1-planet model against a 2-planet model for the TOI-544 system. This system was first characterised in RV by Osborne et al. (2024), who knew TOI-544~b was transiting, but suspected the presence of an outer non-transiting planet TOI-544~c.
In Bayesian inference, the Bayesian evidence (marginal likelihood) \(\mathcal{Z} = \int \mathcal{L}(\theta)\, \pi(\theta)\, d\theta\) quantifies how well a model explains the data, integrating over the full prior volume. Unlike information criteria (AIC, BIC), it naturally penalises unnecessary model complexity through the prior – rather than just through the likelihood and the dimensionality – this is the Bayesian Occam’s razor. In RV we tend to have quite informative priors (e.g. tight priors on \(P\) from transits, or on quasiperiodic GP kernel hyperparameters through photometric studies), so it seems a waste to not use that prior information in our model comparisons by doing AIC or BIC instead.
The problem though is that computing \(\mathcal{Z}\) directly is computationally intractable for high-dimensional problems. In MCMC, we only sample the likelihood and the prior functions, not the evidence (giving us the un-normalised posterior, which is fine for parameter inference, but need it to be normalised for model comparison). Therefore, we will use the harmonic package implementation of the Learned Harmonic Mean Estimator (LHME) which uses normalizing flows to estimate \(\mathcal{Z}\) from our existing posterior MCMC samples. See McEwen et al. (2021) for the theoretical background, but the quick sell is that this is basically free - we already have the MCMC samples from the fitting/parameter inference anyway! We can use any sampling technique we like (such as ravest) and get evidences in five minutes.
This tutorial assumes familiarity with fitting RV models in ravest so we’ll skip over the fine details – see the K2-24 case-study and GP fitting tutorial for more detailed walkthroughs.
from pathlib import Path
import harmonic as hm
import jax
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from ravest.fit import GPFitter
from ravest.gp import GPKernel
from ravest.param import Parameter, Parameterisation
from ravest.prior import HalfNormal, Normal, Uniform
jax.config.update("jax_enable_x64", True)
Data
TOI-544 is an M dwarf observed with HARPS and HARPS-N. The data below have already been prepared: radial velocities are in m/s and times are in BTJD (BJD - 2457000). This data is based on RV data originally from Osborne et al. (2024) | MNRAS 527, 11138–11157.
fpath = Path.cwd() / "example_data" / "TOI-544.csv"
data = pd.read_csv(fpath)
print("Instruments in the dataset:", data["Instrument"].unique())
data.head()
Instruments in the dataset: ['HARPS' 'HARPS-N']
| RV | e_RV | Instrument | BTJD | |
|---|---|---|---|---|
| 0 | 19.4 | 1.7 | HARPS | 2205.63053 |
| 1 | 30.4 | 1.4 | HARPS | 2206.63247 |
| 2 | 21.5 | 1.9 | HARPS | 2207.57568 |
| 3 | 22.7 | 1.3 | HARPS | 2207.66046 |
| 4 | -2.1 | 1.4 | HARPS | 2214.71815 |
plt.figure(figsize=(15, 3.5))
plt.title("TOI-544 radial velocity data")
plt.ylabel("Radial Velocity [m/s]")
plt.xlabel("BTJD (days)")
for inst in data["Instrument"].unique():
mask = data["Instrument"] == inst
plt.errorbar(
data.loc[mask, "BTJD"],
data.loc[mask, "RV"],
yerr=data.loc[mask, "e_RV"],
marker=".",
label=inst,
linestyle="None",
)
plt.legend()
plt.show()
1-Planet Model
We first fit a model with only planet b (\(P = 1.5484\) d) on a circular orbit, using a Quasiperiodic GP to model stellar activity.
We keep the fitting details compact here — see the GP fitting tutorial for a thorough walkthrough of the GP fitting workflow.
time = data["BTJD"].to_numpy(dtype=float)
rv = data["RV"].to_numpy(dtype=float)
e_rv = data["e_RV"].to_numpy(dtype=float)
instruments = data["Instrument"].to_numpy(dtype=str)
t0 = np.median(time)
unique_instruments = sorted(np.unique(instruments).tolist())
print(f"Instruments: {unique_instruments}")
Instruments: ['HARPS', 'HARPS-N']
gpkernel_1pl = GPKernel(kernel_type="Quasiperiodic")
fitter_1pl = GPFitter(
planet_letters=["b"],
parameterisation=Parameterisation("P K secosw sesinw Tc"),
gp_kernel=gpkernel_1pl,
)
fitter_1pl.add_data(time=time, vel=rv, velerr=e_rv, instrument=instruments, t0=t0)
# Planet b: circular orbit (secosw, sesinw fixed to 0)
tc_b = 2459199.0314 - 2457000 # BTJD
# Per-instrument gamma bounds from the data range
rv_range = np.max(rv) - np.min(rv)
g_lower = np.min(rv) - 0.5 * rv_range
g_upper = np.max(rv) + 0.5 * rv_range
params_1pl = {
"P_b": Parameter(1.5484, "d", fixed=False),
"K_b": Parameter(10.0, "m/s", fixed=False),
"secosw_b": Parameter(0.0, "", fixed=True),
"sesinw_b": Parameter(0.0, "", fixed=True),
"Tc_b": Parameter(tc_b, "d", fixed=False),
"gd": Parameter(0.0, "m/s/d", fixed=True),
"gdd": Parameter(0.0, "m/s/d^2", fixed=True),
}
for inst in unique_instruments:
mask = instruments == inst
params_1pl[f"g_{inst}"] = Parameter(np.median(rv[mask]), "m/s", fixed=False)
params_1pl[f"jit_{inst}"] = Parameter(0.0, "m/s", fixed=False)
priors_1pl = {
"P_b": Normal(1.5484, 0.000002),
"K_b": Uniform(0.0, 50.0),
"Tc_b": Normal(tc_b, 0.0007),
}
for inst in unique_instruments:
priors_1pl[f"g_{inst}"] = Uniform(g_lower, g_upper)
priors_1pl[f"jit_{inst}"] = Uniform(0, 30)
hyperparams_1pl = {
"gp_amp": Parameter(1.0, "", fixed=False),
"gp_lambda_e": Parameter(112, "d", fixed=False),
"gp_lambda_p": Parameter(0.519, "", fixed=False),
"gp_period": Parameter(19.343, "d", fixed=False),
}
hyperpriors_1pl = {
"gp_amp": Uniform(0, rv_range),
"gp_lambda_e": Uniform(6, 112 + (3 * 28)),
"gp_lambda_p": Uniform(0.1, 0.519 + (4 * 0.091)),
"gp_period": Uniform(19.343 - 2, 19.343 + 2),
}
fitter_1pl.params = params_1pl
fitter_1pl.priors = priors_1pl
fitter_1pl.hyperparams = hyperparams_1pl
fitter_1pl.hyperpriors = hyperpriors_1pl
print(f"Free parameters: {list(fitter_1pl.free_params_dict.keys())}")
print(f"Free hyperparameters: {list(fitter_1pl.free_hyperparams_names)}")
print(f"Total dimensions: {fitter_1pl.ndim}")
Free parameters: ['P_b', 'K_b', 'Tc_b', 'g_HARPS', 'jit_HARPS', 'g_HARPS-N', 'jit_HARPS-N']
Free hyperparameters: ['gp_amp', 'gp_lambda_e', 'gp_lambda_p', 'gp_period']
Total dimensions: 11
Fit the 1-planet model
For demonstration purposes, we use reduced MCMC settings here to make this notebook run quickly. For real runs you should use way more - perhaps at least 200 walkers!
map_1pl = fitter_1pl.find_map_estimate(method="Powell")
print(f"MAP log-probability: {-map_1pl.fun:.2f}")
MAP parameter results: {'P_b': np.float64(1.5484000062945433), 'K_b': np.float64(1.8352248224195358), 'Tc_b': np.float64(2199.0314135418944), 'g_HARPS': np.float64(5.175554099549803), 'jit_HARPS': np.float64(3.68991752999967), 'g_HARPS-N': np.float64(-13.760005756717021), 'jit_HARPS-N': np.float64(2.4679874150152665)}
MAP hyperparameter results: {'gp_amp': np.float64(7.00084769703356), 'gp_lambda_e': np.float64(69.23181972209414), 'gp_lambda_p': np.float64(0.28989874682622274), 'gp_period': np.float64(19.402144345306436)}
MAP log-probability: -383.87
nwalkers_1pl = 32 # Needs to be at least 2x ndims, but you should use hundreds!
init_1pl = fitter_1pl.generate_initial_walker_positions_random(nwalkers=nwalkers_1pl)
fitter_1pl.run_mcmc(
initial_positions=init_1pl,
nwalkers=nwalkers_1pl,
max_steps=50_000,
check_convergence=True,
convergence_check_interval=1000,
convergence_check_start=15_000,
progress=True,
multiprocessing=True,
)
print("MCMC completed")
INFO:root:Starting MCMC with convergence checks. (Maximum 50000 steps, checking convergence every 1000 steps after iteration 15000)...
INFO:2026-03-30 15:03:25,893:jax._src.xla_bridge:810: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:2026-03-30 15:03:25,917:jax._src.xla_bridge:810: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:2026-03-30 15:03:25,949:jax._src.xla_bridge:810: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:2026-03-30 15:03:25,966:jax._src.xla_bridge:810: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:2026-03-30 15:03:25,997:jax._src.xla_bridge:810: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:2026-03-30 15:03:26,008:jax._src.xla_bridge:810: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:2026-03-30 15:03:26,018:jax._src.xla_bridge:810: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:2026-03-30 15:03:26,019:jax._src.xla_bridge:810: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:2026-03-30 15:03:26,019:jax._src.xla_bridge:810: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:2026-03-30 15:03:26,051:jax._src.xla_bridge:810: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
30%|██▉ | 14998/50000 [03:00<08:10, 71.36it/s]INFO:root:Convergence check: Step 15000: mean(tau)=228.9, max(tau)=569.0
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
32%|███▏ | 15996/50000 [03:15<08:38, 65.60it/s]INFO:root:Convergence check: Step 16000: mean(tau)=230.7, max(tau)=572.8
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
34%|███▍ | 16998/50000 [03:29<07:22, 74.62it/s]INFO:root:Convergence check: Step 17000: mean(tau)=230.7, max(tau)=575.6
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
36%|███▌ | 17996/50000 [03:43<07:04, 75.48it/s]INFO:root:Convergence check: Step 18000: mean(tau)=231.4, max(tau)=577.8
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
38%|███▊ | 18998/50000 [03:57<06:52, 75.13it/s]INFO:root:Convergence check: Step 19000: mean(tau)=231.4, max(tau)=578.9
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
40%|███▉ | 19995/50000 [04:10<06:48, 73.40it/s]INFO:root:Convergence check: Step 20000: mean(tau)=231.4, max(tau)=579.8
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
42%|████▏ | 20993/50000 [04:24<06:20, 76.31it/s]INFO:root:Convergence check: Step 21000: mean(tau)=232.0, max(tau)=581.2
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
44%|████▍ | 21995/50000 [04:38<05:54, 78.90it/s]INFO:root:Convergence check: Step 22000: mean(tau)=231.9, max(tau)=582.6
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
46%|████▌ | 22999/50000 [04:52<05:59, 75.02it/s]INFO:root:Convergence check: Step 23000: mean(tau)=231.3, max(tau)=583.0
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
48%|████▊ | 23999/50000 [05:05<05:40, 76.43it/s]INFO:root:Convergence check: Step 24000: mean(tau)=230.9, max(tau)=583.6
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
50%|████▉ | 24997/50000 [05:19<05:27, 76.25it/s]INFO:root:Convergence check: Step 25000: mean(tau)=230.4, max(tau)=584.1
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
52%|█████▏ | 25996/50000 [05:33<05:08, 77.73it/s]INFO:root:Convergence check: Step 26000: mean(tau)=230.0, max(tau)=584.7
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
54%|█████▍ | 26998/50000 [05:47<05:02, 75.94it/s]INFO:root:Convergence check: Step 27000: mean(tau)=229.4, max(tau)=584.5
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
56%|█████▌ | 27994/50000 [06:00<04:52, 75.14it/s]INFO:root:Convergence check: Step 28000: mean(tau)=229.1, max(tau)=585.1
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
58%|█████▊ | 28999/50000 [06:14<04:36, 75.94it/s]INFO:root:Convergence check: Step 29000: mean(tau)=228.2, max(tau)=585.3
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
60%|█████▉ | 29995/50000 [06:28<04:22, 76.35it/s]INFO:root:Convergence check: Step 30000: mean(tau)=228.3, max(tau)=585.4
INFO:root:Not yet converged (N/50>tau check: True, tau stability check: False)
62%|██████▏ | 30997/50000 [06:42<04:26, 71.35it/s]INFO:root:Convergence check: Step 31000: mean(tau)=227.6, max(tau)=585.6
INFO:root:Not yet converged (N/50>tau check: True, tau stability check: False)
64%|██████▍ | 31995/50000 [06:55<03:49, 78.41it/s]INFO:root:Convergence check: Step 32000: mean(tau)=227.0, max(tau)=585.7
INFO:root:Not yet converged (N/50>tau check: True, tau stability check: False)
66%|██████▌ | 32999/50000 [07:09<03:57, 71.48it/s]INFO:root:Convergence check: Step 33000: mean(tau)=226.9, max(tau)=585.6
INFO:root:Not yet converged (N/50>tau check: True, tau stability check: False)
68%|██████▊ | 33993/50000 [07:24<03:33, 74.93it/s]INFO:root:Convergence check: Step 34000: mean(tau)=227.4, max(tau)=585.7
INFO:root:Not yet converged (N/50>tau check: True, tau stability check: False)
70%|██████▉ | 34997/50000 [07:38<03:30, 71.23it/s]INFO:root:Convergence check: Step 35000: mean(tau)=227.0, max(tau)=585.4
INFO:root:Not yet converged (N/50>tau check: True, tau stability check: False)
72%|███████▏ | 35999/50000 [07:53<03:05, 75.27it/s]INFO:root:Convergence check: Step 36000: mean(tau)=226.8, max(tau)=586.0
INFO:root:Not yet converged (N/50>tau check: True, tau stability check: False)
74%|███████▍ | 36993/50000 [08:07<02:51, 75.76it/s]INFO:root:Convergence check: Step 37000: mean(tau)=226.7, max(tau)=586.1
INFO:root:Not yet converged (N/50>tau check: True, tau stability check: False)
76%|███████▌ | 37998/50000 [08:22<02:43, 73.25it/s]INFO:root:Convergence check: Step 38000: mean(tau)=226.2, max(tau)=585.8
INFO:root:Not yet converged (N/50>tau check: True, tau stability check: False)
78%|███████▊ | 38999/50000 [08:36<02:31, 72.74it/s]INFO:root:Convergence check: Step 39000: mean(tau)=226.3, max(tau)=585.8
INFO:root:Not yet converged (N/50>tau check: True, tau stability check: False)
80%|████████ | 40000/50000 [08:51<02:11, 75.98it/s]INFO:root:Convergence check: Step 40000: mean(tau)=225.7, max(tau)=585.8
INFO:root:Converged at iteration 40000
80%|████████ | 40000/50000 [08:52<02:13, 75.12it/s]
INFO:root:MCMC complete: 40000 steps total
MCMC completed
total_steps_1pl, nwalkers_1pl, ndim_1pl = fitter_1pl.sampler.iteration, fitter_1pl.sampler.nwalkers, fitter_1pl.sampler.ndim
samples_per_walker = 1000 # we want to keep the final 1000 samples from each walker, after thinning by 10
thin_rate = 10
discard_start_1pl = total_steps_1pl - (samples_per_walker * thin_rate)
print(f"nsteps: {total_steps_1pl}, nwalkers: {nwalkers_1pl}, ndim: {ndim_1pl}")
print(f"Discarding first {discard_start_1pl} steps, then thinning by {thin_rate} to get {samples_per_walker} samples per walker.")
print(f"Total samples after discarding and thinning: {(total_steps_1pl - discard_start_1pl) // thin_rate * nwalkers_1pl}")
nsteps: 40000, nwalkers: 32, ndim: 11
Discarding first 30000 steps, then thinning by 10 to get 1000 samples per walker.
Total samples after discarding and thinning: 32000
samples_1pl = fitter_1pl.get_samples_np(
discard_start=discard_start_1pl, thin=thin_rate, flat=False
)
lnprob_1pl = fitter_1pl.get_sampler_lnprob(
discard_start=discard_start_1pl, thin=thin_rate, flat=False
)
print(f"Samples shape: {samples_1pl.shape} (nsteps, nwalkers, ndim)")
print(f"Lnprob shape: {lnprob_1pl.shape} (nsteps, nwalkers)")
Samples shape: (1000, 32, 11) (nsteps, nwalkers, ndim)
Lnprob shape: (1000, 32) (nsteps, nwalkers)
Check the corners to make sure the posteriors look converged and well-behaved
fitter_1pl.plot_corner(discard_start=discard_start_1pl, thin=thin_rate, plot_datapoints=True)
2-Planet Model
Now we fit a model with planet b (as before) plus a candidate planet c at \(P \sim 50.5\) d with free eccentricity. We use a HalfNormal(0.32) eccentricity prior based on the Van Eylen et al. (2019) single-transiting planet distribution.
gpkernel_2pl = GPKernel(kernel_type="Quasiperiodic")
fitter_2pl = GPFitter(
planet_letters=["b", "c"],
parameterisation=Parameterisation("P K secosw sesinw Tc"),
gp_kernel=gpkernel_2pl,
)
fitter_2pl.add_data(time=time, vel=rv, velerr=e_rv, instrument=instruments, t0=t0)
tc_c_lower = 2459205.0 - 2457000 # BTJD
tc_c_upper = 2459225.0 - 2457000 # BTJD
tc_c = (tc_c_lower + tc_c_upper) / 2
params_2pl = {
# Planet b (circular)
"P_b": Parameter(1.5484, "d", fixed=False),
"K_b": Parameter(10.0, "m/s", fixed=False),
"secosw_b": Parameter(0.0, "", fixed=True),
"sesinw_b": Parameter(0.0, "", fixed=True),
"Tc_b": Parameter(tc_b, "d", fixed=False),
# Planet c (eccentric)
"P_c": Parameter(50.5, "d", fixed=False),
"K_c": Parameter(10.0, "m/s", fixed=False),
"secosw_c": Parameter(0.0, "", fixed=False),
"sesinw_c": Parameter(0.0, "", fixed=False),
"Tc_c": Parameter(tc_c, "d", fixed=False),
# Trend (fixed)
"gd": Parameter(0.0, "m/s/d", fixed=True),
"gdd": Parameter(0.0, "m/s/d^2", fixed=True),
}
for inst in unique_instruments:
mask = instruments == inst
params_2pl[f"g_{inst}"] = Parameter(np.median(rv[mask]), "m/s", fixed=False)
params_2pl[f"jit_{inst}"] = Parameter(0.0, "m/s", fixed=False)
priors_2pl = {
# Planet b
"P_b": Normal(1.5484, 0.000002),
"K_b": Uniform(0.0, 50.0),
"Tc_b": Normal(tc_b, 0.0007),
# Planet c
"P_c": Uniform(49.0, 52.0),
"K_c": Uniform(0.0, 50.0),
"e_c": HalfNormal(0.32),
"w_c": Uniform(-np.pi, np.pi),
"Tc_c": Uniform(tc_c_lower, tc_c_upper),
}
for inst in unique_instruments:
priors_2pl[f"g_{inst}"] = Uniform(g_lower, g_upper)
priors_2pl[f"jit_{inst}"] = Uniform(0, 30)
hyperparams_2pl = {
"gp_amp": Parameter(1.0, "", fixed=False),
"gp_lambda_e": Parameter(112, "d", fixed=False),
"gp_lambda_p": Parameter(0.519, "", fixed=False),
"gp_period": Parameter(19.343, "d", fixed=False),
}
hyperpriors_2pl = {
"gp_amp": Uniform(0, rv_range),
"gp_lambda_e": Uniform(6, 112 + (3 * 28)),
"gp_lambda_p": Uniform(0.1, 0.519 + (4 * 0.091)),
"gp_period": Uniform(19.343 - 2, 19.343 + 1),
}
fitter_2pl.params = params_2pl
fitter_2pl.priors = priors_2pl
fitter_2pl.hyperparams = hyperparams_2pl
fitter_2pl.hyperpriors = hyperpriors_2pl
print(f"Free parameters: {list(fitter_2pl.free_params_dict.keys())}")
print(f"Free hyperparameters: {list(fitter_2pl.free_hyperparams_names)}")
print(f"Total dimensions: {fitter_2pl.ndim}")
Free parameters: ['P_b', 'K_b', 'Tc_b', 'P_c', 'K_c', 'secosw_c', 'sesinw_c', 'Tc_c', 'g_HARPS', 'jit_HARPS', 'g_HARPS-N', 'jit_HARPS-N']
Free hyperparameters: ['gp_amp', 'gp_lambda_e', 'gp_lambda_p', 'gp_period']
Total dimensions: 16
Fit the 2-planet model
map_2pl = fitter_2pl.find_map_estimate(method="Powell")
print(f"MAP log-probability: {-map_2pl.fun:.2f}")
MAP parameter results: {'P_b': np.float64(1.5484000246157859), 'K_b': np.float64(2.2295222651118403), 'Tc_b': np.float64(2199.0314296949805), 'P_c': np.float64(50.00893705172687), 'K_c': np.float64(5.4641651390305075), 'secosw_c': np.float64(0.5486256368580782), 'sesinw_c': np.float64(0.0426733327652046), 'Tc_c': np.float64(2212.1792218607093), 'g_HARPS': np.float64(6.1666542224678675), 'jit_HARPS': np.float64(1.5180970050873714), 'g_HARPS-N': np.float64(-17.765161842942295), 'jit_HARPS-N': np.float64(1.761141119197035)}
MAP hyperparameter results: {'gp_amp': np.float64(7.786985106883591), 'gp_lambda_e': np.float64(102.66967216949709), 'gp_lambda_p': np.float64(0.2430817976976905), 'gp_period': np.float64(19.291961308795)}
MAP log-probability: -350.98
nwalkers_2pl = 32 # Needs to be at least 2x ndims, but you should use hundreds!
init_2pl = fitter_2pl.generate_initial_walker_positions_random(nwalkers=nwalkers_2pl)
fitter_2pl.run_mcmc(
initial_positions=init_2pl,
nwalkers=nwalkers_2pl,
max_steps=50_000,
check_convergence=True,
convergence_check_interval=1000,
convergence_check_start=15_000,
progress=True,
multiprocessing=True,
)
print("MCMC completed")
INFO:root:Starting MCMC with convergence checks. (Maximum 50000 steps, checking convergence every 1000 steps after iteration 15000)...
INFO:2026-03-30 15:12:25,117:jax._src.xla_bridge:810: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:2026-03-30 15:12:25,207:jax._src.xla_bridge:810: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:2026-03-30 15:12:25,236:jax._src.xla_bridge:810: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:2026-03-30 15:12:25,253:jax._src.xla_bridge:810: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:2026-03-30 15:12:25,276:jax._src.xla_bridge:810: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:2026-03-30 15:12:25,281:jax._src.xla_bridge:810: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:2026-03-30 15:12:25,324:jax._src.xla_bridge:810: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:2026-03-30 15:12:25,359:jax._src.xla_bridge:810: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:2026-03-30 15:12:25,359:jax._src.xla_bridge:810: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:2026-03-30 15:12:25,363:jax._src.xla_bridge:810: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file)
30%|██▉ | 14993/50000 [03:18<07:37, 76.54it/s]INFO:root:Convergence check: Step 15000: mean(tau)=314.8, max(tau)=723.9
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
32%|███▏ | 15993/50000 [03:31<07:21, 77.09it/s]INFO:root:Convergence check: Step 16000: mean(tau)=321.2, max(tau)=732.2
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
34%|███▍ | 17000/50000 [03:45<07:19, 75.09it/s]INFO:root:Convergence check: Step 17000: mean(tau)=328.7, max(tau)=751.9
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
36%|███▌ | 17998/50000 [03:59<07:20, 72.69it/s]INFO:root:Convergence check: Step 18000: mean(tau)=332.0, max(tau)=756.4
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
38%|███▊ | 18994/50000 [04:13<07:27, 69.27it/s]INFO:root:Convergence check: Step 19000: mean(tau)=332.2, max(tau)=736.9
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
40%|███▉ | 19994/50000 [04:26<06:33, 76.33it/s]INFO:root:Convergence check: Step 20000: mean(tau)=331.7, max(tau)=725.5
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
42%|████▏ | 20997/50000 [04:40<06:22, 75.75it/s]INFO:root:Convergence check: Step 21000: mean(tau)=334.5, max(tau)=734.9
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
44%|████▍ | 21997/50000 [04:54<06:44, 69.31it/s]INFO:root:Convergence check: Step 22000: mean(tau)=336.5, max(tau)=731.1
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
46%|████▌ | 22995/50000 [05:08<05:48, 77.60it/s]INFO:root:Convergence check: Step 23000: mean(tau)=340.0, max(tau)=728.0
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
48%|████▊ | 23995/50000 [05:22<05:39, 76.54it/s]INFO:root:Convergence check: Step 24000: mean(tau)=343.1, max(tau)=729.6
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
50%|████▉ | 24992/50000 [05:36<05:26, 76.56it/s]INFO:root:Convergence check: Step 25000: mean(tau)=347.9, max(tau)=729.2
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
52%|█████▏ | 25999/50000 [05:50<05:11, 77.06it/s]INFO:root:Convergence check: Step 26000: mean(tau)=351.1, max(tau)=729.5
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
54%|█████▍ | 26998/50000 [06:04<05:04, 75.50it/s]INFO:root:Convergence check: Step 27000: mean(tau)=351.1, max(tau)=721.4
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
56%|█████▌ | 27999/50000 [06:18<04:56, 74.29it/s]INFO:root:Convergence check: Step 28000: mean(tau)=352.9, max(tau)=720.6
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
58%|█████▊ | 28995/50000 [06:32<04:40, 74.78it/s]INFO:root:Convergence check: Step 29000: mean(tau)=355.2, max(tau)=719.5
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
60%|█████▉ | 29994/50000 [06:45<04:28, 74.45it/s]INFO:root:Convergence check: Step 30000: mean(tau)=355.9, max(tau)=722.3
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
62%|██████▏ | 30994/50000 [06:59<04:16, 74.18it/s]INFO:root:Convergence check: Step 31000: mean(tau)=356.7, max(tau)=722.6
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
64%|██████▍ | 31999/50000 [07:13<03:52, 77.49it/s]INFO:root:Convergence check: Step 32000: mean(tau)=358.7, max(tau)=724.6
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
66%|██████▌ | 32998/50000 [07:27<03:55, 72.07it/s]INFO:root:Convergence check: Step 33000: mean(tau)=361.7, max(tau)=720.8
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
68%|██████▊ | 33996/50000 [07:42<03:34, 74.76it/s]INFO:root:Convergence check: Step 34000: mean(tau)=363.8, max(tau)=724.8
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
70%|███████ | 35000/50000 [07:57<03:18, 75.50it/s]INFO:root:Convergence check: Step 35000: mean(tau)=366.9, max(tau)=728.6
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
72%|███████▏ | 35998/50000 [08:12<03:11, 73.23it/s]INFO:root:Convergence check: Step 36000: mean(tau)=369.8, max(tau)=731.9
INFO:root:Not yet converged (N/50>tau check: False, tau stability check: False)
74%|███████▍ | 36993/50000 [08:26<02:52, 75.21it/s]INFO:root:Convergence check: Step 37000: mean(tau)=371.1, max(tau)=731.5
INFO:root:Not yet converged (N/50>tau check: True, tau stability check: False)
76%|███████▌ | 37994/50000 [08:41<02:41, 74.52it/s]INFO:root:Convergence check: Step 38000: mean(tau)=369.4, max(tau)=727.7
INFO:root:Not yet converged (N/50>tau check: True, tau stability check: False)
78%|███████▊ | 38998/50000 [08:56<02:23, 76.64it/s]INFO:root:Convergence check: Step 39000: mean(tau)=370.3, max(tau)=729.2
INFO:root:Not yet converged (N/50>tau check: True, tau stability check: False)
80%|███████▉ | 39996/50000 [09:10<02:12, 75.47it/s]INFO:root:Convergence check: Step 40000: mean(tau)=372.7, max(tau)=731.8
INFO:root:Not yet converged (N/50>tau check: True, tau stability check: False)
82%|████████▏ | 40993/50000 [09:25<01:59, 75.19it/s]INFO:root:Convergence check: Step 41000: mean(tau)=374.7, max(tau)=724.3
INFO:root:Not yet converged (N/50>tau check: True, tau stability check: False)
WARNING:root:Approaching max iterations (50000) without convergence! (max tau=724.3, tau stability change=[0.0156478 0.00966202 0.00472649 0.02936193 0.01846816 0.01030093
0.01429811 0.00393237 0.00156828 0.00225264 0.00080257 0.00219052
0.02401969 0.00655273 0.01301511 0.00557646])
84%|████████▍ | 42000/50000 [09:42<02:39, 50.20it/s]INFO:root:Convergence check: Step 42000: mean(tau)=376.6, max(tau)=725.5
INFO:root:Not yet converged (N/50>tau check: True, tau stability check: False)
WARNING:root:Approaching max iterations (50000) without convergence! (max tau=725.5, tau stability change=[0.00425303 0.00697043 0.00122193 0.0077468 0.01924607 0.00166489
0.01299635 0.00887325 0.00744002 0.00290209 0.00642868 0.00126132
0.02356412 0.00686248 0.01765563 0.01120693])
86%|████████▌ | 42998/50000 [09:58<01:38, 71.08it/s]INFO:root:Convergence check: Step 43000: mean(tau)=376.3, max(tau)=720.3
INFO:root:Not yet converged (N/50>tau check: True, tau stability check: False)
WARNING:root:Approaching max iterations (50000) without convergence! (max tau=720.3, tau stability change=[0.00577206 0.0038538 0.00069017 0.01180585 0.00733511 0.00725195
0.00116983 0.00637679 0.00172874 0.00669658 0.00360301 0.01029376
0.0007225 0.01869467 0.00659774 0.00582953])
88%|████████▊ | 44000/50000 [10:14<01:24, 70.61it/s]INFO:root:Convergence check: Step 44000: mean(tau)=377.9, max(tau)=719.5
INFO:root:Not yet converged (N/50>tau check: True, tau stability check: False)
WARNING:root:Approaching max iterations (50000) without convergence! (max tau=719.5, tau stability change=[0.01718438 0.01117355 0.00410733 0.00708313 0.00866909 0.00110055
0.01813813 0.00026693 0.00190634 0.0038081 0.01143116 0.01304897
0.00250748 0.00529952 0.0005265 0.00438942])
90%|█████████ | 45000/50000 [10:30<01:15, 66.45it/s]INFO:root:Convergence check: Step 45000: mean(tau)=377.8, max(tau)=720.6
INFO:root:Converged at iteration 45000
90%|█████████ | 45000/50000 [10:32<01:10, 71.16it/s]
INFO:root:MCMC complete: 45000 steps total
MCMC completed
total_steps_2pl, nwalkers_2pl, ndim_2pl = fitter_2pl.sampler.iteration, fitter_2pl.sampler.nwalkers, fitter_2pl.sampler.ndim
# using the same samples_per_walker=1000 and thin_rate=10 as for 1pl model
discard_start_2pl = total_steps_2pl - (samples_per_walker * thin_rate)
print(f"nsteps: {total_steps_2pl}, nwalkers: {nwalkers_2pl}, ndim: {ndim_2pl}")
print(f"Discarding first {discard_start_2pl} steps, then thinning by {thin_rate} to get {samples_per_walker} samples per walker.")
print(f"Total samples after discarding and thinning: {(total_steps_2pl - discard_start_2pl) // thin_rate * nwalkers_2pl}")
nsteps: 45000, nwalkers: 32, ndim: 16
Discarding first 35000 steps, then thinning by 10 to get 1000 samples per walker.
Total samples after discarding and thinning: 32000
samples_2pl = fitter_2pl.get_samples_np(discard_start=discard_start_2pl, thin=thin_rate, flat=False)
lnprob_2pl = fitter_2pl.get_sampler_lnprob(discard_start=discard_start_2pl, thin=thin_rate, flat=False)
print(f"Samples shape: {samples_2pl.shape} (nsteps, nwalkers, ndim)")
print(f"Lnprob shape: {lnprob_2pl.shape} (nsteps, nwalkers)")
Samples shape: (1000, 32, 16) (nsteps, nwalkers, ndim)
Lnprob shape: (1000, 32) (nsteps, nwalkers)
fitter_2pl.plot_corner(discard_start=discard_start_2pl, thin=thin_rate, plot_datapoints=True)
Comparing the fits
Let’s plot the RV model for both the 1-planet and 2-planet fits using their median values of the posterior parameters (obviously you should always check your posterior distributions rather than just taking these values blindly).
I’ve used custom xlims to zoom in on one clump of HARPS data to make it clearer - you can change these arguments to have a look at the full set of data.
# Build median posterior parameter dicts for both models
samples_flat_1pl = fitter_1pl.get_samples_np(
discard_start=discard_start_1pl, thin=thin_rate, flat=True # flatten the (nsteps, nwalkers, ndim) array into (nsteps*nwalkers, ndim)
)
medians_1pl = np.percentile(samples_flat_1pl, 50, axis=0)
dict_1pl = fitter_1pl.build_params_dict(medians_1pl)
samples_flat_2pl = fitter_2pl.get_samples_np(
discard_start=discard_start_2pl, thin=thin_rate, flat=True # flatten the (nsteps, nwalkers, ndim) array into (nsteps*nwalkers, ndim)
)
medians_2pl = np.percentile(samples_flat_2pl, 50, axis=0)
dict_2pl = fitter_2pl.build_params_dict(medians_2pl)
fitter_1pl.plot_custom_rv(dict_1pl, title="1-Planet Model (median posterior)", n_smooth=5000,
xlim=(2485, 2637), res_xlim=(2485, 2637),
ylim=(-23, 33))
fitter_2pl.plot_custom_rv(dict_2pl, title="2-Planet Model (median posterior)", n_smooth=5000,
xlim=(2485, 2637), res_xlim=(2485, 2637),
ylim=(-23, 33))
At first glance looking at the residuals, the two fits look quite similar, both models seem to do a reasonable job of modelling the data.
Minimum planet masses
Let’s investigate how the mass(es) differ between between the two models. We can compute \(M_p \sin i\) for each model to see how the planet b mass estimate compares, and what the candidate planet c mass looks like. We use the value of \(M_{\star}=0.630 \pm 0.018 M_{\odot}\) from Table 1 of Osborne et al. 2024.
from ravest.model import calculate_mpsini
# Stellar mass from Osborne et al. (2024) Table 1
st_mass = 0.631 # M_sun
st_mass_err = 0.018 # M_sun
# --- 1-planet model: planet b only ---
post_1pl = fitter_1pl.get_mcmc_posterior_dict(discard_start=discard_start_1pl, thin=thin_rate)
n_samples_1pl = len(post_1pl["P_b"])
stellar_masses_1pl = np.random.normal(loc=st_mass, scale=st_mass_err, size=n_samples_1pl)
ecc_b_1pl = post_1pl["secosw_b"]**2 + post_1pl["sesinw_b"]**2 # get e from secosw and sesinw
mpsini_b_1pl = calculate_mpsini(stellar_masses_1pl, post_1pl["P_b"], post_1pl["K_b"], ecc_b_1pl, unit="M_earth")
p16, p50, p84 = np.percentile(mpsini_b_1pl, [16, 50, 84])
print("1-Planet Model:")
print(f" Planet b: Mpsini = {p50:.2f} +{p84 - p50:.2f} -{p50 - p16:.2f} M_Earth")
# --- 2-planet model: planets b and c ---
post_2pl = fitter_2pl.get_mcmc_posterior_dict(discard_start=discard_start_2pl, thin=thin_rate)
n_samples_2pl = len(post_2pl["P_b"])
stellar_masses_2pl = np.random.normal(loc=st_mass, scale=st_mass_err, size=n_samples_2pl)
print("\n2-Planet Model:")
for planet in ["b", "c"]:
ecc = post_2pl[f"secosw_{planet}"]**2 + post_2pl[f"sesinw_{planet}"]**2 # get e from secosw and sesinw
mpsini = calculate_mpsini(stellar_masses_2pl, post_2pl[f"P_{planet}"], post_2pl[f"K_{planet}"], ecc, unit="M_earth")
p16, p50, p84 = np.percentile(mpsini, [16, 50, 84])
print(f" Planet {planet}: Mpsini = {p50:.2f} +{p84 - p50:.2f} -{p50 - p16:.2f} M_Earth")
1-Planet Model:
Planet b: Mpsini = 3.24 +0.52 -0.53 M_Earth
2-Planet Model:
Planet b: Mpsini = 2.96 +0.52 -0.52 M_Earth
Planet c: Mpsini = 21.68 +2.22 -2.23 M_Earth
So while at first glance we have two quite similar fits, we can see that’s two very different results! With the two-planet model, not only do we now have a second planet, but planet b decreases in mass by \(\sim10\)% - that could have a big impact on analyses downstream that use exoplanet masses (e.g. density, bulk compositions).
This motivates using the Learned Harmonic Mean Estimator (LHME) via the harmonic package to compare our models and see which one is preferred.
Bayesian Evidence Estimation with harmonic
Now for the main event. We use harmonic to estimate the Bayesian evidence \(\ln \mathcal{Z}\) for each model from the MCMC posterior samples. The workflow is:
Package the MCMC samples into harmonic’s
ChainsformatSplit into training (25%) and inference (75%) sets
Train a normalizing flow (
RQSplineModel) on the training samples to learn the posterior distributionEstimate the evidence from the inference samples
One important detail: ravest’s get_samples_np(flat=False) returns arrays with shape (nsteps, nwalkers, ndim), but harmonic’s add_chains_3d expects (nchains, nsamples, ndim). We need to transpose axes 0 and 1.
harmonic provides a few diagnostics to check the reliability of the evidence estimate. First, the check_basic_diagnostic command built in to Harmonic. Secondly, we check the kurtosis \(\kappa\) which should be close to 3 (i.e. roughly Gaussian) – anything far away from that indicates an unreliable estimate due to heavy tails. Thirdly, the nu^2/sigma^2 ratio (\(\nu^2/\sigma^2\)) compares the estimated variance of the variance to its theoretical target; values close to the target indicate sufficient effective samples. More information on these in §3.2 of McEwen et al. (2021).
The main caveat is that all of the posteriors need to be unimodal, constrained and well-behaved: because the distribution harmonic makes needs to be constrained within the tails of the posterior distributions (i.e. within your samples). For full details on harmonic, see the documentation and GitHub repository.
1-Planet Model Evidence
# Transpose from (nsteps, nwalkers, ndim) to (nchains, nsamples, ndim)
samples_1pl_hm = np.ascontiguousarray(samples_1pl.transpose(1, 0, 2))
lnprob_1pl_hm = np.ascontiguousarray(lnprob_1pl.T)
print(f"Samples shape for harmonic: {samples_1pl_hm.shape} (nchains, nsamples, ndim)")
print(f"Lnprob shape for harmonic: {lnprob_1pl_hm.shape} (nchains, nsamples)")
chains_1pl = hm.Chains(ndim=fitter_1pl.ndim)
chains_1pl.add_chains_3d(samples=samples_1pl_hm, ln_posterior=lnprob_1pl_hm)
Samples shape for harmonic: (32, 1000, 11) (nchains, nsamples, ndim)
Lnprob shape for harmonic: (32, 1000) (nchains, nsamples)
chains_train_1pl, chains_test_1pl = hm.utils.split_data(
chains_1pl, training_proportion=0.25
)
model_1pl = hm.model.RQSplineModel(
fitter_1pl.ndim,
n_layers=3,
n_bins=8,
hidden_size=[32, 32],
spline_range=(-10.0, 10.0),
standardize=True,
temperature=0.8,
)
model_1pl.fit(chains_train_1pl.samples, epochs=100, batch_size=128, verbose=True)
Training NF: 100%|██████████| 100/100 [00:12<00:00, 7.82it/s]
ev_1pl = hm.Evidence(nchains=chains_test_1pl.nchains, model=model_1pl)
ev_1pl.add_chains(chains_test_1pl)
ln_Z_1pl, ln_Z_1pl_std = ev_1pl.compute_ln_evidence()
check_1pl = ev_1pl.check_basic_diagnostic()
kurtosis_1pl = ev_1pl.kurtosis
n_eff_1pl = ev_1pl.n_eff
nu2_sigma2_1pl = np.exp(0.5 * ev_1pl.ln_evidence_inv_var_var - ev_1pl.ln_evidence_inv_var)
target_nu2_sigma2_1pl = np.sqrt(2.0 / (n_eff_1pl - 1))
print(f"1-Planet Model: ln(Z) = {ln_Z_1pl:.3f} +/- {ln_Z_1pl_std:.3f}")
print(f"Basic diagnostic: {'PASS' if check_1pl else 'FAIL'}")
print(f"Kurtosis: {kurtosis_1pl:.3f} (aim for ~3)")
print(f"nu^2/sigma^2: {nu2_sigma2_1pl:.4f} (target: {target_nu2_sigma2_1pl:.4f})")
print(f"n_eff: {n_eff_1pl:.1f}")
1-Planet Model: ln(Z) = -385.544 +/- -390.184
Basic diagnostic: PASS
Kurtosis: 2.339 (aim for ~3)
nu^2/sigma^2: 0.2336 (target: 0.2949)
n_eff: 24.0
2-Planet Model Evidence
samples_2pl_hm = np.ascontiguousarray(samples_2pl.transpose(1, 0, 2))
lnprob_2pl_hm = np.ascontiguousarray(lnprob_2pl.T)
print(f"Samples shape for harmonic: {samples_2pl_hm.shape} (nchains, nsamples, ndim)")
print(f"Lnprob shape for harmonic: {lnprob_2pl_hm.shape} (nchains, nsamples)")
chains_2pl = hm.Chains(ndim=fitter_2pl.ndim)
chains_2pl.add_chains_3d(samples=samples_2pl_hm, ln_posterior=lnprob_2pl_hm)
Samples shape for harmonic: (32, 1000, 16) (nchains, nsamples, ndim)
Lnprob shape for harmonic: (32, 1000) (nchains, nsamples)
chains_train_2pl, chains_test_2pl = hm.utils.split_data(
chains_2pl, training_proportion=0.25
)
model_2pl = hm.model.RQSplineModel(
fitter_2pl.ndim,
n_layers=3,
n_bins=8,
hidden_size=[32, 32],
spline_range=(-10.0, 10.0),
standardize=True,
temperature=0.8,
)
model_2pl.fit(chains_train_2pl.samples, epochs=100, batch_size=128, verbose=True)
Training NF: 100%|██████████| 100/100 [00:16<00:00, 6.11it/s]
ev_2pl = hm.Evidence(nchains=chains_test_2pl.nchains, model=model_2pl)
ev_2pl.add_chains(chains_test_2pl)
ln_Z_2pl, ln_Z_2pl_std = ev_2pl.compute_ln_evidence()
check_2pl = ev_2pl.check_basic_diagnostic()
kurtosis_2pl = ev_2pl.kurtosis
n_eff_2pl = ev_2pl.n_eff
nu2_sigma2_2pl = np.exp(0.5 * ev_2pl.ln_evidence_inv_var_var - ev_2pl.ln_evidence_inv_var)
target_nu2_sigma2_2pl = np.sqrt(2.0 / (n_eff_2pl - 1))
print(f"2-Planet Model: ln(Z) = {ln_Z_2pl:.3f} +/- {ln_Z_2pl_std:.3f}")
print(f"Basic diagnostic: {'PASS' if check_2pl else 'FAIL'}")
print(f"Kurtosis: {kurtosis_2pl:.3f} (aim for ~3)")
print(f"nu^2/sigma^2: {nu2_sigma2_2pl:.4f} (target: {target_nu2_sigma2_2pl:.4f})")
print(f"n_eff: {n_eff_2pl:.1f}")
2-Planet Model: ln(Z) = -366.428 +/- -370.132
Basic diagnostic: PASS
Kurtosis: 2.077 (aim for ~3)
nu^2/sigma^2: 0.2110 (target: 0.2949)
n_eff: 24.0
Model Comparison
delta_ln_Z = ln_Z_2pl - ln_Z_1pl
print(f"{'Model':<20} {'ln(Z)':>12} {'std':>8}")
print("-" * 42)
print(f"{'1-Planet':<20} {ln_Z_1pl:>12.3f} {ln_Z_1pl_std:>8.3f}")
print(f"{'2-Planet':<20} {ln_Z_2pl:>12.3f} {ln_Z_2pl_std:>8.3f}")
print(f"\nDelta ln(Z) = ln(Z_2pl) - ln(Z_1pl) = {delta_ln_Z:.3f}")
Model ln(Z) std
------------------------------------------
1-Planet -385.544 -390.184
2-Planet -366.428 -370.132
Delta ln(Z) = ln(Z_2pl) - ln(Z_1pl) = 19.116
We can interpret the (log-)Bayes factor \(\Delta\ln{\mathcal{Z}}\) using the scale given in Trotta (2008): inconclusive for \(<1.0\), weak evidence up to \(2.5\), moderate evidence up to \(5.0\), and strong evidence above that.
We find \(\ln \mathcal{Z}_{\text{1pl}} = -385.544\) and \(\ln \mathcal{Z}_{\text{2pl}} = -366.428\), giving \(\Delta \ln \mathcal{Z_{21}} = \ln \mathcal{Z_{2}} - \ln \mathcal{Z_{1}} = 19.1\) – very strong evidence in favour of the 2-planet model! It goes to show that if you aren’t careful and test alternative models, you could end up with a result that looks good at first glance (the one-planet model) when another solution may be preferable, with big implications for the masses (and everything that uses masses downstream.)
Note that unlike AIC or BIC, the Bayesian evidence naturally penalises model complexity through the prior volume. AIC and BIC penalise models only based on the number of free parameters and the likelihood; the priors aren’t involved. A more complex model must fit the data sufficiently better to justify the additional parameters (incorporating the increased prior volume) – the Bayesian Occam’s razor is built in. Obviously though, this means you should have normalised, justified, well-behaved sensible priors for all parameters and hyperparameters, or it may not be a fair comparison!
Comparison with information criteria
As a quick check, we can also compare the models using the classic Bayesian Information Criterion and Akaike Information Criterion (corrected):
chi2_1pl = fitter_1pl.calculate_chi2(dict_1pl)
aicc_1pl = fitter_1pl.calculate_aicc(dict_1pl)
bic_1pl = fitter_1pl.calculate_bic(dict_1pl)
chi2_2pl = fitter_2pl.calculate_chi2(dict_2pl)
aicc_2pl = fitter_2pl.calculate_aicc(dict_2pl)
bic_2pl = fitter_2pl.calculate_bic(dict_2pl)
print(f"{'Metric':<20} {'1-Planet':>12} {'2-Planet':>12}")
print("-" * 46)
print(f"{'Free parameters':<20} {fitter_1pl.ndim:>12} {fitter_2pl.ndim:>12}")
print(f"{'Chi2':<20} {chi2_1pl:>12.2f} {chi2_2pl:>12.2f}")
print(f"{'AICc':<20} {aicc_1pl:>12.2f} {aicc_2pl:>12.2f} (lower is better)")
print(f"{'BIC':<20} {bic_1pl:>12.2f} {bic_2pl:>12.2f} (lower is better)")
print(f"{'ln(Z) [harmonic]':<20} {ln_Z_1pl:>12.3f} {ln_Z_2pl:>12.3f} (higher is better)")
Metric 1-Planet 2-Planet
----------------------------------------------
Free parameters 11 16
Chi2 111.72 110.15
AICc 745.49 700.59 (lower is better)
BIC 773.71 739.90 (lower is better)
ln(Z) [harmonic] -385.544 -366.428 (higher is better)
Summary
We used harmonic to estimate the Bayesian evidence for 1-planet and 2-planet models of TOI-544. The key steps were:
Fit both models with ravest’s
GPFitterto obtain posterior MCMC samplesTranspose the samples from ravest’s
(nsteps, nwalkers, ndim)to harmonic’s(nchains, nsamples, ndim)formatTrain an
RQSplineModelnormalizing flow on a subset of samplesEstimate the evidence from the remaining samples
The 2-planet model is strongly preferred, consistent with both the Bayesian evidence and classical information criteria. This shows the importance of testing your models to ensure you have the most appropriate one. You could consider testing not only the number of planets, but maybe circular vs eccentric orbits, the choices of priors and hyperpriors, whether to use a GP (and different kernels), whether to fit a long-term linear or quadratic trend. Do bear in mind that you need to be comparing the same data in your models each time though.
Further reading:
McEwen et al. (2021) — the learned harmonic mean estimator
Trotta (2008) — Bayesian model comparison review