m-thesis-introduction/fourier/mylib/util.py

"""
Various utilities
"""
import numpy as np

rng = np.random.default_rng()


def phasemod(phase, low=np.pi):
    """
    Modulo phase such that it falls within the
    interval $[-low, 2\pi - low)$.
    """
    return (phase + low) % (2*np.pi) - low

# Alias phase_mod to phasemod
phase_mod = phasemod

def sine_fitfunc(t, amp=1, freq=1, phase=0, baseline=0, t_delay=0):
    """Simple sine wave for fitting purposes"""
    return amp*np.cos( 2*np.pi*freq*(t-t_delay) + phase) + baseline

def sin_delay(f, t, phase=0):
    return sine_fitfunc(t, amp=1, freq=f, phase=phase, baseline=1, t_delay=0)

def sampled_time(sample_rate=1, start=0, end=1, offset=0):
    return offset + np.arange(start, end, 1/sample_rate)

def normalise_sine_params(params):
    params[2] = phase_mod(params[2])
    return params

def noisy_sine_sampling(time, init_params, noise_sigma=1, rng=rng):
    if init_params[2] is None:
        init_params[2] = phasemod(2*np.pi*rng.random())

    samples = sine_fitfunc(time, *init_params)
    noise = rng.normal(0, noise_sigma, size=len(samples))

    return samples, noise

# Alias noisy_sine
noisy_sine = noisy_sine_sampling

def find_nearest(value, array, return_idx=True):
    array = np.asarray(array)
    idx = (np.abs(array - value)).argmin()

    if return_idx:
        return idx
    else:
        return array[idx]