fitting

hmc(params, log_prob_and_grad, num_steps, num_leaps, step_size, comm, key)

Runs Hamilonian Monte Carlo using a leapfrog integrator to approximate Hamilonian dynamics. This is a naive implementaion that will be replaced in the future.

The parallelism model employed here is different that most samplers where each task runs a subset of the chain, instead since the rest of WITCH employs a model where the data is distributed across tasks we do that here as well. In this model the chain evolves simultaneously in all tasks, but only rank 0 actually stores the chain.

Parameters:

Name	Type	Description	Default
params	Array	The initial parameters to start the chain at.	required
log_prob_and_grad	Callable[[Array], tuple[Array, Array]]	Function that returns the log probability of the model and its gradient for a given set of params. This should take params as its first arguments, all other arguments should be fixed ahead of time (ie: using functools.partial). The returned gradient should have shape (len(params),). The compiled sampler is cached on this function (keyed on num_leaps), so reuse the same function object to avoid recompiling and use a new one whenever the model or data change.	required
num_steps	int	The number of steps to run the chain for.	required
num_leaps	int	The number of leapfrog steps to run at each step of the chain.	required
step_size	Array	The step size to use. Should be a scalar jax array. At each leapfrog step the parameters will evolve by step_size*momentum.	required
comm	Comm | Intracomm | NullComm	The MPI comm object to use.	required

Returns:

Name	Type	Description
chain	Array	The chain of samples. Will have shape (num_steps, len(params)) in the rank 0 task. Note that on tasks with rank other than 0 the actual chain is not returned, instead a dummy array of size (0,) is returned.

Source code in witch/fitting.py

