ericteunis/m-thesis-introduction
1
0
Fork 0
mirror of https://gitlab.science.ru.nl/mthesis-edeboone/m-thesis-introduction.git synced 2024-12-21 11:03:37 +01:00
Code Issues Projects Releases Packages Wiki Activity

SNR figure: split script into library and script

Browse source
This commit is contained in:
Eric Teunis de Boone 2022-10-31 18:16:03 +01:00
parent d23f8adff2
commit 007bd7f963
8 changed files with 564 additions and 428 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

227
fourier/04_signal_to_noise.py Executable file
View file

@ -0,0 +1,227 @@
#!/usr/bin/env python3
__doc__ = \
"""
Show the curve for signal-to-noise ratio vs N_samples
"""
import matplotlib.pyplot as plt
import numpy as np
from mylib import *
rng = np.random.default_rng()
def noisy_sine_realisation_snr(
N = 1,
f_sample = 250e6, # Hz
t_length = 1e4 * 1e-9, # s
noise_band = passband(30e6, 80e6),
noise_sigma = 1,
# signal properties
f_sine = 50e6,
signal_band = passband(50e6 - 1e6, 50e6 + 1e6),
sine_amp = 0.2,
sine_offset = 0,
return_ranges_plot = False,
cut_signal_band_from_noise_band = False,
rng=rng
):
"""
Return N signal to noise ratios determined on
N different noise + sine realisations.
"""
N = int(N)
init_params = np.array([sine_amp, f_sine, None, sine_offset])
axs = None
snrs = np.zeros( N )
time = sampled_time(f_sample, end=t_length)
for j in range(N):
samples, noise = noisy_sine_sampling(time, init_params, noise_sigma, rng=rng)
# determine signal to noise
noise_level = bandlevel(noise, f_sample, noise_band)
if cut_signal_band_from_noise_band:
lower_noise_band = passband(noise_band[0], signal_band[0])
upper_noise_band = passband(signal_band[1], noise_band[1])
noise_level = bandlevel(noise, f_sample, lower_noise_band)
noise_level += bandlevel(noise, f_sample, upper_noise_band)
signal_level = bandlevel(samples, f_sample, signal_band)
snrs[j] = np.sqrt(signal_level/noise_level)
# make a nice plot showing what ranges were taken
# and the bandlevels associated with them
if return_ranges_plot and j == 0:
combined_fft, freqs = ft_spectrum(samples+noise, f_sample)
# plot the original signal
if False:
_, ax = plt.subplots()
ax = plot_signal(samples+noise, sample_rate=f_sample/1e6, time_unit='us', ax=ax)
# plot the spectrum
if True:
freq_scaler=1e6
_, axs = plot_combined_spectrum(combined_fft, freqs, freq_scaler=freq_scaler, freq_unit='MHz')
# indicate band ranges and frequency
for ax in axs:
ax.axvline(f_sine/freq_scaler, color='r', alpha=0.4)
ax.axvspan(noise_band[0]/freq_scaler, noise_band[1]/freq_scaler, color='purple', alpha=0.3, label='noiseband')
ax.axvspan(signal_band[0]/freq_scaler, signal_band[1]/freq_scaler, color='orange', alpha=0.3, label='signalband')
# indicate initial phase
axs[1].axhline(init_params[2], color='r', alpha=0.4)
# plot the band levels
levelax = axs[0].twinx()
levelax.set_ylabel("Bandlevel")
levelax.hlines(signal_level, noise_band[0]/freq_scaler, signal_band[1]/freq_scaler, colors=['orange'])
levelax.hlines(noise_level, noise_band[0]/freq_scaler, noise_band[1]/freq_scaler, colors=['purple'])
levelax.set_ylim(bottom=0)
axs[0].legend()
# plot signal_band pass signal
if False:
freqs = np.fft.fftfreq(len(samples), 1/f_sample)
bandmask = bandpass_mask(freqs, band=signal_band)
fft = np.fft.fft(samples)
fft[ ~bandmask ] = 0
bandpassed_samples = np.fft.ifft(fft)
_, ax3 = plt.subplots()
ax3 = plot_signal(bandpassed_samples, sample_rate=f_sample/1e6, time_unit='us', ax=ax3)
ax3.set_title("Bandpassed Signal")
return snrs, axs
if __name__ == "__main__":
from argparse import ArgumentParser
from myscriptlib import save_all_figs_to_path_or_show
rng = np.random.default_rng(1)
parser = ArgumentParser(description=__doc__)
parser.add_argument("fname", metavar="path/to/figure[/]", nargs="?", help="Location for generated figure, will append __file__ if a directory. If not supplied, figure is shown.")
args = parser.parse_args()
default_extensions = ['.pdf', '.png']
if args.fname == 'none':
args.fname = None
###
t_lengths = np.linspace(1e3, 5e4, 5)* 1e-9 # s
N = 10e1
fs_sine = [33.3e6, 50e6, 73.3e6] # Hz
fs_sample = [250e6, 500e6] # Hz
if False:
# show t_length and fs_sample really don't care
fs_iter = [ (fs_sample[0], f_sine, t_lengths) for f_sine in fs_sine ]
fs_iter2 = [ (fs_sample[1], f_sine, t_lengths/2) for f_sine in fs_sine ]
fs_iter += fs_iter2
del fs_iter2
else:
fs_iter = [ (f_sample, f_sine, t_lengths) for f_sample in fs_sample for f_sine in fs_sine ]
if False:
f_sine = fs_sine[0]
f_sample = fs_sample[0]
N = 1 # Note: keep this low, N figures will be displayed!
N_t_length = 10
for t_length in t_lengths[-N_t_length-1:-1]:
snrs = np.zeros( int(N))
for i in range(int(N)):
delta_f = 1/t_length
signal_band = passband(f_sine- 3*delta_f, f_sine + 3*delta_f)
noise_band = passband(30e6, 80e6)
snrs[i], axs = noisy_sine_realisation_snr(
N=1,
t_length=t_length,
f_sample=f_sample,
# signal properties
f_sine = fs_sine[0],
sine_amp = 1,
noise_sigma = 1,
noise_band = noise_band,
signal_band = signal_band,
return_ranges_plot=True,
rng=rng,
)
axs[0].set_title("SNR: {}, N:{}".format(snrs[i], t_length*f_sample))
axs[0].set_xlim(
(f_sine - 20*delta_f)/1e6,
(f_sine + 20*delta_f)/1e6
)
print(snrs, "M:",np.mean(snrs))
plt.show(block=False)
else:
#original code
my_snrs = np.zeros( (len(fs_iter), len(t_lengths), int(N)) )
for i, (f_sample, f_sine, t_lengths) in enumerate( fs_iter ):
for k, t_length in enumerate(t_lengths):
return_ranges_plot = ((k==0) and True) or ( (k==(len(t_lengths)-1)) and True) and i < 1
delta_f = 1/t_length
signal_band = passband(f_sine- 3*delta_f, f_sine + 3*delta_f)
noise_band=passband(30e6, 80e6)
my_snrs[i,k], axs = noisy_sine_realisation_snr(
N=N,
t_length=t_length,
f_sample = f_sample,
# signal properties
f_sine = f_sine,
sine_amp = 1,
noise_sigma = 1,
noise_band = noise_band,
signal_band = signal_band,
return_ranges_plot=return_ranges_plot,
rng=rng
)
if return_ranges_plot:
ranges_axs = axs
# plot the snrs
fig, axs2 = plt.subplots()
axs2.set_xlabel("$N = T*f_s$")
axs2.set_ylabel("SNR")
for i, (f_sample, f_sine, t_lengths) in enumerate(fs_iter):
# plot the means
l = axs2.plot(t_lengths*f_sample, np.mean(my_snrs[i], axis=-1), marker='*', ls='none', label='f:{}MHz, fs:{}MHz'.format(f_sine/1e6, f_sample/1e6), markeredgecolor='black')
color = l[0].get_color()
for k, t_length in enumerate(t_lengths):
t_length = np.repeat(t_length * f_sample, my_snrs.shape[-1])
axs2.plot(t_length, my_snrs[i,k], ls='none', color=color, marker='o', alpha=max(0.01, 1/my_snrs.shape[-1]))
axs2.legend()
### Save or show figures
save_all_figs_to_path_or_show(args.fname, default_basename=__file__, default_extensions=default_extensions)

