[useful_functions.py] added contrast_sorting function

code/useful_functions.py

@@ -154,6 +154,59 @@ def calculate_integral(freq, power, point, delta = 2.5):
     local_mean = np.mean([l_integral, r_integral])
     return integral, local_mean, p_power
 
+def contrast_sorting(sams, con_1 = 20, con_2 = 10, con_3 = 5, stim_count = 3, stim_dur = 2):
+    '''
+    sorts the sams into three contrasts
+
+    Parameters
+    ----------
+    sams : ReproRuns
+        The sams to be sorted.
+    con_1 : int, optional
+        the first contrast. The default is 20.
+    con_2 : int, optional
+        the second contrast. The default is 10.
+    con_3 : int, optional
+        the third contrast. The default is 5.
+    stim_count : int, optional
+        the amount of stimuli per sam in a good sam. The default is 3.
+    stim_dur : int, optional
+        The stimulus duration. The default is 2.
+
+    Returns
+    -------
+    contrast_sams : dictionary
+        A dictionary containing all sams sorted to the contrasts.
+
+    '''
+    # dictionary for the contrasts
+    contrast_sams = {con_1 : [],
+                     con_2 : [],
+                     con_3 : []}
+    # loop over all sams
+    for sam in sams:
+        # get the contrast
+        avg_dur, contrast, _, _, _, _, _ = sam_data(sam)
+        # check for valid trails
+        if np.isnan(contrast):
+            continue
+        elif sam.stimulus_count < stim_count: #aborted trials
+            continue
+        elif avg_dur < (stim_dur * 0.8):
+            continue
+        else:
+            contrast = int(contrast) # get integer of contrast
+            # sort them accordingly
+            if contrast == con_1:
+                contrast_sams[con_1].append(sam)
+            elif contrast == con_2:
+                contrast_sams[con_2].append(sam)
+            elif contrast == con_3:
+                contrast_sams[con_3].append(sam)
+            else:
+                continue
+    return contrast_sams
+
 def extract_stim_data(stimulus):
     '''
     extracts all necessary metadata for each stimulus