def hmc(
    params: jax.Array,
    log_prob_and_grad: Callable[[jax.Array], tuple[jax.Array, jax.Array]],
    num_steps: int,
    num_leaps: int,
    step_size: jax.Array,
    comm: MPI.Comm | MPI.Intracomm | NullComm,
    key: jax.Array,
) -> tuple[jax.Array, jax.Array]:
    """
    Runs Hamilonian Monte Carlo using a leapfrog integrator to approximate Hamilonian dynamics.
    This is a naive implementaion that will be replaced in the future.

    The parallelism model employed here is different that most samplers where each task runs
    a subset of the chain, instead since the rest of WITCH employs a model where the data is
    distributed across tasks we do that here as well.
    In this model the chain evolves simultaneously in all tasks,
    but only rank 0 actually stores the chain.

    Parameters
    ----------
    params : jax.Array
        The initial parameters to start the chain at.
    log_prob_and_grad : Callable[[jax.Array], tuple[jax.Array, jax.Array]]
        Function that returns the log probability of the model and its gradient
        for a given set of params. This should take `params` as its
        first arguments, all other arguments should be fixed ahead of time
        (ie: using `functools.partial`). The returned gradient should have
        shape `(len(params),)`.
        The compiled sampler is cached on this function (keyed on `num_leaps`),
        so reuse the same function object to avoid recompiling and use a new
        one whenever the model or data change.
    num_steps : int
        The number of steps to run the chain for.
    num_leaps : int
        The number of leapfrog steps to run at each step of the chain.
    step_size : jax.Array
        The step size to use.
        Should be a scalar jax array.
        At each leapfrog step the parameters will evolve by `step_size`*`momentum`.
    comm : MPI.Comm | MPI.Intracomm | NullComm
        The MPI comm object to use.

    Returns
    -------
    chain : jax.Array
        The chain of samples.
        Will have shape `(num_steps, len(params))` in the rank 0 task.
        Note that on tasks with rank other than 0 the actual
        chain is not returned, instead a dummy array of size `(0,)` is
        returned.
    """
    rank = comm.Get_rank()
    vnorm = jax.vmap(
        partial(jax.random.normal, shape=params[0].shape, dtype=params.dtype)
    )
    npar = len(params)
    ones = jnp.ones(npar, dtype=bool)

    def _leap(_, args):
        params, momentum, _, grad, step_size = args
        momentum = momentum.at[:].add(0.5 * step_size * grad)  # kick
        params = params.at[:].add(step_size * momentum)  # drift
        # The gradient at the end of one leap is reused for the next first kick
        log_prob, grad = log_prob_and_grad(params)
        momentum = momentum.at[:].add(0.5 * step_size * grad)  # kick

        return params, momentum, log_prob, grad, step_size

    @jax.jit
    def _sample(key, params, log_prob, grad, step_size):
        mpi4jax.barrier(comm=comm)
        key = mpi4jax.bcast(key, 0, comm=comm)
        key, uniform_key = jax.random.split(key, 2)

        # generate random momentum
        momentum = vnorm(jax.random.split(key, npar))
        new_params, new_momentum, new_log_prob, new_grad, _ = jax.lax.fori_loop(
            0, num_leaps, _leap, (params, momentum, log_prob, grad, step_size)
        )

        # MH correction
        dpe = new_log_prob - log_prob
        dke = -0.5 * (jnp.sum(new_momentum**2) - jnp.sum(momentum**2))
        log_accept = dke + dpe
        accept_prob = jnp.minimum(jnp.exp(log_accept), 1)
        accept = jax.random.uniform(uniform_key) < accept_prob
        params = jax.lax.select(accept * ones, new_params, params)
        grad = jax.lax.select(accept * ones, new_grad, grad)
        log_prob = jax.lax.select(accept, new_log_prob, log_prob)

        return key, params, log_prob, grad, accept_prob

    # Cache on the likelihood function so that closures are never stale
    sample_cache = log_prob_and_grad.__dict__.setdefault("_hmc_sample", {})
    c_sample = sample_cache.get((num_leaps, npar), None)
    log_prob, grad = log_prob_and_grad(params)
    if c_sample is None:
        if rank == 0:
            print("Compiling MC sample function. This can take a few minutes!")
        t0 = time.time()
        _ = _sample(key.copy(), params.copy(), log_prob, grad, step_size.copy())
        jax.block_until_ready(_)
        c_sample = _sample
        t1 = time.time()
        if rank == 0:
            print(f"Compiled MC sample function in {t1-t0} s")
        sample_cache[(num_leaps, npar)] = c_sample
    else:
        if rank == 0:
            print("Using cached sample function")

    if rank == 0:
        chain = jnp.zeros((num_steps, len(params)))
        accept_prob = jnp.zeros(num_steps)
    else:
        chain = jnp.zeros(0)
        accept_prob = jnp.zeros(0)
    for i in tqdm(range(num_steps), disable=(rank != 0)):
        key, params, log_prob, grad, prob = c_sample(
            key, params, log_prob, grad, step_size
        )
        if rank == 0:
            chain, accept_prob = _update_chain(chain, accept_prob, params, prob, i)
    if rank == 0:
        print(f"Accepted {accept_prob.mean():.2%} of samples")
        sys.stdout.flush()
    return chain, accept_prob

invsafe(matrix, thresh=1e-16)

Safe SVD based psuedo-inversion of the matrix. This zeros out modes that are too small when inverting. Use with caution in cases where you really care about what the inverse is.

Parameters:

Name	Type	Description	Default
matrix	Array	The matrix to invert. Should be a (n, n) array.	required
thresh	float	Threshold at which to zero out a mode.	1e-14

Returns:

Name	Type	Description
invmat	Array	The inverted matrix. Same shape as matrix.

Source code in witch/fitting.py