8
fourier/mylib/__init__.py Normal file
View file

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

44
fourier/mylib/fft.py Normal file
View file

@ -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:]

81
fourier/mylib/passband.py Normal file
View file

@ -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

95
fourier/mylib/plotting.py Normal file
View file

@ -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

31
fourier/mylib/util.py Normal file
View file

@ -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

78
fourier/myscriptlib/__init__.py Normal file
View file

@ -0,0 +1,78 @@
"""
Functions for easy script writing
mostly for simpler figure saving from cli
"""
import matplotlib.pyplot as plt
import numpy as np
import os.path as path
from argparse import ArgumentParser
def ArgumentParserWithFigure(*args, **kwargs):
parser = ArgumentParser(*args, **kwargs)
parser.add_argument("fname", metavar="path/to/figure[/]", nargs="?", help="Location for generated figure, will append __file__ if a directory. If not supplied, figure is shown.")
return parser
def save_all_figs_to_path_or_show(fnames, figs=None, default_basename=None, default_extensions=['.pdf', '.png']):
"""
Save all figures to fnames.
If fnames is empty, simply call plt.show()
"""
if not fnames:
# empty list, False, None
plt.show()
return
if figs is None:
figs = [plt.figure(i) for i in plt.get_fignums()]
default_basename = path.basename(default_basename)
# singular value
if isinstance(fnames, (str, True)):
fnames = [fnames]
if len(fnames) == len(figs):
fnames_list = zip(figs, fnames, False)
elif len(fnames) == 1:
tmp_fname = fnames[0] #needed for generator
fnames_list = ( (fig, tmp_fname, len(figs) > 1) for fig in figs)
else:
# outer product magic
fnames_list = ( (fig,fname, False) for fname in fnames for fig in figs )
del fnames
# format fnames
pad_width = max(2, int(np.floor(np.log10(len(figs))+1)))
fig_fnames = []
for fig, fnames, append_num in fnames_list:
if not hasattr(fnames, '__len__') or isinstance(fnames, str):
# single name
fnames = [fnames]
new_fnames = []
for fname in fnames:
if path.isdir(fname):
if default_basename is not None:
fname = path.join(fname, path.splitext(default_basename)[0]) # leave off extension
elif hasattr(fig, 'basefilename'):
fname = path.join(fname, path.splitext(fig.basefilename)[0]) # leave off extension
if append_num is True:
fname += ("_fig{:0"+str(pad_width)+"d}").format(fig.number)
if not path.splitext(fname)[1]: # no extension
for ext in default_extensions:
new_fnames.append(fname+ext)
else:
new_fnames.append(fname)
fig_fnames.append(new_fnames)
# save files
for fnames, fig in zip(fig_fnames, figs):
for fname in fnames:
fig.savefig(fname, transparent=True)

