line_tracking_of_fish_movement/do_calc_eod_times.py

108 lines
4.6 KiB
Python

import sys
import os
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt
from thunderfish.dataloader import open_data
from thunderfish.eodanalysis import eod_waveform
from IPython import embed
def load_tracker_data(foldername):
'''
Load data
:param foldername: where are the data saved
:return: ident_v: identity vector
power: power vector
freq_v: frequency vector
timeidx: time vector in indices
realtime: time vector in relative real time to the beginning
'''
ident_v = np.load('../../../data/' + foldername + '/all_ident_v.npy', allow_pickle=True)
power = np.load('../../../data/' + foldername + '/all_sign_v.npy', allow_pickle=True)
freq_v = np.load('../../../data/' + foldername + '/all_fund_v.npy', allow_pickle=True)
timeidx = np.load('../../../data/' + foldername + '/all_idx_v.npy', allow_pickle=True)
realtime = np.load('../../../data/' + foldername + '/all_times.npy', allow_pickle=True)
aifl = np.load('../data/' + foldername + '/aifl2.npy', allow_pickle=True)
return ident_v, power, freq_v, timeidx, realtime, aifl
def extract_times_freq_for_plotting(aifl, ident_v, freq_v, power_v, timeidx, realtime, fish_number):
t = []
f = []
time_res = np.diff(realtime)[0]
dt = 10 # IN SECONDS
for channel in tqdm(range(16)):
for id in aifl[channel,fish_number][~np.isnan(aifl[channel, fish_number])]:
i_df_power_channel = []
itr = np.arange(len(ident_v[channel]))[ident_v[channel] == id]
for i0 in itr[np.array(np.round(np.linspace(0, len(itr)*.9-1, 100)), dtype=int)]:
freqs_of_interest = freq_v[channel][(np.arange(len(ident_v[channel])) >= i0) &
(np.arange(len(ident_v[channel])) <= i0 + dt / time_res) &
(ident_v[channel] == id)]
df = np.max(freqs_of_interest) - np.min(freqs_of_interest)
mean_f = np.mean(freqs_of_interest)
# if df <= 0.5:
sign_of_interest = power_v[channel][(np.arange(len(ident_v[channel])) >= i0) &
(np.arange(len(ident_v[channel])) <= i0 + dt / time_res) &
(ident_v[channel] == id)]
max_channel = np.argmax(sign_of_interest, axis=1)
power = list(map(lambda x: np.max(x), sign_of_interest))
i_df_power_channel.append([i0, mean_f, np.min(power), max_channel[0], df])
i0, mean_f, power, c, df = i_df_power_channel[np.nanargmax(np.array(i_df_power_channel)[:, 4])]
t.append(realtime[timeidx[channel][i0]])
f.append(mean_f)
return t, f
def main(filename):
'''
Extracts the time points and corresponding frequencies where the EOD waveforms should be calculated
:param filename:
:return: all_t: list; all time points of each fish where the eod should be extracted
all_f: list; all frequencies of each fish where the eod should be extracted
all_fish: list; fish_number of the aifl
'''
# load data
ident_v, power_v, freq_v, timeidx, realtime, aifl = load_tracker_data(filename)
fish_in_aifl = list(np.unique(np.where(~np.isnan(aifl[:, :, 0]))[1]))
all_t = []
all_f = []
all_fish = []
for fish_number in fish_in_aifl:
print(fish_number)
# if 400 <= all_q10[fish_number, 2] <= 650 and power_means[fish_number] >= -90.0:
times, freqs = extract_times_freq_for_plotting(aifl, ident_v, freq_v, power_v, timeidx, realtime, fish_number)
all_t.append(times)
all_f.append(freqs)
all_fish.append(fish_number)
return all_t, all_f, all_fish
if __name__ == '__main__':
for index, filename_idx in enumerate([0, 2, 3, 5]):
filename = sorted(os.listdir('../../../data/mount_data/sanmartin/softgrid_1x16/'))[filename_idx]
print('new file: '+ filename)
all_t, all_f, all_fish = main(filename)
#############################################################################################################
# save data
if filename not in os.listdir('../data/'):
os.mkdir('../data/' + filename)
np.save('../data/' + filename + '/eod_times_new_new.npy', all_t)
np.save('../data/' + filename + '/eod_freq_new_new.npy', all_f)
np.save('../data/' + filename + '/eod_fishnumber_new_new.npy', all_fish)