walkthrough.py

# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.16.7
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Spectrometer playground
# This notebook/script uses simulated noise to provide a testbed for the package's features.

# %%
import tempfile
from functools import partial

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import scipy as sc

from python_spectrometer import daq, Spectrometer


# %% [markdown]
# Matplotlib backend, choose `qt` for a separate window.

# %%
# %matplotlib qt
# # %matplotlib widget

# %% [markdown]
# Define an interesting noise spectrum:
# %%
def spectrum(f, A=1e-2, exp=1.5, add_colored=True, add_50hz=False,
             baseline=1e-5, npeaks=None, **_):
    f = np.asarray(f)
    S = np.zeros_like(f)

    if add_colored:
        with np.errstate(divide='ignore'):
            S += A / f ** exp

    if add_50hz:
        # sophisticated algorithm!
        harmonics = abs(f % 50) < np.diff(f).mean()
        idx, = harmonics.nonzero()
        p = sc.stats.beta.sf(np.linspace(0, 1, idx.size, endpoint=False),
                             5, 2)
        idx = rng.choice(idx, size=(min(10, idx.size) if npeaks is None else npeaks),
                         replace=False, p=p/p.sum())
        S[(idx,)] += 5e0 * A / (10 * f[f.nonzero()].min()) ** rng.random(size=idx.size)

    S += baseline
    return S



# %%
SEED = 1
mpl.style.use('qutil.plotting.presentation_notex')
rng = np.random.default_rng(SEED)

fig, ax = plt.subplots(layout='constrained')
ax.loglog(f := np.linspace(0, 1234.5, 501), spectrum(f, add_50hz=True))
ax.set_xlabel('$f$ (Hz)')
ax.set_ylabel(r'$S(f)$ (V$^2/\sqrt{\mathrm{Hz}}$)')

# %% [markdown]
# ## Sequential spectrum acquisition
# The `daq.simulator` module currently provides three different classes:
#  - `QoptColoredNoise`, which just generates noise according to a given power spectral density
#  - `DemodulatorQoptColoredNoise` uses `QoptColoredNoise` for noise generation but mimicks a Lock-in amplifier by demodulating (and optionally modulating) the noise.
#  - `MonochromaticNoise` generates a single-frequency tone. Useful for analyzing effects of spectral windows etc.
#
# We can instantiate the two necessary objects like so:

# %%
qopt_daq = daq.simulator.QoptColoredNoise(spectrum)
speck = Spectrometer(qopt_daq, savepath=tempfile.mkdtemp(),
                     # By default acquisition is asynchronous,
                     # set this to avoid race conditions
                     blocking_acquisition=True,
                     plot_style='qutil.plotting.presentation_notex',
                     figure_kw=dict(layout='constrained'), 
                     legend_kw=dict(loc='lower left'))

# %% [markdown]
# By default, only the main spectum plot is shown (scaled as an amplitude spectral density, the square root of the PSD). Changing plot settings can be done either in the constructor or interactively using the property setters.
#
# The main method used to acquire new data is `Spectrometer.take()`. Keyword arguments to this method are passed through to `DAQ.setup()`, `DAQ.acquire()`, and all processing functions downstream.

# %%
speck.take('my first spectrum', add_50hz=True, f_min=0.1)
plt.show()  # needed in notebook+widget for repeated cell execution.

# %% [markdown]
# Plot the design spectrum:

# %%
ln, = speck.ax[0].loglog(speck[0]['f_processed'], spectrum(speck[0]['f_processed']) ** .5,
                         color='tab:red', ls='--', zorder=10)
# %% [markdown]
# Looks like the spectrum estimation worked, gelukkig!
# We can index the `Spectrometer` object using the same identifiers shown in the plot to see the internally stored data.
# This includes all data acquisition settings, the comment, the filepath to where the file is stored (which is relative to `speck.savepath` by default) as well as the actual data.

# %%
speck['my first spectrum']

# %% [markdown]
# Now show the timetrace of the data underlying the previously acquired spectrum:

# %%
ln.remove()
speck.plot_timetrace = True

# %% [markdown]
# Now show the cumulative spectrum:

# %%
speck.plot_cumulative = True
speck.plot_cumulative_normalized = False

# %% [markdown]
# You can also play back the timetrace as audio:

# %%
speck.play_sound = True
speck.play(0)

# %% [markdown]
# Take some more data, but remove the cumulative plot again

# %%
speck.plot_cumulative = False
speck.take('my second spectrum', nperseg=4096, f_max=1e3)
speck.take('my third spectrum', n_avg=5, f_min=2e3/4096, f_max=1e3)
speck.take('my fourth spectrum', n_seg=10, nperseg=4096, f_max=1e3)

# %% [markdown]
# Hide data we are not interested in anymore

# %%
speck.hide('my first spectrum', 1)

# %% [markdown]
# Whoops, didn't mean to hide the second one