@jax.jit
def invsafe(matrix: jax.Array, thresh: float = 1e-16) -> jax.Array:
    """
    Safe SVD based psuedo-inversion of the matrix.
    This zeros out modes that are too small when inverting.
    Use with caution in cases where you really care about what the inverse is.

    Parameters
    ----------
    matrix : jax.Array
        The matrix to invert.
        Should be a `(n, n)` array.
    thresh : float, default: 1e-14
        Threshold at which to zero out a mode.

    Returns
    -------
    invmat: jax.Array
        The inverted matrix.
        Same shape as `matrix`.
    """
    u, s, v = jnp.linalg.svd(matrix, False)
    s_inv = jnp.array(jnp.where(jnp.abs(s) <= thresh * jnp.max(s), 0, 1 / s))

    return jnp.dot(jnp.transpose(v), jnp.dot(jnp.diag(s_inv), jnp.transpose(u)))

invscale(matrix, thresh=1e-16, do_invsafe=False)

Invert and rescale a matrix by the diagonal. This uses invsafe for the inversion.

Parameters:

Name	Type	Description	Default
matrix	Array	The matrix to invert and sxane. Should be a (n, n) array.	required
thresh	float	Threshold for invsafe. See that function for more info.	1e-14

Returns:

Name	Type	Description
invmat	Array	The inverted and rescaled matrix. Same shape as matrix.

Source code in witch/fitting.py

@partial(jax.jit, static_argnums=(2,))
def invscale(matrix: jax.Array, thresh: float = 1e-16, do_invsafe=False) -> jax.Array:
    """
    Invert and rescale a matrix by the diagonal.
    This uses `invsafe` for the inversion.

    Parameters
    ----------
    matrix : jax.Array
        The matrix to invert and sxane.
        Should be a `(n, n)` array.
    thresh : float, default: 1e-14
        Threshold for `invsafe`.
        See that function for more info.

    Returns
    -------
    invmat: jax.Array
        The inverted and rescaled matrix.
        Same shape as `matrix`.
    """
    diag = jnp.diag(matrix)
    vec = jnp.array(jnp.where(diag != 0, 1.0 / jnp.sqrt(jnp.abs(diag)), thresh))
    mm = jnp.outer(vec, vec)

    if do_invsafe:
        return mm * invsafe(mm * matrix, thresh)
    return mm * jnp.linalg.inv(mm * matrix)

run_lmfit(metamodel, maxiter=10, chitol=1e-05)

Fit a set of models to datasets jointly. This uses a modified Levenberg–Marquardt fitter with flat priors. This function is MPI aware.

Parameters:

Name	Type	Description	Default
metamodel	MetaModel	The MetaModel that describes the models and datasets.	required
maxiter	int	The maximum number of iterations to fit.	10
chitol	float	The delta chisq to use as the convergence criteria.	1e-5

Returns:

Name	Type	Description
models	tuple[Model, ...]	Model with the final set of fit parameters, errors, and chisq.
final_iter	int	The number of iterations the fitter ran for.
delta_chisq	float	The final delta chisq.
curve	Array	The final covariance matrix.

Source code in witch/fitting.py

