Merge branch 'main' of https://github.com/mancusolab/susiepca

docs/conf.py

@@ -148,9 +148,19 @@
 python_transform_type_annotations_pep585 = False
 python_transform_type_annotations_pep604 = True
 python_transform_type_annotations_concise_literal = True
+python_strip_self_type_annotations = True
+
+python_type_aliases = {
+    "jnp.ndarray": "jax.numpy.ndarray",
+}
+
+python_module_names_to_strip_from_xrefs = [
+    "jax._src.numpy.lax_numpy",
+]
 
 # fix namedtuple attribute types
 napoleon_use_ivar = True
+napoleon_google_docstring = True
 
 python_apigen_rst_prolog = """
 .. default-role:: py:obj

susiepca/infer.py

@@ -284,7 +284,7 @@ def update_tau(X, params):
 
 
 #
-def compute_elbo(X, params) -> ELBOResults:
+def compute_elbo(X: jnp.ndarray, params: ModelParams) -> ELBOResults:
     """Create function to compute evidence lower bound (ELBO)
 
     Args:
@@ -350,7 +350,7 @@ def compute_pip(params: ModelParams) -> jnp.ndarray:
         params: instance of infered parameters
 
     Returns:
-        Array of posterior inclusion probabilities (PIPs) for each of
+        jnp.ndarray: Array of posterior inclusion probabilities (PIPs) for each of
         `K x P` factor, feature combinations
     """
 
@@ -366,8 +366,8 @@ def compute_pve(params: ModelParams) -> jnp.ndarray:
         params: instance of infered parameters
 
     Returns:
-        Array of length `K` that contains percent of variance explained by each
-        factor (PVE)
+        jnp.ndarray: Array of length `K` that contains percent of variance
+        explained by each factor (PVE)
     """
 
     n_dim, z_dim = params.mu_z.shape
@@ -501,7 +501,7 @@ def susie_pca(
             iteration
 
     Returns:
-        A :py:obj:`SuSiEPCAResults` instance that has member variables for learned
+        :py:obj:`SuSiEPCAResults`: tuple that has member variables for learned
         parameters (:py:obj:`ModelParams`), evidence lower bound (ELBO) results
         (:py:obj:`ELBOResults`) from the last iteration, the percent of variance
         explained (PVE) for each of the `K` factors (:py:obj:`jax.numpy.ndarray`),

susiepca/sim.py

@@ -41,13 +41,11 @@ def generate_sim(
         p_dim: Number of feature in the data
         z_dim: Number of Latent dimensions
         effect_size: The effect size of features contributing to the factor.
-                      The default is 1.
+                      (default = 1).
 
     Returns:
-        Z: Simulated factors (N by K)
-        W: Simulated factor loadings (K by P)
-        X: Simulated data (N by P)
-
+        SimulatedData: Tuple that contains simulated factors (`N x K`),
+        W (factor loadings (`K x P`), and data X (data (`N x P`).
     """
     rng_key = random.PRNGKey(seed)
     rng_key, z_key, b_key, obs_key = random.split(rng_key, 4)