428
fourier/signal_to_noise.py
View file

@ -1,428 +0,0 @@
#!/usr/bin/env python3
__doc__ = \
"""
Show the curve for signal-to-noise ratio vs N_samples
"""
from collections import namedtuple
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import numpy as np
import scipy.fftpack as ft
rng = np.random.default_rng()
passband = namedtuple("Band", ['low', 'high'], defaults=[0, np.inf])
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:]
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
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 save_all_figs_to_path(fnames, figs=None, default_basename=__file__, default_extensions=['.pdf', '.png']):
if figs is None:
figs = [plt.figure(i) for i in plt.get_fignums()]
default_basename = path.basename(default_basename)
# singular value
if isinstance(fnames, (str, True)):
fnames = [fnames]
if len(fnames) == len(figs):
fnames_list = zip(figs, fnames, False)
elif len(fnames) == 1:
tmp_fname = fnames[0] #needed for generator
fnames_list = ( (fig, tmp_fname, len(figs) > 1) for fig in figs)
else:
# outer product magic
fnames_list = ( (fig,fname, False) for fname in fnames for fig in figs )
del fnames
# format fnames
pad_width = max(2, int(np.floor(np.log10(len(figs))+1)))
fig_fnames = []
for fig, fnames, append_num in fnames_list:
if not hasattr(fnames, '__len__') or isinstance(fnames, str):
# single name
fnames = [fnames]
new_fnames = []
for fname in fnames:
if path.isdir(fname):
fname = path.join(fname, path.splitext(default_basename)[0]) # leave off extension
if append_num is True:
fname += ("_fig{:0"+str(pad_width)+"d}").format(fig.number)
if not path.splitext(fname)[1]: # no extension
for ext in default_extensions:
new_fnames.append(fname+ext)
else:
new_fnames.append(fname)
fig_fnames.append(new_fnames)
# save files
for fnames, fig in zip(fig_fnames, figs):
for fname in fnames:
fig.savefig(fname, transparent=True)
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 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 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
def main(
N = 1,
f_sample = 250e6, # Hz
t_length = 1e4 * 1e-9, # s
noise_band = passband(30e6, 80e6),
noise_sigma = 1,
# signal properties
f_sine = 50e6,
signal_band = passband(50e6 - 1e6, 50e6 + 1e6),
sine_amp = 0.2,
sine_offset = 0,
return_ranges_plot = False,
cut_signal_band_from_noise_band = False
):
N = int(N)
init_params = np.array([sine_amp, f_sine, None, sine_offset])
axs = None
snrs = np.zeros( N )
time = sampled_time(f_sample, end=t_length)
for j in range(N):
samples, noise = noisy_sine_sampling(time, init_params, noise_sigma)
# determine signal to noise
noise_level = bandlevel(noise, f_sample, noise_band)
if cut_signal_band_from_noise_band:
lower_noise_band = passband(noise_band[0], signal_band[0])
upper_noise_band = passband(signal_band[1], noise_band[1])
noise_level = bandlevel(noise, f_sample, lower_noise_band)
noise_level += bandlevel(noise, f_sample, upper_noise_band)
signal_level = bandlevel(samples, f_sample, signal_band)
snrs[j] = np.sqrt(signal_level/noise_level)
# make a nice plot showing what ranges were taken
# and the bandlevels associated with them
if return_ranges_plot and j == 0:
combined_fft, freqs = ft_spectrum(samples+noise, f_sample)
# plot the original signal
if False:
_, ax = plt.subplots()
ax = plot_signal(samples+noise, sample_rate=f_sample/1e6, time_unit='us', ax=ax)
# plot the spectrum
if True:
freq_scaler=1e6
_, axs = plot_combined_spectrum(combined_fft, freqs, freq_scaler=freq_scaler, freq_unit='MHz')
# indicate band ranges and frequency
for ax in axs:
ax.axvline(f_sine/freq_scaler, color='r', alpha=0.4)
ax.axvspan(noise_band[0]/freq_scaler, noise_band[1]/freq_scaler, color='purple', alpha=0.3, label='noiseband')
ax.axvspan(signal_band[0]/freq_scaler, signal_band[1]/freq_scaler, color='orange', alpha=0.3, label='signalband')
# indicate initial phase
axs[1].axhline(init_params[2], color='r', alpha=0.4)
# plot the band levels
levelax = axs[0].twinx()
levelax.set_ylabel("Bandlevel")
levelax.hlines(signal_level, noise_band[0]/freq_scaler, signal_band[1]/freq_scaler, colors=['orange'])
levelax.hlines(noise_level, noise_band[0]/freq_scaler, noise_band[1]/freq_scaler, colors=['purple'])
levelax.set_ylim(bottom=0)
axs[0].legend()
# plot signal_band pass signal
if False:
freqs = np.fft.fftfreq(len(samples), 1/f_sample)
bandmask = bandpass_mask(freqs, band=signal_band)
fft = np.fft.fft(samples)
fft[ ~bandmask ] = 0
bandpassed_samples = np.fft.ifft(fft)
_, ax3 = plt.subplots()
ax3 = plot_signal(bandpassed_samples, sample_rate=f_sample/1e6, time_unit='us', ax=ax3)
ax3.set_title("Bandpassed Signal")
return snrs, axs
if __name__ == "__main__":
from argparse import ArgumentParser
import os.path as path
rng = np.random.default_rng(1)
parser = ArgumentParser(description=__doc__)
parser.add_argument("fname", metavar="path/to/figure[/]", nargs="?", help="Location for generated figure, will append __file__ if a directory. If not supplied, figure is shown.")
args = parser.parse_args()
default_extensions = ['.pdf', '.png']
if args.fname == 'none':
args.fname = None
###
t_lengths = np.linspace(1e3, 5e4)* 1e-9 # s
N = 10e1
f_sine = 53.3e6 # Hz
f_sample = 250e6 # Hz
if False:
N = 1 # Note: keep this low, N figures will be displayed!
N_t_length = 10
for t_length in t_lengths[-N_t_length-1:-1]:
snrs = np.zeros( int(N))
for i in range(int(N)):
delta_f = 1/t_length
snrs[i], axs = main(
N=1,
t_length=t_length,
f_sample=f_sample,
# signal properties
f_sine = f_sine,
sine_amp = 1,
noise_sigma = 1,
noise_band = passband(30e6, 80e6),
signal_band = passband(f_sine- 3*delta_f, f_sine + 3*delta_f),
return_ranges_plot=True
)
axs[0].set_title("SNR: {}, N:{}".format(snrs[i], t_length*f_sample))
axs[0].set_xlim(
(f_sine - 20*delta_f)/1e6,
(f_sine + 20*delta_f)/1e6
)
print(snrs, "M:",np.mean(snrs))
plt.show(block=False)
else:
#original code
my_snrs = np.zeros( (len(t_lengths), int(N)) )
for j, t_length in enumerate(t_lengths):
return_ranges_plot = ((j==0) and True) or ( (j==(len(t_lengths)-1)) and True)
delta_f = 1/t_length
my_snrs[j], axs = main(
N=N,
t_length=t_length,
f_sample = f_sample,
# signal properties
f_sine = f_sine,
sine_amp = 1,
noise_sigma = 1,
noise_band = passband(30e6, 80e6),
signal_band = passband(f_sine- 3*delta_f, f_sine + 3*delta_f),
return_ranges_plot=return_ranges_plot,
)
if return_ranges_plot:
ranges_axs = axs
fig, axs2 = plt.subplots()
axs2.set_xlabel("N = T*$f_s$")
axs2.set_ylabel("SNR")
for j, t_length in enumerate(t_lengths):
t_length = t_length * f_sample
axs2.plot(np.repeat(t_length, my_snrs.shape[1]), my_snrs[j], ls='none', color='blue', marker='o', alpha=max(0.01, 1/my_snrs.shape[1]))
# plot the means
axs2.plot(t_lengths*f_sample, np.mean(my_snrs, axis=-1), color='green', marker='*', ls='none')
### Save or show figures
if not args.fname:
# empty list, False, None
plt.show()
else:
save_all_figs_to_path(args.fname, default_basename=__file__)