Merge pull request #221 from elybrand/central_frequency

[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_peak_oscillation

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_peak_oscillation(sig_ap, fs, freq, bw, height):
+    """Returns a time series whose full power spectrum consists of an aperiodic component
+    and a gaussian peak at ``freq`` with standard deviation ``bw`` and relative ``height``.
+
+    Parameters
+    -----------
+    sig_ap : 1d array
+        The timeseries of the aperiodic component.
+    fs: float
+        Sampling rate of ``sig_ap``.
+    freq: float
+        Central frequency for the gaussian peak in Hz.
+    bw: float
+        Bandwidth, or standard deviation, of gaussian peak in Hz.
+    height: float
+        Relative height in log_10(Hz) over the aperiodic component of the gaussian peak at the
+        central ``freq``.
+
+    Returns
+    -------
+    sig: 1d array
+        Time series with desired power spectrum.
+
+    Notes
+    -----
+    The periodic component of the signal will be sinusoidal.
+
+    Examples
+    --------
+    Simulate an aperiodic component with exponent -2 superimposed over
+    an oscillatory component with central frequency 20.
+
+    >>> from neurodsp.sim import sim_powerlaw
+    >>> fs = 500
+    >>> sig_ap = sim_powerlaw(n_seconds=10, fs=fs)
+    >>> sig = sim_peak_oscillation(sig_ap, fs=fs, freq=20, bw=5, height=7)
+    """
+
+    sig_len = len(sig_ap)
+    times = create_times(sig_len/fs, 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_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([height * np.exp(-(lin_freq - freq)**2/(2*bw**2)) for lin_freq 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*lin_freq*times), 2)**2 for lin_freq 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,10 +2,12 @@
 
 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 *
+from neurodsp.sim import sim_powerlaw
+from neurodsp.spectral import compute_spectrum
 
 ###################################################################################################
 ###################################################################################################
@@ -29,3 +31,16 @@ def test_sim_combined():
     variances = [0.5, 1]
     with raises(ValueError):
         out = sim_combined(N_SECONDS, FS, simulations, variances)
+
+def test_sim_peak_oscillation():
+
+    sig_ap = sim_powerlaw(N_SECONDS, FS)
+    sig = sim_peak_oscillation(sig_ap, FS, FREQ1, bw=5, height=10)
+
+    check_sim_output(sig)
+
+    # Ensure the peak is at or close (+/- 5hz) to FREQ1
+    _, powers_ap = compute_spectrum(sig_ap, FS)
+    _, powers = compute_spectrum(sig, FS)
+
+    assert abs(np.argmax(powers-powers_ap) - FREQ1) < 5