moving in molevol and sequences.py

matsengrp · Jun 10, 2024 · ea24a89 · ea24a89
1 parent 8a3ddc3
commit ea24a89
Show file tree

Hide file tree

Showing 4 changed files with 597 additions and 12 deletions.
diff --git a/netam/dnsm.py b/netam/dnsm.py
@@ -24,14 +24,7 @@
 
 from tqdm import tqdm
 
-from epam.torch_common import optimize_branch_length
 from epam.models import WrappedBinaryMutSel
-import epam.molevol as molevol
-import epam.sequences as sequences
-from epam.sequences import (
-    aa_subs_indicator_tensor_of,
-    translate_sequences,
-)
 
 from netam.common import (
     MAX_AMBIG_AA_IDX,
@@ -42,7 +35,12 @@
 )
 import netam.framework as framework
 from netam.hyper_burrito import HyperBurrito
-
+import netam.molevol as molevol
+import netam.sequences as sequences
+from netam.sequences import (
+    aa_subs_indicator_tensor_of,
+    translate_sequences,
+)
 
 class DNSMDataset(Dataset):
     def __init__(
@@ -357,7 +355,7 @@ def _find_optimal_branch_length(
         )
         if type(starting_branch_length) == torch.Tensor:
             starting_branch_length = starting_branch_length.detach().item()
-        return optimize_branch_length(
+        return molevol.optimize_branch_length(
             log_pcp_probability, starting_branch_length, **optimization_kwargs
         )
 

diff --git a/netam/framework.py b/netam/framework.py
@@ -25,8 +25,7 @@
     VRC01_NT_SEQ,
 )
 from netam import models
-
-from epam.torch_common import optimize_branch_length
+from netam import molevol
 
 
 def encode_mut_pos_and_base(parent, child, site_count=None):
@@ -905,7 +904,7 @@ def log_pcp_probability(log_branch_length):
             rate_loss = self.bce_loss(mut_prob_masked, mutation_indicator_masked)
             return -rate_loss
 
-        return optimize_branch_length(
+        return molevol.optimize_branch_length(
             log_pcp_probability,
             starting_branch_length.double().item(),
             **optimization_kwargs,