[MNT] - Simulations refactors

neurodsp/sim/aperiodic.py

@@ -168,21 +168,26 @@ def sim_knee(n_seconds, fs, chi1, chi2, knee):
     times = create_times(n_seconds, fs)
     n_samples = compute_nsamples(n_seconds, fs)
 
-    # Create frequencies for the power spectrum, which will be freqs of the summed cosines
-    freqs = np.linspace(0, fs/2, num=int(n_samples//2 + 1), endpoint=True)
-
-    # Drop the DC component
+    # Create frequencies for the power spectrum and drop the DC component
+    #   These frequencies are used to create the cosines to sum
+    freqs = np.linspace(0, fs / 2, num=int(n_samples // 2 + 1), endpoint=True)
     freqs = freqs[1:]
 
-    # Map the frequencies under the (square root) Lorentzian
-    #   This will give us the amplitude coefficients for the sinusoids
-    cosine_coeffs = np.array([np.sqrt(1 / (freq ** -chi1 * (freq ** (-chi2 - chi1) + knee))) \
-        for freq in freqs])
+    # Define sub-function for computing the amplitude coefficients for the cosines
+    #   This computes the square root of the Lorentzian
+    def _coscoef(freq):
+        return np.sqrt(1 / (freq ** -chi1 * (freq ** (-chi2 - chi1) + knee)))
+
+    # Define & vectorize sub-function for computing the set of cosines, adding a random phase shift
+    def _create_cosines(freq):
+        return _coscoef(freq) * np.cos(2 * np.pi * freq * times + 2 * np.pi * np.random.rand())
+    vect_create_cosines = np.vectorize(_create_cosines, signature='()->(n)')
+
+    # Create the set of cosines
+    cosines = vect_create_cosines(freqs)
 
-    # Add sinusoids with a random phase shift
-    sig = np.sum(np.array([cosine_coeffs[ell] * \
-                          np.cos(2 * np.pi * freq * times + 2 * np.pi * np.random.rand()) \
-                 for ell, freq in enumerate(freqs)]), axis=0)
+    # Create the final signal by summing the cosines
+    sig = np.sum(cosines, axis=0)
 
     return sig
 

neurodsp/sim/combined.py

@@ -122,41 +122,46 @@ def sim_peak_oscillation(sig_ap, fs, freq, bw, height):
 
     >>> from neurodsp.sim import sim_powerlaw
     >>> fs = 500
-    >>> sig_ap = sim_powerlaw(n_seconds=10, fs=fs)
+    >>> sig_ap = sim_powerlaw(n_seconds=10, fs=fs, exponent=-2.0)
     >>> 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 FFT since the signal is real
+    # Compute the Fourier transform of the aperiodic signal
+    #   We extract the first half of the frequencies from the FFT, 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
+    # Create the corresponding frequency vector, which is used to create the cosines to sum
     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])
+    # Define sub-function for computing the relative height above the aperiodic power spectrum
+    def _hght(f_val):
+        return height * np.exp(-(f_val - freq) ** 2 / (2 * bw ** 2))
 
-    # Build an array of the sum of squares of the cosines to use in the amplitude calculation
-    cosine_norms = np.array([norm(np.cos(2 * np.pi * lin_freq * times), 2) ** 2 \
-        for lin_freq in freqs])
+    # Define sub-function for computing the sum of squares of the cosines
+    def _cosn(f_val):
+        return norm(np.cos(2 * np.pi * f_val * times), 2) ** 2
 
-    # 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])])
+    # Define sub-function for computing the amplitude coefficients
+    def _coef(f_val, fft):
+        return (-np.real(fft) + np.sqrt(np.real(fft) ** 2 + (10 ** _hght(f_val) - 1) * \
+            np.abs(fft) ** 2)) / _cosn(f_val)
 
-    # 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)
+    # Define & vectorize sub-function to create the set of cosines for the periodic component
+    #   Cosines are created with their respective coefficients & a random phase shift
+    def _create_cosines(f_val, fft):
+        return _coef(f_val, fft) * np.cos(2 * np.pi * f_val * times + 2 * np.pi * np.random.rand())
+    vect_create_cosines = np.vectorize(_create_cosines, signature='(),()->(n)')
 
+    # Create the set of cosines, defined from the frequency and FFT values
+    cosines = vect_create_cosines(freqs, sig_ap_hat)
+
+    # Create the periodic component by summing the cosines
+    sig_periodic = np.sum(cosines, axis=0)
+
+    # Create the combined signal by summing periodic & aperiodic
     sig = sig_ap + sig_periodic
 
     return sig