SNR figure: split script into library and script

fourier/mylib/__init__.py

@@ -0,0 +1,8 @@
+"""
+Module to automatically load another (local) module.
+"""
+
+from .passband import *
+from .fft import *
+from .plotting import *
+from .util import *

fourier/mylib/fft.py

@@ -0,0 +1,44 @@
+"""
+Simple FFT stuff
+"""
+
+import numpy as np
+import scipy.fftpack as ft
+
+def get_freq_spec(val,dt):
+    """From earsim/tools.py"""
+    fval = np.fft.fft(val)[:len(val)//2]
+    freq =  np.fft.fftfreq(len(val),dt)[:len(val)//2]
+    return fval, freq
+
+
+def ft_spectrum( signal, sample_rate=1, ftfunc=None, freqfunc=None, mask_bias=False, normalise_amplitude=False):
+    """Return a FT of $signal$, with corresponding frequencies"""
+
+    if True:
+        return get_freq_spec(signal, 1/sample_rate)
+
+    n_samples = len(signal)
+
+    if ftfunc is None:
+        real_signal = np.isrealobj(signal)
+        if False and real_signal:
+            ftfunc = ft.rfft
+            freqfunc = ft.rfftfreq
+        else:
+            ftfunc = ft.fft
+            freqfunc = ft.fftfreq
+
+    if freqfunc is None:
+        freqfunc = ft.fftfreq
+
+    normalisation = 2/len(signal) if normalise_amplitude else 1
+
+    spectrum = normalisation * ftfunc(signal)
+    freqs = freqfunc(n_samples, 1/sample_rate)
+
+    if not mask_bias:
+        return spectrum, freqs
+    else:
+        return spectrum[1:], freqs[1:]
+

fourier/mylib/passband.py

@@ -0,0 +1,81 @@
+
+import numpy as np
+from collections import namedtuple
+
+from .fft import ft_spectrum
+
+class passband(namedtuple("passband", ['low', 'high'], defaults=[0, np.inf])):
+    """
+    Band for a bandpass filter.
+    It encapsulates a tuple.
+    """
+
+    def size():
+        return bandsize(self)
+
+    def freq_mask(frequencies):
+        return bandpass_mask(frequencies, self)
+
+    def signal_level(samples, samplerate, normalise_bandsize=True, **ft_kwargs):
+
+        return bandlevel(samples, samplerate, self, normalise_bandsize, **ft_kwargs)
+
+    def filter_samples(samples, samplerate, **ft_kwargs):
+        """
+        Bandpass the samples with this passband.
+        This is a hard filter.
+        """
+        fft, freqs = ft_spectrum(samples, samplerate, **ft_kwargs)
+
+        fft[ ~ self.freq_mask(freqs) ] = 0
+
+        return irfft(fft)
+
+
+def bandpass_samples(samples, samplerate, band=passband(), **ft_kwargs):
+        """
+        Bandpass the samples with this passband.
+        This is a hard filter.
+        """
+        fft, freqs = ft_spectrum(samples, samplerate, **ft_kwargs)
+
+        fft[ ~ self.freq_mask(freqs) ] = 0
+
+        return np.fft.irfft(fft)
+
+def bandpass_mask(freqs, band=passband()):
+    low_pass = abs(freqs) <= band[1]
+    high_pass = abs(freqs) >= band[0]
+
+    return low_pass & high_pass
+
+def bandsize(band = passband()):
+    return band[1] - band[0]
+
+def bandlevel(samples, samplerate=1, band=passband(), normalise_bandsize=True, **ft_kwargs):
+    fft, freqs = ft_spectrum(samples, samplerate, **ft_kwargs)
+
+    bandmask = bandpass_mask(freqs, band=band)
+
+    if normalise_bandsize:
+        bins = np.count_nonzero(bandmask, axis=-1)
+    else:
+        bins = 1
+
+    level = np.sum(np.abs(fft[bandmask])**2)
+
+    return level/bins
+
+def signal_to_noise( samplerate, samples, noise, signal_band, noise_band=None):
+    if noise_band is None:
+        noise_band = sample_band
+
+    if noise is None:
+        noise = samples
+
+    noise_level = bandlevel(noise, samplerate, noise_band)
+
+    signal_level = bandlevel(samples, samplerate, signal_band)
+
+    return (signal_level/noise_level)**0.5
+

fourier/mylib/plotting.py

@@ -0,0 +1,95 @@
+"""
+Functions to simplify plotting
+"""
+import matplotlib.pyplot as plt
+import matplotlib.gridspec as gridspec
+import numpy as np
+
+def plot_spectrum( spectrum, freqs, plot_complex=False, plot_power=False, plot_amplitude=None, ax=None, freq_unit="Hz", freq_scaler=1):
+    """ Plot a signal's spectrum on an Axis object"""
+    plot_amplitude = plot_amplitude or (not plot_power and not plot_complex)
+    alpha = 1
+
+    if ax is None:
+        ax = plt.gca()
+
+    ax.set_title("Spectrum")
+    ax.set_xlabel("f" + (" ["+freq_unit+"]" if freq_unit else "" ))
+    ylabel = ""
+    if plot_amplitude or plot_complex:
+        ylabel = "Amplitude"
+    if plot_power:
+        if ylabel:
+            ylabel += "|"
+        ylabel += "Power"
+    ax.set_ylabel(ylabel)
+
+    if plot_complex:
+        alpha = 0.5
+        ax.plot(freqs/freq_scaler, np.real(spectrum), '.-', label='Real', alpha=alpha)
+        ax.plot(freqs/freq_scaler, np.imag(spectrum), '.-', label='Imag', alpha=alpha)
+
+    if plot_power:
+        ax.plot(freqs/freq_scaler, np.abs(spectrum)**2, '.-', label='Power', alpha=alpha)
+
+    if plot_amplitude:
+        ax.plot(freqs/freq_scaler, np.abs(spectrum), '.-', label='Abs', alpha=alpha)
+
+    ax.legend()
+
+    return ax
+
+def plot_phase( spectrum, freqs, ylim_epsilon=0.5, ax=None, freq_unit="Hz", freq_scaler=1):
+    if ax is None:
+        ax = plt.gca()
+
+    ax.set_ylabel("Phase")
+    ax.set_xlabel("f" + (" ["+freq_unit+"]" if freq_unit else "" ))
+
+    ax.plot(freqs/freq_scaler, np.angle(spectrum), '.-')
+    ax.set_ylim(-1*np.pi - ylim_epsilon, np.pi + ylim_epsilon)
+
+    return ax
+
+def plot_signal( signal, sample_rate = 1, ax=None, time=None, time_unit="s", **kwargs):
+    if ax is None:
+        ax = plt.gca()
+
+    if time is None:
+        time = np.arange(len(signal))/sample_rate
+
+    ax.set_title("Signal")
+    ax.set_xlabel("t" + (" ["+time_unit+"]" if time_unit else "" ))
+    ax.set_ylabel("A(t)")
+
+    ax.plot(time, signal, **kwargs)
+
+    return ax
+
+def plot_combined_spectrum(spectrum, freqs,
+                           spectrum_kwargs={}, fig=None, gs=None, freq_scaler=1, freq_unit="Hz"):
+    """Plot both the frequencies and phase in one figure."""
+
+    # configure plotting layout
+    if fig is None:
+        fig = plt.figure(figsize=(8, 16))
+
+    if gs is None:
+        gs = gridspec.GridSpec(2, 1, figure=fig, height_ratios=[3,1], hspace=0)
+
+    ax1 = fig.add_subplot(gs[:-1, -1])
+    ax2 = fig.add_subplot(gs[-1, -1], sharex=ax1)
+
+    axes = np.array([ax1, ax2])
+
+    # plot the spectrum
+    plot_spectrum(spectrum, freqs, ax=ax1, freq_scaler=freq_scaler, freq_unit=freq_unit, **spectrum_kwargs)
+
+    # plot the phase
+    plot_phase(spectrum, freqs, ax=ax2, freq_scaler=freq_scaler, freq_unit=freq_unit)
+
+    ax1.xaxis.tick_top()
+    [label.set_visible(False) for label in ax1.get_xticklabels()]
+
+    return fig, axes
+

fourier/mylib/util.py

@@ -0,0 +1,31 @@
+"""
+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
+
+def sine_fitfunc(t, amp=1, freq=1, phase=0, off=0):
+    """Simple sine wave for fitting purposes"""
+    return amp*np.sin( 2*np.pi*freq*t + phase) + off
+
+def sampled_time(sample_rate=1, start=0, end=1, offset=0):
+    return offset + np.arange(start, end, 1/sample_rate)
+
+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
+