sodetlib 0.6.1rc1__py3-none-any.whl

sodetlib/operations/uxm_relock.py

@@ -0,0 +1,406 @@
+import time
+import numpy as np
+import sodetlib as sdl
+from sodetlib.operations.tracking import relock_tracking_setup
+from sodetlib.operations import uxm_setup, bias_steps
+
+import matplotlib.pyplot as plt
+import os
+if not os.environ.get('NO_PYSMURF', False):
+    from pysmurf.client.base.smurf_control import SmurfControl
+
+
+
+@sdl.set_action()
+def reload_tune(S, cfg, bands, setup_notches=False,
+                new_master_assignment=False, tunefile=None, update_cfg=True):
+    """
+    Reloads an existing tune, runs setup-notches and serial grad descent
+    and eta scan.
+
+    Args
+    -----
+    S : SmurfControl
+        Pysmurf instance
+    cfg : DetConfig
+        Det config instance
+    bands : list, optional
+        List of bands to run
+    setup_notches : bool
+        Whether to run setup_notches
+    new_master_assignment : bool
+        Whether to create a new master assignment
+    tunefile : str
+        Tunefile to load. Defaults to the tunefile in the device cfg.
+    update_cfg : bool
+        If true, will update the device cfg with the tunefile created by
+        setup notches.
+    """
+
+    if tunefile is None:
+        tunefile = cfg.dev.exp['tunefile']
+
+    S.load_tune(tunefile)
+
+    for band in bands:
+        sdl.stop_point(S)
+        bcfg = cfg.dev.bands[band]
+        S.set_att_uc(band, bcfg['uc_att'])
+        S.set_att_dc(band, bcfg['dc_att'])
+        sdl.pub_ocs_log(S, f"Relocking tune: Band {band}")
+        if setup_notches:
+            S.log(f"Setup notches, new_master_assignment={new_master_assignment}")
+            S.setup_notches(band, new_master_assignment=new_master_assignment)
+        else:
+            S.relock(band)
+
+    if setup_notches and update_cfg:
+        # Update tunefile
+        cfg.dev.update_experiment({'tunefile': S.tune_file}, update_file=True)
+
+    run_grad_descent_and_eta_scan(S, cfg, bands=bands, update_tune=False)
+
+    return True, None
+
+
+@sdl.set_action()
+def run_grad_descent_and_eta_scan(
+    S, cfg, bands=None, update_tune=False, force_run=False, max_iters=None,
+    gain=None): 
+    """
+    This function runs serial gradient and eta scan for each band.
+    Critically, it pulls in gradient descent tune parameters from the device
+    config and uses them to setup the grad descent operation before running.
+
+    Args
+    ----
+    S : SmurfControl
+        Pysmurf instance
+    cfg : DetConfig
+        Det config instance
+    bands : list, optional
+        List of bands to run on. Defaults to all 8
+    update_tune : bool
+        If this is set to True, the new resonance frequency and eta parameters
+        will be loaded into the smurf tune, and a new tunefile will be written
+        based on the new measurements.
+    force_run : bool
+        If True, will reset the ``etaScanInProgress`` variable to force it
+        to re-run. This might be necessary if serial gradient descent failed out.
+    """
+
+    if bands is None:
+        bands = cfg.dev.exp['active_bands']
+    bands = np.atleast_1d(bands)
+
+    for b in bands:
+        sdl.stop_point(S)
+        in_progress_reg = S._cryo_root(b) + 'etaScanInProgress'
+        if S._caget(in_progress_reg):
+            if force_run:
+                S._caput(in_progress_reg, 0)
+            else:
+                raise RuntimeError(
+                    "Failed during grad descent -- scan already in progress"
+                )
+
+        bcfg = cfg.dev.bands[b]
+
+        if max_iters is None:
+            max_iters = bcfg['gradientDescentMaxIters']
+        if gain is None:
+            gain = bcfg['gradientDescentGain']
+
+        S.set_gradient_descent_step_hz(b, bcfg['gradientDescentStepHz'])
+        S.set_gradient_descent_max_iters(b, max_iters)
+        S.set_gradient_descent_gain(b, gain)
+        S.set_gradient_descent_converge_hz(b, bcfg['gradientDescentConvergeHz'])
+        S.set_gradient_descent_beta(b, bcfg['gradientDescentBeta'])
+
+        S.log(f"Running grad descent and eta scan on band {b}")
+
+        init_scale_array = S.get_amplitude_scale_array(b)
+        init_center_freq_array = S.get_center_frequency_array(b)
+        S.run_serial_gradient_descent(b)
+
+        if (S.get_amplitude_scale_array(b) != init_scale_array).any():
+            S.log("Grad descent errored out! Resetting state")
+            S.set_amplitude_scale_array(b, init_scale_array)
+            S.set_center_frequency_array(b, init_center_freq_array)
+            S._caput(in_progress_reg, 0)
+            raise RuntimeError(
+                "Failed during grad descent -- check smurf-streamer logs"
+            )
+
+        S.run_serial_eta_scan(b)
+
+        if update_tune:
+            # This basically does the inverse of what's in the relock function
+            band_center_mhz = S.get_band_center_mhz(b)
+            subband_offset = S.get_tone_frequency_offset_mhz(b)
+            center_freq_array = S.get_center_frequency_array(b)
+            eta_phase_array = S.get_eta_phase_array(b)
+            eta_mag_array = S.get_eta_mag_array(b)
+            res_freqs = band_center_mhz + subband_offset + center_freq_array
+
+            for res in S.freq_resp[b]['resonances'].values():
+                ch = res['channel']
+                f0 = res['freq']
+                # Pysmurf does this to ignore channels with bad freqs so I 
+                # guess we will too
+                for ll, hh in S._bad_mask:
+                    if ll < f0 < hh:
+                        ch = -1
+
+                res['freq'] = res_freqs[ch]
+                res['offset'] = center_freq_array[ch]
+                res['eta_phase'] = eta_phase_array[ch]
+                res['eta_scaled'] = eta_mag_array[ch]
+
+    if update_tune:
+        S.log("Saving tune!")
+        S.save_tune()
+        cfg.dev.exp['tunefile'] = S.tune_file
+        cfg.dev.update_file()
+
+
+def get_full_band_sweep(S, cfg, band, chan):
+    """
+    This runs full_band_ampl_sweep to get the resonator response around a single
+    channel. This is useful for debugging bad channels.
+
+    Args
+    -----
+    S : SmurfControl
+        Pysmurf instance
+    cfg : DetConfig
+        Det Config instance
+    band : int
+        smurf-band of channel
+    channel : int
+        smurf-channel of channel
+
+    Returns
+    ----------
+    fs : np.ndarray
+        Array of frequencies that were swept over
+    resp : np.ndarray
+        Complex transmission across the frequency range
+    """
+    sb = S.get_subband_from_channel(band, chan)
+    tone_power = cfg.dev.bands[band]['tone_power']
+    center_freq_array = S.get_center_frequency_array(band)
+    amp_scale_array = S.get_amplitude_scale_array(band)
+    eta_phase_array = S.get_eta_phase_array(band)
+    eta_mag_array = S.get_eta_mag_array(band)
+    fb_enable = S.get_feedback_enable(band)
+    fb_enable_arr = S.get_feedback_enable_array(band)
+
+    try:
+        fs, resp = S.full_band_ampl_sweep(band, [sb], tone_power, 2, n_step=256)
+    finally:
+        S.set_center_frequency_array(band, center_freq_array)
+        S.set_amplitude_scale_array(band, amp_scale_array)
+        S.set_eta_phase_array(band, eta_phase_array)
+        S.set_eta_mag_array(band, eta_mag_array)
+        S.set_feedback_enable(band, fb_enable)
+        S.set_feedback_enable_array(band, fb_enable_arr)
+    return fs, resp
+
+
+def plot_channel_resonance(S, cfg, band, chan):
+    """
+    Measures and plots resonator properties for a single smurf channel.
+    This will perform a scan around the specified channel, and will plot the
+    resonance dip along with where smurf thinks the resonance frequency is.
+    This is very useful for debugging why a channel isn't reading out properly.
+
+    Args
+    -----
+    S : (SmurfControl)
+        pysmurf instance
+    cfg : (DetConfig)
+        DetConfig instance
+    band : (int)
+        smurf band number
+    chan : (int)
+        smurf channel number
+    """
+    res_freq = S.channel_to_freq(band, chan)
+    fs, resp = get_full_band_sweep(S, cfg, band, chan)
+
+    fs = fs.ravel() + S.get_band_center_mhz(band)
+    resp = resp.ravel()
+    fs = fs[np.abs(resp) > 0]
+    resp = resp[np.abs(resp) > 0]
+
+    fig, axes = plt.subplots(2, 1, figsize=(10, 10))
+
+    # Res-dip plot
+    ax = axes[0]
+    ax.plot(fs, np.abs(resp))
+    ax = ax.twinx()
+    ax.plot(fs, np.unwrap(np.angle(resp)), color='C1')
+    ax.axvline(res_freq, color='grey', ls='--')
+
+    # Circ plot
+    eta_phase = S.get_eta_phase_array(band)[chan]
+    eta = np.exp(1.0j * eta_phase * (2*np.pi) / 360)
+    Q = np.real(resp * eta)
+    I = -np.imag(resp * eta)
+
+    ax = axes[1]
+    ax.axvline(0, color='grey')
+    ax.axhline(0, color='grey')
+    ax.scatter(I, Q, c=np.abs(resp))
+    ax.scatter(np.real(resp), np.imag(resp), c=np.abs(resp), alpha=0.3)
+
+    return fig, axes
+
+
+@sdl.set_action()
+def uxm_relock(
+    S, cfg, bands=None, show_plots=False,
+    setup_notches=False, new_master_assignment=False, reset_rate_khz=None,
+    nphi0=None, skip_setup_amps=False):
+    """
+    Relocks resonators by running the following steps:
+
+        1. Reset state (all off, disable waveform, etc.)
+        2. Set amps and check tolerance
+        3. load tune, setup_notches, serial grad descent and eta scan
+        4. Tracking setup
+        5. Measure noise
+
+    Args
+    -----
+    S : SmurfControl
+        Pysmurf instance
+    cfg : DetConfig
+        Det config instance
+    bands : list, optional
+        List of bands to run on. Defaults to all 8
+    setup_notches : bool
+        If True, will run setup notches instead of relocking
+    new_master_assignment : bool
+        Whether to run setup_notches with new_master_assignment.
+        Defaults to False
+    reset_rate_khz : float, optional
+        Flux Ramp Reset Rate to set (kHz), defaults to the value in the dev cfg
+    nphi0 : int, optional
+        Number of phi0's to ramp through. Defaults to the value that was used
+        during setup.
+    show_plots : bool
+        If True will show plots
+    skip_setup_amps : bool
+        If True will skip amplifier setup.
+
+    Returns
+    --------
+    summary : dict
+        Dictionary containing a summary of all run operations.
+    """
+    if bands is None:
+        bands = cfg.dev.exp['active_bands']
+    bands = np.atleast_1d(bands)
+
+    exp = cfg.dev.exp
+
+    #############################################################
+    # 1. Reset to known state
+    #############################################################
+    S.all_off()  # Turn off Flux ramp, tones, and biases
+    S.set_rtm_arb_waveform_enable(0)
+    S.set_filter_disable(0)
+    S.set_downsample_factor(exp['downsample_factor'])
+    if exp['coupling_mode'] == 'dc':
+        S.set_mode_dc()
+    else:
+        S.set_mode_ac()
+
+    for band in bands:
+        band_cfg = cfg.dev.bands[band]
+        S.set_synthesis_scale(band, exp['synthesis_scale'])
+        S.set_att_uc(band, band_cfg['uc_att'])
+        S.set_att_dc(band, band_cfg['dc_att'])
+        S.set_band_delay_us(band, band_cfg['band_delay_us'])
+        S.amplitude_scale[band] = band_cfg['tone_power']
+
+    summary = {}
+    summary['timestamps'] = []
+    #############################################################
+    # 2. Setup amps
+    #############################################################
+    sdl.stop_point(S)
+    if not skip_setup_amps:
+        summary['timestamps'].append(('setup_amps', time.time()))
+        sdl.set_session_data(S, 'timestamps', summary['timestamps'])
+        success, summary['amps'] = uxm_setup.setup_amps(S, cfg)
+
+        if not success:
+            return False, summary
+    else:
+        print("Skipping amp setup")
+
+    #############################################################
+    # 3. Load tune
+    #############################################################
+    sdl.stop_point(S)
+    summary['timestamps'].append(('load_tune', time.time()))
+    sdl.set_session_data(S, 'timestamps', summary['timestamps'])
+    success, summary['reload_tune'] = reload_tune(
+        S, cfg, bands, setup_notches=setup_notches,
+        new_master_assignment=new_master_assignment)
+
+    sdl.set_session_data(S, 'reload_tune', summary['reload_tune'])
+    if not success:
+        return False, summary
+
+    #############################################################
+    # 4. Tracking Setup
+    #############################################################
+    sdl.stop_point(S)
+    summary['timestamps'].append(('tracking_setup', time.time()))
+    sdl.set_session_data(S, 'timestamps', summary['timestamps'])
+
+    # Ensure feedback is enabled
+    for b in bands:
+        S.set_feedback_enable(b, 1)
+
+    tr = relock_tracking_setup(
+        S, cfg, bands, show_plots=show_plots,
+        reset_rate_khz=reset_rate_khz, nphi0=nphi0
+    )
+    summary['tracking_setup_results'] = tr
+
+    #############################################################
+    # 5. Noise
+    #############################################################
+    sdl.stop_point(S)
+    summary['timestamps'].append(('noise', time.time()))
+    sdl.set_session_data(S, 'timestamps', summary['timestamps'])
+    am, summary['noise'] = sdl.noise.take_noise(S, cfg, 30,
+                                                show_plot=show_plots,
+                                                save_plot=True,
+                                                g3_tag='post_relock')
+    summary['noise']['am'] = am
+
+    sdl.set_session_data(S, 'noise', {
+        'band_medians': summary['noise']['band_medians']
+    })
+
+    #############################################################
+    # 6. Bias group mapping
+    #############################################################
+    sdl.stop_point(S)
+    summary['timestamps'].append(('bg_map', time.time()))
+    sdl.set_session_data(S, 'timestamps', summary['timestamps'])
+
+    bsa = bias_steps.take_bgmap(S, cfg,
+                                show_plots=show_plots)
+    summary['bg_map'] = bsa
+
+    summary['timestamps'].append(('end', time.time()))
+    sdl.set_session_data(S, 'timestamps', summary['timestamps'])
+
+    return True, summary