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.28942711e-06]
[-7.57241205e-07 8.70506149e-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:35<00:00, 83.86it/s]
State([[-1.96822071e+00 1.76008758e+00 5.16908388e-01 -6.16566764e-03
-6.91729744e-03 -2.08150685e-03]
[-2.05912774e+00 1.80495978e+00 5.23776581e-01 -8.88490895e-03
-5.45571174e-03 -1.86267435e-03]
[-2.33746858e+00 1.94242048e+00 5.44814483e-01 1.08505499e-02
-1.49813626e-02 -3.32194792e-03]
[-2.02424189e+00 1.78774947e+00 5.21143992e-01 -2.43636696e-03
-8.71191275e-03 -2.35744145e-03]
[-1.94433684e+00 1.74827812e+00 5.15107936e-01 -1.01119693e-02
-4.97741261e-03 -1.78914516e-03]
[-2.18225127e+00 1.86576577e+00 5.33079723e-01 -8.45587578e-04
-9.32623021e-03 -2.45384843e-03]
[-2.31090005e+00 1.92930402e+00 5.42805031e-01 1.39124294e-02
-1.65377358e-02 -3.56018394e-03]
[-1.84877441e+00 1.70110284e+00 5.07882995e-01 -1.48046542e-02
-2.74683095e-03 -1.44081864e-03]
[-1.91494612e+00 1.73377534e+00 5.12883994e-01 -1.26650596e-02
-3.73678395e-03 -1.59415304e-03]
[-2.37567667e+00 1.96127077e+00 5.47699900e-01 1.68759961e-02
-1.79418944e-02 -3.77699542e-03]
[-2.20644652e+00 1.87771702e+00 5.34906884e-01 6.69292822e-03
-1.30556672e-02 -3.02377724e-03]
[-2.07592317e+00 1.81325549e+00 5.25052792e-01 -3.35531474e-03
-8.19450181e-03 -2.28160771e-03]
[-2.13566644e+00 1.84275458e+00 5.29553309e-01 -3.10901890e-03
-8.25027354e-03 -2.28956320e-03]
[-2.02843875e+00 1.78979986e+00 5.21445078e-01 -1.04662473e-02
-4.70338166e-03 -1.74170989e-03]
[-1.85925917e+00 1.70628901e+00 5.08679428e-01 -1.00618179e-02
-5.09906241e-03 -1.80267536e-03]
[-2.08169147e+00 1.81611592e+00 5.25477853e-01 -8.21647528e-04
-9.45468268e-03 -2.47269487e-03]
[-2.27191261e+00 1.91005125e+00 5.39855696e-01 1.05297458e-02
-1.48985261e-02 -3.30887036e-03]
[-1.93474893e+00 1.74356234e+00 5.14380488e-01 -8.70882233e-03
-5.68722793e-03 -1.89531385e-03]
[-2.04732011e+00 1.79913437e+00 5.22885893e-01 -9.49516838e-03
-5.16902669e-03 -1.81588506e-03]
[-2.14732547e+00 1.84850675e+00 5.30438013e-01 -2.87833985e-03
-8.35022545e-03 -2.30465403e-03]
[-2.00992064e+00 1.78068028e+00 5.20059185e-01 -6.80738269e-03
-6.55264811e-03 -2.02650283e-03]
[-1.97063601e+00 1.76128346e+00 5.17082742e-01 -6.91817103e-03
-6.53844310e-03 -2.02328330e-03]
[-1.94915868e+00 1.75066381e+00 5.15468143e-01 -9.25179484e-03
-5.39828204e-03 -1.85059309e-03]
[-2.07657281e+00 1.81358843e+00 5.25093483e-01 -1.87100395e-03
-8.93703729e-03 -2.39200582e-03]
[-2.40247988e+00 1.97451911e+00 5.49722621e-01 1.63922923e-02
-1.76684769e-02 -3.73476231e-03]
[-2.07398907e+00 1.81231496e+00 5.24904179e-01 -9.20992050e-04
-9.41486663e-03 -2.46750088e-03]
[-2.10847182e+00 1.82932770e+00 5.27494573e-01 -4.64985209e-03
-7.51294601e-03 -2.17639228e-03]
[-1.96254928e+00 1.75728893e+00 5.16482909e-01 -7.67602865e-03
-6.17297954e-03 -1.96896813e-03]
[-2.41585601e+00 1.98110971e+00 5.50720799e-01 1.45470587e-02
-1.67325828e-02 -3.58706278e-03]
[-1.92501842e+00 1.73875630e+00 5.13640044e-01 -8.53221830e-03
-5.79027082e-03 -1.90910001e-03]
[-2.05877608e+00 1.80479525e+00 5.23753028e-01 -7.71795684e-03
-6.03888726e-03 -1.95135942e-03]
[-2.01983214e+00 1.78557716e+00 5.20800641e-01 -2.79165198e-03
-8.54337177e-03 -2.33213322e-03]], log_prob=[-1.19979528 -3.84520422 -6.81261065 -2.45972481 -8.1471832 -3.39452986
-4.47665361 -1.91957865 -1.46007216 -8.78583821 -2.37657684 -2.75830232
-3.03551128 -6.821811 -3.89392403 -1.63271288 -3.17950212 -2.07684559
-3.83926129 -2.7507804 -1.25251194 -4.34338908 -1.50715187 -0.57740064
-4.30133866 -1.71881732 -5.30581834 -0.97396473 -5.58051141 -3.72617898
-4.27630643 -4.24591076], blobs=None, random_state=('MT19937', array([2724096517, 1893056562, 3803867166, 774504455, 3845294866,
1499388596, 151782859, 2670781967, 415446228, 2498140160,
3724110279, 1077785034, 2955547244, 4109581407, 763361584,
1869461708, 1101153246, 1163735629, 3891517902, 142817592,
2855173452, 102927600, 2594680300, 1407444533, 2121478398,
2621400654, 1242796601, 1404506295, 194007921, 753952106,
280608449, 2189235990, 3140616207, 893145539, 999358929,
1697024349, 4082581176, 389738188, 1466317789, 3628843259,
3841824295, 2786791246, 2995082531, 3876138939, 4200102978,
1914111299, 3763898801, 248036533, 192966135, 788520892,
858953896, 1508317539, 3946878993, 2974896184, 1654273646,
3224934390, 2134270537, 4270785549, 1394962561, 324820768,
3846238584, 1105881865, 2505983590, 3582476641, 718076105,
1093786117, 769753218, 2859241595, 398609831, 1732334576,
2692203869, 1009270795, 2816776966, 1905575356, 3157724965,
1324607100, 1704464103, 3626343281, 504667210, 691508037,
1041588818, 2335484711, 3676706668, 2949831297, 3542376142,
920615681, 3750441174, 2487827270, 676590208, 2642494002,
3058977839, 1098052273, 3637807408, 1889008602, 3196991081,
4166513638, 3498630446, 2245230441, 1532403964, 3279168758,
2902507880, 1249352614, 2403927218, 3414631852, 3000308251,
3428757020, 3465830014, 2503958360, 2745268469, 1527965035,
1639704128, 1903257226, 628097221, 2193271122, 693171246,
3911299897, 2078924191, 3291625791, 3650390221, 896055415,
2947442082, 188997243, 363148380, 2368092741, 793796504,
819510019, 1166040437, 381341126, 986781311, 3145476137,
1283367363, 3093569359, 4013175021, 4069713511, 535979315,
820124380, 2738805444, 997908974, 3121914995, 4054221742,
1257173061, 4111996714, 1516502472, 1741377291, 2211489857,
2453862168, 3931274988, 1987389108, 4218229271, 4140075286,
2150000900, 2015167886, 188021492, 1083512349, 1732478784,
3130505729, 2699600197, 127812934, 400939901, 1939899653,
3767641323, 1469144850, 2631639866, 1707293606, 3807042633,
903140625, 354131720, 3904271655, 4222731761, 4243955695,
1979424119, 2128226245, 768773983, 376319895, 4002114714,
1806070451, 1554090949, 3461583635, 3027837863, 2405655021,
2576336677, 1099118085, 1631072359, 1559359324, 2052607724,
616122838, 3118546122, 361817154, 2442800777, 1324154012,
694261312, 1238085406, 4068112750, 1749283947, 772259256,
3017732952, 1944115519, 3233954025, 867469169, 989430678,
1873491847, 427358404, 4092534901, 3631697970, 3880162316,
2356512256, 1426705827, 1347239755, 2656970955, 136473072,
3960796213, 1612467953, 2401218534, 3467277806, 1049352910,
4158269526, 2173090115, 3682566382, 2745000818, 2548720212,
2964506916, 1213221474, 1982379704, 4218980926, 4126521040,
117751430, 1682816074, 608553340, 1047514793, 597779300,
2988681681, 1452336647, 1687232188, 1580972059, 1677819049,
3455459686, 438847133, 3980795801, 2128761215, 2748304668,
2440991282, 1588368552, 3864292846, 802634004, 103656925,
134802906, 2432913881, 3707033781, 1276884587, 982203932,
1633048098, 682170102, 2371795572, 4048219142, 1697912285,
3943754334, 1149656285, 641983766, 3926143178, 2772063517,
1024689625, 606886537, 2471886433, 2647130997, 1508910397,
3771525387, 4191287885, 4069792223, 4034222222, 3884115584,
4025219644, 2712849831, 4036421075, 3320813223, 1710639379,
682662872, 2642892004, 3128230568, 3129173082, 5813376,
4238382866, 833724802, 4223384870, 1794160842, 3189399729,
877410280, 4210339408, 1575697724, 4235512066, 1347210661,
2360315341, 3670320161, 996103834, 1764575294, 2848859949,
3992478259, 4059906755, 3135375477, 550215042, 2689149630,
57350636, 1814131793, 2185024924, 229741504, 697093330,
2104810009, 1401051836, 2761982987, 151423827, 1108226996,
2702850874, 4011103988, 822285666, 1139535660, 1568304987,
305708569, 231917218, 4245730868, 3733986488, 1111397228,
2827772848, 263131221, 1492018352, 1125047230, 1558187581,
4094545261, 1168612376, 3451860116, 1263393880, 2313570488,
2648731200, 644169300, 3505200065, 1252748055, 1921833541,
2863668351, 4175589076, 2766007843, 3453550363, 2948265950,
3031519107, 1830598583, 484616744, 298374041, 2200973136,
1389855369, 1312998666, 1713670912, 93565877, 2728325620,
3793728814, 2758903484, 3985598225, 320508485, 1774955148,
2302887683, 1688879531, 1593574381, 2970379844, 3521011060,
2597137079, 3831022772, 2737071359, 3676518208, 1046144127,
3286271404, 2729328595, 215742586, 4181249444, 2697947919,
1732974683, 1561854938, 663069082, 1192819124, 3910626652,
475771488, 545789183, 169230429, 2185612942, 3334854244,
4213974090, 3163672778, 694012019, 440857599, 3703873473,
22286190, 1857496612, 1756809401, 3999174810, 1536642371,
3372073365, 1139037829, 2412516072, 1237080439, 973432518,
2482046915, 1396465147, 1639783310, 1615554306, 2638767887,
945064951, 1620581719, 3692500700, 916345108, 3230023829,
2962569177, 2705935517, 2794838524, 2963229971, 1882546623,
2474966302, 1879643543, 1312252551, 1565054159, 890042088,
4207139455, 2298518147, 1279495278, 976816329, 3042950377,
364721428, 828231976, 4249743977, 3336246382, 1010707076,
221890524, 1912913706, 337479643, 2464111046, 2008380067,
2070904335, 1246907064, 4239667027, 3299552602, 512038677,
539519068, 1178821660, 1119982741, 2652695029, 918873588,
4272106825, 1701062164, 1796290299, 1748165810, 2034517970,
1958057138, 4277133363, 756798729, 2483644986, 3399961883,
4157844448, 2744717689, 1073287509, 1454704012, 1902170895,
2634944329, 2453214156, 736991643, 3271528807, 1885378954,
2479154824, 980438070, 4265021605, 3752530378, 3736824317,
1757520855, 1866934910, 1970978498, 1166836960, 3353243028,
332192049, 1600719988, 2844302399, 1578028479, 3509498031,
90155550, 4029022352, 722404085, 393750119, 1257768179,
694554301, 1509647, 2722731603, 3428968056, 2718823378,
3152475512, 371404350, 1991218301, 2242837899, 1117466656,
534907184, 938807292, 713799731, 3489065028, 2620021297,
478754530, 3970179860, 1599532102, 727265397, 376301630,
3019287872, 20000885, 4016692395, 824626612, 181669284,
2965986015, 1053328108, 3359915423, 1069126300, 1068430316,
2486048097, 3544396146, 1044903483, 1178989710, 2940985091,
3220642481, 2177664896, 4175279852, 3582894093, 266930068,
1741217127, 2918337428, 569322765, 1801147033, 1872290446,
2126829718, 2133928174, 2120639860, 1199994500, 2338337955,
2769756217, 1172536183, 3090437907, 1393083532, 210333880,
2311531116, 3059696637, 1700211371, 1279214501, 365974977,
1682818973, 1735481437, 2070220390, 679946300, 2525074498,
3656971002, 2651420342, 2454262674, 3092341336, 2943389204,
1882316960, 3875711151, 868827502, 832927346, 1183533271,
582794812, 2879439193, 4127024575, 195696361, 541676672,
4214666265, 3657473421, 3840042753, 520600148, 4109448195,
1887066305, 834271711, 2811559537, 2049249908, 1962556484,
4123641756, 3563663963, 3093611643, 616036765, 3401179841,
4153525013, 36472873, 2410821933, 1286109411, 411835227,
366336027, 2251465342, 2340604530, 1139486379, 2635527480,
645033653, 2892329924, 983101017, 3229237240, 1166867485,
213384789, 388546522, 361823361, 3851036786, 2019619381,
2669974580, 2296532796, 696083494, 2290853879, 161584884,
2214552461, 337656556, 1100932266, 1002479921, 1838949706,
3868573082, 1300561838, 3078462166, 1087241853, 3814145971,
3134722087, 3545255766, 3661848442, 2277461603, 451493515,
890876291, 3836172322, 89468544, 159912009, 1968537069,
4035399868, 3392695765, 2966127745, 3676035110], dtype=uint32), 601, 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.51it/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, 76.62it/s]
acceptance fraction: 0.518
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 [01:01<00:00, 48.54it/s]
acceptance fraction: 0.516
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()
