Docs: simplify README examples and scripts

README.md

@@ -46,11 +46,11 @@ Examples below require `matplotlib`:
 python -m pip install matplotlib
 ```
 
-### Example 1 — Multi-taper spectrum of a signal
+### Example 1 — Single sinusoid: signal + multitaper spectrum + spectrogram
 Run:
 
 ```bash
-python examples/readme_examples/example1_multitaper_spectrum.py
+python examples/readme_examples/example1_multitaper_and_spectrogram.py
 ```
 
 ```python
@@ -61,51 +61,53 @@ from pathlib import Path
 
 import matplotlib.pyplot as plt
 import numpy as np
+from scipy.signal import spectrogram
 
 from nstat.compat.matlab import SignalObj
 
-rng = np.random.default_rng(0)
 fs_hz = 1000.0
 dt = 1.0 / fs_hz
 duration_s = 2.0
+f0_hz = 10.0
 time = np.arange(0.0, duration_s, dt, dtype=float)
 
-signal = (
-    1.0 * np.sin(2.0 * np.pi * 10.0 * time)
-    + 0.6 * np.sin(2.0 * np.pi * 40.0 * time + 0.3)
-    + 0.2 * np.sin(2.0 * np.pi * 75.0 * time)
-    + 0.12 * rng.standard_normal(time.size)
-)
-
-sig_obj = SignalObj(time=time, data=signal, name="synthetic_signal", units="a.u.")
+signal = np.sin(2.0 * np.pi * f0_hz * time)
+sig_obj = SignalObj(time=time, data=signal, name="sine_signal", units="a.u.")
 freq_hz, psd = sig_obj.MTMspectrum()
+f_spec, t_spec, sxx = spectrogram(signal, fs=fs_hz, nperseg=256, noverlap=224, scaling="density", mode="psd")
 
-fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(7.0, 5.0), sharex=False)
+fig, axes = plt.subplots(3, 1, figsize=(7.5, 7.5))
 preview_mask = time <= 1.0
-ax1.plot(time[preview_mask], signal[preview_mask], color="black", linewidth=1.0)
-ax1.set_xlabel("time (s)")
-ax1.set_ylabel("amplitude")
-ax1.set_title("Synthetic signal (first 1 s)")
-ax2.plot(freq_hz, psd, color="tab:blue", linewidth=1.2)
-ax2.set_xlim(0.0, 150.0)
-ax2.set_xlabel("frequency (Hz)")
-ax2.set_ylabel("power spectral density")
-ax2.set_title("Multi-taper spectrum")
+axes[0].plot(time[preview_mask], signal[preview_mask], color="tab:blue", linewidth=1.4)
+axes[0].set_title("Signal (10 Hz sinusoid)")
+axes[0].set_xlabel("time (s)")
+axes[0].set_ylabel("amplitude")
+axes[1].plot(freq_hz, psd, color="tab:orange", linewidth=1.2)
+axes[1].set_xlim(0.0, 100.0)
+axes[1].set_title("Multi-taper spectrum")
+axes[1].set_xlabel("frequency (Hz)")
+axes[1].set_ylabel("PSD")
+im = axes[2].pcolormesh(t_spec, f_spec, sxx, shading="auto", cmap="magma")
+axes[2].set_ylim(0.0, 100.0)
+axes[2].set_title("Spectrogram")
+axes[2].set_xlabel("time (s)")
+axes[2].set_ylabel("frequency (Hz)")
+fig.colorbar(im, ax=axes[2], pad=0.01, label="PSD")
 fig.tight_layout()
 
 out_dir = Path("examples/readme_examples/images")
 out_dir.mkdir(parents=True, exist_ok=True)
-fig.savefig(out_dir / "readme_example1_multitaper_spectrum.png", dpi=180)
+fig.savefig(out_dir / "readme_example1_multitaper_and_spectrogram.png", dpi=180)
 ```
 
 **Expected output**
-![Multi-taper spectrum](examples/readme_examples/images/readme_example1_multitaper_spectrum.png)
+![Multitaper and spectrogram](examples/readme_examples/images/readme_example1_multitaper_and_spectrogram.png)
 
-### Example 2 — Simulate a spike train from a time-varying CIF
+### Example 2 — Time-varying CIF over 10 seconds (single-frequency sinusoid)
 Run:
 
 ```bash
-python examples/readme_examples/example2_simulate_cif_spiketrain.py
+python examples/readme_examples/example2_simulate_cif_spiketrain_10s.py
 ```
 
 ```python
