Source code for prsctrl.scripts.lock_in_benchmark

import matplotlib.pyplot as plt
import numpy as np
import time
import os

from devctrl.utility.file_io import listdir_abs, listdir_filter
from devctrl.devices import LockInAmp



def benchmark_lock_in_buffer(lockin: LockInAmp, nvalues: int, nchann: 2, srat=128) -> float:
    whats = ["R", "theta", "ADC1"]
    lockin.buffer_setup(whats[:nchann], nvalues, srat)
    lockin.buffer_start_fill()
    timeout = 2 * nvalues/srat
    while timeout > 0:
        if lockin.buffer_is_done(): break
        timeout -= 1
        time.sleep(1)
    if timeout < 0: raise RuntimeError(f"Timed out")
    t_start = time.time()
    data = lockin.buffer_get_data()
    t_stop = time.time()
    return t_stop - t_start




def full_benchmark_lock_in_buffer(lockin: LockInAmp, nchann:int, report_file_path: str|None=None):
    from devctrl.devices.lock_in.impl import sr830, model7260
    # these are just to calculate useful number of points
    measurement_time_s = 30
    srats = [16, 32, 64, 128, 256, 512]
    n_points = [measurement_time_s * srat for srat in srats]
    benchmarks = []
    for i, n in enumerate(n_points):
        try:
            t = benchmark_lock_in_buffer(lockin, nvalues=n, nchann=nchann, srat=128)
            benchmarks.append((n, t))
        except Exception as e:
            print(f"ERROR at n = {n}:", e)
    if report_file_path is not None:
        try:
            import math
            with open(report_file_path, "w") as file:
                file.write("# Lock-In-Amplifier data transmission benchmark\n")
                file.write(f"# lock-in: {lockin.get_device_name()}\n")
                file.write(f"# number of channels: {nchann}\n")
                file.write(f"n_points,time [s]")
                for n, t in benchmarks:
                    file.write(f"\n{n},{math.ceil(t)}")
        except Exception as e:
            print(e)
    return benchmarks




def plot_benchmarks(search_dir, regex_filter, plot_name="lock-in-benchmark.pdf"):
    benchmarks = listdir_filter(search_dir, regex_filter)
    fig, ax = plt.subplots()
    for b in benchmarks:
        with open(b, "r") as file:
            lines = file.readlines()
        lockin_name = lines[1].split(":")[1].strip(" \n")
        nchann = lines[2].split(":")[1].strip(" \n")
        data = np.array([list(map(float, lines[i].strip("\n").split(","))) for i in range(4, len(lines))])
        ax.plot(data[:,0], data[:,1], label=f"{lockin_name}, $n_\\text{{CH}}={nchann}$")
    ax.set_xlabel("$n_\\text{points}$")
    ax.set_ylabel("$t$ [s]")
    sax = ax.secondary_xaxis("top", (lambda n: n/30, lambda n: n*30))
    sax.set_xlabel(f"Sample rate for $t=30\\,$s")
    sax.grid()
    ax.legend()
    fig.suptitle(f"Lock-In-Amplifier data transmission benchmark")
    fig.tight_layout()
    if plot_name:
        if not os.path.isabs(plot_name):
            plot_name = os.path.join(search_dir, plot_name)
        fig.savefig(plot_name)
    return fig, ax