diff --git a/event_time_analysis.py b/event_time_analysis.py index 3601250..31d573d 100644 --- a/event_time_analysis.py +++ b/event_time_analysis.py @@ -2,11 +2,16 @@ import os import sys import argparse import numpy as np -import cupy as cp +try: + import cupy as cp +except ImportError: + import numpy as cp + import matplotlib.pyplot as plt import matplotlib.gridspec as gridspec import pandas as pd from IPython import embed +from tqdm import tqdm def load_and_converete_boris_events(trial_path, recording, sr): @@ -86,30 +91,76 @@ def kde(event_dt, max_dt = 60): plt.plot(conv_t, conv_array) -def permulation_kde(event_dt, repetitions = 100, max_dt = 60): +def permulation_kde(select_event_dt, repetitions = 2000, max_dt = 60, max_mem_use_GB = 1): + def chunk_permutation(select_event_dt, conv_tt, n_chuck, max_jitter, kernal_w, kernal_h): + # array.shape = (120, 100, 15486) = (len(conv_t), repetitions, len(event_dt)) + # event_dt_perm = cp.tile(event_dt, (len(conv_t), repetitions, 1)) + event_dt_perm = cp.tile(select_event_dt, (len(conv_tt), n_chuck, 1)) + jitter = np.random.uniform(-max_jitter, max_jitter, size=(event_dt_perm.shape[1], event_dt_perm.shape[2])) + jitter = np.expand_dims(jitter, axis=0) + + event_dt_perm += jitter + # conv_t_perm = cp.tile(conv_tt, (1, repetitions, len(event_dt))) + + gauss_3d = cp.exp(-((conv_tt - event_dt_perm) / kernal_w) ** 2 / 2) * kernal_h + kde_3d = cp.sum(gauss_3d, axis = 2).transpose() + + try: + kde_3d_numpy = cp.asnumpy(kde_3d) + del event_dt_perm, gauss_3d, kde_3d + return kde_3d_numpy + + except AttributeError: + del event_dt_perm, gauss_3d + return kde_3d + embed() quit() - kernal_w = 1 kernal_h = 0.2 + max_jitter = 30 + select_event_dt = event_dt[np.abs(event_dt) <= max_dt + max_jitter*2] + conv_t = cp.arange(-max_dt, max_dt, 1) conv_tt = cp.reshape(conv_t, (len(conv_t), 1, 1)) - # array.shape = (120, 100, 15486) = (len(conv_t), repetitions, len(event_dt)) - event_dt_perm = cp.tile(event_dt, (len(conv_t), repetitions, 1)) - # conv_t_perm = cp.tile(conv_tt, (1, repetitions, len(event_dt))) + chunk_size = int(np.floor(max_mem_use_GB / (select_event_dt.nbytes * conv_t.size / 1e9))) + chunk_collector =[] + + # for _ in range(repetitions // chunk_size): + for _ in range(3): + chunk_boot_KDE = chunk_permutation(select_event_dt, conv_tt, chunk_size, max_jitter, kernal_w, kernal_h) + chunk_collector.extend(chunk_boot_KDE) + # # array.shape = (120, 100, 15486) = (len(conv_t), repetitions, len(event_dt)) + # # event_dt_perm = cp.tile(event_dt, (len(conv_t), repetitions, 1)) + # event_dt_perm = cp.tile(event_dt, (len(conv_t), chunk_size, 1)) + # jitter = np.random.uniform(-max_jitter, max_jitter, size=(event_dt_perm.shape[1], event_dt_perm.shape[2])) + # jitter = np.expand_dims(jitter, axis=0) + # + # event_dt_perm += jitter + # # conv_t_perm = cp.tile(conv_tt, (1, repetitions, len(event_dt))) + # + # gauss_3d = cp.exp(-((conv_tt - event_dt_perm) / kernal_w) ** 2 / 2) * kernal_h + # kde_3d = cp.sum(gauss_3d, axis = 2).transpose() + # try: + # kde_3d_numpy = cp.asnumpy(kde_3d) + # chunk_collector.extend(kde_3d_numpy) + # except AttributeError: + # chunk_collector.extend(kde_3d) + # del event_dt_perm, gauss_3d, kde_3d + chunk_boot_KDE = chunk_permutation(select_event_dt, conv_tt, repetitions % chunk_size, max_jitter, kernal_w, kernal_h) + chunk_collector.extend(chunk_boot_KDE) + + - gauss_3d = cp.exp(-((conv_tt - event_dt_perm) / kernal_w) ** 2 / 2) * kernal_h - kde_3d = cp.sum(gauss_3d, axis = 2).transpose() - kde_3d_numpy = cp.asnumpy(kde_3d) def main(base_path): trial_summary = pd.read_csv('trial_summary.csv', index_col=0) lose_chrips_centered_on_ag_off_t = [] - for index, trial in trial_summary.iterrows(): + for index, trial in tqdm(trial_summary.iterrows()): trial_path = os.path.join(base_path, trial['recording']) if trial['group'] < 5: @@ -128,12 +179,12 @@ def main(base_path): load_and_converete_boris_events(trial_path, trial['recording'], sr=20_000) ### communication - got_chirps = False - if os.path.exists(os.path.join(trial_path, 'chirp_times_cnn.npy')): - chirp_t = np.load(os.path.join(trial_path, 'chirp_times_cnn.npy')) - chirp_ids = np.load(os.path.join(trial_path, 'chirp_ids_cnn.npy')) - chirp_times = [chirp_t[chirp_ids == trial['win_ID']], chirp_t[chirp_ids == trial['lose_ID']]] - got_chirps = True + if not os.path.exists(os.path.join(trial_path, 'chirp_times_cnn.npy')): + continue + chirp_t = np.load(os.path.join(trial_path, 'chirp_times_cnn.npy')) + chirp_ids = np.load(os.path.join(trial_path, 'chirp_ids_cnn.npy')) + chirp_times = [chirp_t[chirp_ids == trial['win_ID']], chirp_t[chirp_ids == trial['lose_ID']]] + rise_idx = np.load(os.path.join(trial_path, 'analysis', 'rise_idx.npy'))[::sorter] rise_idx_int = [np.array(rise_idx[i][~np.isnan(rise_idx[i])], dtype=int) for i in range(len(rise_idx))]