diff --git a/neurodsp/sim/periodic.py b/neurodsp/sim/periodic.py index 470d73f2..9472cb2c 100644 --- a/neurodsp/sim/periodic.py +++ b/neurodsp/sim/periodic.py @@ -347,14 +347,21 @@ def make_bursts(n_seconds, fs, is_oscillating, cycle): n_samples_cycle = len(cycle) burst_sig = np.zeros([n_samples]) - for sig_ind, is_osc in zip(range(0, n_samples, n_samples_cycle), is_oscillating): + + for ind, sig_ind in enumerate(range(0, n_samples, n_samples_cycle)): + + is_osc = is_oscillating[ind] # If set as an oscillating cycle, add cycle to signal # The sample check is to check there are enough samples left to add a full cycle # If there are not, this skipps the add, leaving zeros instead of adding part of a cycle - if is_osc and sig_ind + n_samples_cycle < n_samples: + if is_osc and sig_ind + n_samples_cycle + 1 < n_samples: burst_sig[sig_ind:sig_ind+n_samples_cycle] = cycle + # Complete the last cycle of a burst + if is_osc and ind < len(is_oscillating) - 1 and not is_oscillating[ind+1]: + burst_sig[sig_ind+n_samples_cycle] = cycle[0] + return burst_sig diff --git a/neurodsp/tests/sim/test_periodic.py b/neurodsp/tests/sim/test_periodic.py index cc91fb09..de6e5848 100644 --- a/neurodsp/tests/sim/test_periodic.py +++ b/neurodsp/tests/sim/test_periodic.py @@ -101,11 +101,14 @@ def test_sim_damped_oscillation(): def test_make_bursts(): is_osc = np.array([False, False, True, True, False, True, False, True, True, False]) - cycle = np.ones([10]) + cycle = np.ones([100]) sig = make_bursts(N_SECONDS, FS, is_osc, cycle) check_sim_output(sig) + # Three is added since each of the three continuous bursts are fully completed + assert sig.sum() == (is_osc.sum() * 100) + 3 + # Test make bursts with uneven division of signal and cycle divisions # In this test, there aren't enough samples in the signal to add last cycle is_osc = np.array([False, True, True])