[ENH] - Simulation of Central Frequency in the Time Domain

neurodsp/sim/__init__.py

@@ -6,4 +6,4 @@
 from .aperiodic import sim_powerlaw, sim_random_walk, sim_synaptic_current, sim_poisson_pop
 from .cycles import sim_cycle
 from .transients import sim_synaptic_kernel
-from .combined import sim_combined
+from .combined import sim_combined, sim_central_freq

neurodsp/sim/combined.py

@@ -3,9 +3,11 @@
 from itertools import repeat
 
 import numpy as np
+from scipy.linalg import norm
 
 from neurodsp.sim.info import get_sim_func
 from neurodsp.utils.decorators import normalize
+from neurodsp.utils.data import create_times
 
 ###################################################################################################
 ###################################################################################################
@@ -82,3 +84,81 @@ def sim_combined(n_seconds, fs, components, component_variances=1):
     sig = np.sum(sig_components, axis=0)
 
     return sig
+
+@normalize
+def sim_central_freq(n_seconds, fs, chi, central_freq, bw, ht):
+    """
+    Returns a time series whose full power spectrum consists of a power law with exponent chi
+    and a gaussian peak at central_freq with standard deviation bw and relative height ht.
+
+
+    Parameters
+    -----------
+    n_seconds: float
+        Number of seconds elapsed in the time series.
+    fs: float
+        Sampling rate.
+    chi: float
+        Power law exponent.
+    central_freq: float
+        Central frequency for the gaussian peak in Hz.
+    bw: float
+        Bandwidth, or standard deviation, of gaussian peak in Hz.
+    ht: float
+        Relative height in log_10(Hz) over the aperiodic component of the gaussian peak at central_freq.
+
+    Returns
+    -------
+    sig: 1d array
+        Time series with desired power spectrum.
+
+    Examples
+    --------
+    Simulate aperiodic noise with exponent -2 superimposed over an oscillatory component with
+    central frequency 20.
+
+    >>> sig = sim_gauss_peak(n_seconds=50, fs=500, chi=-2, central_freq=20, bw=5, ht=7)
+    """
+
+    times = create_times(n_seconds, fs)
+
+    # Construct the aperiodic component and compute its Fourier transform. Only use the
+    # first half of the frequencies from the Fourier transform since the signal is real.
+    sig_components = {'sim_powerlaw': {'exponent': chi}}
+    sig_ap = sim_combined(n_seconds, fs, sig_components)
+    sig_len = sig_ap.shape[0]
+
+    sig_ap_hat = np.fft.fft(sig_ap)[0:(sig_len//2+1)]
+
+    # Create the range of frequencies that appear in the power spectrum since these
+    # will be the frequencies in the cosines we sum below
+    freqs = np.linspace(0, fs/2, num=sig_len//2 + 1, endpoint=True)
+
+    # Construct the array of relative heights above the aperiodic power spectrum
+    rel_heights = np.array([ ht * np.exp(-(f - central_freq)**2/(2*bw**2)) for f in freqs])
+
+    # Build an array of the sum of squares of the cosines as they appear in the calculation of the
+    # amplitudes
+    cosine_norms = np.array([
+                          norm(
+                            np.cos(2*np.pi*f*times), 2
+                          )**2 for f in freqs
+                        ])
+
+    # Build an array of the amplitude coefficients
+    cosine_coeffs = np.array([
+                    (-np.real(sig_ap_hat[ell]) + np.sqrt(np.real(sig_ap_hat[ell])**2 + (10**rel_heights[ell] - 1)*np.abs(sig_ap_hat[ell])**2))/cosine_norms[ell]
+                    for ell in range(cosine_norms.shape[0])]
+                )
+
+    # Add cosines with the respective coefficients and with a random phase shift for each one
+    sig_periodic = np.sum(
+                    np.array(
+                        [cosine_coeffs[ell]*np.cos(2*np.pi*freqs[ell]*times + 2*np.pi*np.random.rand()) for ell in range(cosine_norms.shape[0])]
+                    ),
+                    axis=0
+                    )
+
+    sig = sig_ap + sig_periodic
+
+    return sig

neurodsp/tests/sim/test_combined.py

@@ -2,7 +2,7 @@
 
 from pytest import raises
 
-from neurodsp.tests.settings import FS, N_SECONDS, FREQ1, FREQ2
+from neurodsp.tests.settings import FS, N_SECONDS, FREQ1, FREQ2, EXP1
 from neurodsp.tests.tutils import check_sim_output
 
 from neurodsp.sim.combined import *
@@ -29,3 +29,8 @@ def test_sim_combined():
     variances = [0.5, 1]
     with raises(ValueError):
         out = sim_combined(N_SECONDS, FS, simulations, variances)
+
+def test_sim_central_freq():
+
+    sig = sim_central_freq(N_SECONDS, FS, EXP1, FREQ1, bw=5, ht=5)
+    check_sim_output(sig)