Update mutagenizer.py

evanseitz · Apr 3, 2024 · 9820ef8 · 9820ef8
1 parent 6090263
commit 9820ef8
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/squid/mutagenizer.py b/squid/mutagenizer.py
@@ -146,9 +146,12 @@ def swap_elements(x, t):
         # mutagenize each sequence based on number of mutations; i.e., samples from alphabet pool
         one_hot = np.zeros(shape=(num_sim, L, A))
         for i, num_mut in enumerate(tqdm(num_muts, desc="Mutagenesis")):
-            options_list = choice(alphabet_pool, size=num_mut, replace=True) # sample 'num_mut' characters from alphabet_pool with replacement
-            mut_seq = ''.join(swap_elements(seq, options_list))
-            one_hot[i,:,:] = seq2twohot(mut_seq)
+            if i == 0:
+                one_hot[i,:,:] = seq2twohot(''.join(seq))
+            else:
+                options_list = choice(alphabet_pool, size=num_mut, replace=True) # sample 'num_mut' characters from alphabet_pool with replacement
+                mut_seq = ''.join(swap_elements(seq, options_list))
+                one_hot[i,:,:] = seq2twohot(mut_seq)
 
         return one_hot