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()