From 20ea559a8b239a6378e5a8adcdc2dd70fe808834 Mon Sep 17 00:00:00 2001
From: Jared Galloway <jaredgalloway07@gmail.com>
Date: Mon, 8 Jul 2024 12:37:23 -0700
Subject: [PATCH] linting and documentation

---
 multidms/model.py            | 24 ++++++++++++++----------
 multidms/model_collection.py |  5 ++++-
 tests/test_data.py           |  2 +-
 3 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/multidms/model.py b/multidms/model.py
index 55571d3..d348b11 100644
--- a/multidms/model.py
+++ b/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
 
diff --git a/multidms/model_collection.py b/multidms/model_collection.py
index 5032e2b..b3f84f4 100644
--- a/multidms/model_collection.py
+++ b/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
diff --git a/tests/test_data.py b/tests/test_data.py
index 33b1061..9d6b156 100644
--- a/tests/test_data.py
+++ b/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)