Individual-level inference¶

def infer_sushie(
    Xs: Sequence[ArrayLike],
    ys: Sequence[ArrayLike],
    covar: Sequence[ArrayLike] | None = None,
    L: int = 10,
    no_scale: bool = False,
    no_regress: bool = False,
    no_update: bool = False,
    pi: ArrayLike | None = None,
    resid_var: utils.ListFloatOrNone = None,
    effect_var: utils.ListFloatOrNone = None,
    rho: utils.ListFloatOrNone = None,
    max_iter: int = 500,
    min_tol: float = 1e-4,
    threshold: float = 0.95,
    purity: float = 0.5,
    purity_method: str = "weighted",
    max_select: int = 250,
    min_snps: int = 100,
    no_reorder: bool = False,
    seed: int = 12345,
) -> SushieResult:
    """The main inference function for running SuShiE.

    Args:
        Xs: Genotype data for multiple ancestries.
        ys: Phenotype data for multiple ancestries.
        covar: Covariate data for multiple ancestries.
        L: Inferred number of eQTLs for the gene.
        no_scale: Do not scale the genotype and phenotype. Default is to scale.
        no_regress: Do not regress covariates on genotypes. Default is to regress.
        no_update: Do not update the effect size prior. Default is to update.
        pi: The probability prior for one SNP to be causal ($\\pi$ in [model description](../model.md)).
            Default is $1$ over the number of SNPs by specifying it as ``None``.
        resid_var: Prior residual variance ($\\sigma^2_e$ in [model description](../model.md)).
            Default is $0.001$ by specifying it as ``None``.
        effect_var: Prior causal effect size variance ($\\sigma^2_{i,b}$ in [model description](../model.md)).
            Default is $0.001$ by specifying it as ``None``.
        rho: Prior effect size correlation ($\\rho$ in [model description](../model.md)).
            Default is $0.1$ by specifying it as ``None``.
        max_iter: The maximum iteration for optimization. Default is $500$.
        min_tol: The convergence tolerance. Default is $10^{-4}$.
        threshold: The credible set threshold. Default is $0.95$.
        purity: The minimum pairwise correlation across SNPs to be eligible as output credible set.
            Default is $0.5$.
        purity_method: The method to compute purity across ancestries. Default is ``weighted``.
        max_select: The maximum number of selected SNPs to compute purity. Default is $250$.
        min_snps: The minimum number of SNPs to fine-map. Default is $100$.
        no_reorder: Do not re-order single effects based on Frobenius norm of effect size covariance prior.
            Default is to re-order.
        seed: The randomization seed for selecting SNPs in the credible set to compute purity. Default is $12345$.

    Returns:
        `SushieResult`: A SuShiE result object that contains prior (`Prior`),
        posterior (`Posterior`), ``cs``, ``pip``, ``elbo``, and ``elbo_increase``.

    Example:
        Basic usage with two-ancestry data:

        ```python
        import numpy as np
        from sushie.infer import infer_sushie

        # Generate example data for 2 ancestries
        # Ancestry 1: 100 samples, 500 SNPs
        X1 = np.random.randn(100, 500)
        y1 = np.random.randn(100)

        # Ancestry 2: 150 samples, 500 SNPs
        X2 = np.random.randn(150, 500)
        y2 = np.random.randn(150)

        # Run SuShiE fine-mapping
        result = infer_sushie(Xs=[X1, X2], ys=[y1, y2], L=5)

        # Access results
        print(result.pip_all)  # Posterior inclusion probabilities
        print(result.cs)       # Credible sets
        ```

    """
    Xs = [jnp.asarray(X) for X in Xs]
    ys = [jnp.asarray(y) for y in ys]
    covar = None if covar is None else [jnp.asarray(c) for c in covar]
    pi_array = None if pi is None else jnp.asarray(pi)

    if len(Xs) == len(ys):
        n_pop = len(Xs)
    else:
        raise ValueError(f"The number of geno ({len(Xs)}) and pheno ({len(ys)}) data does not match. Check your input.")

    if covar is not None and len(covar) != n_pop:
        raise ValueError(
            f"The number of covariate ({len(covar)}) and geno ({n_pop}) data does not match. Check your input."
        )

    # check x and y have the same sample size
    for idx in range(n_pop):
        if Xs[idx].shape[0] != ys[idx].shape[0]:
            raise ValueError(
                f"Ancestry {idx + 1}: The sample size of geno ({Xs[idx].shape[0]}) "
                + f"and pheno ({ys[idx].shape[0]}) data does not match. Check your input."
            )

    # check each ancestry has the same number of SNPs
    for idx in range(1, n_pop):
        if Xs[idx - 1].shape[1] != Xs[idx].shape[1]:
            raise ValueError(
                f"Ancestry {idx} and ancestry {idx} do not have "
                + f"the same number of SNPs ({Xs[idx - 1].shape[1]} vs {Xs[idx].shape[1]})."
            )

    if L <= 0:
        raise ValueError(f"Inferred L ({L}) is invalid, choose a positive L.")

    if min_tol > 0.1:
        log.logger.warning(f"Minimum intolerance ({min_tol}) is greater than 0.1. Inference may not be accurate.")

    if not 0 < threshold < 1:
        raise ValueError(
            f"Credible set PIP threshold ({threshold}) must be greater than 0 and less than 1."
            + " Specify a valid value using '--threshold' for command-line usage"
            + " or 'threshold=' for in-Python function calls."
        )

    if not 0 < purity < 1:
        raise ValueError(
            f"Purity threshold ({purity}) must be greater than or equal to 0 and less than 1. "
            + " Specify a valid value using '--purity' for command-line usage"
            + " or 'purity=' for in-Python function calls."
        )

    if max_select <= 0:
        raise ValueError(
            "The maximum selected number of SNPs for purity must be greater than 0."
            + " Specify a valid value using '--max-select' for command-line usage"
            + " or 'max_select=' for in-Python function calls."
        )

    if min_snps <= 0:
        raise ValueError(
            "The minimum number of SNPs to fine-map must be greater than 0."
            + " Specify a valid value using '--min-snps' for command-line usage"
            + " or 'min_snps=' for in-Python function calls."
        )

    _, n_snps = Xs[0].shape

    if pi_array is None:
        pi_array = jnp.ones(n_snps) / float(n_snps)
    else:
        if not (pi_array > 0).all():
            raise ValueError("Prior probability/weights must be all positive values.")

        if pi_array.shape[0] != Xs[0].shape[1]:
            raise ValueError(
                f"Prior probability/weights ({pi_array.shape[0]}) does not match the number of SNPs ({Xs[0].shape[1]})."
            )

        if jnp.sum(pi_array) > 1:
            log.logger.debug("Prior probability/weights sum is not equal to 1. Will normalize to sum to 1.")
            pi_array = pi_array.astype(float) / jnp.sum(pi_array)

    # Regress covariates before centering/scaling so the model sees residualized arrays.
    if covar is not None:
        for idx in range(n_pop):
            Xs[idx], ys[idx] = utils.regress_covar(Xs[idx], ys[idx], covar[idx], no_regress)

    # Center and optionally scale each ancestry before padding to a shared size.
    for idx in range(n_pop):
        Xs[idx] -= jnp.mean(Xs[idx], axis=0)
        ys[idx] -= jnp.mean(ys[idx])
        if not no_scale:
            Xs[idx] /= jnp.std(Xs[idx], axis=0)
            ys[idx] /= jnp.std(ys[idx])

        ys[idx] = jnp.squeeze(ys[idx])

    if resid_var is None:
        resid_var_array = jnp.array([jnp.var(y, ddof=1) for y in ys])
    else:
        if len(resid_var) != n_pop:
            raise ValueError(
                f"Number of specified residual prior ({len(resid_var)}) does not match ancestry number ({n_pop})."
            )
        resid_var_array = jnp.array([float(i) for i in resid_var])
        if jnp.any(resid_var_array <= 0):
            raise ValueError(f"The input of residual prior ({resid_var}) is invalid (<0). Check your input.")

    if min_snps < L:
        raise ValueError(
            f"The number of minimum common SNPs across ancestries ({min_snps}) is less than inferred L ({L})."
            + " Specify a larger value using '--min-snps' for command-line usage"
            + " or 'min_snps=' for in-Python function calls."
        )

    if n_snps < min_snps:
        raise ValueError(
            f"The number of common SNPs across ancestries ({n_snps}) is less than minimum common"
            + " number of SNPs (100) specified."
            + " Users can specify a smaller value using '--min-snps' for command-line usage"
            + " or 'min_snps=' for in-Python function calls."
        )

    param_effect_var = effect_var
    if effect_var is None:
        effect_var = [1e-3] * n_pop
    else:
        if len(effect_var) != n_pop:
            raise ValueError(
                f"Number of specified effect prior ({len(effect_var)}) does not match ancestry number ({n_pop})."
            )
        effect_var = [float(i) for i in effect_var]
        if jnp.any(jnp.array(effect_var) <= 0):
            raise ValueError(f"The effect size prior variance ({effect_var}) must be positive.")

    exp_num_rho = math.comb(n_pop, 2)
    param_rho = rho
    if rho is None:
        rho = [0.1] * exp_num_rho
    else:
        if n_pop == 1:
            log.logger.debug("Running single-ancestry SuShiE. The '--rho' parameter is specified but will be ignored.")

        if (len(rho) != exp_num_rho) and n_pop != 1:
            raise ValueError(
                f"Number of specified rho ({len(rho)}) does not match expected" + f" number {exp_num_rho}.",
            )
        rho = [float(i) for i in rho]
        # double-check the if it's invalid rho
        if jnp.any(jnp.abs(jnp.array(rho)) >= 1):
            raise ValueError(f"Effect size prior correlation ({rho}) must be between -1 and 1 (inclusive).")

    effect_covar = jnp.diag(jnp.array(effect_var))
    ct = 0
    for col in range(n_pop):
        for row in range(1, n_pop):
            if col < row:
                _two_sd = jnp.sqrt(effect_var[row] * effect_var[col])
                effect_covar = effect_covar.at[row, col].set(rho[ct] * _two_sd)
                effect_covar = effect_covar.at[col, row].set(rho[ct] * _two_sd)
                ct += 1

    if no_update:
        # if we specify no_update and rho, we want to keep rho through iterations and update variance
        if param_effect_var is None and param_rho is not None and n_pop != 1:
            prior_adjustor = _PriorAdjustor(
                times=jnp.eye(n_pop),
                plus=effect_covar - jnp.diag(jnp.diag(effect_covar)),
            )

            log.logger.info("No updates on the prior effect correlation rho while updating prior effect variance.")
        # if we specify no_update and effect_covar, we want to keep variance through iterations, and update rho
        elif param_effect_var is not None and param_rho is None and n_pop != 1:
            prior_adjustor = _PriorAdjustor(
                times=jnp.ones((n_pop, n_pop)) - jnp.eye(n_pop),
                plus=effect_covar * jnp.eye(n_pop),
            )
            log.logger.info("No updates on the prior effect variance while updating prior effect correlation rho.")
        # if we (do not specify effect_covar and rho) or (specify both effect_covar and rho)
        # nothing is updated through iterations
        else:
            prior_adjustor = _PriorAdjustor(times=jnp.zeros((n_pop, n_pop)), plus=effect_covar)
            log.logger.info("No updates on the prior effect size variance/covariance matrix.")
    else:
        prior_adjustor = _PriorAdjustor(times=jnp.ones((n_pop, n_pop)), plus=jnp.zeros((n_pop, n_pop)))

    # define:
    # k is ancestry
    # n is sample size
    # p is SNP
    # l is the number of effects

    priors = Prior(
        pi=pi_array,
        resid_var=resid_var_array[:, jnp.newaxis],
        effect_covar=jnp.array([effect_covar] * L),
    )

    posteriors = Posterior(
        alpha=jnp.zeros((L, n_snps)),
        post_mean=jnp.zeros((L, n_snps, n_pop)),
        post_mean_sq=jnp.zeros((L, n_snps, n_pop, n_pop)),
        weighted_sum_covar=jnp.zeros((L, n_pop, n_pop)),
        kl=jnp.zeros((L,)),
        log_bf=jnp.zeros((L, n_snps)),
    )

    opt_v_func = _EMOptFunc() if not no_update else _NoopOptFunc()

    # Pad ancestries to a common sample dimension, then stack once for JIT kernels.
    sample_sizes = jnp.array([X.shape[0] for X in Xs])
    max_sample_size = int(jnp.max(sample_sizes))
    for idx in range(n_pop):
        pad_rows = max_sample_size - int(sample_sizes[idx])
        Xs[idx] = jnp.pad(Xs[idx], ((0, pad_rows), (0, 0)), "constant")
        ys[idx] = jnp.pad(ys[idx], (0, pad_rows), "constant")

    ns = sample_sizes[:, jnp.newaxis]
    Xs_array = jnp.stack(Xs)
    ys_array = jnp.stack(ys)
    XtXs = jnp.sum(Xs_array**2, axis=1)

    elbo_tracker = jnp.array([-jnp.inf])
    elbo_increase = True
    decimal_digit = len(str(min_tol)) - str(min_tol).find(".") - 1
    for o_iter in range(max_iter):
        log.logger.debug(f"Starting optimization iteration {o_iter + 1}.")

        prev_priors = priors
        prev_posteriors = posteriors

        priors, posteriors, elbo_cur = _update_effects(
            Xs_array,
            ys_array,
            XtXs,
            ns,
            priors,
            posteriors,
            prior_adjustor,
            opt_v_func,
        )

        log.logger.debug(f"Iteration {o_iter + 1} finished.")

        elbo_last = elbo_tracker[o_iter]
        elbo_tracker = jnp.append(elbo_tracker, elbo_cur)
        elbo_increase = bool(jnp.logical_or(elbo_cur >= elbo_last, jnp.isclose(elbo_cur, elbo_last, atol=1e-8)))

        if not elbo_increase:
            log.logger.warning(
                f"Optimization finished after {o_iter + 1} iterations."
                + f" ELBO decreased. Final ELBO score: {elbo_cur}. Return last iteration's results."
                + " It can be precision issue,"
                + " and adding 'import jax; jax.config.update('jax_enable_x64', True)' may fix it."
                + " If this issue keeps rising for many genes, contact the developers."
            )
            priors = prev_priors
            posteriors = prev_posteriors
            break

        if jnp.abs(elbo_cur - elbo_last) < min_tol:
            log.logger.info(
                f"Optimization concludes after {o_iter + 1} iterations. Final ELBO score: {elbo_cur:.{decimal_digit}f}."
                + f" Reach minimum tolerance threshold {min_tol}.",
            )
            break

        if o_iter + 1 == max_iter:
            log.logger.info(
                f"Optimization concludes after {o_iter + 1} iterations. Final ELBO score: {elbo_cur:.{decimal_digit}f}."
                + f" Reach maximum iteration threshold {max_iter}.",
            )

    l_order = jnp.arange(L)
    if not no_reorder:
        log.logger.debug("Reordering effects based on Frobenius norm of effect size covariance prior.")
        priors, posteriors, l_order = _reorder_l(priors, posteriors)

    log.logger.debug("Computing credible sets.")

    cs, full_alphas, pip_all, pip_cs = make_cs(
        posteriors.alpha,
        posteriors.log_bf,
        ns,
        Xs_array,
        None,
        threshold,
        purity,
        purity_method,
        max_select,
        seed,
    )

    log.logger.debug("Inference and credible set computation complete. Beginning to write results.")

    return SushieResult(
        priors,
        posteriors,
        pip_all,
        pip_cs,
        cs,
        full_alphas,
        ns,
        elbo_tracker,
        elbo_increase,
        l_order,
    )

