fixes ants module version, and also style

brainsss/brain_utils.py

@@ -4,18 +4,126 @@
 import nibabel as nib
 import os
 import time
+from scipy.signal import butter, filtfilt, freqz
 
-def extract_traces(fictrac, stim_times, pre_window, post_window):
+def extract_traces(fictrac, stim_times, pre_window, post_window, val=None):
     traces = []
     for i in range(len(stim_times)):
-        trace = fictrac['Z'][stim_times[i]-pre_window:stim_times[i]+post_window]
+        if val != None:
+            trace = fictrac[val][stim_times[i]-pre_window:stim_times[i]+post_window]
+        else:
+            trace = fictrac[stim_times[i]-pre_window:stim_times[i]+post_window]
         if len(trace) == pre_window + post_window: # this handles fictrac that crashed or was aborted or some bullshit
             traces.append(trace)
     traces = np.asarray(traces)
     mean_trace = np.mean(traces,axis=0)
     sem_trace = scipy.stats.sem(traces,axis=0)
     return traces, mean_trace, sem_trace
 
+def get_vals_for_analysis(tp_width_sec=0, trial_time_sec=0, before_stim_sec=0, pre_trial_window_sec=0, pre_trial_size_sec=0, post_trial_window_sec=0, post_trial_size_sec=0, thresh=0, fr=0):
+
+    stim_end_sec = before_stim_sec + trial_time_sec
+
+    stim_idx = int(before_stim_sec/tp_width_sec)
+    stim_end_idx = int(stim_end_sec/tp_width_sec)
+    pre_trial_window_idx = int(pre_trial_window_sec/tp_width_sec)
+
+    # first val i use to make delta to compare change in vel
+    first_val_end_idx = int(stim_idx - pre_trial_window_idx)
+    first_val_start_idx = int(first_val_end_idx - (pre_trial_size_sec/tp_width_sec)) #only needed if you're averaging across a window, otherwise just start val is used
+
+    # second val used to make delta
+    second_val_start_idx = int(stim_idx + post_trial_size_sec/tp_width_sec)
+    second_val_end_idx = int(second_val_start_idx + post_trial_window_sec/tp_width_sec)
+    return stim_idx, stim_end_idx, pre_trial_window_idx, first_val_end_idx, first_val_start_idx, second_val_end_idx, second_val_start_idx
+
+def butter_lowpass(cutoff, fs, order=5):
+    return butter(order, cutoff, fs=fs, btype='low', analog=False)
+
+def butter_lowpass_filter(data, cutoff, fs, order=5):
+    b, a = butter_lowpass(cutoff, fs, order=order)
+    y = filtfilt(b, a, data, method="gust")
+    return y
+
+def apply_butter_lowpass(behavior_traces, stim_idx, fr):
+    # Filter requirements.
+    order = 4
+    fs = fr      # sample rate, Hz
+    cutoff = 3  # desired cutoff frequency of the filter, Hz
+
+    # Filter the data, and plot both the original and filtered signals.
+    lpf_behavior = np.zeros((behavior_traces.shape[0], behavior_traces.shape[1] - stim_idx))
+
+    for i in range(behavior_traces.shape[0]):
+        lpf_behavior[i, :] = butter_lowpass_filter(behavior_traces[i,stim_idx:], cutoff, fs, order)
+
+    return lpf_behavior
+
+def get_speed_change(behavior, stim_idx, first_val_start_idx, first_val_end_idx, second_val_start_idx, second_val_end_idx):
+    pre_tsi= first_val_start_idx-stim_idx
+    pre_tei = first_val_end_idx-stim_idx
+    post_tsi = second_val_start_idx-stim_idx
+    post_tei = second_val_end_idx-stim_idx
+
+    if first_val_start_idx!=first_val_end_idx:
+        first_speed = np.mean(behavior[:,pre_tsi:pre_tei], axis = 1)
+    else:
+        first_val_start_idx = int(first_val_start_idx)
+        first_speed = np.asarray(behavior[:,pre_tsi])
+
+    if second_val_start_idx!=second_val_end_idx:
+        second_speed = np.mean(behavior[:,post_tsi:post_tei], axis = 1)
+    else: 
+        second_speed = np.asarray(behavior[:,post_tsi])
+    speed_change = second_speed-first_speed
+    return speed_change
+
+def get_trials(speed_change, thresh):
+    decrease_trials = speed_change <-thresh
+    decrease_idx = np.where(decrease_trials)
+    increase_trials = speed_change>thresh
+    increase_idx = np.where(increase_trials)
+    flat_trials = np.logical_and(speed_change>-thresh, speed_change<thresh)
+    flat_idx = np.where(flat_trials)
+    return increase_trials, increase_idx, decrease_trials, decrease_idx, flat_trials, flat_idx
+
+def separate_traces(behavior_struct, value_struct):
+
+    increase_total= {}
+    increase_trial_num = 0
+    decrease_total= {}
+    decrease_trial_num = 0
+    flat_total= {}
+    flat_trial_num = 0
+
+    stim_idx, stim_end_idx, pre_trial_window_idx, first_val_end_idx, first_val_start_idx, second_val_end_idx, second_val_start_idx = get_vals_for_analysis(**value_struct)
+
+    for key in behavior_struct:
+        traces = behavior_struct[key]
+        lpf_behavior = apply_butter_lowpass(traces, stim_idx, value_struct["fr"])
+        speed_change = get_speed_change(lpf_behavior, stim_idx, first_val_start_idx, first_val_end_idx, second_val_start_idx, second_val_end_idx)
+        increase_trials,increase_idx, decrease_trials, decrease_idx, flat_trials, flat_idx = get_trials(speed_change, value_struct["thresh"])
+        increase_traces = traces[increase_trials]
+        increase_trial_num += np.shape(increase_traces)[0]
+        increase_total[key]={}
+        increase_total[key]['traces'] = increase_traces
+        increase_total[key]['idx'] = increase_idx
+        decrease_traces = traces[decrease_trials]
+        decrease_trial_num += np.shape(decrease_traces)[0]
+        decrease_total[key]={}
+        decrease_total[key]['traces'] = decrease_traces
+        decrease_total[key]['idx'] = decrease_idx
+        flat_traces = traces[flat_trials]
+        flat_trial_num += np.shape(flat_traces)[0]
+        flat_total[key]={}
+        flat_total[key]['traces'] = flat_traces
+        flat_total[key]['idx'] = flat_idx
+
+    print(f"Increase trial number: {increase_trial_num}")
+    print(f"Decrease trial number: {decrease_trial_num}")
+    print(f"Flat trial number: {flat_trial_num}")
+    return increase_total, decrease_total, flat_total
+
 def get_visually_evoked_turns(traces, mean_turn, start, stop, r_thresh, av_thresh, stim_times, expected_direction):
     ### this will flip the sign of the trace to get the correct av_thresh comparison
     if expected_direction == 'pos':

