Sine fitting figure showing SNR vs residuals

fourier/mylib/passband.py

@@ -66,9 +66,9 @@ def bandpower(samples, samplerate=1, band=passband(), normalise_bandsize=True, *
 
     return power/bins
 
-def signal_to_noise( samplerate, samples, noise, signal_band, noise_band=None):
+def signal_to_noise(samples, noise, samplerate=1, signal_band=passband(), noise_band=None):
     if noise_band is None:
-        noise_band = sample_band
+        noise_band = signal_band
 
     if noise is None:
         noise = samples

fourier/mylib/plotting.py

@@ -1,10 +1,14 @@
 """
 Functions to simplify plotting of fourier spectra
 """
+# vim: fdm=indent ts=4
+
 import matplotlib.pyplot as plt
 import matplotlib.gridspec as gridspec
 import numpy as np
 
+from .fft import ft_spectrum
+
 def plot_spectrum(
         spectrum, freqs,
         plot_complex=False, plot_power=False, plot_amplitude=None,
@@ -49,6 +53,7 @@ def plot_phase(
         spectrum, freqs,
         ylim_epsilon=0.5, ax=None, grid=True,
         freq_unit="Hz", freq_scaler=1, xlabel='Frequency',
+        major_divider=2, minor_divider=12,
         **plot_kwargs
     ):
     """
@@ -65,14 +70,21 @@ def plot_phase(
     ax.plot(freqs/freq_scaler, np.angle(spectrum), '.-', **plot_kwargs)
     ax.set_ylim(-1*np.pi - ylim_epsilon, np.pi + ylim_epsilon)
 
-    major_tick = Multiple(2, np.pi, '\pi')
-    minor_tick = Multiple(12, np.pi, '\pi')
-    ax.yaxis.set_major_locator(major_tick.locator())
-    ax.yaxis.set_major_formatter(major_tick.formatter())
-    ax.yaxis.set_minor_locator(minor_tick.locator())
+    axis_pi_ticker(ax.yaxis, major_divider=major_divider, minor_divider=minor_divider)
 
     return ax
 
+def axis_pi_ticker(axis, major_divider=2, minor_divider=12):
+
+    major_tick = Multiple(major_divider, np.pi, '\pi')
+    minor_tick = Multiple(minor_divider, np.pi, '\pi')
+
+    axis.set_major_locator(major_tick.locator())
+    axis.set_major_formatter(major_tick.formatter())
+    axis.set_minor_locator(minor_tick.locator())
+
+    return axis
+
 def plot_signal(
         signal, sample_rate = 1,
         time=None, ax=None,
@@ -139,6 +151,43 @@ def plot_combined_spectrum(
 
     return fig, axes
 
+def plot_signal_and_spectrum(
+        signal, sample_rate=1, title=None,
+        signal_kwargs={}, ft_kwargs={},
+        spectrum_kwargs={}, phase_kwargs={'major_divider':1, 'minor_divider': 6},
+        **phase_spectrum_kwargs
+        ):
+    """
+    Create a figure showing both the signal and the combined spectrum.
+    """
+    fig = plt.figure(figsize=(16, 4))
+
+    if title:
+        fig.suptitle(title)
+
+    # setup plot layout
+    gs0 = gridspec.GridSpec(1, 2, figure=fig)
+    gs00 = gs0[0].subgridspec(1, 1)
+    gs01 = gs0[1].subgridspec(2, 1, height_ratios=[3,1], hspace=0)
+
+    # plot the signal
+    ax1 = fig.add_subplot(gs00[0, 0])
+    plot_signal(signal, sample_rate, ax=ax1, **signal_kwargs)
+
+    # plot spectrum
+    signal_fft, freqs = ft_spectrum(signal, sample_rate, **ft_kwargs)
+    _, (ax2, ax3) = plot_combined_spectrum(
+            signal_fft, freqs,
+            fig=fig, gs=gs01,
+            spectrum_kwargs=spectrum_kwargs, phase_kwargs=phase_kwargs,
+            **phase_spectrum_kwargs
+            )
+
+    # return the axes
+    axes = np.array([ax1, ax2, ax3])
+
+    return fig, axes
+
 def multiple_formatter(denominator=2, number=np.pi, latex='\pi'):
     """
     From https://stackoverflow.com/a/53586826

fourier/mylib/util.py

@@ -13,13 +13,23 @@ def phasemod(phase, low=np.pi):
     """
     return (phase + low) % (2*np.pi) - low
 
-def sine_fitfunc(t, amp=1, freq=1, phase=0, off=0):
+# Alias phase_mod to phasemod
+phase_mod = phasemod
+
+def sine_fitfunc(t, amp=1, freq=1, phase=0, off=0, t_delay=0):
     """Simple sine wave for fitting purposes"""
-    return amp*np.cos( 2*np.pi*freq*t + phase) + off
+    return amp*np.cos( 2*np.pi*freq*(t-t_delay) + phase) + off
+
+def sin_delay(f, t, phase=0):
+    return sine_fitfunc(t, amp=1, freq=f, phase=phase, off=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())
@@ -29,3 +39,15 @@ def noisy_sine_sampling(time, init_params, noise_sigma=1, rng=rng):
 
     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]
+