def make_cs(
    alpha: ArrayLike,
    log_bf: ArrayLike,
    ns: ArrayLike,
    Xs: ArrayLike | None = None,
    lds: ArrayLike | None = None,
    threshold: float = 0.9,
    purity: float = 0.5,
    purity_method: str = "weighted",
    max_select: int = 500,
    seed: int = 12345,
) -> tuple[pl.DataFrame, pl.DataFrame, Array, Array]:
    """The function to compute the credible sets.

    Args:
        alpha: $L \\times p$ matrix that contains posterior probability for SNP to be causal
            (i.e., $\\alpha$ in [model description](../model.md)).
        log_bf: $L \\times p$ matrix that contains log Bayes factor for each SNP in each effect.
        Xs: Genotype data for multiple ancestries. It cannot be None if lds is None.
        lds: LD matrix for multiple ancestries. It cannot be None if Xs is None.
        ns: Sample size for each ancestry.
        threshold: The credible set threshold.
        purity: The minimum pairwise correlation across SNPs to be eligible as output credible set.
        purity_method: The method to compute purity across ancestries.
        max_select: The maximum number of selected SNPs to compute purity.
        seed: The randomization seed for selecting SNPs in the credible set to compute purity.

    Returns:
        `Tuple[pl.DataFrame, pl.DataFrame, Array, Array]`: A tuple of
            - credible set (`pl.DataFrame`) after pruning for purity,
            - full credible set (`pl.DataFrame`) before pruning for purity,
            - PIPs (`Array`) across $L$ credible sets,
            - PIPs (`Array`) across credible sets that are not pruned. An array of zero if all credible sets
                are pruned.

    Example:
        Compute credible sets from a SuShiE posterior:

        ```python
        import numpy as np

        from sushie.infer import infer_sushie, make_cs

        X1 = np.random.randn(100, 500)
        X2 = np.random.randn(150, 500)
        y1 = np.random.randn(100)
        y2 = np.random.randn(150)
        result = infer_sushie(Xs=[X1, X2], ys=[y1, y2], L=5)
        cs, full_cs, pip_all, pip_cs = make_cs(
            alpha=result.posteriors.alpha,
            log_bf=result.posteriors.log_bf,
            ns=np.array([X1.shape[0], X2.shape[0]]),
            Xs=[X1, X2],
            threshold=0.9,
            purity=0.5,
        )
        ```

    """
    if Xs is None and lds is None:
        raise ValueError("Both Xs and lds are None. Please specify at least one of them.")

    rng_key = random.PRNGKey(seed)
    alpha = jnp.asarray(alpha)
    log_bf = jnp.asarray(log_bf)
    ns = jnp.asarray(ns)
    Xs = None if Xs is None else jnp.asarray(Xs)
    lds = None if lds is None else jnp.asarray(lds)
    n_l, n_snp = alpha.shape

    cs_frames = []
    full_alphas = pl.DataFrame({"SNPIndex": np.arange(n_snp, dtype=np.int64)})

    for ldx in range(n_l):
        # Select original SNP indices by descending alpha.
        sorted_idx = jnp.argsort(-alpha[ldx])
        sorted_alpha = alpha[ldx, sorted_idx]
        c_alpha = jnp.cumsum(sorted_alpha)
        n_row = int(jnp.sum(c_alpha < threshold))

        # Include the first SNP that reaches the requested cumulative-alpha threshold.
        n_selected = min(n_row + 1, n_snp)
        select_idx = jnp.arange(n_selected)

        snp_idx = sorted_idx[select_idx].astype("int64")

        # output CS Index is 1-based
        tmp_cs = pl.DataFrame(
            {
                "CSIndex": np.repeat(ldx + 1, len(select_idx)),
                "SNPIndex": np.asarray(snp_idx),
                "alpha": np.asarray(sorted_alpha[select_idx]),
                "c_alpha": np.asarray(c_alpha[select_idx]),
            }
        )

        # Estimate purity on a bounded subset so very large credible sets do not
        # dominate runtime.
        if len(snp_idx) > max_select:
            snp_idx = random.choice(rng_key, snp_idx, shape=(max_select,), replace=False)

        if Xs is not None:
            ld_Xs = Xs[:, :, snp_idx]
            ld = jnp.einsum("ijk,ijm->ikm", ld_Xs, ld_Xs) / ns[:, jnp.newaxis]
        elif lds is not None:
            ld = lds[:, snp_idx, :][:, :, snp_idx]

        min_abs_corr = jnp.min(jnp.abs(ld), axis=(1, 2))
        if purity_method == "weighted":
            ancestry_weight = jnp.squeeze(ns / jnp.sum(ns))
            avg_corr = jnp.sum(min_abs_corr * ancestry_weight)
        elif purity_method == "max":
            avg_corr = jnp.max(min_abs_corr)
        elif purity_method == "min":
            avg_corr = jnp.min(min_abs_corr)
        else:
            raise ValueError(f"Invalid purity method {purity_method}. Choose from 'weighted', 'max', or 'min'.")

        in_cs = jnp.zeros(n_snp, dtype=int).at[sorted_idx[select_idx]].set(1)
        kept = bool(avg_corr > purity)
        full_alphas = full_alphas.with_columns(
            pl.Series(f"alpha_l{ldx + 1}", np.asarray(alpha[ldx])),
            pl.Series(f"in_cs_l{ldx + 1}", np.asarray(in_cs)),
            pl.lit(float(avg_corr)).alias(f"purity_l{ldx + 1}"),
            pl.lit(int(kept)).alias(f"kept_l{ldx + 1}"),
            pl.Series(f"log_bf_l{ldx + 1}", np.asarray(log_bf[ldx, :])),
        )

        if kept:
            cs_frames.append(tmp_cs)

    pip_all = utils.make_pip(alpha)
    cs = (
        pl.concat(cs_frames)
        if len(cs_frames) != 0
        else pl.DataFrame(
            schema={
                "CSIndex": pl.Int64,
                "SNPIndex": pl.Int64,
                "alpha": pl.Float64,
                "c_alpha": pl.Float64,
            }
        )
    )

    # CSIndex is now 1-based
    kept_effects = cs["CSIndex"].unique().to_numpy().astype(int) - 1
    pip_cs = utils.make_pip(alpha[kept_effects])

    n_snp_cs = cs["SNPIndex"].to_numpy().astype(int)
    n_snp_cs_unique = jnp.unique(n_snp_cs)

    if len(n_snp_cs) != len(n_snp_cs_unique):
        log.logger.warning(
            "Same SNPs appear in different credible set, which is very unusual."
            + " You may want to check this gene in details."
        )

    cs = cs.with_columns(
        pl.Series("pip_all", np.asarray(pip_all[n_snp_cs])),
        pl.Series("pip_cs", np.asarray(pip_cs[n_snp_cs])),
    )
    full_alphas = full_alphas.with_columns(
        pl.Series("pip_all", np.asarray(pip_all)),
        pl.Series("pip_cs", np.asarray(pip_cs)),
    )

    log.logger.info(
        f"{cs['CSIndex'].n_unique()} out of {n_l} credible sets remain after pruning based on purity ({purity})."
        + " For detailed results, specify --alphas."
    )

    return cs, full_alphas, pip_all, pip_cs