def run_lmfit(
    metamodel: MetaModel,
    maxiter: int = 10,
    chitol: float = 1e-5,
) -> tuple[MetaModel, int, jax.Array]:
    """
    Fit a set of models to datasets jointly.
    This uses a modified Levenberg–Marquardt fitter with flat priors.
    This function is MPI aware.

    Parameters
    ----------
    metamodel : MetaModel
        The MetaModel that describes the models and datasets.
    maxiter : int, default: 10
        The maximum number of iterations to fit.
    chitol : float, default: 1e-5
        The delta chisq to use as the convergence criteria.

    Returns
    -------
    models : tuple[Model, ...]
        Model with the final set of fit parameters, errors, and chisq.
    final_iter : int
        The number of iterations the fitter ran for.
    delta_chisq : float
        The final delta chisq.
    curve : jax.Array
        The final covariance matrix.
    """
    zero = jnp.array(0.0, jnp.float32)
    chitol = jnp.float32(chitol)
    tf = np.where(np.array(metamodel.to_fit))[0]

    # Every rank has to make the same choice here since the chunked path reduces
    # with mpi4py and the jitted path with mpi4jax. In a joint fit the map ranks
    # follow the TOD ranks onto the chunked path (one chunk per map).
    local_multiple_tods = any(
        dataset.mode == "tod" and len(dataset.datavec.tods) > 1
        for dataset in metamodel.datasets
    )
    if isinstance(metamodel.global_comm, NullComm):
        use_tod_chunks = local_multiple_tods
    else:
        use_tod_chunks = (
            metamodel.global_comm.allreduce(int(local_multiple_tods), op=MPI.MAX) > 0
        )
    if use_tod_chunks:
        return _run_lmfit_tod_chunks(metamodel, maxiter, chitol)

    @jax.jit
    def _cond_func(val):
        i, delta_chisq, lmd, *_ = val
        iterbool = jax.lax.lt(i, maxiter)
        chisqbool = jax.lax.ge(delta_chisq, chitol) + jax.lax.gt(lmd, zero)
        return iterbool * chisqbool

    @jax.jit
    def _body_func(val):
        i, delta_chisq, lmd, metamodel, curve, grad = val
        curve_use = curve.at[:].add(lmd * jnp.diag(jnp.diag(curve)))
        # Get the step
        step = jnp.dot(
            invscale(curve_use.at[tf, :].get().at[:, tf].get(), do_invsafe=True),
            grad.at[tf].get(),
        )
        new_pars, to_fit = _prior_pars_fit(
            metamodel.priors,
            metamodel.parameters.at[tf].add(step),
            jnp.array(metamodel.to_fit),
        )
        # Get errs
        full_cov = invscale(curve_use, do_invsafe=True)
        errs = jnp.where(to_fit, jnp.sqrt(jnp.diag(full_cov)), 0)
        cov = jnp.where(to_fit, full_cov, 0)
        # Now lets get an updated model
        new_metamodel = copy(metamodel).update(
            pars=new_pars, errs=errs, cov=cov, chisq=metamodel.chisq
        )
        new_chisq, new_grad, new_curve = joint_objective(
            new_metamodel, True, True, True
        )
        new_metamodel = copy(new_metamodel).update(
            pars=new_pars, errs=errs, cov=cov, chisq=new_chisq
        )

        new_delta_chisq = jnp.astype(metamodel.chisq - new_metamodel.chisq, jnp.float32)
        metamodel, grad, curve, delta_chisq, lmd = jax.lax.cond(
            new_delta_chisq > 0,
            _success,
            _failure,
            metamodel,
            new_metamodel,
            grad,
            new_grad,
            curve,
            new_curve,
            delta_chisq,
            new_delta_chisq,
            lmd,
        )
        stalled = _stalled(new_delta_chisq, chitol)
        lmd = jnp.where(stalled, zero, lmd)
        delta_chisq = jnp.where(stalled, 0.0, delta_chisq)

        return (i + 1, delta_chisq, lmd, metamodel, curve, grad)

    pars, _ = _prior_pars_fit(
        metamodel.priors, metamodel.parameters, jnp.array(metamodel.to_fit)
    )
    # Recompute chisq at the starting pars so that it matches the current noise model,
    # the stored value may be from before the noise was reestimated
    chisq, grad, curve = joint_objective(metamodel, True, True, True)
    jax.block_until_ready((chisq, grad, curve))
    metamodel = metamodel.update(
        pars=pars, errs=metamodel.errs, cov=metamodel.cov, chisq=chisq
    )
    i, delta_chisq, lmd, metamodel, *_ = jax.lax.while_loop(
        _cond_func,
        _body_func,
        (0, jnp.astype(jnp.inf, jnp.float32), zero.copy(), metamodel, curve, grad),
    )

    return metamodel, i, delta_chisq, lmd

run_mcmc(metamodel, num_steps=5000, num_leaps=10, step_size=jnp.array(0.02), sample_which=-1, burn_in=0.1, max_tries=20)

