Lib: add simple util functions

lib/__init__.py

@@ -1,6 +1,7 @@
 from . import signals
 from . import location
 from . import sampling
+from .plotting import *
 from .util import *
 
 

lib/plotting.py

@@ -0,0 +1,25 @@
+"""
+Routines to assist in plotting
+"""
+
+def annotate_width(ax, name, x1, x2, y, text_kw={}, arrow_kw={}):
+    default_arrow_kw = dict(
+            xy = (x1, y),
+            xytext = (x2,y),
+            arrowprops = dict(
+                arrowstyle="<->",
+                shrinkA=False,
+                shrinkB=False
+            ),
+    )
+
+    default_text_kw = dict(
+            va='bottom',
+            ha='center',
+            xy=((x1+x2)/2, y)
+    )
+
+    an1 = ax.annotate("", **{**default_arrow_kw, **arrow_kw})
+    an2 = ax.annotate(name, **{**default_text_kw, **text_kw})
+
+    return [an1, an2]

lib/util.py

@@ -3,6 +3,7 @@ Various useful utilities (duh)
 """
 
 import numpy as np
+import scipy.fft as ft
 
 def sampled_time(sample_rate=1, start=0, end=1, offset=0):
     return offset + np.arange(start, end, 1/sample_rate)
@@ -36,3 +37,75 @@ def detect_edges(threshold, data, rising=True, falling=False):
         mask |= (data[:-1] > threshold) & (data[1:] < threshold)
 
     return np.flatnonzero(mask)+1
+
+def sin_delay(f, t, t_delay=0, phase=0):
+    return np.sin( 2*np.pi*f*(t - t_delay) + phase )
+
+def time2phase(time, frequency=1):
+    return 2*np.pi*frequency*time
+
+def phase2time(phase, frequency=1):
+    return phase/(2*np.pi*frequency)
+
+def time_roll(a, samplerate, time_shift, *roll_args, int_func=lambda x: np.rint(x).astype(int), **roll_kwargs):
+    """
+    Like np.roll, but use samplerate and time_shift to approximate
+    the offset to roll.
+    """
+    shift = int_func(time_shift*samplerate)
+    return np.roll(a, shift, *roll_args, **roll_kwargs)
+
+### signal generation
+def fft_bandpass(signal, band, samplerate):
+    """
+    Simple bandpassing function employing a FFT.
+
+    Parameters
+    ----------
+    signal : arraylike
+    band : tuple(low, high)
+        Frequencies for bandpassing
+    samplerate : float
+    """
+    signal = np.asarray(signal)
+
+    fft = ft.rfft(signal)
+    freqs = ft.rfftfreq(signal.size, 1/samplerate)
+    fft[(freqs < band[0]) | (freqs > band[1])] = 0
+
+    return ft.irfft(fft, signal.size), (fft, freqs)
+
+def deltapeak(timelength=1e3, samplerate=1, offset=None, peaklength=1):
+    """
+    Generate a series of zeroes with a deltapeak.
+
+    If offset is not specified, it puts it at a random location.
+
+    Note: the series is regarded as periodic.
+
+    Parameters
+    ----------
+    timelength : float
+    samplerate : float
+    offset : float or tuple(float, float)
+        Start of the peak
+    peaklength : int
+        Length of the peak
+    """
+
+    N_samples = int(timelength * samplerate)
+    if offset is None:
+        offset = (None,None)
+
+    if isinstance(offset, (tuple, list)):
+        offset_min = 0 if offset[0] is None else offset[0]
+        offset_max = N_samples if offset[-1] is None else offset[-1]
+
+        offset = (np.random.random(1)*(offset_max - offset_min)+offset_min).astype(int) % N_samples
+
+    position = (offset + np.arange(0, peaklength)).astype(int) % N_samples
+
+    signal = np.zeros(N_samples)
+    signal[position] = 1
+
+    return signal, position