sodetlib 0.6.1rc1__py3-none-any.whl

sodetlib/resonator_fitting.py

@@ -0,0 +1,508 @@
+'''
+Given a complex s21 sweep, returns the data fit to the resonator model and the
+resonator parameters and provides tools for plotting the fit results.
+
+Based on equation 11 from Kahlil et al. and adapted from Columbia KIDs open
+source analysis code.
+'''
+import numpy as np
+from lmfit import Model
+import matplotlib.pyplot as plt
+from matplotlib.gridspec import GridSpec
+
+def linear_resonator(f, f_0, Q, Q_e_real, Q_e_imag):
+    '''
+    Function for a resonator with asymmetry parameterized by the imaginary
+    part of ``Q_e``. The real part of ``Q_e`` is what we typically refer to as
+    the coupled Q, ``Q_c``.
+    '''
+    Q_e = Q_e_real + 1j*Q_e_imag
+    return (1 - (Q * Q_e**(-1) /(1 + 2j * Q * (f - f_0) / f_0) ) )
+
+def cable_delay(f, delay, phi, f_min):
+    '''
+    Function implements a time delay (phase variation linear with frequency).
+    '''
+    return np.exp(1j * (-2 * np.pi * (f - f_min) * delay + phi))
+
+def general_cable(f, delay, phi, f_min, A_mag, A_slope):
+    '''
+    Function implements a time delay (phase variation linear with frequency) and
+    attenuation slope characterizing a background RF cable transfer function.
+    '''
+    phase_term =  cable_delay(f,delay,phi,f_min)
+    magnitude_term = ((f-f_min)*A_slope + 1)* A_mag
+    return magnitude_term*phase_term
+
+def resonator_cable(f, f_0, Q, Q_e_real, Q_e_imag, delay, phi, f_min, A_mag,
+                    A_slope):
+    '''
+    Function that includes asymmetric resonator (``linear_resonator``) and cable
+    transfer functions (``general_cable``). Which most closely matches our full
+    measured transfer function.
+    '''
+    resonator_term = linear_resonator(f, f_0, Q, Q_e_real, Q_e_imag)
+    cable_term = general_cable(f, delay, phi, f_min, A_mag, A_slope)
+    return cable_term*resonator_term
+
+def full_fit(freqs, real, imag):
+    '''
+    Fitting function that takes in frequency and real and imaginary parts of the
+    transmission of a resonator (needs to trimmed down to only data for a
+    single resonator) and returns fitted parameters to the ``resonator_cable``
+    model.
+
+    Args
+    ----
+    freqs : float ndarray
+        Frequencies that line up with complex transmission data.
+    real : float ndarray
+        Real part of resonator complex transmission to be fit
+    imag : float ndarray
+        Imaginary part of resonator complex transmission to be fit.
+
+    Returns
+    -------
+    result : (`lmfit.Model.ModelResult`)
+        This is a class of lmfit that contains all of the fitted parameters as
+        well as a number of other pieces of data and metadata related to the fit
+        and some helper functions for plotting and manipulating the data in the
+        object.
+    '''
+    s21_complex = np.vectorize(complex)(real, imag)
+
+    #set our initial guesses
+    argmin_s21 = np.abs(s21_complex).argmin()
+    fmin = freqs.min()
+    fmax = freqs.max()
+    f_0_guess = freqs[argmin_s21]
+    Q_min = 0.1 * (f_0_guess / (fmax - fmin))
+    delta_f = np.diff(freqs)
+    min_delta_f = delta_f[delta_f > 0].min()
+    Q_max = f_0_guess / min_delta_f
+    Q_guess = np.sqrt(Q_min * Q_max)
+    s21_min = np.abs(s21_complex[argmin_s21])
+    s21_max = np.abs(s21_complex).max()
+    Q_e_real_guess = Q_guess / (1 - s21_min / s21_max)
+    A_slope, A_offset = np.polyfit(freqs - fmin, np.abs(s21_complex), 1)
+    A_mag = A_offset
+    A_mag_slope = A_slope / A_mag
+    phi_slope, phi_offset = np.polyfit(freqs - fmin, np.unwrap(np.angle(s21_complex)), 1)
+    delay = -phi_slope / (2 * np.pi)
+
+    #make our model
+    totalmodel = Model(resonator_cable)
+    params = totalmodel.make_params(f_0=f_0_guess,
+                                Q=Q_guess,
+                                Q_e_real=Q_e_real_guess,
+                                Q_e_imag=0,
+                                delay=delay,
+                                phi=phi_offset,
+                                f_min=fmin,
+                                A_mag=A_mag,
+                                A_slope=A_mag_slope)
+    #set some bounds
+    params['f_0'].set(min=freqs.min(), max=freqs.max())
+    params['Q'].set(min=Q_min, max=Q_max)
+    params['Q_e_real'].set(min=1, max=1e7)
+    params['Q_e_imag'].set(min=-1e7, max=1e7)
+    params['phi'].set(min=phi_offset-np.pi, max=phi_offset+np.pi)
+
+    #fit it
+    result = totalmodel.fit(s21_complex, params, f=freqs)
+    return result
+
+def get_qi(Q, Q_e_real, Q_e_imag):
+    '''
+    Function for deriving the internal quality factor from the fitted quality
+    factors (Q and Qc).
+
+    Args
+    ----
+    Q : float
+        Total resonator quality factor output parameter of ``full_fit``
+    Q_e_real : float
+        Resonator coupled quality factor output parameter of ``full_fit``
+    Returns
+    -------
+    Qi : float
+        Resonator internal quality factor.
+    '''
+    Qi = (Q**-1 - np.real((Q_e_real+1j*Q_e_imag)**-1))**-1
+    return Qi
+
+def get_br(Q, f_0):
+    '''
+    Function for deriving the resonator bandwidth from the fit results.
+
+    Args
+    ----
+    Q : float
+        Total resonator quality factor output parameter of ``full_fit``
+    f_0 : float
+        Resonance frequency output parameter of ``full_fit`` in Hz.
+    Returns
+    -------
+    br : float
+        Resonator bandwidth in Hz.
+    '''
+    br = f_0/Q
+    return br
+
+def reduced_chi_squared(ydata, ymod, n_param=9, sd=None):
+    '''
+    Reduced chi squared in lmfit does not return something reasonable so this
+    is a handwritten function to calculate it since we want standard deviation
+    to be the complex error.
+
+    Args
+    ----
+    ydata : float ndarray
+        complex data to calculate reduced chi squared on.
+    ymod : float ndarray
+        model fit result of ydata at same x-locations as ydata is sampled.
+    n_param : int
+        Number of parameters in the fit.
+    sd : float ndarray
+        standard deviation of data
+    Returns
+    -------
+    br : float
+        Resonator bandwidth in Hz.
+    '''
+    if sd is None:
+        sdr = np.std(np.real(ydata))
+        sdi = np.std(np.imag(ydata))
+        sd = sdr + 1j*sdi
+    chisq = np.sum((np.real(ydata) - np.real(ymod))**2/((np.real(sd))**2)) +\
+            np.sum((np.imag(ydata) - np.imag(ymod))**2/((np.imag(sd))**2))
+    nu=2*ydata.size-n_param     #multiply  the usual by 2 since complex
+    red_chisq = chisq/nu
+    return chisq, red_chisq
+
+def fit_tune(tunefile):
+    """
+    Automated fitting of resonator parameters from one tuning file.
+
+    Args
+    ----
+    tunefile : str, filepath
+        Full directories of one tunning file.
+    Returns
+    -------
+    dres : dict
+        a dictionary containing all of the fit results for all resonances in the provided tunefile. The keys are organized as::
+        
+          {
+            tunefile: str, filepath - full directories of one tunning file
+            band: 1d int array with shape (ndets) - smurf band for each detector
+            channels : 1d int array with shape (ndets) - Smurf channel for each
+                       detector. If unassigned, set to -1
+            res_index: 1d int array with shape (ndets) -- Index of the resonator within each band
+            res_freqs: 1d int array of resonance freqs (as is found in the tunefile)
+            model_params : dictionary for 9 parameters of resonator_cable 
+                f_0 : 1d array with shape (ndets)
+                Q : 1d array with shape (ndets)
+                Q_e_real : 1d array with shape (ndets)
+                Q_e_imag: 1d array with shape (ndets)
+                delay: 1d array with shape (ndets)
+                phi: 1d array with shape (ndets)
+                f_min: 1d array with shape (ndets)
+                A_mag: 1d array with shape (ndets)
+                A_slope: 1d array with shape (ndets)
+            derived_params : dictionary for derived parameters
+                Qi: 1d array with shape (ndets)
+                br: 1d array with shape (ndets)
+                depth: 1d array with shape (ndets)
+            find_freq_ctime : 1d string array with shape (ndets) - ctime find_freq was taken
+            S21: 2d array (ndets, nsamps) of measured S21
+            scan_freqs: 2d array (ndets, ndsamps) of freqs used for S21 scan
+            chi2 : 2d float array of shape (ndets x 2) - chi-squared goodness of fit
+            sid : str, session ID
+          }
+    """
+    dres = {}
+    data = np.load(tunefile,allow_pickle=True).item()
+
+    dres['tunefile'] = tunefile
+    model_params = {'f_0': [], 'Q': [], 'Q_e_real' : [], 'Q_e_imag': [],
+                   'delay': [], 'phi': [], 'f_min': [], 'A_mag': [],
+                    'A_slope': []}
+    derived_params = {'Qi': [], 'br': [], 'depth': []}
+    bands = []
+    channels = []
+    res_index = []
+    chi2 = []
+    sid = []
+    find_freq_ctime = []
+    S21s = []
+    scan_freqs = []
+    res_freqs = []
+    
+    for band in sorted(list(data.keys())):
+        if 'resonances' in list(data[band].keys()):
+            for idx in list(data[band]['resonances'].keys()):
+                scan=data[band]['resonances'][idx]
+                f=scan['freq_eta_scan']
+                if (band > 3) & (np.mean(f) > 6000):
+                    f-=2000
+                S21=scan['resp_eta_scan']
+                result=full_fit(f,S21.real,S21.imag)
+
+                # Need to check if this plays well with being in a dict/pickling
+                S21_mod = result.best_fit.real+1j*result.best_fit.imag                
+                
+                bands.append(band)
+                channels.append(scan['channel'])
+                res_index.append(idx)
+                res_freqs.append(scan['freq'])
+                chi2.append(reduced_chi_squared(S21, S21_mod))
+                find_freq_ctime.append(data[band]['find_freq']['timestamp'][0])
+                S21s.append(S21)
+                scan_freqs.append(f)
+
+                model_params['f_0'].append(result.best_values['f_0'])
+                model_params['Q'].append(result.best_values['Q'])
+                model_params['Q_e_real'].append(result.best_values['Q_e_real'])
+                model_params['Q_e_imag'].append(result.best_values['Q_e_imag'])
+                model_params['delay'].append(result.best_values['delay'])
+                model_params['phi'].append(result.best_values['phi'])
+                model_params['f_min'].append(result.best_values['f_min'])
+                model_params['A_mag'].append(result.best_values['A_mag'])
+                model_params['A_slope'].append(result.best_values['A_slope'])
+                
+                Qi = get_qi(result.best_values['Q'], result.best_values['Q_e_real'],
+                            result.best_values['Q_e_imag'])
+                br = get_br(result.best_values['Q'], result.best_values['f_0'])
+                depth = np.abs(S21_mod).max()/np.abs(S21_mod).min()
+                derived_params['Qi'].append(Qi)
+                derived_params['br'].append(br) 
+                derived_params['depth'].append(depth)
+                
+    bands = np.array(bands)
+    channels = np.array(channels)
+    res_index = np.array(res_index)
+    S21s = np.array(S21s)
+    scan_freqs = np.array(scan_freqs)
+    chi2 = np.array(chi2)
+    sid = np.array(sid)
+    for par in model_params:
+        model_params[par] = np.array(model_params[par])
+    for par in derived_params:
+        derived_params[par] = np.array(derived_params[par])
+
+    dres['bands'] = bands
+    dres['channels'] = channels
+    dres['res_index'] = res_index
+    dres['res_freqs'] = res_freqs
+    dres['sid'] = sid
+    dres['S21'] = S21s
+    dres['scan_freqs'] = scan_freqs
+    dres['chi2'] = chi2
+    dres['model_params'] = model_params
+    dres['derived_params'] = derived_params
+    return dres
+
+def get_resfit_plot_txt(resfit_dict, band, rix):
+    '''
+    Function to assemble some key fit information out of the fit dictionary
+    into a text block for adding to plots.
+
+    Args
+    ----
+    resfit_dict : dict
+        Dictionary with fit results output from ``fit_tune``
+    band : int
+        Smurf band of channel to get plot text for.
+    rix : int
+        Resonator index in tunefile (sorted by frequency order of setup_notches
+        channels) of channel to get plot text for.
+
+    Returns
+    -------
+    text : str
+        text block to add to resonator fit channel plot.
+    '''
+#     Get index of band, rix pair provided by user
+    idx = np.where(
+        (resfit_dict['res_index'] == rix)
+        & (resfit_dict['bands'] == band)
+    )[0][0]
+    # idx = int(np.where((resfit_dict['res_index'][:,0] == band) \
+    #              & (resfit_dict['res_index'][:,1] == rix))[0])
+            
+    mparams = resfit_dict['model_params']
+    dparams = resfit_dict['derived_params']
+    channel = resfit_dict['channels'][idx]
+    chi2 = resfit_dict['chi2'][idx]
+    if channel == -1:
+        text = f'Channel unassigned'
+    else:
+        text = f'Band {band} Channel {channel}'
+    text += '\n$f_r$: '+ f"{np.round(mparams['f_0'][idx],1)} MHz"
+    text += '\n$Q_i$: '+ f"{int(dparams['Qi'][idx])}"
+    text += '\nBW: '+ f"{int(dparams['br'][idx]*1e3)} kHz"
+#     text += '\nfit $\chi^2$: ' + f"{np.round(chi2,3)}"
+    return text
+
+def plot_channel_fit(fit_dict, idx):
+    '''
+    Function for plotting single channel eta_scan data from a tunefile with
+    a fit to an asymmetric resonator model ``resonator_cable``.
+
+    Args
+    ----
+    fit_dict : dict
+        fit results dictionary from ``fit_tune``
+    band : int
+        smurf band of resonator to plot
+    channel : int
+        smurf channel of resonator to plot
+    '''
+
+    rix = fit_dict['res_index'][idx]
+    band = fit_dict['bands'][idx]
+    
+    freqs = fit_dict['scan_freqs'][idx]
+    res_freq = fit_dict['res_freqs'][idx]
+    freqs_plot = 1e3*(freqs - res_freq)
+    resp_data = fit_dict['S21'][idx]
+    
+    params = {}
+    for p_opt in fit_dict['model_params']:
+        params[p_opt] = fit_dict['model_params'][p_opt][idx]
+    
+    resp_model = resonator_cable(freqs,**params)
+    fr_idx = np.argmin(np.abs(freqs - fit_dict['model_params']['f_0'][idx]))
+
+    fig = plt.figure(figsize = (12,6),constrained_layout=True)
+    gs = GridSpec(2, 6, figure=fig)
+    ax1 = fig.add_subplot(gs[:1,0:2])
+    ax1.plot(freqs_plot,20*np.log10(np.abs(resp_data)),
+             'C0o',label = 'Data')
+    ax1.plot(freqs_plot,20*np.log10(np.abs(resp_model)),
+             'C1-',label = 'Fit')
+    ax1.plot(freqs_plot[fr_idx],20*np.log10(np.abs(resp_data[fr_idx])),
+             'rx',ms = 12,label = '$f_r$ - fit')
+    ax1.set_xlabel('Offset $(f-f_{min})$ [kHz]')
+    ax1.set_ylabel('$|S_{21}|$ [dB]')
+    ax1.legend(loc = 'lower left',fontsize = 12)
+    ax2 = fig.add_subplot(gs[1:,0:2])
+    ax2.plot(freqs_plot,np.rad2deg(np.unwrap(np.angle(resp_data))),'C0o')
+    ax2.plot(freqs_plot,np.rad2deg(np.unwrap(np.angle(resp_model))),'C1-')
+    ax2.plot(freqs_plot[fr_idx],np.rad2deg(np.unwrap(np.angle(resp_data))[fr_idx]),
+             'rx',ms = 12)
+    ax2.set_xlabel('Offset $(f-f_r)$ [kHz]')
+    ax2.set_ylabel('Phase $S_{21}$ [$^{\\circ}$]')
+    ax3 = fig.add_subplot(gs[:,2:])
+    ax3.plot(np.real(resp_data),np.imag(resp_data),'C0o')
+    ax3.plot(np.real(resp_model),np.imag(resp_model),'C1-')
+    ax3.plot(np.real(resp_model[fr_idx]),np.imag(resp_model[fr_idx]),
+             'rx',ms = 12)
+    ax3.set_xlabel('Re($S_{21}$) [I]')
+    ax3.set_ylabel('Im($S_{21}$) [Q]')
+    mrange = 1.1*np.max(np.concatenate((np.abs(np.real(resp_data)),
+                                   np.abs(np.imag(resp_data)),
+                                   np.abs(np.real(resp_model)),
+                                   np.abs(np.imag(resp_model)))))
+    ax3.set_xlim(-mrange,mrange)
+    ax3.set_ylim(-mrange,mrange)
+    ax3.axhline(0,color = 'k')
+    ax3.axvline(0,color = 'k')
+
+    restext = get_resfit_plot_txt(fit_dict, band, rix)
+    ax3.text(0.025,0.05,restext,
+             bbox=dict(facecolor='wheat',
+                       alpha=0.5,
+                      boxstyle="round",),
+            transform=ax3.transAxes)
+    return fig, np.array([ax1, ax2, ax3])
+
+def plot_fit_summary(fit_dict, plot_style=None, quantile=0.98):
+    '''
+    Function for plotting full wafer eta_scan data from a tunefile.
+        Plots distributions of Qi, dip depth, bandwidth, and frequency spacing
+        for all resonators within specified quantile.
+
+    Args
+    ----
+    fit_dict : dict
+        fit results dictionary from ``fit_tune``
+    plot_style : dict
+        keyword arguments to pass to the histogram plotting for formatting.
+    quantile: float
+        inner X percent to graph
+    '''
+    
+    Qis = fit_dict['derived_params']['Qi']
+    depths = fit_dict['derived_params']['depth']
+    bws = fit_dict['derived_params']['br']
+    frs = fit_dict['model_params']['f_0']
+    seps = np.diff(frs)
+
+    Qi_quant = Qis[(Qis>=np.nanquantile(Qis, 1-quantile)) \
+                   & (Qis<=np.nanquantile(Qis, quantile))]
+    depth_quant = depths[(depths>=np.nanquantile(depths, 1-quantile)) \
+                   & (depths<=np.nanquantile(depths, quantile))]
+    bw_quant = bws[(bws>=np.nanquantile(bws, 1-quantile)) \
+                   & (bws<=np.nanquantile(bws, quantile))]
+    seps_quant = seps[(seps>=np.nanquantile(seps, 1-quantile)) \
+                   & (seps<=np.nanquantile(seps, quantile))]
+    
+    fig, axes = plt.subplots(2, 2, figsize=(16, 8))
+
+    if plot_style is None:
+        plot_style = {'bins': 30,
+                      'color': 'gold',
+                      'alpha': 0.5,
+                      'edgecolor': 'orange',
+                      'lw': 2}
+    #Qi plot
+    ax = axes[0, 0]
+    ax.hist(Qi_quant/1e5, **plot_style)
+    Qi_med = np.median(Qi_quant/1e5)
+    ax.axvline(np.median(Qi_quant/1e5),color = 'purple',
+                label = f'Median: {np.round(Qi_med,2)} $\\times10^5$')
+    ax.legend(loc = 'upper right')
+    ax.set_xlabel('$Q_i\\times10^5$')
+    ax.set_ylabel('Counts')
+
+    #Dip depth plot
+    ax = axes[0, 1]
+    ax.hist(20*np.log10(depth_quant),**plot_style)
+    dep_med = np.median(20*np.log10(depth_quant))
+    ax.axvline(dep_med,color = 'purple',
+                label = f'Median: {np.round(dep_med,2)} dB')
+    ax.legend(loc = 'upper right')
+    ax.set_xlabel('Dip Depth [dB]')
+    ax.set_ylabel('Counts')
+
+    #Bandwidth plot
+    ax = axes[1, 0]
+    ax.hist(bw_quant*1e3, **plot_style)
+    bw_med = np.median(bw_quant*1e3)
+    ax.axvline(bw_med,color = 'purple',
+                label = f'Median: {np.round(bw_med,2)} kHz')
+    ax.legend(loc = 'upper right')
+    ax.set_xlabel('Bandwidth [kHZ]')
+    ax.set_ylabel('Counts')
+
+    #Frequency Separation plot
+    ax = axes[1, 1]
+    mean, std = np.nanmean(seps_quant), np.std(seps_quant)
+    nstd = 2
+    rng = (
+        max(np.min(seps_quant), mean - nstd * std),
+        min(np.max(seps_quant), mean + nstd * std),
+    )
+
+    ax.hist(seps, range=rng, **plot_style)
+    sep_med = np.median(seps_quant)
+    ax.axvline(sep_med,color = 'purple',
+                label = f'Median: {np.round(sep_med,2)} MHz')
+    ax.legend(loc = 'upper right')
+    ax.set_xlabel('Resonator Separation [Mhz]')
+    ax.set_ylabel('Counts')
+
+    plt.tight_layout()
+    return fig, axes