The main inference function for SuSiE PCA.

Parameters:

X: Array | ndarray | bool_ | number | bool | int | float | complex | JAXSparse¶

Input data. Should be an array-like

z_dim: int¶

Latent factor dimension (int; K)

l_dim: int¶

Number of single-effects comprising each factor (int; L)

A: Array | ndarray | bool_ | number | bool | int | float | complex | None = None¶

Annotation matrix to use in parameterized-prior mode. If not None, leading dimension should match the feature dimension of X.

tau: float = 1.0¶

initial value of residual precision (default = 1)

standardize: bool = False¶

Whether to center and scale the input data with mean 0 and variance 1 (default = False)

init: 'pca' | 'random' = 'pca'¶

How to initialize the variational mean parameters for latent factors. Either “pca” or “random” (default = “pca”)

seed: int = 0¶

Seed for “random” initialization (int)

max_iter: int = 200¶

Maximum number of iterations for inference (int)

tol: float = 0.001¶

Numerical tolerance for ELBO convergence (float)

verbose: bool = True¶

Flag to indicate displaying log information (ELBO value) in each iteration

Returns:

tuple that has member variables for learned parameters (ModelParams), evidence lower bound (ELBO) results (ELBOResults) from the last iteration, the percent of variance explained (PVE) for each of the K factors (jax.numpy.ndarray), the posterior inclusion probabilities (PIPs) for each of the K factors and P features (jax.numpy.ndarray).

Return type:

SuSiEPCAResults

Raises:

• ValueError – Invalid l_dim or z_dim values. Invalid initialization scheme.

• Data X contains inf or nan. If annotation matrix A is not None, raises –

• if A contains inf, nan or does not match feature dimension with X. –