116 lines
6.9 KiB
Python
116 lines
6.9 KiB
Python
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() |