brainsss/visual.py

@@ -107,7 +107,7 @@ def get_stimulus_metadata(vision_path, printlog=None):
 		try:
 			## if no key error it must be a visprotocol metadata file ##
 			fly_ids = file['Flies']
-			metadata = parse_visprotocol_metadata(file)
+			metadata = parse_visprotocol_metadata(file, printlog)
 		except KeyError:
 			## if keyerror it is a visual_stimulation metadata file ##
 			metadata = parse_visual_stimulation_metadata(file)
@@ -124,7 +124,7 @@ def get_stimulus_metadata(vision_path, printlog=None):
 		return metadata['stim_ids'], metadata['angles']
 		printlog('Could not get visual metadata.')
 
-def parse_visprotocol_metadata(file):
+def parse_visprotocol_metadata(file, printlog):
 	### loop over flies and series to find the one that has many stim presentations (others were aborted)
 	# note it is critical each fly has their own .h5 file saved
 	found_a_full_series = False

scripts/background_subtraction.py

@@ -0,0 +1,161 @@
+import os
+import sys
+import numpy as np
+import argparse
+import subprocess
+import json
+import nibabel as nib
+import brainsss
+import h5py
+from scipy import signal
+import datetime
+from skimage import io
+import matplotlib.pyplot as plt
+from time import time
+from time import strftime
+from time import sleep
+class BgRemover3D:
+
+    def __init__(self, img_path, save_path, half_wid=25):
+        self.path = img_path
+        self.save = save_path
+        self.half_wid = half_wid
+        # self.img shoud have dimension x, y, z, t here, x is along the line scan direction
+        if '.nii' in img_path:
+            self.img = np.asarray(nib.load(img_path).get_data().squeeze(), dtype='float32')
+        elif '.h5' in img_path:
+            with h5py.File(img_path, "r") as f:
+                self.img = f['data'][:]
+        else:
+            print('format not supported!')
+        self.make_savedir()
+
+    def make_savedir(self):
+        working_dir = self.save
+        saving_dir = os.path.join(working_dir, 'background_subtraction')
+        if not os.path.exists(saving_dir):
+            os.mkdir(saving_dir)
+        self.saving_dir = saving_dir
+        self.file_head = os.path.basename(self.path).split('.')[0]
+
+    def draw_bg(self):
+        half_wid = self.half_wid
+        wid = 2*half_wid
+        kernel = np.ones(wid)/wid
+        template = np.mean(self.img, axis=-1)
+        template = np.moveaxis(template, (0, 2), (2, 0))
+        bg_ind = []
+        for patch in template:
+            bg_ind_tmp = []
+            for line in patch:
+                tmp = np.convolve(line, kernel, 'valid')
+                bg_center = np.argmin(tmp) + half_wid
+                bg_ind_tmp.append([bg_center-half_wid, bg_center+half_wid])
+            bg_ind.append(bg_ind_tmp)
+
+        self.bg_ind = bg_ind
+
+    def show_bg(self):
+        bg_ind = self.bg_ind
+        show_bg = np.mean(self.img, axis=-1)
+        mv = np.round(np.max(show_bg))
+        show_bg = np.moveaxis(show_bg, (0, 2), (2, 0))
+        for i in range(show_bg.shape[0]):
+            for j in range(show_bg.shape[1]):
+                show_bg[i, j, bg_ind[i][j][0]:bg_ind[i][j][1]] = mv
+
+        save_name = os.path.join(self.saving_dir, self.file_head+'_bg_selection.tif')
+        io.imsave(save_name, np.round(show_bg).astype('int16'))
+
+
+    def remove_bg(self, offset=300):
+        bg_ind = self.bg_ind
+        img = np.moveaxis(self.img, (0,1,2,3), (3,1,2,0))
+        out = np.zeros_like(img)
+        for ind_y in range(img.shape[1]):
+            for ind_z in range(img.shape[2]):
+                patch = img[:, ind_y, ind_z, :]
+                bg_patch = img[:, ind_y, ind_z, bg_ind[ind_z][ind_y][0]:bg_ind[ind_z][ind_y][1]]
+                bg = bg_patch.mean(axis=-1)
+                patch = patch-bg[None].T
+                out[:, ind_y, ind_z, :] = patch
+        self.out = np.moveaxis(out, (0,1,2,3), (3,1,2,0))
+
+    def show_spectrum(self, fs=180):
+        half_wid = 5
+        half_y = 15 
+        kernel2d = np.ones((half_wid*2, half_y*2))/(4*half_wid*half_y)
+        conv_template = signal.convolve2d(self.img.mean(-1).mean(-1), kernel2d, boundary='symm', mode='valid')
+        test_x, test_y = np.unravel_index(np.argmin(conv_template), conv_template.shape)
+        test_x += half_wid
+        test_y += half_y
+
+        test_patch = self.img[test_x-half_wid:test_x+half_wid, test_y-half_y:test_y+half_y, :, :]
+        test = test_patch.mean(axis=(0,1))
+        test = test.flatten(order='F') 
+        test = (test-test.mean())/test.std()
+        f, Pxx_den = signal.periodogram(test, fs)
+        plt.semilogy(f, Pxx_den)
+        plt.ylim([1e-7, 1000])
+        plt.savefig(os.path.join(self.saving_dir, self.file_head+'_before_removal.png'))
+        plt.close()
+
+        test_patch = self.out[test_x-half_wid:test_x+half_wid, test_y-half_y:test_y+half_y, :, :]
+        test = test_patch.mean(axis=(0,1))
+        test = test.flatten(order='F') 
+        test = (test-test.mean())/test.std()
+        f, Pxx_den = signal.periodogram(test, fs)
+        plt.semilogy(f, Pxx_den)
+        plt.ylim([1e-7, 1000])
+        plt.savefig(os.path.join(self.saving_dir, self.file_head+'_after_removal.png'))
+        plt.close()
+
+    def save_out(self):
+        assert self.img.shape == self.out.shape
+        if '.nii' in self.path:
+            save_name = os.path.join(self.saving_dir, self.file_head+'_cleaned.nii')
+            nib.Nifti1Image(np.round(self.out).astype('int16'), np.eye(4)).to_filename(save_name)
+        elif '.h5' in self.path:
+            save_name = os.path.join(self.saving_dir, self.file_head+'.h5')
+            with h5py.File(save_name, "w") as data_file:
+                data_file.create_dataset("data", data=np.round(self.out).astype('int16'))
+
+
+def main(args):
+
+    load_directory = args['load_directory']
+    save_directory = args['save_directory']
+    brain_file = args['brain_file']
+
+    full_load_path = os.path.join(load_directory, brain_file)
+
+    #####################
+    ### SETUP LOGGING ###
+    #####################
+
+    width = 120
+    logfile = args['logfile']
+    printlog = getattr(brainsss.Printlog(logfile=logfile), 'print_to_log')
+
+    ##############################
+    ### BACKGROUND SUBTRACTION ###
+    ##############################
+
+    printlog("Beginning BACKGROUND SUBTRACTION")
+    with h5py.File(full_load_path, 'r') as hf:
+        data = hf['data'][:]
+        dims = np.shape(data)
+
+        printlog("Data shape is {}".format(dims))  
+        br = BgRemover3D(full_load_path, save_directory, half_wid=5)
+
+        br.draw_bg()
+        br.show_bg()
+        br.remove_bg()
+        br.show_spectrum(fs=180)
+        br.save_out()
+
+    printlog("background subtraction done")
+
+if __name__ == '__main__':
+    main(json.loads(sys.argv[1]))

