Adding per-base inference

netam/common.py

@@ -7,8 +7,10 @@
 import torch.optim as optim
 from torch import nn, Tensor
 
+BIG = 1e9
 SMALL_PROB = 1e-6
 BASES = ["A", "C", "G", "T"]
+BASES_AND_N_TO_INDEX = {"A": 0, "C": 1, "G": 2, "T": 3, "N": 4}
 AA_STR_SORTED = "ACDEFGHIKLMNPQRSTVWY"
 AA_STR_SORTED_AMBIG = AA_STR_SORTED + "X"
 MAX_AMBIG_AA_IDX = len(AA_STR_SORTED_AMBIG) - 1
@@ -53,6 +55,10 @@ def mask_tensor_of(seq_str, length=None):
     return mask
 
 
+def informative_site_count(seq_str):
+    return sum(c != "N" for c in seq_str)
+
+
 def clamp_probability(x: Tensor) -> Tensor:
     return torch.clamp(x, min=SMALL_PROB, max=(1.0 - SMALL_PROB))
 

netam/dnsm.py

@@ -239,10 +239,8 @@ def train_test_datasets_of_pcp_df(pcp_df, train_frac=0.8, branch_length_multipli
 
 
 class DNSMBurrito(framework.Burrito):
-    def __init__(self, *args, device=pick_device(), **kwargs):
+    def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.device = device
-        self.model.to(self.device)
         self.wrapped_model = WrappedBinaryMutSel(self.model, weights_directory=None)
 
     def load_branch_lengths(self, in_csv_prefix):
@@ -401,11 +399,11 @@ def burrito_of_model(
         l2_regularization_coeff=1e-6,
         verbose=False,
     ):
+        model.to(device)
         burrito = DNSMBurrito(
             self.train_dataset,
             self.val_dataset,
             model,
-            device=device,
             batch_size=batch_size,
             learning_rate=learning_rate,
             min_learning_rate=min_learning_rate,

netam/experiment.py

@@ -38,14 +38,14 @@ def build_model_instances(self, prename):
                 filter_count=14,
                 dropout_prob=0.1,
             ),
-            f"{prename}_cnn_med_orig": models.CNNModel(
+            f"{prename}_cnn_med": models.CNNModel(
                 kmer_length=3,
-                kernel_size=11,
-                embedding_dim=9,
-                filter_count=9,
-                dropout_prob=0.1,
+                kernel_size=9,
+                embedding_dim=7,
+                filter_count=16,
+                dropout_prob=0.2,
             ),
-            f"{prename}_cnn_med": models.CNNModel(
+            f"{prename}_ind_rscnn_med": models.IndepRSCNNModel(
                 kmer_length=3,
                 kernel_size=9,
                 embedding_dim=7,
@@ -135,7 +135,8 @@ def train_or_load(
 
         our_burrito_params = deepcopy(self.burrito_params)
         our_burrito_params.update(training_params)
-        burrito = framework.SHMBurrito(
+        burrito_class = framework.burrito_class_of_model(model)
+        burrito = burrito_class(
             train_dataset, val_dataset, model, verbose=False, **our_burrito_params
         )
         train_history = burrito.multi_train(epochs=self.epochs)
@@ -211,7 +212,8 @@ def train_experiment_df(self, experiment_df, pretrained_dir="../pretrained"):
 
     def calculate_loss(self, model, dataset):
         model.eval()
-        burrito = framework.SHMBurrito(
+        burrito_class = framework.burrito_class_of_model(model)
+        burrito = burrito_class(
             dataset, dataset, model, verbose=False, **self.burrito_params
         )
         loss = burrito.evaluate()