@@ -121,40 +123,42 @@ from nstat.compat.matlab import CIF, Covariate
 
 np.random.seed(0)
 dt = 0.001
-duration_s = 2.0
-time = np.arange(0.0, duration_s + 0.5 * dt, dt, dtype=float)
+duration_s = 10.0
+t = np.arange(0.0, duration_s + dt, dt, dtype=float)
 
-lambda_t = 12.0 + 5.5 * np.sin(2.0 * np.pi * 2.0 * time) + 1.5 * np.cos(2.0 * np.pi * 6.0 * time)
-lambda_t = np.clip(lambda_t, 0.1, None)
+f_hz = 0.5
+baseline_hz = 15.0
+amp_hz = 10.0
+lam = np.clip(baseline_hz + amp_hz * np.sin(2.0 * np.pi * f_hz * t), 0.2, None)
 
-lambda_cov = Covariate(time=time, data=lambda_t, name="Lambda(t)", units="spikes/s", labels=["lambda"])
-coll = CIF.simulateCIFByThinningFromLambda(lambda_cov, 1, dt)
-spike_times = coll.getNST(0).spike_times
+lambda_cov = Covariate(time=t, data=lam, name="Lambda", units="spikes/s", labels=["lambda"])
+spikes = CIF.simulateCIFByThinningFromLambda(lambda_cov, 1, dt)
+spike_times = spikes.getNST(0).spike_times
 
-fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(7.0, 5.0), sharex=True, gridspec_kw={"height_ratios": [2.0, 1.0]})
-ax1.plot(time, lambda_t, color="tab:blue", linewidth=1.4)
+fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8.0, 4.8), sharex=True, gridspec_kw={"height_ratios": [2.0, 1.0]})
+ax1.plot(t, lam, color="tab:blue", linewidth=1.3)
 ax1.set_ylabel("rate (spikes/s)")
-ax1.set_title("Time-varying CIF")
+ax1.set_title("Time-varying CIF over 10 s")
 ax2.vlines(spike_times, 0.0, 1.0, color="black", linewidth=0.8)
 ax2.set_ylim(0.0, 1.0)
+ax2.set_yticks([])
 ax2.set_xlabel("time (s)")
-ax2.set_ylabel("spikes")
 ax2.set_title("Simulated spike train")
 fig.tight_layout()
 
 out_dir = Path("examples/readme_examples/images")
 out_dir.mkdir(parents=True, exist_ok=True)
-fig.savefig(out_dir / "readme_example2_simulate_cif_spiketrain.png", dpi=180)
+fig.savefig(out_dir / "readme_example2_simulate_cif_spiketrain_10s.png", dpi=180)
 ```
 
 **Expected output**
-![CIF spike train simulation](examples/readme_examples/images/readme_example2_simulate_cif_spiketrain.png)
+![CIF spike train simulation](examples/readme_examples/images/readme_example2_simulate_cif_spiketrain_10s.png)
 
-### Example 3 — Spike-train raster (collection, nstColl.plot)
+### Example 3 — Spike train collection raster from Example 2
 Run:
 
 ```bash
-python examples/readme_examples/example3_spike_train_raster_nstcoll.py
+python examples/readme_examples/example3_nstcoll_raster_from_example2.py
 ```
 
 ```python
@@ -170,32 +174,34 @@ from nstat.compat.matlab import CIF, Covariate
 
 np.random.seed(0)
 dt = 0.001
-duration_s = 2.0
+duration_s = 10.0
 n_units = 20
-time = np.arange(0.0, duration_s + 0.5 * dt, dt, dtype=float)
+t = np.arange(0.0, duration_s + dt, dt, dtype=float)
 
-lambda_t = 9.0 + 4.0 * np.sin(2.0 * np.pi * 1.5 * time) + 2.0 * np.sin(2.0 * np.pi * 4.0 * time + 0.25)
-lambda_t = np.clip(lambda_t, 0.1, None)
+f_hz = 0.5
+baseline_hz = 15.0
+amp_hz = 10.0
+lam = np.clip(baseline_hz + amp_hz * np.sin(2.0 * np.pi * f_hz * t), 0.2, None)
 
-lambda_cov = Covariate(time=time, data=lambda_t, name="Lambda(t)", units="spikes/s", labels=["lambda"])
+lambda_cov = Covariate(time=t, data=lam, name="Lambda", units="spikes/s", labels=["lambda"])
 coll = CIF.simulateCIFByThinningFromLambda(lambda_cov, n_units, dt)
 
