m-thesis-introduction/lib/util.py

150 lines
3.8 KiB
Python

"""
Various useful utilities (duh)
"""
import numpy as np
try:
import scipy.fft as ft
except ImportError:
import numpy.fft as ft
class MethodMappingProxy():
def __init__(self, *elements):
self.elements = elements
def __getattr__(self, name):
# Test name exists on all elements
for el in self.elements:
attr = getattr(el, name)
if not callable(attr):
raise AttributeError("Attribute `{name}` is not callable on `{el}`.")
def mapping_func(*args, **kwargs):
return [ getattr(el, name)(*args, **kwargs) for el in self.elements ]
return mapping_func
def sampled_time(sample_rate=1, start=0, end=1, offset=0):
return offset + np.arange(start, end, 1/sample_rate)
def rot_vector(phi1=0.12345):
"""
Return a unit vector rotated by phi radians.
"""
unit = np.array([
phi1,
phi1 - np.pi/2
])
return np.cos(unit)
def detect_edges(threshold, data, rising=True, falling=False):
"""
Detect rising/falling edges in data, returning the indices
of the detected edges.
https://stackoverflow.com/a/50365462
"""
mask = np.full(len(data)-1, False)
if rising:
mask |= (data[:-1] < threshold) & (data[1:] > threshold)
if falling:
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 phase_modulo(phase, low=np.pi):
"""
Modulo phase such that it falls within the
interval $[-low, 2\pi - low)$.
"""
return (phase + low) % (2*np.pi) - low
def time_roll(a, samplerate, time_shift, sample_shift=0, 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 + sample_shift)
return np.roll(a, shift, **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, rng=None):
"""
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]
if 0 < offset_min < 1:
offset_min *= N_samples
if 0 < offset_max < 1:
offset_max *= N_samples
if rng is not None:
rand = rng.random(1)
else:
rand = np.random.random(1)
rand = np.asarray(rand)
offset = (rand * (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