linting and documentation

multidms/model.py

@@ -88,6 +88,10 @@ class Model:
     init_theta_bias : float
         Initialize the bias parameter :math:`\theta_{\text{bias}}` of
         a two parameter epistatic model (Sigmoid or Softplus).
+    init_beta_variance : float
+        Beta parameters are initialized by sampling from
+        a normal distribution. This parameter specifies the
+        variance of the distribution being sampled.
     n_hidden_units : int or None
         If using :func:`multidms.biophysical.nn_global_epistasis`
         as the epistatic model, this is the number of hidden units
@@ -215,7 +219,7 @@ def __init__(
         n_hidden_units=5,
         init_theta_scale=5.0,
         init_theta_bias=-5.0,
-        init_beta_variance=0.0,  # TODO document and test
+        init_beta_variance=0.0,
         name=None,
     ):
         """See class docstring."""
@@ -603,9 +607,9 @@ def get_mutations_df(
 
         for condition in self.data.conditions:
             single_mut_binary = self.data.single_mut_encodings[condition]
-            mutations_df[f"predicted_func_score_{condition}"] = (
-                self.phenotype_frombinary(single_mut_binary, condition=condition)
-            )
+            mutations_df[
+                f"predicted_func_score_{condition}"
+            ] = self.phenotype_frombinary(single_mut_binary, condition=condition)
 
             if phenotype_as_effect:
                 wt_func_score = self.wildtype_df.loc[condition, "predicted_func_score"]
@@ -862,9 +866,9 @@ def add_phenotypes_to_df(
             if phenotype_as_effect:
                 latent_predictions -= wildtype_df.loc[condition, "predicted_latent"]
             latent_predictions[nan_variant_indices] = onp.nan
-            ret.loc[condition_df.index.values, latent_phenotype_col] = (
-                latent_predictions
-            )
+            ret.loc[
+                condition_df.index.values, latent_phenotype_col
+            ] = latent_predictions
 
             # func_score predictions on binary variants, X
             phenotype_predictions = onp.array(
@@ -876,9 +880,9 @@ def add_phenotypes_to_df(
                     condition, "predicted_func_score"
                 ]
             phenotype_predictions[nan_variant_indices] = onp.nan
-            ret.loc[condition_df.index.values, observed_phenotype_col] = (
-                phenotype_predictions
-            )
+            ret.loc[
+                condition_df.index.values, observed_phenotype_col
+            ] = phenotype_predictions
 
         return ret
 

multidms/model_collection.py

@@ -54,7 +54,7 @@ def fit_one_model(
     # gamma_corrected=False,  # GAMMA
     init_theta_scale=6.5,
     init_theta_bias=-3.5,
-    init_beta_variance=1.0,  # TODO document
+    init_beta_variance=1.0,
     n_hidden_units=5,
     lower_bound=None,
     PRNGKey=0,
@@ -80,6 +80,9 @@ def fit_one_model(
         The scale to use for initializing the model parameters. The default is 6.5.
     init_theta_bias : float, optional
         The bias to use for initializing the model parameters. The default is -3.5.
+    init_beta_variance : float, optional
+        The variance to use for initializing the model's beta parameters from a normal
+        distribution. The default is 1.0.
     n_hidden_units : int, optional
         The number of hidden units to use in the neural network model. The default is 5.
     lower_bound : float, optional

tests/test_data.py

@@ -302,7 +302,7 @@ def test_linear_model_multi_cond_fit_simple():
         reference="a",
         assert_site_integrity=False,
     )
-    assert np.all(data.bundle_idxs["a"] == False)
+    assert np.all(data.bundle_idxs["a"] is False)
     model = multidms.Model(data, multidms.biophysical.identity_activation, PRNGKey=23)
     model.fit(maxiter=2, warn_unconverged=False)