Run Hamilonian Monte Carlo to estimate the posterior for our model. Currently this function only support flat priors, but more will be supported down the line. In order to ensure accuracy of the noise model used, it is reccomended that you run at least one round of fit_tods followed by noise reestimation before this function.

This is MPI aware. Eventually this will be replaced with something more jaxy.

Parameters:

Name	Type	Description	Default
metamodel	MetaModel	The MetaModel that describes the models and datasets.	required
num_steps	int	The number of steps to run MCMC for.	5000
num_leaps	int	The number of leapfrog steps to take at each sample.	10
step_size	Array	The step size to use in the leapfrog algorithm. This should be tuned to get an acceptance fraction of ~.65.	.02
sample_which	int	Sets which parameters to sample. If this is >= 0 then we will sample which ever parameters were fit in that round of fitting. If this is -1 then we will sample which ever parameters were fit in the last round of fitting. If this is -2 then any parameters that were ever fit will be sampled. If this is <= -3 or >= model.n_rounds then all parameters are sampled.	-1,
burn_in	float	Fractional burn-in period of samples to discard	0.1
max_tries	int	Number of tries to tune step size that will be attemted. If 0 no tuning will be run.	10

Returns:

Name	Type	Description
metamodel	MetaModel	The metamodel with MCMC estimated parameters and errors. The parameters are estimated as the mean of the samples. The errors are estimated as the standard deviation. This also has the chi-squared of the estimated parameters.
flat_samples	Array	Array of samples from running MCMC.

Source code in witch/fitting.py