-fig, ax = plt.subplots(figsize=(7.0, 4.5))
+fig, ax = plt.subplots(figsize=(8.0, 4.8))
 plt.sca(ax)
 coll.plot()
 ax.set_xlabel("time (s)")
 ax.set_ylabel("unit index")
-ax.set_title("Spike-train raster (nstColl.plot)")
+ax.set_title("Spike-train collection raster (nstColl.plot)")
 ax.set_ylim(0.5, n_units + 0.5)
 fig.tight_layout()
 
 out_dir = Path("examples/readme_examples/images")
 out_dir.mkdir(parents=True, exist_ok=True)
-fig.savefig(out_dir / "readme_example3_spike_train_raster_nstcoll.png", dpi=180)
+fig.savefig(out_dir / "readme_example3_nstcoll_raster.png", dpi=180)
 ```
 
 **Expected output**
-![Spike train raster](examples/readme_examples/images/readme_example3_spike_train_raster_nstcoll.png)
+![Spike train raster](examples/readme_examples/images/readme_example3_nstcoll_raster.png)
 
 ## Examples and notebooks
 

examples/readme_examples/example1_multitaper_and_spectrogram.py

@@ -0,0 +1,80 @@
+import matplotlib
+matplotlib.use("Agg")
+
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy.signal import spectrogram
+
+from nstat.compat.matlab import SignalObj
+
+
+def _fallback_multitaper_psd(signal: np.ndarray, fs_hz: float) -> tuple[np.ndarray, np.ndarray]:
+    from scipy.signal.windows import dpss
+
+    n = signal.size
+    tapers = dpss(n, NW=3.5, Kmax=4, sym=False)
+    centered = signal - np.mean(signal)
+    tapered = tapers * centered[np.newaxis, :]
+    fft_vals = np.fft.rfft(tapered, axis=1)
+    psd = np.mean(np.abs(fft_vals) ** 2, axis=0) / (fs_hz * n)
+    freq = np.fft.rfftfreq(n, d=1.0 / fs_hz)
+    return freq, psd
+
+
+def main() -> None:
+    fs_hz = 1000.0
+    duration_s = 2.0
+    f0_hz = 10.0
+    dt = 1.0 / fs_hz
+    time = np.arange(0.0, duration_s, dt, dtype=float)
+
+    signal = np.sin(2.0 * np.pi * f0_hz * time)
+    sig_obj = SignalObj(time=time, data=signal, name="sine_signal", units="a.u.")
+
+    try:
+        freq_hz, psd = sig_obj.MTMspectrum()
+    except Exception:
+        freq_hz, psd = _fallback_multitaper_psd(signal, fs_hz)
+
+    f_spec, t_spec, sxx = spectrogram(
+        signal,
+        fs=fs_hz,
+        nperseg=256,
+        noverlap=224,
+        scaling="density",
+        mode="psd",
+    )
+
+    fig, axes = plt.subplots(3, 1, figsize=(7.5, 7.5))
+
+    preview = time <= 1.0
+    axes[0].plot(time[preview], signal[preview], color="tab:blue", linewidth=1.4)
+    axes[0].set_title("Signal (10 Hz sinusoid)")
+    axes[0].set_xlabel("time (s)")
+    axes[0].set_ylabel("amplitude")
+
+    axes[1].plot(freq_hz, psd, color="tab:orange", linewidth=1.2)
+    axes[1].set_xlim(0.0, 100.0)
+    axes[1].set_title("Multi-taper spectrum")
+    axes[1].set_xlabel("frequency (Hz)")
+    axes[1].set_ylabel("PSD")
+
+    im = axes[2].pcolormesh(t_spec, f_spec, sxx, shading="auto", cmap="magma")
+    axes[2].set_ylim(0.0, 100.0)
+    axes[2].set_title("Spectrogram")
+    axes[2].set_xlabel("time (s)")
+    axes[2].set_ylabel("frequency (Hz)")
+    fig.colorbar(im, ax=axes[2], pad=0.01, label="PSD")
+
+    fig.tight_layout()
+
+    out_dir = Path(__file__).resolve().parent / "images"
+    out_dir.mkdir(parents=True, exist_ok=True)
+    fig.savefig(out_dir / "readme_example1_multitaper_and_spectrogram.png", dpi=180)
+    plt.close(fig)
+
+
+if __name__ == "__main__":
+    main()