Summary-level inference¶

def infer_sushie_ss(
    lds: Sequence[ArrayLike],
    ns: ArrayLike,
    zs: Sequence[ArrayLike],
    L: int = 10,
    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,
) -> infer.SushieResult:
    """The main inference function for running SuShiE.

    Args:
        lds: LD matrix for multiple ancestries.
        zs: molQTL scan z scores for multiple ancestries.
        ns: Sample size for each ancestry.
        L: Inferred number of eQTLs for the gene.
        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 summary statistics:

        ```python
        import numpy as np
        from sushie.infer_ss import infer_sushie_ss

        # Generate example data for 2 ancestries, 500 SNPs
        n_snps = 500

        # LD matrices (correlation matrices)
        LD1 = np.eye(n_snps)  # Identity for simplicity
        LD2 = np.eye(n_snps)

        # Z-scores from GWAS
        z1 = np.random.randn(n_snps)
        z2 = np.random.randn(n_snps)

        # Sample sizes
        ns = np.array([1000, 1500])

        # Run SuShiE fine-mapping with summary statistics
        result = infer_sushie_ss(
            lds=[LD1, LD2],
            zs=[z1, z2],
            ns=ns,
            L=5,
        )

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

    """
    ns = jnp.asarray(ns)
    if ns.ndim == 1:
        ns = ns[:, jnp.newaxis]
    elif ns.ndim != 2 or ns.shape[1] != 1:
        raise ValueError("Sample sizes must be a vector or a single-column matrix. Check your input.")

    lds = [jnp.asarray(ld) for ld in lds]
    pi_array = None if pi is None else jnp.asarray(pi)

    n_pop = ns.shape[0]

    if len(lds) != n_pop:
        raise ValueError(f"The number of LD matrices ({len(lds)}) does not match the number of ancestries ({n_pop}).")

    if not all(ld.shape == lds[0].shape for ld in lds):
        raise ValueError("LD matrices do not have the same shape. Check your input.")

    if lds[0].shape[0] != lds[0].shape[1]:
        raise ValueError("LD matrices are not square matrices. Check your input.")

    if zs is None:
        raise ValueError("Z scores are not provided. Check your input.")

    zs = [jnp.asarray(z) for z in zs]
    if len(zs) != n_pop:
        raise ValueError(f"The number of Z scores ({len(zs)}) does not match the number of ancestries ({n_pop}).")

    if not all(z.shape == zs[0].shape for z in zs):
        raise ValueError("Z scores across ancestries do not have the same shape. Check your input.")

    if zs[0].shape[0] != lds[0].shape[0]:
        raise ValueError("Z scores do not have the same number of SNPs as the LD matrices. Check your input.")

    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 is invalid. Choose a positive integer.")

    n_snps = lds[0].shape[0]

    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] != lds[0].shape[1]:
            raise ValueError(
                f"Prior probability/weights ({pi_array.shape[0]}) does not match "
                + f"the number of SNPs ({lds[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)

    if resid_var is None:
        resid_var_array = jnp.ones(n_pop)
    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 = infer._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 = infer._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 = infer._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 = infer._PriorAdjustor(times=jnp.ones((n_pop, n_pop)), plus=jnp.zeros((n_pop, n_pop)))

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

    posteriors = infer.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 = infer._EMOptFunc() if not no_update else infer._NoopOptFunc()

    # Stack once after validation so summary-stat kernels see canonical arrays.
    zs_array = jnp.stack(zs)
    lds_array = jnp.stack(lds)
    sigma2 = ns / (ns + zs_array**2)
    Xtys = jnp.sqrt(ns) * jnp.sqrt(sigma2) * zs_array
    XtXs = ns[:, :, jnp.newaxis] * lds_array

    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_ss(
            Xtys,
            XtXs,
            ns,
            priors,
            posteriors,
            prior_adjustor,
            opt_v_func,
        )
        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)))

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

        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 = infer._reorder_l(priors, posteriors)

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

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

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

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