import os import numpy as np import scipy.signal as sig import matplotlib.pyplot as plt from chirp_stimulation import create_chirp from util import despine figure_folder = "figures" def get_signals(eodfs, condition, contrast, chirp_size, chirp_duration, chirp_amplitude_dip, chirp_times, duration, dt): if not isinstance(condition, str) or ("self" not in condition and "other" not in condition): raise ValueError("Condition argument must be either 'self' or 'other'!") if not isinstance(eodfs, dict) or (not "self" in eodfs.keys() or not "other" in eodfs.keys()): raise ValueError("EOFs must be a dict containing 'self' and 'other' fish's eod frequency!") time = np.arange(0.0, duration, dt) non_chirper_freq = eodfs["self"] if condition == "other" else eodfs["other"] non_chirper_signal = np.sin(non_chirper_freq * time * 2 * np.pi) non_chirper_freq_profile = np.ones(time.shape) * non_chirper_freq chirper_freq = eodfs["other"] if condition == "other" else eodfs["self"] _, chirper_signal, _, chirper_freq_profile = create_chirp(eodf=chirper_freq, chirpsize=chirp_size, chirpduration=chirp_duration, ampl_reduction=chirp_amplitude_dip, chirptimes=chirp_times, duration=duration, dt=dt) other_ampl = contrast/100 if condition == "self": self_signal = chirper_signal self_freq = chirper_freq_profile other_signal = non_chirper_signal * other_ampl other_freq = non_chirper_freq_profile else: self_signal = non_chirper_signal self_freq = non_chirper_freq_profile other_signal = chirper_signal * other_ampl other_freq = chirper_freq_profile return time, self_signal, self_freq, other_signal, other_freq if __name__ == "__main__": eod_frequencies = {"self": 600, "other": 620} # Hz, eod frequencies of the two fish, 'self' is the one that we "record" from eod_contrasts = [20, 10, 5, 2.5, 1.25, 0.625, 0.3125] # %, strength of 'other' relative to 'self' chirp_size = 100 # Hz, frequency excursion chirp_duration = 0.015 # s, chirp duration chirp_amplitude_dip = 0.05 # %, amplitude drop during chirp chirp_frequency = 10 # Hz, how often does the fish chirp total_duration = 0.5 # s, total duration of simulation dt = 0.00001 # s, stepsize of the simulation chirp_times = np.arange(0.125+chirp_duration, 0.125 + total_duration - chirp_duration, 1./chirp_frequency) grid_shape = (5 + len(eod_contrasts) - 1, 7) conditions = ["other", "self"] fig = plt.figure(figsize=(4.5, 5.5)) for i, condition in enumerate(conditions): time, self_signal, self_freq, other_signal, other_freq = get_signals(eod_frequencies, condition, eod_contrasts[0], chirp_size, chirp_duration, chirp_amplitude_dip, chirp_times, total_duration + 0.25, dt) plot_time = time[(time >= 0.125) & (time < total_duration + 0.125)] - 0.125 ax = plt.subplot2grid(grid_shape, (0, i * 3 + i * 1), rowspan=2, colspan=3, fig=fig) ax.plot(plot_time, self_freq[(time >= 0.125) & (time < total_duration + 0.125)], color="#ff7f0e", label="%iHz" % eod_frequencies["self"]) ax.plot(plot_time, other_freq[(time >= 0.125) & (time < total_duration + 0.125)], color="#1f77b4", label="%iHz" % eod_frequencies["other"]) if i == 0: ax.text(1.15 * plot_time[-1], eod_frequencies["self"], "self", color="#ff7f0e", va="center", ha="left", fontsize=9) ax.text(1.15 * plot_time[-1], eod_frequencies["other"], "other", color="#1f77b4", va="center", ha="left", fontsize=9) ax.text(-0.05 * plot_time[-1], eod_frequencies["self"], "%iHz" % eod_frequencies["self"], color="#ff7f0e", va="center", ha="right", fontsize=9) ax.text(-0.05 * plot_time[-1], eod_frequencies["other"], "%iHz" % eod_frequencies["other"], color="#1f77b4", va="center", ha="right", fontsize=9) despine(ax, spines=["top", "bottom", "left", "right"]) ax = plt.subplot2grid(grid_shape, (3, i * 3 + i * 1), rowspan=2, colspan=3, fig=fig) combined = self_signal + other_signal plot_combined = combined[(time >= 0.125) & (time < total_duration + 0.125)] am = np.abs(sig.hilbert(combined)) plot_am = am[(time >= 0.125) & (time < total_duration + 0.125)] ax.plot(plot_time, plot_combined, color="#2ca02c", label="self + other") ax.plot(plot_time, plot_am, color="#d62728", label="am") ax.set_ylim([-1.25, 1.25]) if i == 0: ax.text(1.25 * plot_time[-1], np.mean(combined), "contrast=\n20%",color="#d62728", va="center", ha="center", fontsize=9) ax.text(-0.05 * plot_time[-1], np.mean(am), "am", color="#d62728", va="center", ha="right", fontsize=9) ax.text(-0.05 * plot_time[-1], np.mean(combined), "self+\nother", color="#2ca02c", va="center", ha="right", fontsize=9) despine(ax, spines=["top", "bottom", "left", "right"]) for j, contrast in enumerate(eod_contrasts[1:]): time, self_signal, self_freq, other_signal, other_freq = get_signals(eod_frequencies, condition, contrast, chirp_size, chirp_duration, chirp_amplitude_dip, chirp_times, total_duration + 0.25, dt) combined = self_signal + other_signal am = np.abs(sig.hilbert(combined)) plot_time = time[(time >= 0.125) & (time < total_duration + 0.125)] - 0.125 plot_combined = combined[(time >= 0.125) & (time < total_duration + 0.125)] plot_am = am[(time >= 0.125) & (time < total_duration + 0.125)] ax = plt.subplot2grid(grid_shape, (5 + j, i * 3 + i * 1), rowspan=1, colspan=3) ax.plot(plot_time, plot_am, color="#d62728", label="am") ax.text(1.25 * plot_time[-1], np.mean(am), "%.2f" % contrast, color="#d62728", va="center", ha="center", fontsize=9) ax.set_ylim([0.8, 1.2]) if j == len(eod_contrasts)-2: despine(ax, spines=["top", "left", "right"]) ax.set_xticks(np.arange(0.0, total_duration + 0.001, 0.25)) ax.set_xticklabels(np.arange(0.0, total_duration * 1000+1, 250), fontsize=7) ax.set_xlabel("time [ms]", fontsize=9) else: despine(ax, spines=["top", "bottom", "left", "right"]) fig.subplots_adjust(left=0.1, bottom=0.1, top=0.99, right=0.99) plt.savefig(os.path.join(figure_folder, "Chirp_induced_AMs.pdf")) plt.close()