MCMC Orbit Fitting#
The maximum likelihood tutorial showed how to find the MAP orbit and compute likelihood gradients. Here we go one step further and sample the full posterior using emcee, an affine-invariant ensemble sampler.
The second part of this notebook adds progressively more observations and overlays the resulting posteriors — a demonstration of how quickly the orbit uncertainty shrinks once a second opposition’s worth of data is included.
Note that we don’t do any of the stuff that you really should do when using MCMC, like checking for convergence, discarding burn-in, etc. This is just a quick demo of how to set up the sampler and visualize the results: when using MCMC for real, you should be much more careful about these things.
import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
import numpy as np
import astropy.units as u
import corner
import emcee
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
from astropy.coordinates import SkyCoord
from astropy.time import Time
from astroquery.jplhorizons import Horizons
from jorbit import Observations, Particle
from jorbit.data.constants import SPEED_OF_LIGHT
from jorbit.utils.states import CartesianState
Part 1: Single-opposition MCMC#
We will use the same fake dataset as the maximum likelihood tutorial: nine astrometric observations of (274301) Wikipedia spread over three nights in January 2025, queried from JPL Horizons.
nights = [Time("2025-01-01 07:00"), Time("2025-01-02 07:00"), Time("2025-01-05 07:00")]
times = []
for n in nights:
times.extend([n + i * 1 * u.hour for i in range(3)])
times = Time(times)
obj = Horizons(id="274301", location="695@399", epochs=times.utc.jd)
pts = obj.ephemerides(extra_precision=True, quantities="1")
coords = SkyCoord(pts["RA"], pts["DEC"], unit=(u.deg, u.deg))
times = Time(pts["datetime_jd"], format="jd", scale="utc")
obs = Observations(
observed_coordinates=coords,
times=times,
observatories="kitt peak",
astrometric_uncertainties=1 * u.arcsec,
)
For the initial state we pull the true barycentric state vector from Horizons. In practice you would start from a max_likelihood result; here we start at the truth to keep the notebook self-contained.
obj = Horizons(id="274301", location="500@0", epochs=times.tdb.jd[0])
vecs = obj.vectors(refplane="earth")
true_x0 = jnp.array([vecs["x"], vecs["y"], vecs["z"]]).T[0]
true_v0 = jnp.array([vecs["vx"], vecs["vy"], vecs["vz"]]).T[0]
p = Particle(
x=true_x0,
v=true_v0,
time=times[0],
name="274301 Wikipedia",
observations=obs,
)
The static integration path#
When a Particle is created with observations, jorbit automatically builds a static_residuals function by precomputing the IAS15 step sizes and perturber (planet + asteroid) positions along the reference orbit. The precomputed data are frozen into a JIT-compiled closure, so each subsequent likelihood evaluation avoids rerunning the adaptive integrator and avoids re-querying the ephemeris — making it much faster for the thousands of calls that MCMC demands.
We can verify it is working by checking that the true-orbit residuals are near zero:
# static_residuals returns (n_obs, 2) tangent-plane residuals in arcseconds
print(p.static_residuals(p.cartesian_state))
[[-3.32660383e-07 -4.91174805e-07]
[ 8.18208454e-07 -1.30179371e-06]
[-5.41303605e-07 -1.09084419e-06]
[-7.47107233e-07 -1.28947281e-06]
[-7.57241205e-07 8.70551844e-07]
[ 7.41829122e-07 9.18007346e-07]
[ 9.92006244e-07 1.02686012e-06]
[-1.70074093e-07 -3.05681937e-07]
[-4.09222656e-07 -1.75906677e-07]]
Defining the log-probability for emcee#
emcee samples in a flat parameter space. We use the six Cartesian components [x, y, z, vx, vy, vz] in barycentric ICRS (AU and AU/day). The log-likelihood is a simple chi-squared built from static_residuals:
$$\ln p(\theta) = -\frac{1}{2} \sum_{i} \left(\frac{\xi_i^2 + \eta_i^2}{\sigma^2}\right)$$
where $\xi_i, \eta_i$ are the tangent-plane residuals in arcseconds and $\sigma = 1$ arcsec.
A few implementation notes:
time=jnp.array(0.0)is correct because jorbit stores all JAX-visible times as offsets from the particle’s reference epoch (t_ref), so the initial state is always at offset 0.static_residualsis alreadyjax.jit-compiled; the first call triggers compilation (~10–30 s), after which evaluations are very fast.
def log_prob(params):
state = CartesianState(
x=jnp.array(params[:3]).reshape(1, 3),
v=jnp.array(params[3:]).reshape(1, 3),
acceleration_func_kwargs={"c2": SPEED_OF_LIGHT**2},
# state is at p's epoch (offset 0 from t_ref_jd anchor)
relative_time=jnp.array(0.0),
time_reference=p._t_ref_jd,
)
residuals = p.static_residuals(state) # (n_obs, 2) arcsec
return float(-0.5 * jnp.sum(residuals**2)) # chi-sq, sigma = 1 arcsec
Running the sampler#
We initialize 32 walkers in a tight ball around the known MAP and run for 3000 steps. With 1 arcsec uncertainties and only nine observations over a short arc, the posterior is broad — especially along the radial direction.
ndim = 6
nwalkers = 32
p0_map = np.concatenate([true_x0, true_v0])
# Start walkers within ~km (position) and ~mm/s (velocity) of the MAP.
# emcee's stretch move will carry them out to explore the full posterior width.
scale = np.array([1e-7, 1e-7, 1e-7, 1e-10, 1e-10, 1e-10])
initial_pos = p0_map + scale * np.random.default_rng(42).standard_normal(
(nwalkers, ndim)
)
sampler = emcee.EnsembleSampler(nwalkers, ndim, log_prob)
sampler.run_mcmc(initial_pos, 3000, progress=True)
100%|██████████| 3000/3000 [00:34<00:00, 85.84it/s]
State([[-2.08978826e+00 1.82011952e+00 5.26095313e-01 -1.10061930e-03
-9.30708705e-03 -2.45303267e-03]
[-2.18392455e+00 1.86660083e+00 5.33215187e-01 5.87143616e-03
-1.26749460e-02 -2.97032908e-03]
[-2.08289901e+00 1.81671003e+00 5.25569690e-01 -5.43744787e-03
-7.15173242e-03 -2.11949069e-03]
[-2.22937633e+00 1.88904618e+00 5.36643445e-01 6.68907009e-03
-1.30299843e-02 -3.02332289e-03]
[-2.06121914e+00 1.80599447e+00 5.23925707e-01 -4.72783969e-03
-7.52486614e-03 -2.17596362e-03]
[-2.02904939e+00 1.79011457e+00 5.21506092e-01 -8.00309084e-03
-5.93226989e-03 -1.93577188e-03]
[-1.93048128e+00 1.74143518e+00 5.14048168e-01 -1.11916884e-02
-4.45303508e-03 -1.70266066e-03]
[-1.97966378e+00 1.76572578e+00 5.17769823e-01 -9.78402559e-03
-5.09791848e-03 -1.80459715e-03]
[-2.41253643e+00 1.97948217e+00 5.50473587e-01 2.12250249e-02
-2.00657601e-02 -4.09699811e-03]
[-2.16497737e+00 1.85724572e+00 5.31777416e-01 2.98267498e-03
-1.12551467e-02 -2.74998944e-03]
[-2.08332801e+00 1.81692441e+00 5.25607790e-01 -4.59296940e-03
-7.57449329e-03 -2.18806456e-03]
[-1.95369574e+00 1.75290223e+00 5.15799365e-01 -1.10014345e-02
-4.52060561e-03 -1.71314212e-03]
[-2.01089608e+00 1.78114516e+00 5.20135801e-01 -6.23031632e-03
-6.83403896e-03 -2.07221392e-03]
[-2.00030867e+00 1.77592230e+00 5.19328847e-01 -8.21766106e-03
-5.85792433e-03 -1.92143129e-03]
[-2.01499433e+00 1.78318204e+00 5.20438363e-01 -5.42428538e-03
-7.23556045e-03 -2.13224826e-03]
[-2.11087327e+00 1.83051930e+00 5.27679713e-01 -1.76826451e-03
-8.94644247e-03 -2.39603700e-03]
[-2.10441255e+00 1.82731577e+00 5.27200881e-01 -4.59576076e-03
-7.54203798e-03 -2.18170594e-03]
[-1.99228143e+00 1.77195279e+00 5.18728454e-01 -1.08193839e-02
-4.56945504e-03 -1.72383448e-03]
[-2.03814420e+00 1.79461464e+00 5.22185962e-01 -5.31377016e-03
-7.26224770e-03 -2.13309029e-03]
[-2.24592599e+00 1.89722690e+00 5.37891262e-01 8.59825849e-03
-1.39646987e-02 -3.16597199e-03]
[-2.19388289e+00 1.87152063e+00 5.33957677e-01 8.74575451e-03
-1.40958768e-02 -3.18393543e-03]
[-1.85071013e+00 1.70206418e+00 5.08033953e-01 -1.20470172e-02
-4.11872720e-03 -1.65372244e-03]
[-1.93990888e+00 1.74610291e+00 5.14762094e-01 -1.19699550e-02
-4.05538736e-03 -1.64169757e-03]
[-1.94850237e+00 1.75034629e+00 5.15416626e-01 -1.14694163e-02
-4.29890569e-03 -1.68051560e-03]
[-1.83270707e+00 1.69317941e+00 5.06663740e-01 -1.01368463e-02
-5.09188986e-03 -1.79810418e-03]
[-1.85150918e+00 1.70244947e+00 5.08086364e-01 -1.79863242e-02
-1.15670294e-03 -1.19866340e-03]
[-2.08920883e+00 1.81982999e+00 5.26053023e-01 -2.60850550e-03
-8.55888462e-03 -2.33662991e-03]
[-2.15205617e+00 1.85086462e+00 5.30796404e-01 3.79418311e-03
-1.16752366e-02 -2.81107223e-03]
[-2.30477512e+00 1.92627403e+00 5.42340250e-01 1.34706444e-02
-1.63255748e-02 -3.52950056e-03]
[-2.30712993e+00 1.92741776e+00 5.42527418e-01 1.14939283e-02
-1.53313036e-02 -3.38009273e-03]
[-2.13619309e+00 1.84302581e+00 5.29592850e-01 4.25314595e-03
-1.19183405e-02 -2.85051610e-03]
[-2.02641166e+00 1.78881331e+00 5.21297664e-01 -8.30603283e-03
-5.78320678e-03 -1.90889091e-03]], log_prob=[-3.10435584 -3.37917277 -1.95289263 -2.97404434 -3.18955745 -2.03474378
-5.44116693 -1.18623692 -7.12608186 -2.15957104 -4.23938842 -3.45227049
-2.51670789 -0.9941759 -1.0766713 -2.62139151 -3.3042199 -2.86103771
-3.32272451 -6.18536087 -3.20239559 -2.76566704 -1.81347025 -2.07886224
-6.3530006 -1.96115356 -3.97403263 -1.94720924 -4.05672142 -7.56669215
-4.70274893 -2.05821431], blobs=None, random_state=('MT19937', array([1721897736, 1392670817, 1513426754, 685118538, 774076185,
2404707027, 689868471, 2290839921, 267563483, 255116882,
120897149, 1274296145, 705115367, 3486589526, 2906713487,
1126558692, 3841673932, 1575663244, 2412718865, 1296728012,
1201342382, 2032609585, 610725578, 1996035226, 3517265496,
2852241109, 1439871326, 1075224205, 688560265, 3400055555,
3388872703, 1871218817, 644289709, 2705617848, 2150810167,
2450362926, 1048803018, 1802379049, 1336895588, 3170025130,
3797837131, 3732478692, 2259719293, 1029515396, 193978684,
3987615522, 2654067906, 4107602161, 2617550130, 909449520,
695386005, 1086033827, 3427548778, 2001113446, 3000137827,
2799889367, 2690765705, 1658640094, 1327518450, 2809270019,
1880482437, 4257064517, 749440864, 443208907, 3069724261,
346549571, 534537528, 447578998, 629437273, 212915477,
2696187280, 1447805657, 4206355025, 2426733149, 3883471318,
3895025406, 29368805, 1688264796, 3474860481, 762319105,
1412046081, 855831558, 3293421805, 824533778, 785344109,
3840602828, 3808242394, 2162673717, 3863947439, 2202514996,
3582336733, 1227482232, 3832654345, 1785874610, 1069324617,
4250278522, 2415822392, 2277945932, 158070946, 1303409982,
2549342572, 1586731907, 4003978224, 2924397734, 1244127165,
1957525753, 1906532513, 185508046, 1347787456, 3626965933,
1197036682, 2634916834, 3418801072, 2311258527, 2481938257,
2932280272, 4218295926, 3779547546, 2814180156, 4198652694,
1847183973, 3146703539, 2017292911, 196780988, 952239503,
193350678, 3064995194, 1378137832, 2469843530, 2802651146,
3470710494, 779057821, 2971993392, 2827947888, 3117281561,
2982220754, 1306267795, 476842360, 2883159476, 971456061,
182462552, 3276670552, 117769155, 3932754325, 694227625,
2962651921, 3771741368, 3592642749, 693899243, 1019786445,
3222080326, 1348148716, 1787594004, 3287593246, 4280177238,
528947878, 1191301373, 3972293108, 2540771604, 3243518928,
1814303059, 1325883459, 700419377, 1493498018, 1556547703,
3286796295, 1432567702, 241181142, 2437408899, 726637865,
3891846003, 1658641808, 4142098870, 576904390, 4072802654,
1970241606, 298517350, 334413106, 177747513, 3951210418,
546111079, 4241147305, 2908774973, 325744029, 3893609620,
199146614, 1367684459, 2302929816, 2677476211, 4025387153,
3069863656, 1892523491, 2518776250, 2545075979, 3332531438,
4226698231, 3955382844, 3728663202, 2399361451, 3291556211,
899875588, 2438579988, 2084678237, 160984070, 2978047928,
2692875909, 2533092833, 2733697461, 1254850879, 662981193,
2772895395, 374985472, 672676883, 2299041753, 3121926432,
3787456543, 4252535604, 1012731780, 3317759788, 2723054539,
1010318775, 348792684, 4292978466, 2777601577, 4183805450,
1928222013, 823934012, 737707596, 2783918995, 4114455777,
1624555686, 997700756, 3527853686, 1859857820, 2913439217,
1281559523, 456175527, 3704139956, 409033676, 2956360929,
245664670, 2191420477, 1745776477, 3520282100, 3145146509,
3889295482, 2245238201, 2494345745, 2634005806, 774979563,
3525090028, 713156348, 756049597, 3660969087, 1166462231,
93482492, 2078978288, 1871147094, 3107027264, 129063480,
3000857962, 717669977, 317050737, 235344973, 854942574,
3396387958, 1212667673, 3455411873, 2457410339, 2938011152,
4011292575, 2351253154, 2278909998, 2585380189, 549121698,
3535324008, 3132679102, 832298991, 3159008566, 83969205,
1838739460, 2695184571, 91940273, 360233498, 2140983341,
1566032317, 965326227, 2829098135, 541908511, 2760893891,
3180929842, 908571193, 518570224, 1721529024, 4275391790,
1436349781, 1878237548, 1039479505, 3043882123, 1901380256,
3277940364, 1659209931, 799581344, 2636743437, 3281463455,
3754380570, 816788728, 856507427, 2322649437, 996016981,
980930181, 1561615397, 1523573894, 3077893458, 3924171326,
2542096907, 1063784461, 4184181596, 1870057907, 1279692126,
1250858594, 389093763, 2515433330, 381571913, 90699579,
2499629627, 3008626557, 2309680851, 1650404939, 2111111483,
40693872, 2247696513, 2339127107, 105477612, 583772360,
2415207099, 2020379292, 2938120109, 2138449530, 1523098436,
319421871, 4009706490, 1986946059, 1469796054, 2726731854,
90260545, 3258401266, 1355957111, 3437903694, 1260748477,
441317837, 3712341283, 2212326368, 735942724, 383251442,
3801367995, 1694893976, 2435764645, 1530330365, 1584752235,
670034690, 816054498, 3455298099, 4141060950, 693684003,
1829530, 1849076789, 540849949, 2764574576, 2083734920,
2084366476, 677555175, 4183731894, 677003258, 729171372,
4123197511, 1383441735, 3814206208, 514224621, 1669273281,
1379821672, 208539301, 2186797767, 140029747, 753044318,
1432433587, 2298348433, 3158726227, 1757163803, 838724695,
3186502536, 3400261271, 104764375, 185858268, 1725961445,
1656677327, 3162338370, 3670510563, 488666026, 1792661616,
990937193, 16143456, 2530882511, 839248130, 1129902597,
3111752597, 123451504, 3265834755, 628305743, 920591824,
1455811610, 2455090419, 2750987505, 1726681682, 4112927918,
2007590136, 2592475145, 737810650, 163367193, 771828312,
486658475, 1156032362, 1113982831, 3462823906, 254536750,
3266897786, 3547871930, 74410972, 1959311010, 3162794394,
2370416998, 251739625, 4041992932, 4279591267, 2092791194,
1264367608, 1923435095, 1849860782, 738040418, 279172367,
1479587492, 3980875339, 1400162297, 3477375791, 2888556462,
1938724310, 2925731907, 2952447644, 749677657, 955476327,
2609255129, 3066406185, 3404372910, 2424351706, 2832984871,
4133311099, 3711207900, 2073884768, 3736214975, 1539158255,
343068047, 2044495407, 1805545639, 457365259, 3942482873,
4036539979, 3591243904, 1892466995, 1000563381, 1652053074,
212720810, 919473573, 3659370878, 1900545781, 1566538171,
3093477559, 3846349088, 1239338809, 946608967, 913989607,
304806526, 2583676799, 3536330581, 2511067461, 1200314613,
2942955617, 947330408, 1146236843, 893564277, 3272040223,
1689748038, 2403890284, 881726695, 3049840239, 1968956285,
3582608774, 3218686575, 3414679261, 2979240480, 296340847,
3217934287, 7314455, 1340664697, 495244686, 100464398,
4200742524, 1161114938, 1322209289, 2573878217, 2267993704,
1599035339, 3466082171, 3997181152, 2055068894, 1363565904,
1981026262, 3330018758, 2465224790, 1351042068, 31754617,
3677829956, 76957095, 2217916510, 1646531285, 3903490287,
789338927, 4150557953, 3578231671, 3870715817, 932354240,
2296485689, 2356955963, 1750972034, 3905008758, 1797680322,
99496383, 2915668987, 3893641845, 1204780972, 3069298258,
2722411348, 1354517184, 3252036627, 2983895518, 3190246646,
3651426590, 2719728852, 2905874738, 1881033320, 1201299361,
2581249964, 1142549443, 1582291805, 2801658584, 719595117,
939097498, 605147777, 796069673, 1472478298, 2110742649,
854046072, 1946053984, 2248851136, 2358667417, 2739569868,
2117395995, 612592489, 3980606795, 518355212, 4146452904,
1228672438, 4231637157, 3108268127, 1160698712, 737207189,
558680789, 2811106746, 1144897811, 2475290132, 2888332916,
2448771031, 1719042324, 2992526026, 2350170003, 167727961,
3043518736, 1841637237, 2684250085, 3187303150, 15325695,
3397473511, 2394271465, 258787432, 929943540, 1110661353,
3486711375, 3256495909, 2647111814, 2243510876, 2554303534,
1332980497, 3282357446, 3110979453, 2631363703, 2422146091,
4007525331, 3476030646, 2097372865, 2083942222, 1799950195,
1895130353, 3683766641, 2245227228, 833178533, 2724818867,
3455677066, 961484751, 1670740130, 3740686851, 2243778255,
750154863, 618160395, 1252354852, 3475799494], dtype=uint32), 39, 0, 0.0))
flat_samples = sampler.get_chain(discard=500, thin=5, flat=True)
print(f"Acceptance fraction: {np.mean(sampler.acceptance_fraction):.3f}")
print(f"Samples retained: {len(flat_samples)}")
Acceptance fraction: 0.516
Samples retained: 16000
Corner plot#
We plot the samples directly in the Cartesian state vector used by the sampler: barycentric ICRS position in AU and velocity in AU/day.
cartesian_labels = [
"$x$ (AU)",
"$y$ (AU)",
"$z$ (AU)",
"$v_x$ (AU/day)",
"$v_y$ (AU/day)",
"$v_z$ (AU/day)",
]
fig = corner.corner(
flat_samples,
labels=cartesian_labels,
show_titles=True,
title_fmt=".3g",
)
plt.suptitle(
"Posterior: 9 obs, 3 nights (single opposition)",
y=1.01,
fontsize=12,
)
plt.show()
Part 2: How the posterior shrinks with more data#
Short arcs within a single opposition leave parts of the Cartesian state vector poorly constrained, especially the line-of-sight position and velocity components. Adding nights from the same opposition helps, but adding data from a second opposition (a year or so later) is far more powerful: the Earth has completed nearly one full orbit while the asteroid has moved appreciably, providing strong leverage on the heliocentric distance and motion.
Below we run three MCMC chains:
The baseline: three nights from January 2025
Extended single opposition: seven nights in January 2025
Two oppositions: the seven January 2025 nights plus three nights in October 2025
All chains share the same true starting state and the same sampler settings.
Run 1 — baseline: 3 nights, single opposition#
print("Run 1: 3 nights (Jan 1, 2, 5 2025)")
samples1 = run_mcmc(obs)
Run 1: 3 nights (Jan 1, 2, 5 2025)
100%|██████████| 3000/3000 [00:34<00:00, 86.26it/s]
acceptance fraction: 0.516
Run 2 — extended arc: 7 nights, single opposition#
We add four more nights in January 2025 and use the + operator to combine them with the baseline observations.
nights_extra = [Time(f"2025-01-{d:02d} 07:00") for d in [8, 12, 17, 20]]
times_extra = []
for n in nights_extra:
times_extra.extend([n + i * 1 * u.hour for i in range(3)])
times_extra = Time(times_extra)
obj = Horizons(id="274301", location="695@399", epochs=times_extra.utc.jd)
pts = obj.ephemerides(extra_precision=True, quantities="1")
coords_extra = SkyCoord(pts["RA"], pts["DEC"], unit=(u.deg, u.deg))
times_extra = Time(pts["datetime_jd"], format="jd", scale="utc")
obs_extra = Observations(
observed_coordinates=coords_extra,
times=times_extra,
observatories="kitt peak",
astrometric_uncertainties=1 * u.arcsec,
)
# Combine with the baseline using the + operator
obs_ext = obs + obs_extra
print(f"Run 2: {len(obs_ext.ra)} observations over 7 nights")
samples2 = run_mcmc(obs_ext)
Run 2: 21 observations over 7 nights
100%|██████████| 3000/3000 [00:39<00:00, 75.32it/s]
acceptance fraction: 0.520
Run 3 — two oppositions: Jan 2025 + Oct 2025#
Adding observations from a second epoch roughly nine months later extends the arc to cover a substantial fraction of the asteroid’s orbital period, strongly constraining the Cartesian state vector.
nights_opp2 = [Time(f"2025-10-{d:02d} 07:00") for d in [1, 5, 10]]
times_opp2 = []
for n in nights_opp2:
times_opp2.extend([n + i * 1 * u.hour for i in range(3)])
times_opp2 = Time(times_opp2)
obj = Horizons(id="274301", location="695@399", epochs=times_opp2.utc.jd)
pts = obj.ephemerides(extra_precision=True, quantities="1")
coords_opp2 = SkyCoord(pts["RA"], pts["DEC"], unit=(u.deg, u.deg))
times_opp2 = Time(pts["datetime_jd"], format="jd", scale="utc")
obs_opp2 = Observations(
observed_coordinates=coords_opp2,
times=times_opp2,
observatories="kitt peak",
astrometric_uncertainties=1 * u.arcsec,
)
obs_all = obs_ext + obs_opp2
print(f"Run 3: {len(obs_all.ra)} observations across two epochs")
samples3 = run_mcmc(obs_all)
Run 3: 30 observations across two epochs
100%|██████████| 3000/3000 [00:59<00:00, 50.33it/s]
acceptance fraction: 0.521
Overlaying the three posteriors#
We show only the Cartesian position subspace $(x, y, z)$ for readability and overlay 68% and 95% credible contours for each run.
colors = ["C0", "C1", "C2"]
labels_runs = [
"3 nights (1 epoch, Jan 2025)",
"7 nights (1 epoch, Jan 2025)",
"10 nights (2 epochs, Jan + Oct 2025)",
]
corner_kwargs = dict(
plot_density=False,
plot_datapoints=False,
levels=[0.68, 0.95],
no_fill_contours=False,
contourf_kwargs={"alpha": 0.25},
)
fig = corner.corner(
samples1[:, :3],
labels=cartesian_labels[:3],
color=colors[0],
**corner_kwargs,
)
corner.corner(samples2[:, :3], fig=fig, color=colors[1], **corner_kwargs)
corner.corner(samples3[:, :3], fig=fig, color=colors[2], **corner_kwargs)
handles = [
mpatches.Patch(color=c, alpha=0.7, label=l) for c, l in zip(colors, labels_runs)
]
fig.axes[0].legend(handles=handles, loc="lower left", frameon=False, fontsize=9)
plt.suptitle("Posterior shrinkage with additional data", y=1.02, fontsize=13)
plt.show()
Part 3: Batched likelihoods with System#
Everything above evaluated the likelihood one walker at a time: emcee’s default mode calls log_prob separately for each walker, and each call runs Particle.static_residuals on a single state. For an affine-invariant ensemble this leaves performance on the table — the walkers are independent evaluations of the same forward model and could be computed together.
When you pass an Observations object to a System, jorbit builds a batched forward model that scores a whole (P, 6) stack of candidate states in one vectorized call, exposed as System.loglike (alongside System.residuals, System.chi2, and System.model_radec). This is exactly the interface emcee’s vectorize=True mode expects: each step, the active half-ensemble of walkers is handed to log_prob as a single (nwalkers/2, ndim) array.
A few things to keep in mind:
The component particles of a
Systemare dynamically independent — they could be entirely different asteroids. Supplyingobservationsmeans every state in the batch is scored against the same shared data. Here we exploit that to treat the emcee walkers as the batch axis: one templateSystem, an arbitrary(P, 6)batch of candidate states per call.System.loglikeuses the dynamic IAS15 dense path (a single bounded-arc buffer, no host-side stitching), not the static path from Part 1. For a gradient-free sampler like emcee that distinction is invisible; the two paths agree to sub-mas. The static path remains the route to take when you need reverse-mode gradients (e.g. for HMC/NUTS).If the observation arc is long enough to exceed one dense buffer, the affected states return
-inf(the schedule is shared across the batch, so truncation is batch-wide) — the sampler simply rejects them rather than crashing. For arcs that long, fall back toSystem.ephemeris.
We reuse the same nine-observation baseline from Part 1 so we can check the batched path lands on the same posterior.
from jorbit import System
# Reuse the Part 1 particle (built from `obs`) and wrap it in a System carrying the
# same observations. The batched forward model is built at construction time; no
# explicit "compile" step and no JAX boilerplate are needed.
system = System(
particles=[p],
observations=obs,
gravity="default solar system",
)
# System.loglike maps a (P, 6) batch of candidate states -> (P,) log-likelihoods,
# which is exactly the array signature emcee's vectorize=True mode expects. The
# whole hand-rolled CartesianState / SPEED_OF_LIGHT / offset dance from Part 1 is
# gone: you hand loglike a plain (P, 6) array of [x, y, z, vx, vy, vz] at t_ref.
def log_prob_batched(coords):
return np.asarray(system.loglike(jnp.asarray(coords)))
sampler_batched = emcee.EnsembleSampler(
nwalkers, ndim, log_prob_batched, vectorize=True
)
sampler_batched.run_mcmc(initial_pos, 3000, progress=True)
flat_batched = sampler_batched.get_chain(discard=500, thin=5, flat=True)
print(f"Acceptance fraction: {np.mean(sampler_batched.acceptance_fraction):.3f}")
100%|██████████| 3000/3000 [00:31<00:00, 95.57it/s]
Acceptance fraction: 0.516
Cross-check against Part 1#
The batched System path and the per-walker static_residuals path condition on the same nine observations, so their posteriors should agree to within sampling noise. (They will not be bit-for-bit identical: System.loglike uses the full per-observation covariance normalization and the dynamic dense integrator, whereas Part 1 used an isotropic 1 arcsec chi-square evaluated on the static path. These differ only by an additive constant in the log-probability, which does not affect the sampled distribution.) Overlaying the two confirms they land on the same posterior:
fig = corner.corner(
flat_samples,
labels=cartesian_labels,
color="C0",
plot_density=False,
plot_datapoints=False,
levels=[0.68, 0.95],
)
corner.corner(
flat_batched,
fig=fig,
color="C3",
plot_density=False,
plot_datapoints=False,
levels=[0.68, 0.95],
)
handles = [
mpatches.Patch(color="C0", label=r"per-walker $\mathtt{static\_residuals}$"),
mpatches.Patch(color="C3", label=r"batched $\mathtt{System.loglike}$"),
]
fig.legend(handles=handles, loc="upper right", frameon=False, fontsize=10)
plt.suptitle("Same posterior, two evaluation paths", y=1.02, fontsize=13)
plt.show()
