diff --git a/chirp_ams.py b/chirp_ams.py new file mode 100644 index 0000000..8660dbe --- /dev/null +++ b/chirp_ams.py @@ -0,0 +1,122 @@ +import numpy as np +import scipy.signal as sig +import matplotlib.pyplot as plt + +from chirp_stimulation import create_chirp +from IPython import embed + +def despine(axis, spines=None, hide_ticks=True): + + def hide_spine(spine): + spine.set_visible(False) + + for spine in axis.spines.keys(): + if spines is not None: + if spine in spines: + hide_spine(axis.spines[spine]) + else: + hide_spine(axis.spines[spine]) + if hide_ticks: + axis.xaxis.set_ticks([]) + axis.yaxis.set_ticks([]) + + +def get_signals(eodfs, condition, contrast, c_size, c_duration, c_ampl_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=c_size, chirpduration=c_duration, + ampl_reduction=c_ampl_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, 4.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("times [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.show() \ No newline at end of file