sodetlib 0.6.1rc1__py3-none-any.whl

sodetlib/quality_control.py

@@ -0,0 +1,415 @@
+import sodetlib as sdl
+import numpy as np
+import time
+from scipy import signal
+from tqdm.auto import trange
+from pysmurf.client.base.smurf_control import SmurfControl
+import matplotlib.pyplot as plt
+
+def check_packet_loss(Ss, cfgs, dur=10, fr_khz=4, nchans=2000, slots=None, downsample_factor=1):
+    """
+    Takes a short G3 Stream on multiple slots simultaneously and checks for
+    dropped samples. This function is strange since it requires simultaneous
+    streaming on multiple slots to properly test, so it doesn't follow the
+    standard sodetlib function / data format.
+
+    Args
+    -----
+    Ss : dict[SmurfController]1
+        Dict of pysmurf instances where the key is the slot-number
+    cfgs : dict[DetConfig]
+        Dict of DetConfigs where the key is the slot number
+    dur : float
+        Duration of data stream (sec)
+    fr_khz : float
+        Frequency of FR rate (khz)
+    nchans : int
+        Number of channels to stream
+    slots : list
+        Which slots to stream data on. If None, will stream on all slots in the
+        Ss object.
+    downsample_factor : int
+        Downsample factor to use
+
+    Returns
+    --------
+    ams : dict[AxisManagers]
+        Dict of axis managers indexed by slot-number
+    res : dict
+        Dict where the key is the slot number, and the values are dicts
+        containing frame counters, number of dropped frames, etc.
+    """
+    if slots is None:
+        slots = Ss.keys()
+
+    for s in slots:
+        S = Ss[s]
+        S.flux_ramp_setup(fr_khz, 0.4, band=0)
+        sdl.stream_g3_on(
+            S, channel_mask=np.arange(nchans), downsample_factor=downsample_factor,
+            subtype='check_packet_loss'
+        )
+
+    time.sleep(dur)
+    
+    sids = {}
+    for s in slots:
+        sids[s] = sdl.stream_g3_off(Ss[s])
+
+    ams = {}
+    for s, sid in sids.items():
+        ams[s] = sdl.load_session(cfgs[s].stream_id, sid, no_signal=True)
+
+    res = {}
+    for s, am in ams.items():
+        dropped_samps = np.sum(np.diff(am.primary['FrameCounter']) - 1)
+        total_samps = len(am.primary['FrameCounter'])
+        res[s] = {
+            'sid': sids[s],
+            'meta': sdl.get_metadata(Ss[s], cfgs[s]),
+            'frame_counter': am.primary['FrameCounter'],
+            'dropped_samples': dropped_samps,
+            'dropped_frac': dropped_samps / total_samps,
+        }
+
+    return ams, res
+
+@sdl.set_action()
+def measure_bias_line_resistances(
+    S: SmurfControl, cfg, vstep=0.001, bgs=None, sleep_time=2.0):
+    """
+    Function to measure the bias line resistance and high-low-current-ratio for
+    each bias group. This needs to be run with the smurf hooked up to the
+    cryostat and the detectors superconducting, and a bgmap should exist.
+
+    Args
+    -------
+    S : SmurfControl
+        Pysmurf instance
+    cfg : DetConfig
+        Det Config instance
+    vstep : float
+        Voltage step size (in low-current-mode volts)
+    bgs : list
+        Bias lines to measure. Will default to active bias lines
+    sleep_time : float
+        Time to wait at each step.
+    """
+    if bgs is None:
+        bgs = cfg.dev.exp['active_bgs']
+    bgs = np.atleast_1d(bgs)
+    bgmap = np.load(cfg.dev.exp['bgmap_file'], allow_pickle=True).item()
+    bg_m = np.asarray(sum([bgmap['bgmap'] == bg for bg in bgs]), dtype=bool)
+
+    vbias = S.get_tes_bias_bipolar_array()
+    vb_low = vbias.copy()
+    vb_low[bgs] = 0
+    vb_high = vbias.copy()
+    vb_high[bgs] = vstep
+    segs = []
+
+    S.set_tes_bias_bipolar_array(vb_low)
+    sdl.set_current_mode(S, bgs, 0, const_current=False)
+
+    def take_step(bias_arr, sleep_time, wait_time=0.2):
+        S.set_tes_bias_bipolar_array(bias_arr)
+        time.sleep(wait_time)
+        t0 = time.time()
+        time.sleep(sleep_time)
+        t1 = time.time()
+        return (t0, t1)
+
+    sdl.stream_g3_on(S, subtype='measure_bias_line_resistance')
+    time.sleep(0.5)
+
+    segs.append(take_step(vb_low, sleep_time, wait_time=0.5))
+    segs.append(take_step(vb_high, sleep_time, wait_time=0.5))
+
+    S.set_tes_bias_bipolar_array(vb_low)
+    time.sleep(0.5)
+    sdl.set_current_mode(S, bgs, 1, const_current=False)
+
+    segs.append(take_step(vb_low, sleep_time, wait_time=0.05))
+    segs.append(take_step(vb_high, sleep_time, wait_time=0.05))
+
+    sid = sdl.stream_g3_off(S)
+
+    am = sdl.load_session(cfg.stream_id, sid)
+    ts = am.timestamps
+    sigs = []
+    for (t0, t1) in segs:
+        m = (t0 < ts) & (ts < t1)
+        sigs.append(np.mean(am.signal[bg_m][:, m], axis=1) * S.pA_per_phi0 / (2*np.pi))
+
+    Rbl_low = vstep / (np.abs(sigs[1] - sigs[0]) * 1e-12)
+    Rbl_high = vstep / (np.abs(sigs[3] - sigs[2]) * 1e-12)
+    high_low_ratio = Rbl_low / Rbl_high
+
+    cfg.dev.exp['bias_line_resistance'] = np.nanmedian(Rbl_low)
+    cfg.dev.exp['high_low_current_ratio'] = np.nanmedian(high_low_ratio)
+    cfg.dev.update_file()
+
+    path = sdl.make_filename(S, 'measure_bias_line_info')
+    data = {
+        'Rbl_low_all': Rbl_low,
+        'Rbl_high_all': Rbl_high,
+        'high_low_ratio_all': high_low_ratio,
+        'bias_line_resistance': np.nanmedian(Rbl_low),
+        'high_current_mode_resistance': np.nanmedian(Rbl_high),
+        'high_low_ratio': np.nanmedian(high_low_ratio),
+        'sid': sid,
+        'vstep': vstep,
+        'bgs': bgs,
+        'meta': sdl.get_metadata(S, cfg),
+        'segs': segs,
+        'sigs': sigs,
+        'path': path,
+    }
+    np.save(path, data, allow_pickle=True)
+    S.pub.register_file(path, 'bias_line_resistances', format='npy')
+
+    return am, data
+
+def setup_fixed_tones(S, cfg, tones_per_band=256, bands=None, jitter=0.5,
+                      tone_power=None):
+    """
+    Enables many fixed tones across a selection of bands.
+
+    Args
+    ----------
+    S : SmurfControl
+        Pysmurf instance
+    cfg : DetConfig
+        Det config instance
+    tones_per_band : int
+        Number of fixed tones to create in each band
+    bands : int, list[int]
+        Bands to set fixed tones in. Defaults to all 8.
+    jitter : float
+        Noise [Mhz] to add to the center freq so that fixed tones are not
+        equispaced.
+    tone_power : int
+        Tone power of fixed tones.
+    """
+    if bands is None:
+        bands = np.arange(8)
+    bands = np.atleast_1d(bands)
+
+    chans_per_band = S.get_number_channels()
+    for band in bands:
+        S.log(f"Setting fixed tones for band {band}")
+        if tone_power is None:
+            tone_power = cfg.dev.bands[band]['tone_power']
+        
+        sbs = np.linspace(0, chans_per_band, tones_per_band, dtype=int, endpoint=False)
+        asa = np.zeros_like(S.get_amplitude_scale_array(band))
+        asa[sbs] = tone_power
+        S.set_amplitude_scale_array(band, asa)
+        S.set_center_frequency_array(band, np.random.uniform(-jitter/2, jitter/2, chans_per_band))
+        S.set_feedback_enable_array(band, np.zeros(chans_per_band, dtype=int))
+
+def get_noise_dBcHz(S, band, chan, nsamp=2**20, nperseg=2**16,
+                    noise_freq=30e3, noise_bw=100):
+    """
+    Takes debug data and measures I/Q noise in dBc/Hz.
+
+    Args
+    -----
+    S : SmurfControl
+        Pysmurf instance
+    band : int
+        Smurf band
+    chan : int
+        Smurf chan
+    nsamp : int
+        Number of samples to take
+    nperseg : int
+        Nperseg to use when creating the psd
+    noise_freq : float
+        Freq to measure the readout noise at.
+    noise_bw : float
+        Frequency bandwidth over which the noise median will be taken
+
+    Returns
+    ---------
+    noise_i : float
+        Noise of the I data stream in dBc/Hz
+    noise_q : float
+        Noise of the Q data stream in dBc/Hz
+    """
+    fsamp = S.get_channel_frequency_mhz() * 1e6
+    
+    sig_i, sig_q, _ = S.take_debug_data(band, channel=chan, rf_iq=True, nsamp=nsamp)
+    datfile = S.get_streamdatawriter_datafile()
+
+    fs, pxxi = signal.welch(sig_i, fs=fsamp, nperseg=nperseg)
+    fs, pxxq = signal.welch(sig_q, fs=fsamp, nperseg=nperseg)
+
+    magfac = np.mean(sig_q)**2 + np.mean(sig_i)**2
+    pxxi_dbc = 10. * np.log10(pxxi/magfac)
+    pxxq_dbc = 10. * np.log10(pxxq/magfac)
+
+    f0, f1 = noise_freq - noise_bw/2, noise_freq + noise_bw/2
+    m = (f0 < fs) & (fs < f1)
+    noise_i = np.nanmedian(pxxi_dbc[m])
+    noise_q = np.nanmedian(pxxq_dbc[m])
+
+    return noise_i, noise_q, datfile
+    
+
+def plot_fixed_tone_loopback(res):
+    """Plot results from fixed_tone_loopback"""
+    fig, ax = plt.subplots()
+    fig.patch.set_facecolor('white')
+
+    ax.plot(res['freqs'], res['noise_q'], '.', alpha=0.8)
+    ax.plot(res['freqs'], res['noise_i'], '.', alpha=0.8)
+    txt = '\n'.join([
+        f"num tones: {len(res['ft_freqs_all'])}",
+        f"tone power: {res['tone_power']}",
+        f"UC att: {res['att_uc']}",
+        f"DC att: {res['att_dc']}",
+    ])
+    ax.text(0.05, 0.8, txt, transform=ax.transAxes,
+            bbox=dict(fc='white', alpha=0.6))
+    ax.set_xlabel("Freq [MHz]")
+    ax.set_ylabel("Noise [dBc/Hz]")
+    crate_id = res['meta']['crate_id']
+    slot = res['meta']['slot']
+    ax.set_title(f"Crate {crate_id}, Slot {slot}")
+
+    return fig, ax
+
+
+def fixed_tone_loopback(
+        S, cfg, bands=None, tones_per_band=256, meas_chans_per_band=5,
+        setup_tones=False, tone_power=12, show_pb=True, noise_freq=30e3,
+        noise_bw=100, att_uc=None, att_dc=None):
+    """
+    Runs QC test to check noise levels across band with many fixed tones enabled.
+    
+    Args
+    ------
+    S : SmurfControl
+        Pysmurf instance
+    cfg : DetConfig
+        Det config instance
+    bands : int, list[int]
+        Bands to run on. Default is all 8
+    tones_per_band : int
+        Number of fixed tones to enable per band. This defaults to 256, which
+        is 2048 tones total with all 8 bands enabled.
+    meas_chans_per_band : int
+        Number of channels per band to measure readout noise
+    setup_fixed_tones : bool
+        If true, will set up fixed tones across specified bands. If false,
+        this will skip the setup and assume tones are already set up.
+    tone_power : int
+        Tone power to use
+    show_pb : bool
+        If True will show a progress bar
+    noise_freq : float
+        Target frequency to measure the readout noise
+    noise_bw : float
+        Frequency bandwidth over which the noise median will be taken
+    att_uc : int
+        UC atten. If not set, will use current uc atten.
+    att_dc : int
+        DC atten. If not set, will use current dc atten
+    """
+    if bands is None:
+        bands = np.arange(8)
+    bands = np.atleast_1d(bands)
+
+    if att_uc is None:
+        att_uc = S.get_att_uc(bands[0])
+    else:
+        for b in bands:
+            S.set_att_uc(b, att_uc)
+
+    if att_dc is None:
+        att_dc = S.get_att_dc(bands[0])
+    else:
+        for b in bands:
+            S.set_att_dc(b, att_dc)
+
+    if setup_tones:
+        setup_fixed_tones(S, cfg, tones_per_band=tones_per_band, bands=bands,
+                          tone_power=tone_power)
+    else: # Just update the tone power
+        S.log(f"Setting tone power to {tone_power} for bands {bands}...")
+        for band in bands:
+            asa = S.get_amplitude_scale_array(band)
+            if tone_power in np.unique(asa):
+                continue
+            asa[asa != 0] = tone_power
+            S.set_amplitude_scale_array(band, asa)
+    
+    meas_bands = []
+    meas_chans = []
+    meas_freqs = []
+    ft_chans_all = []
+    ft_freqs_all = []
+    ft_bands_all = []
+    S.log("Finding fixed tones and meas_channels")
+    for band in bands:
+        freqs = S.get_center_frequency_array(band) \
+                 + S.get_tone_frequency_offset_mhz(band) \
+                 + S.get_band_center_mhz(band)
+
+        ft_chans = np.where(S.get_amplitude_scale_array(band))[0]
+        freqs = freqs[ft_chans]
+        sort_idx = np.argsort(freqs)
+        meas_idx = np.unique(np.round(np.linspace(
+            0, len(ft_chans) - 1, meas_chans_per_band)).astype(int))
+        # We want to sort first so meas_chans are evenly distributed across
+        # freq space
+        meas_chans.append(ft_chans[sort_idx][meas_idx])
+        meas_bands.append([band for _ in meas_idx])
+        meas_freqs.append(freqs[sort_idx][meas_idx])
+        ft_chans_all.append(ft_chans)
+        ft_bands_all.append([band for _ in ft_chans])
+        ft_freqs_all.append(freqs)
+    meas_bands = np.hstack(meas_bands)
+    meas_chans = np.hstack(meas_chans)
+    meas_freqs = np.hstack(meas_freqs)
+    ft_chans_all = np.hstack(ft_chans_all)
+    ft_bands_all = np.hstack(ft_freqs_all)
+    ft_freqs_all = np.hstack(ft_freqs_all)
+
+    noise_i = np.full_like(meas_freqs, np.nan)
+    noise_q = np.full_like(meas_freqs, np.nan)
+    datfiles = []
+    for i in trange(len(meas_bands), disable=not(show_pb)):
+        b, c = meas_bands[i], meas_chans[i]
+        S.log(f"Band {b}, Chan {c}")
+        try:
+            noise_i[i], noise_q[i], _ = get_noise_dBcHz(
+                S, b, c, noise_freq=noise_freq,
+                noise_bw=noise_bw)
+        except IndexError as e:
+            S.log(f"Take Data failed...\n{e}")
+            S.log("Skipping channel")
+        datfiles.append(S.get_streamdatawriter_datafile())
+
+    fname = sdl.make_filename(S, 'fixed_tone_loopback.npy')
+    res = dict(
+        meta=sdl.get_metadata(S, cfg),
+        bands=meas_bands, channels=meas_chans, freqs=meas_freqs,
+        noise_i=noise_i, noise_q=noise_q,
+        ft_bands_all=ft_bands_all, ft_channels_all=ft_chans_all,
+        ft_freqs_all=ft_freqs_all, noise_freq=noise_freq,
+        att_uc=att_uc, att_dc=att_dc, tone_power=tone_power,
+        datfiles=datfiles
+    )
+    np.save(fname, res, allow_pickle=True)
+    S.pub.register_file(fname, 'loopback', format='npy')
+
+
+    fname = sdl.make_filename(S, 'fixed_tone_loopback.png', plot=True)
+    fig, _ = plot_fixed_tone_loopback(res)
+    fig.savefig(fname)
+    S.pub.register_file(fname, 'loopback', format='png', plot=True)
+
+    return res
+