# %%
speck.show('my second spectrum')

# %% [markdown]
# Investigate peak RMS amplitude by switching the *scaling* from a spectral density to a spectrum (see the [SciPy docs](https://docs.scipy.org/doc/scipy/tutorial/signal.html#tutorial-spectralanalysis) for example).

# %%
speck.plot_timetrace = False
speck.plot_density = False

# %% [markdown]
# Now compare the powers relative to the fourth spectrum

# %%
speck.set_reference_spectrum('my fourth spectrum')
speck.plot_dB_scale = True
speck.plot_amplitude = False
speck.leg.set_loc('upper left')

# %% [markdown]
# Once done with spectral analysis, we might want to store the current state on disk and either pick up where we left off later or just inspect data from a different PC. This is easily done using `speck.serialize_to_disk()` and `speck.recall_from_disk()`.

# %%
speck.plot_dB_scale = False
speck.serialize_to_disk('foobar')

# %%
speck_loaded = Spectrometer.recall_from_disk(speck.savepath / 'foobar')

# %% [markdown]
# Alternatively, you can also just add spectra saved on disk to compare them.

# %%
speck_fresh = Spectrometer()
speck_fresh.add_spectrum_from_file(speck.savepath / speck[0]['filepath'])
speck_fresh.add_spectrum_from_file(speck.savepath / speck[3]['filepath'])

# %% [markdown]
# ## Lock-in DAQs
# The `DemodulatorQoptColoredNoise` simluates a Lock-in amplifier by multiplying the noise trace by
# $$
# \text{IQ}(t) = \sqrt{2}\exp(-i\omega t),
# $$
# and subsequently low-pass filtering with an RC filter with $3\,\text{dB}$ cutoff frequency `f_max` and order `filter_order`. Finally, the signal can also be modulated by $\Re(\text{IQ})$ to mimick the output signal of the Lock-in amplifier.
#
# The demodulated signal is complex with $R(t) = X(t) + iY(t)$ and its Fourier spectrum hence complex (two-sided). Thus, the `Spectrometer` plots negative frequencies by default. Furthermore, modulation with frequency $\omega$ shifts the entire spectrum by $\omega$, meaning the DC component of the resulting spectrum is actually the component at frequency $\omega$. This can be shifted back by setting `Spectrometer.plot_absolute_frequencies` to `True`.
#
# See [this resource](https://www.zhinst.com/europe/en/resources/principles-of-lock-in-detection) for more details on Lock-in amplification.

# %%
demod_daq = daq.simulator.DemodulatorQoptColoredNoise(
    partial(spectrum, add_50hz=True)
)
speck = Spectrometer(demod_daq, savepath=tempfile.mkdtemp(),
                     blocking_acquisition=True, 
                     plot_absolute_frequencies=False,
                     plot_style='qutil.plotting.presentation_notex',
                     figure_kw=dict(layout='constrained'))

# %%
speck.take(n_avg=5, freq=789, fs=13.4e3, nperseg=1 << 12, delay=True,
           modulate_signal=False)
plt.show()  # necessary if using notebook widget

# %% [markdown]
# With a different demodulation frequency, the $1/f$ peak moves.

# %%
speck.take(n_avg=5, freq=13, fs=13.4e3, nperseg=1 << 12, delay=True,
           modulate_signal=False)
plt.show()  # necessary if using notebook widget

# %% [markdown]
# We can also turn on modulation of the signal, that is, simulate as if the signal were originating from the Lockin output.

# %%
speck.hide(0, 1)
speck.take('modulated', n_avg=5, freq=77, fs=13.4e3, nperseg=1 << 12, delay=True,
           modulate_signal=True)
speck.take('unmodulated', n_avg=5, freq=77, fs=13.4e3, nperseg=1 << 12, delay=True,
           modulate_signal=False)

# %% [markdown]
# Do not plot negative part of the spectrum, and enable absolute frequencies, ie., do not reference to the carrier frequency.
# Note that in this case, all spectra are plotted for their own reference frequencies even if they might have been acquired for a different carrier frequency.

# %%
speck.plot_negative_frequencies = False
speck.plot_absolute_frequencies = True
speck.ax[0].set_xscale('log')

# %% [markdown]
# ## Live spectrum acquisition
# For noise hunting, it is often convenient to continuously acquire noise spectra while changing things around the setup. For this, use the `live_view` method.
#
# Note: the plot seems to have some issues in `widget` mode, use `qt` for a more reliable plot.

# %%
view, = speck.live_view(11, fs=13.4e3, nperseg=1 << 12, delay=True,
                        in_process=False)

# %% [markdown]
# If oscilloscope mode is active, another view opens. We can also run the figures in a separate process instead of thread by setting `in_process=True`.

# %%
speck.plot_timetrace = True
views = speck.live_view(11, fs=13.4e3, nperseg=1 << 12, delay=True,
                        in_process=True)

# %%
plt.close('all')