From 2328e025c77b9e480cdbc884f3853e37fd89d036 Mon Sep 17 00:00:00 2001
From: Jan Benda <jan.benda@uni-tuebingen.de>
Date: Mon, 26 Jan 2026 19:20:30 +0100
Subject: [PATCH] added stimuli

---
 twobeats.py | 35 +++++++++++++++++++++++++++++++++--
 1 file changed, 33 insertions(+), 2 deletions(-)

diff --git a/twobeats.py b/twobeats.py
index 95bfc80..6447c9a 100644
--- a/twobeats.py
+++ b/twobeats.py
@@ -78,6 +78,35 @@ def align_spikes(spikes, period):
     return spikes
 
 
+def plot_symbols(ax, s):
+    ax.show_spines('')
+
+
+def plot_stimulus(ax, s, tmax, eodf, f1, f2, c=0.1):
+    time = np.arange(0, tmax, 0.0001)
+    eod = np.cos(2*np.pi*eodf*time)
+    am = np.ones(len(time))
+    ams = {}
+    label = '$f_{EOD}$'
+    if f1 is not None:
+        am += c*np.cos(2*np.pi*f1*time)
+        ams = s.lsF01
+        label += r' \& $f_1$'
+    if f2 is not None:
+        am += c*np.cos(2*np.pi*f2*time)
+        ams = s.lsF02
+        label += r' \& $f_2$'
+    if f1 is not None and f2 is not None:
+        ams = s.lsF012
+    ax.show_spines('')
+    ax.plot(1000*time, am*eod, **s.lsStim)
+    if len(ams) > 0:
+        ax.plot(1000*time, am, **ams)
+    ax.set_xlim(0, 1000*tmax)
+    ax.set_ylim(-1.02 - 2*c, 1.02 + 2*c)
+    ax.text(0, 1.1, label, transform=ax.transAxes)
+
+
 def plot_raster(ax, s, spikes, tmin, tmax):
     spikes_ms = [1000*(s[(s > tmin) & (s < tmax)] - tmin) for s in spikes]
     ax.show_spines('')
@@ -140,7 +169,6 @@ if __name__ == '__main__':
     print(f'    Df2  = {df2:.1f}Hz')
     print(f'    {len(spikes)} trials')
 
-
     s = plot_style()
     fig, axs = plt.subplots(5, 4, cmsize=(s.plot_width, 0.6*s.plot_width),
                             height_ratios=[1, 2, 1, 3, 6])
@@ -149,14 +177,17 @@ if __name__ == '__main__':
     tmax = 0.2
     twins = [[-t0, 0], [t1, t1 + t2], [0, t1], [t1 + t2, t1 + t2 + t12]]
     freqs = [eodf, df2, df1, df2]
+    stim_freqs = [[None, None], [df2, None], [None, df1], [df1, df2]]
     for i in range(axs.shape[1]):
         tstart, tend = twins[i]
+        plot_symbols(axs[0, i], s)
+        plot_stimulus(axs[1, i], s, tmax - tmin, eodf, *stim_freqs[i])
         sub_spikes = [times[(times >= tstart) & (times <= tend)] - tstart for times in spikes]
+        plot_psd(axs[4, i], s, sub_spikes, tend - tstart, fmax)
         print(f'align spikes for frequency {freqs[i]:.0f}Hz:')
         sub_spikes = align_spikes(sub_spikes, abs(1/freqs[i]))
         plot_raster(axs[2, i], s, sub_spikes, tmin, tmax)
         plot_rate(axs[3, i], s, sub_spikes, tmin, tmax)
-        plot_psd(axs[4, i], s, sub_spikes, tend - tstart, fmax)
     #fig.savefig()
     plt.show()
     print()