In silico mixtures can now be made

coa_detection/in_silico_mixture.py

@@ -3,7 +3,7 @@ from pathlib import Path
 import numpy as np
 import random
 import matplotlib.pyplot as plt
-
+from pyabf.abfWriter import writeABF1
 
 class InSilicoGenerator:
     def __init__(self, file_list):
@@ -32,35 +32,48 @@ class InSilicoGenerator:
             elif 'A6' in file.name:
                 self.coa6.append(AbfData(file))
 
-    def fill_squiggle(self, coa4_ratio, coa5_ratio, coa6_ratio, unfiltered=True):
+    def fill_squiggle(self, coa4_count, coa5_count, coa6_count, unfiltered=True,
+                      insert_prob=0.95):
         """
-        :param coa4_ratio: number of coa4 events
-        :param coa5_ratio: number of coa5 events
-        :param coa6_ratio: number of coa6 events
+        :param coa4_count: number of coa4 events
+        :param coa5_count: number of coa5 events
+        :param coa6_count: number of coa6 events
         """
         current_index = 0
-        # Just start with some noise first
+        # Keep track of how many coa4,5,6 are inserted into the squiggle and what the target count is
+        count_dict = {4: [0, coa4_count],
+                      5: [0, coa5_count],
+                      6: [0, coa6_count]}
         while current_index < len(self.squiggle):
             fragment_width = random.randint(100, 3000)
             coa_type = random.choice(self.all_coas)
-            fragment_id = 1  # TODO base this on coa type that is chosen
             coa_object = random.choice(coa_type)
-            positive_example = random.choice([True, False])
+            fragment_id = int(coa_object.abf_fn.stem[-1])
+            if count_dict[fragment_id][0] < count_dict[fragment_id][1]:
+                positive_example = random.random() > insert_prob
+            else:
+                positive_example = False
+
             if current_index + fragment_width > self.squiggle_length:
                 # Edge case: always fill final bit of squiggle with noise
                 positive_example = False
                 fragment_width = self.squiggle_length - current_index
             if positive_example:
                 squiggle_fragment = coa_object.get_pos(fragment_width,
-                                                       unfiltered=unfiltered)  # TODO this might not match up with real data
+                                                       unfiltered=unfiltered)
+                count_dict[fragment_id][0] += 1
             else:
                 squiggle_fragment = coa_object.get_neg(fragment_width,
-                                                       unfiltered=unfiltered)
+                                                       nb_neg=1,
+                                                       unfiltered=unfiltered)[0]
                 fragment_id *= -1
             end_index = current_index + fragment_width
             self.squiggle[current_index:end_index] = squiggle_fragment
             self.ground_truth_vector[current_index:end_index] = fragment_id
             current_index += fragment_width
+        print('cOA type: [number of events in in silico squiggle, target number]')
+        print(count_dict)
+        print('If it is too far off, change insert_prob')
         return self.squiggle, self.ground_truth_vector
 
 
@@ -69,6 +82,8 @@ if __name__ == '__main__':
              Path(r"/mnt/c/Users/benno/Downloads/zooi/+120mV_cOA5.abf"),
              Path(r"/mnt/c/Users/benno/Downloads/zooi/+120mV_cOA6.abf")]
     generator = InSilicoGenerator(files)
-    simulated_squiggle, ground_truth = generator.fill_squiggle(1, 2, 3)
+    simulated_squiggle, ground_truth = generator.fill_squiggle(30, 30, 30)
     plt.plot(simulated_squiggle)
     plt.show()
+    # simulated_squiggle = np.reshape(simulated_squiggle, (-1, 1))
+    # writeABF1(simulated_squiggle, r'/mnt/c/Users/benno/PycharmProjects/baseless/coa_detection/output_test.abf', 50000)