Eliminate use of deprecated scheduler epoch param, preclude lambda in…

…finite loop
broadinstitute · Jun 19, 2020 · 41d87e3 · 41d87e3
1 parent 5134191
commit 41d87e3
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 7 deletions.
diff --git a/cellbender/remove_background/infer.py b/cellbender/remove_background/infer.py
@@ -1013,9 +1013,12 @@ def _lambda_binary_search_given_fpr(self,
 
         # Travel in one direction until the direction of FPR error changes.
         lam_limit = lam_mult_init
-        while (lam_limit < 500) and (residual.sign() == initial_residual_sign):
+        i = 0
+        while ((lam_limit < 500) and (lam_limit > 0.1) and
+                (residual.sign() == initial_residual_sign) and (i < max_iterations)):
+
             lam_limit = lam_limit * (initial_residual_sign * 2).exp().item()
-            print(f'{lam_limit:.4f}')
+            print(f'{lam_limit}')
 
             # Calculate an expected false positive rate for this lam_mult value.
             expected_fpr = self._calculate_expected_fpr_given_lambda_mult(
@@ -1026,6 +1029,7 @@ def _lambda_binary_search_given_fpr(self,
                 alpha_est=alpha_est)
 
             residual = fpr - expected_fpr
+            i = i + 1  # one dataset had this go into an infinite loop, taking lam_limit -> 0
 
         # Define the values that bracket the correct answer.
         lam_mult_bracket = np.sort(np.array([lam_mult_init, lam_limit]))

diff --git a/cellbender/remove_background/train.py b/cellbender/remove_background/train.py
@@ -138,11 +138,13 @@ def run_training(model: RemoveBackgroundPyroModel,
                 t = time.time()
 
             # Train, and keep track of training loss.
-            if epoch >= epoch_start_cooldown:
-                # Co-opt the scheduler: make it think we're on epoch 0.
-                total_epoch_loss_train = train_epoch(svi, train_loader, epoch=0)
-            else:
-                total_epoch_loss_train = train_epoch(svi, train_loader)
+            # if epoch >= epoch_start_cooldown:
+            #     # Co-opt the scheduler: make it think we're on epoch 0.
+            #     total_epoch_loss_train = train_epoch(svi, train_loader, epoch=0)  # TODO: deprecated: change this
+            # else:
+            #     total_epoch_loss_train = train_epoch(svi, train_loader)
+
+            total_epoch_loss_train = train_epoch(svi, train_loader)
 
             train_elbo.append(-total_epoch_loss_train)
             model.loss['train']['epoch'].append(epoch)