70 lines
1.5 KiB
Python
70 lines
1.5 KiB
Python
import numpy as np
|
|
from numba import jit
|
|
|
|
|
|
@jit(nopython=True)
|
|
def square(time, freq, jitter_SD, jitter, cycle_num):
|
|
"""
|
|
|
|
A scipy like square wave with jitter
|
|
Parameters
|
|
----------
|
|
time : timestamp, in seconds
|
|
|
|
freq : frequency in Hz
|
|
|
|
jitter_SD : standard deviation of jitter, in seconds
|
|
|
|
jitter : gaussian noise with `jitter_SD` as its standard deviation
|
|
|
|
cycle_num : time // period = cycle_num
|
|
----------
|
|
|
|
"""
|
|
|
|
period = 1/(freq)
|
|
quarter = (period / 4)
|
|
out = 0
|
|
|
|
# A T/4 shift is applied to the following to match scipy square fn
|
|
# ┌──────┐
|
|
# ───┘ └────
|
|
# T/4 3T/4
|
|
nth_cycle, t = np.divmod(time + period / 4, period)
|
|
if t >= (quarter + jitter) and t <= (3*quarter + jitter):
|
|
out = 1
|
|
else:
|
|
out = 0
|
|
|
|
# update jitter every cycle
|
|
if nth_cycle != cycle_num:
|
|
jitter = np.random.normal(0, jitter_SD)
|
|
cycle_num = nth_cycle
|
|
|
|
return out, cycle_num, jitter
|
|
|
|
|
|
@jit(nopython=True)
|
|
def square_arr(time, freq, jitter_SD):
|
|
"""
|
|
|
|
A scipy like square wave with jitter
|
|
Parameters
|
|
----------
|
|
time : numpy array
|
|
|
|
freq : frequency in Hz
|
|
|
|
jitter_SD : standard deviation of jitter, in seconds
|
|
----------
|
|
|
|
"""
|
|
wave = np.zeros(len(time))
|
|
jitter = np.random.normal(0, jitter_SD)
|
|
cycle_num = 0
|
|
|
|
for i, t in enumerate(time):
|
|
wave[i], cycle_num, jitter = square(t, freq, jitter_SD, jitter, cycle_num)
|
|
|
|
return wave
|