def run_mcmc(
    metamodel: MetaModel,
    num_steps: int = 5000,
    num_leaps: int = 10,
    step_size: jax.Array = jnp.array(0.02),
    sample_which: int = -1,
    burn_in: float = 0.1,
    max_tries: int = 20,
) -> tuple[MetaModel, jax.Array]:
    """
    Run Hamilonian Monte Carlo to estimate the posterior for our model.
    Currently this function only support flat priors, but more will be supported
    down the line. In order to ensure accuracy of the noise model used, it is
    reccomended that you run at least one round of `fit_tods` followed by noise
    reestimation before this function.

    This is MPI aware.
    Eventually this will be replaced with something more jaxy.

    Parameters
    ----------
    metamodel : MetaModel
        The MetaModel that describes the models and datasets.
    num_steps : int, default: 5000
        The number of steps to run MCMC for.
    num_leaps : int, default: 10
        The number of leapfrog steps to take at each sample.
    step_size : lax.Array, default: .02
        The step size to use in the leapfrog algorithm.
        This should be tuned to get an acceptance fraction of ~.65.
    sample_which : int, default: -1,
        Sets which parameters to sample.
        If this is >= 0 then we will sample which ever parameters were
        fit in that round of fitting.
        If this is -1 then we will sample which ever parameters were fit
        in the last round of fitting.
        If this is -2 then any parameters that were ever fit will be sampled.
        If this is <= -3 or >= `model.n_rounds` then all parameters are sampled.
    burn_in : float, default: 0.1
        Fractional burn-in period of samples to discard
    max_tries : int, default: 10
        Number of tries to tune step size that will be attemted.
        If 0 no tuning will be run.

    Returns
    -------
    metamodel : MetaModel
        The metamodel with MCMC estimated parameters and errors.
        The parameters are estimated as the mean of the samples.
        The errors are estimated as the standard deviation.
        This also has the chi-squared of the estimated parameters.
    flat_samples : jax.Array
        Array of samples from running MCMC.

    """
    global_comm = metamodel.global_comm
    mpi4jax.barrier(comm=global_comm)
    rank = global_comm.Get_rank()
    step_size = jnp.array(step_size)

    if burn_in >= 1.0 or burn_in < 0:
        raise ValueError("Error: burn_in must be in range [0, 1)")

    if sample_which >= 0 and sample_which < metamodel.n_rounds:
        metamodel.set_round(sample_which)
        to_fit = metamodel.to_fit
    elif sample_which == -1:
        to_fit = metamodel.to_fit
    elif sample_which == -2:
        to_fit = metamodel.to_fit_ever
    else:
        to_fit = jnp.ones_like(metamodel.to_fit_ever, dtype=bool)
    to_fit = jnp.array(to_fit)
    metamodel = metamodel.add_round(jnp.array(to_fit))

    init_pars = jnp.array(metamodel.parameters)
    init_errs = jnp.zeros_like(init_pars)
    init_cov = jnp.zeros((len(init_pars), len(init_pars)))
    final_pars = init_pars.copy()
    final_errs = init_errs.copy()

    prior_l, prior_u = metamodel.priors
    scale = (jnp.abs(prior_l) + jnp.abs(prior_u)) / 2.0
    scale = jnp.where(jnp.isfinite(scale), scale, init_pars)
    scale = jnp.where(scale == 0, 1, scale)
    init_pars = init_pars.at[:].multiply(1.0 / scale)

    @jax.jit
    def _log_prob_and_grad(pars, metamodel=metamodel, init_pars=init_pars):
        full_pars = init_pars.at[to_fit].set(pars)
        full_pars = full_pars.at[:].multiply(scale)
        _, in_bounds = _prior_pars_fit(
            metamodel.priors, full_pars, jnp.array(metamodel.to_fit)
        )
        log_prior = jnp.sum(
            jnp.where(
                jnp.array(jnp.where(metamodel.to_fit, in_bounds, True)), 0, -1 * jnp.inf
            )
        )
        temp_metamodel = copy(metamodel).update(
            full_pars, init_errs, init_cov, jnp.array(0.0)
        )
        chisq, grad, _ = joint_objective(temp_metamodel, True, True, False)
        log_like = -0.5 * chisq + log_prior
        grad = grad.at[:].multiply(1.0 / scale)
        log_like_grad = grad.at[to_fit].get().ravel()
        log_like_grad = log_like_grad.at[:].add(log_prior)
        return log_like, log_like_grad

    def _get_step_size(step_size, key, comm, max_tries):
        mpi4jax.barrier(comm=global_comm)

        def _can_interp(probs):
            probs = jnp.array(probs)
            msk = jnp.isfinite(probs) * probs > 0
            if jnp.sum(msk) < 2:
                return False
            dprob = (probs[msk] - 0.63) / 0.13
            for n in range(1, 4):
                upper, lower = (
                    jnp.sum((dprob <= n) * (dprob >= 0)),
                    jnp.sum((dprob >= -1 * n) * (dprob < 0)),
                )
                if upper >= 1 and lower >= 1 and upper + lower >= n + 1:
                    return True
            return False

        def _hmc(n_steps, step_size):
            _, chain_probs = hmc(
                init_pars.at[to_fit].get().ravel(),
                _log_prob_and_grad,
                num_steps=n_steps,
                num_leaps=num_leaps,
                step_size=step_size,
                comm=comm,
                key=key,
            )
            prob = 0.0
            if rank == 0:
                prob = jnp.mean(chain_probs)
            prob = jnp.array(mpi4jax.bcast(prob, 0, comm=comm))

            return prob

        if max_tries == 0:
            return step_size
        step_fac = 1
        n_steps = 1
        i = 0
        step_sizes = []
        probs = []
        for i in range(max_tries):
            mpi4jax.barrier(comm=global_comm)
            step_size *= step_fac
            if rank == 0:
                print(f"Trying step size {step_size}")
            prob = _hmc(n_steps, step_size)
            step_sizes += [step_size]
            probs += [prob]

            if prob == 0 or not jnp.isfinite(prob):
                step_fac = 1e-1
                n_steps = 1
            elif prob > 0.6 and prob < 0.66:
                break
            elif _can_interp(probs):
                break
            elif prob <= 0.1:
                step_fac = 1 / 5
                n_steps = 3
            elif prob >= 0.9:
                step_fac = 5
                n_steps = 3
            elif prob <= 0.5:
                step_fac = 1 / 3
                n_steps = 5
            elif prob >= 0.76:
                step_fac = 3
                n_steps = 5
            elif prob > 0.5 and prob < 0.63:
                step_fac = 1 / 1.5
                n_steps = 10
            elif prob < 0.76 and prob > 0.63:
                step_fac = 1 / 1.5
                n_steps = 10
            else:
                raise ValueError(f"Unclear what to do with probability {prob}")
        if probs[-1] > 0.60 and probs[-1] < 0.66:
            if rank == 0:
                print(
                    f"Final step size of {step_size} with an accept probability of {probs[-1]}, took {i} tries to find."
                )
                sys.stdout.flush()
            return step_sizes[-1]
        mpi4jax.barrier(comm=global_comm)
        probs = jnp.array(probs)
        step_sizes = jnp.array(step_sizes)
        msk = jnp.isfinite(probs) & (probs > 0)
        srt = jnp.argsort(probs[msk])
        if jnp.sum(msk) == 0:
            raise ValueError(
                "No finite non-zero step sizes found! Manual intervention is needed!"
            )
        step_size_interp = jnp.interp(
            jnp.array([0.63]), probs[msk][srt], step_sizes[msk][srt]
        )
        step_size = step_size_interp[0].item()
        prob = _hmc(5, step_size)
        if rank == 0:
            print(
                f"Interpolated to get a step size of {step_size} with an acceptance probability {prob:.2%}"
            )
            sys.stdout.flush()
        return step_size

    key = jax.random.PRNGKey(0)
    key = jnp.array(mpi4jax.bcast(key, 0, comm=global_comm))

    step_size = _get_step_size(
        step_size=step_size,
        key=key,
        comm=global_comm,
        max_tries=max_tries,
    )

    chain, _ = hmc(
        init_pars.at[to_fit].get().ravel(),
        _log_prob_and_grad,
        num_steps=num_steps,
        num_leaps=num_leaps,
        step_size=step_size,
        comm=global_comm,
        key=key,
    )
    flat_samples = chain.at[:].multiply(scale.at[to_fit].get())

    fit_idx = np.flatnonzero(np.array(to_fit))
    final_cov = jnp.zeros((len(final_pars), len(final_pars)))
    if rank == 0:
        burn_in = int(num_steps * burn_in)
        flat_samples = flat_samples[burn_in:]
        final_pars = final_pars.at[to_fit].set(jnp.median(flat_samples, axis=0).ravel())
        final_errs = final_errs.at[to_fit].set(jnp.std(flat_samples, axis=0).ravel())
        # The sample covariance only covers the fit parameters, embed it in the full matrix
        final_cov = final_cov.at[fit_idx[:, None], fit_idx[None, :]].set(
            jnp.atleast_2d(jnp.cov(flat_samples, rowvar=False))
        )
    final_pars = jnp.array(mpi4jax.bcast(final_pars, 0, comm=global_comm))
    final_errs = jnp.array(mpi4jax.bcast(final_errs, 0, comm=global_comm))
    final_cov = jnp.array(mpi4jax.bcast(final_cov, 0, comm=global_comm))
    # Placeholder cov and chisq so we can evaluate the chisq at the final parameters
    metamodel = metamodel.update(
        final_pars.block_until_ready(),
        final_errs.block_until_ready(),
        jnp.zeros_like(final_cov),
        jnp.array(0.0),
    )
    chisq, *_ = joint_objective(metamodel, True, False, False)
    metamodel = metamodel.update(final_pars, final_errs, final_cov, chisq)
    jax.block_until_ready(metamodel)

    return metamodel, flat_samples