def _compute_posterior(
    rTZDinv: Array,
    inv_shat2: Array,
    priors: Prior,
    posteriors: Posterior,
    l_iter: int,
) -> tuple[Prior, Posterior]:
    n_snps, n_pop, _ = inv_shat2.shape

    # prior_covar is kxk
    prior_covar = priors.effect_covar[l_iter]
    # post_covar is pxkxk
    post_covar = jnp.linalg.inv(inv_shat2 + jnp.linalg.inv(prior_covar))

    # dim m = dim k for the next two lines
    post_mean = jnp.einsum("pkm,pm->pk", post_covar, rTZDinv)
    post_mean_sq = post_covar + jnp.einsum("pk,pm->pkm", post_mean, post_mean)

    # compute the ABF in the original susie paper
    # which is equivalent to the inverse posterior density at 0
    log_bf = -1 * stats.multivariate_normal.logpdf(
        # origianlly it was logpdf(jnp.zeros((n_snps, n_pop)), post_mean, post_covar)
        # but to match what our math derivation in the paper, we change it to following
        # it will not change the result as N(0; 1, 1) is the same as N(1; 0, 1) in terms of density value
        post_mean,
        jnp.zeros((n_snps, n_pop)),
        post_covar,
    )

    alpha = nn.softmax(jnp.log(priors.pi) + log_bf)

    weighted_post_mean = post_mean * alpha[:, jnp.newaxis]
    weighted_post_mean_sq = post_mean_sq * alpha[:, jnp.newaxis, jnp.newaxis]
    # this is also the prior in our E step
    weighted_sum_covar = jnp.sum(weighted_post_mean_sq, axis=0)
    kl_alpha = _kl_categorical(alpha, priors.pi)
    kl_betas = alpha @ _kl_mvn(post_mean, post_covar, jnp.zeros_like(post_mean), prior_covar)

    priors = priors._replace(effect_covar=priors.effect_covar.at[l_iter].set(weighted_sum_covar))

    posteriors = posteriors._replace(
        alpha=posteriors.alpha.at[l_iter].set(alpha),
        post_mean=posteriors.post_mean.at[l_iter].set(weighted_post_mean),
        post_mean_sq=posteriors.post_mean_sq.at[l_iter].set(weighted_post_mean_sq),
        weighted_sum_covar=posteriors.weighted_sum_covar.at[l_iter].set(weighted_sum_covar),
        kl=posteriors.kl.at[l_iter].set(kl_alpha + kl_betas),
        log_bf=posteriors.log_bf.at[l_iter].set(log_bf),
    )

    return priors, posteriors

def _reorder_l(priors: Prior, posteriors: Posterior) -> tuple[Prior, Posterior, Array]:

    frob_norm = jnp.sum(jnp.linalg.svd(posteriors.weighted_sum_covar, compute_uv=False), axis=1)

    # we want to reorder them based on the Frobenius norm
    l_order = jnp.argsort(-frob_norm)

    # priors effect_covar
    priors = priors._replace(effect_covar=priors.effect_covar[l_order])

    posteriors = posteriors._replace(
        alpha=posteriors.alpha[l_order],
        post_mean=posteriors.post_mean[l_order],
        post_mean_sq=posteriors.post_mean_sq[l_order],
        weighted_sum_covar=posteriors.weighted_sum_covar[l_order],
        kl=posteriors.kl[l_order],
        log_bf=posteriors.log_bf[l_order],
    )

    return priors, posteriors, l_order