scripts/butter_highpass.py

@@ -0,0 +1,77 @@
+import os
+import sys
+import numpy as np
+import json
+import brainsss
+import h5py
+import datetime
+from time import time
+from scipy.signal import butter, filtfilt, freqz
+
+def main(args):
+
+    load_directory = args['load_directory']
+    save_directory = args['save_directory']
+    brain_file = args['brain_file']
+
+    full_load_path = os.path.join(load_directory, brain_file)
+    save_file = os.path.join(save_directory, brain_file.split('.')[0] + '_butter_highpass.h5')
+
+    #####################
+    ### SETUP LOGGING ###
+    #####################
+
+    width = 120
+    logfile = args['logfile']
+    printlog = getattr(brainsss.Printlog(logfile=logfile), 'print_to_log')
+
+
+    def butter_highpass(cutoff, fs, order=5):
+        return butter(order, cutoff, fs=fs, btype='high', analog=False)
+
+    def butter_highpass_filter(data, cutoff, fs, order=5):
+        b, a = butter_highpass(cutoff, fs, order=order)
+        y = filtfilt(b, a, data)
+        return y
+
+    def apply_butter_highpass(data, z, cutoff, order, fs):
+        # Get the filter coefficients so we can check its frequency response.
+        b, a = butter_highpass(cutoff, fs, order)
+        hpf_data = butter_highpass_filter(data[:,:,z, :], cutoff, fs, order)
+        return hpf_data
+
+    #################
+    ### HIGH PASS ###
+    #################
+
+    printlog("Beginning butter high pass")
+
+    #filter requirements 
+    order = 2     
+    fs = 1.8      # sample rate, Hz
+    cutoff = 0.01  # desired cutoff frequency of the filter, Hz
+
+    with h5py.File(full_load_path, 'r') as hf:
+        data = hf['data'][:] 
+        dims = np.shape(data)
+        printlog("Data shape is {}".format(dims))
+
+        hpf_total = []
+        for z in range(dims[-2]):
+            hpf_data = apply_butter_highpass(data, z, cutoff, order, fs)
+            hpf_total.append(hpf_data)
+        hpf_total = np.array(hpf_total)
+        hpf_total = np.transpose(hpf_total, (1,2,0,3))
+        dims_hpf = np.shape(hpf_total)
+        printlog("High Pass Filter Data shape is {}".format(dims_hpf))
+
+        ### Save ###
+
+        #save numpy matrix as .h5 file
+        with h5py.File(save_file, 'w') as hf:
+            hf.create_dataset('data', data=hpf_total, dtype='float32')
+
+    printlog("Butter high pass done")
+
+if __name__ == '__main__':
+    main(json.loads(sys.argv[1]))