<commit_before>"""Functions for fitting hanger-style resonators.

These are resonators with an inverted Lorentzian lineshape (i.e. 0 on resonance, 1 off resonance).
"""

import warnings

import lmfit as lf
import numpy as np
import scipy.signal as sps
import scipy.special as spc


def hanger_fit(paramsVec, res, residual=True, **kwargs):
    """Return complex S21 of hanger resonator model or, if data is specified, a residual.

    Parameters
    ----------
    params : list-like
        A an ``lmfit.Parameters`` object containing (df, f0, qc, qi, gain0, gain1, gain2, pgain0, pgain1, pgain2)
    res : scraps.Resonator object
        A Resonator object.
    residual : bool
        Whether to return a residual (True) or to return the model calcuated at the frequencies present in res (False).

    Keyword Arguments
    -----------------
    freqs : list-like
        A list of frequency points at which to calculate the model. Only used if `residual=False`

    remove_baseline : bool
        Whether or not to remove the baseline during calculation (i.e. ignore pgain and gain polynomials). Default is False.

    only_baseline: bool
        Whether or not to calculate and return only the baseline. Default is False.

    Returns
    -------
    model or (model-data)/eps : ``numpy.array``
        If residual=True is specified, the return is the residuals weighted by the uncertainties. If residual=False, the return is the model
        values calculated at the frequency points. The returned array is in the form
        ``I + Q`` or ``residualI + residualQ``.

    """
    # Check if the paramsVec looks like a lmfit params object. If so, unpack to list
    if hasattr(paramsVec, "valuesdict"):
        paramsDict = paramsVec.valuesdict()
        paramsVec = [value for value in paramsDict.values()]

    # intrinsic resonator parameters
    df = paramsVec[0]  # frequency shift due to mismatched impedances
    f0 = paramsVec[1]  # resonant frequency
    qc = paramsVec[2]  # coupling Q
    qi = paramsVec[3]  # internal Q

    # 0th, 1st, and 2nd terms in a taylor series to handle magnitude gain different than 1
    gain0 = paramsVec[4]
    gain1 = paramsVec[5]
    gain2 = paramsVec[6]

    # 0th and 1st terms in a taylor series to handle phase gain different than 1
    pgain0 = paramsVec[7]
    pgain1 = paramsVec[8]

    # Need to add an if statement here to keep compatibility with older version
    # Will be deprecated in the future
    if len(paramsVec) == 12:
        pgain2 = paramsVec[9]

        # Voltage offset at mixer output. Not needed for VNA
        Ioffset = paramsVec[10]
        Qoffset = paramsVec[11]
    else:
        warnings.warn(
            "Warning: new model also fits for quadratic phase. Setting pgain2 = 0. \
                    \nIf using Resonator.do_lmfit() pass kwarg: pgain2_vary = False for legacy behavior."
        )
        pgain2 = 0

        # Voltage offset at mixer output. Not needed for VNA
        Ioffset = paramsVec[9]
        Qoffset = paramsVec[10]

    if residual == False:
        # Use a custom vector of frequencies
        freqs = np.array(kwargs.pop("freqs", res.freq))
    else:
        # Grab a shortcut to the resonant frequency
        freqs = res.freq

    # Repackage resonator data in 1D vector form
    data = np.concatenate((res.I, res.Q), axis=0)

    if (res.sigmaI is not None) and (res.sigmaQ is not None):
        cmplxSigma = np.concatenate((res.sigmaI, res.sigmaQ), axis=0)
    else:
        cmplxSigma = None

    # Make everything referenced to the shifted, unitless, reduced frequency
    fs = f0 + df
    ff = (freqs - fs) / fs

    # Except for the gain, which should reference the file midpoint
    # This is important because the baseline coefs shouldn't drift
    # around with changes in f0 due to power or temperature

    # Of course, this philosophy goes out the window if different sweeps have
    # different ranges.
    fm = freqs[int(np.round((len(freqs) - 1) / 2.0))]
    ffm = (freqs - fm) / fm

    # Calculate the total Q_0
    q0 = 1.0 / (1.0 / qi + 1.0 / qc)

    # Calculate magnitude and phase gain
    gain = gain0 + gain1 * ffm + 0.5 * gain2 * ffm ** 2
    pgain = np.exp(1j * (pgain0 + pgain1 * ffm + 0.5 * pgain2 * ffm ** 2))

    # Allow for voltage offset of I and Q
    offset = Ioffset + 1j * Qoffset

    remove_baseline = kwargs.pop("remove_baseline", False)

    if remove_baseline:
        # Set the baseline vector to a bunch of ones, with zero imaginary part
        total_gain = np.ones_like(freqs, dtype=np.complex)
    else:
        # Actually calculate the baseline
        total_gain = gain * pgain

    only_baseline = kwargs.pop("only_baseline", False)

    if only_baseline:
        # Only return the baseline
        modelCmplx = total_gain
    else:
        # Calculate model from params at each point in freqs
        modelCmplx = (
            total_gain * (1.0 / qi + 1j * 2.0 * (ff + df / fs)) / (1.0 / q0 + 1j * 2.0 * ff)
            + offset
        )

    # Package complex data in 1D vector form
    modelI = np.real(modelCmplx)
    modelQ = np.imag(modelCmplx)
    model = np.concatenate((modelI, modelQ), axis=0)

    # Calculate eps from stdev of first 10 pts of data if not supplied
    if cmplxSigma is None and residual == True:
        epsI = np.std(sps.detrend(res.I[0:10]))
        epsQ = np.std(sps.detrend(res.Q[0:10]))
        cmplxSigma = np.concatenate((np.full_like(res.I, epsI), np.full_like(res.Q, epsQ)))

    # Return model or residual
    if residual == True:
        return (model - data) / cmplxSigma
    else:
        return model


def hanger_params(res, **kwargs):
    """Initialize fitting parameters used by the hanger_fit function.

    Parameters
    ----------
    res : ``scraps.Resonator`` object
        The object you want to calculate parameter guesses for.
<commit_msg>Slight refactor to make the fundamental model available as an import<commit_after>"""Functions for fitting hanger-style resonators.

These are resonators with an inverted Lorentzian lineshape (i.e. 0 on resonance, 1 off resonance).
"""

import warnings

import lmfit as lf
import numpy as np
import scipy.signal as sps
import scipy.special as spc


def cmplx_hanger(freqs, f0, df, qc, q0):
    """Assymetric function that fits a hanger or notch type inverted Lorentzian."""
    ft = f0 + df
    x = (freqs - ft) / ft
    dx = df / ft

    qi = q0 * qc / (qc - q0)

    s21 = (1.0 / qi + 1j * 2.0 * (x + dx)) / (1.0 / q0 + 1j * 2.0 * x)
    return s21


def hanger_fit(paramsVec, res, residual=True, **kwargs):
    """Return complex S21 of hanger resonator model or, if data is specified, a residual.

    Parameters
    ----------
    params : list-like
        A an ``lmfit.Parameters`` object containing (df, f0, qc, qi, gain0, gain1, gain2, pgain0, pgain1, pgain2)
    res : scraps.Resonator object
        A Resonator object.
    residual : bool
        Whether to return a residual (True) or to return the model calcuated at the frequencies present in res (False).

    Keyword Arguments
    -----------------
    freqs : list-like
        A list of frequency points at which to calculate the model. Only used if `residual=False`

    remove_baseline : bool
        Whether or not to remove the baseline during calculation (i.e. ignore pgain and gain polynomials). Default is False.

    only_baseline: bool
        Whether or not to calculate and return only the baseline. Default is False.

    Returns
    -------
    model or (model-data)/eps : ``numpy.array``
        If residual=True is specified, the return is the residuals weighted by the uncertainties. If residual=False, the return is the model
        values calculated at the frequency points. The returned array is in the form
        ``I + Q`` or ``residualI + residualQ``.

    """
    # Check if the paramsVec looks like a lmfit params object. If so, unpack to list
    if hasattr(paramsVec, "valuesdict"):
        paramsDict = paramsVec.valuesdict()
        paramsVec = [value for value in paramsDict.values()]

    # intrinsic resonator parameters
    df = paramsVec[0]  # frequency shift due to mismatched impedances
    f0 = paramsVec[1]  # resonant frequency
    qc = paramsVec[2]  # coupling Q
    qi = paramsVec[3]  # internal Q

    # 0th, 1st, and 2nd terms in a taylor series to handle magnitude gain different than 1
    gain0 = paramsVec[4]
    gain1 = paramsVec[5]
    gain2 = paramsVec[6]

    # 0th and 1st terms in a taylor series to handle phase gain different than 1
    pgain0 = paramsVec[7]
    pgain1 = paramsVec[8]

    # Need to add an if statement here to keep compatibility with older version
    # Will be deprecated in the future
    if len(paramsVec) == 12:
        pgain2 = paramsVec[9]

        # Voltage offset at mixer output. Not needed for VNA
        Ioffset = paramsVec[10]
        Qoffset = paramsVec[11]
    else:
        warnings.warn(
            "Warning: new model also fits for quadratic phase. Setting pgain2 = 0. \
                    \nIf using Resonator.do_lmfit() pass kwarg: pgain2_vary = False for legacy behavior."
        )
        pgain2 = 0

        # Voltage offset at mixer output. Not needed for VNA
        Ioffset = paramsVec[9]
        Qoffset = paramsVec[10]

    if residual == False:
        # Use a custom vector of frequencies
        freqs = np.array(kwargs.pop("freqs", res.freq))
    else:
        # Grab a shortcut to the resonant frequency
        freqs = res.freq

    # Repackage resonator data in 1D vector form
    data = np.concatenate((res.I, res.Q), axis=0)

    if (res.sigmaI is not None) and (res.sigmaQ is not None):
        cmplxSigma = np.concatenate((res.sigmaI, res.sigmaQ), axis=0)
    else:
        cmplxSigma = None

    # Gain should reference the file midpoint
    # This is important because the baseline coefs shouldn't drift
    # around with changes in f0 due to power or temperature

    # Of course, this philosophy goes out the window if different sweeps have
    # different ranges.
    fm = freqs[int(np.round((len(freqs) - 1) / 2.0))]
    ffm = (freqs - fm) / fm

    # Calculate the total Q_0
    q0 = 1.0 / (1.0 / qi + 1.0 / qc)

    # Calculate magnitude and phase gain
    gain = gain0 + gain1 * ffm + 0.5 * gain2 * ffm ** 2
    pgain = np.exp(1j * (pgain0 + pgain1 * ffm + 0.5 * pgain2 * ffm ** 2))

    # Allow for voltage offset of I and Q
    offset = Ioffset + 1j * Qoffset

    remove_baseline = kwargs.pop("remove_baseline", False)

    if remove_baseline:
        # Set the baseline vector to a bunch of ones, with zero imaginary part
        total_gain = np.ones_like(freqs, dtype=np.complex)
    else:
        # Actually calculate the baseline
        total_gain = gain * pgain

    only_baseline = kwargs.pop("only_baseline", False)

    if only_baseline:
        # Only return the baseline
        modelCmplx = total_gain
    else:
        # Calculate model from params at each point in freqs
        modelCmplx = total_gain * cmplx_hanger(freqs, f0, df, qc, q0) + offset

    # Package complex data in 1D vector form
    modelI = np.real(modelCmplx)
    modelQ = np.imag(modelCmplx)
    model = np.concatenate((modelI, modelQ), axis=0)

    # Calculate eps from stdev of first 10 pts of data if not supplied
    if cmplxSigma is None and residual == True:
        epsI = np.std(sps.detrend(res.I[0:10]))
        epsQ = np.std(sps.detrend(res.Q[0:10]))
        cmplxSigma = np.concatenate((np.full_like(res.I, epsI), np.full_like(res.Q, epsQ)))

    # Return model or residual
    if residual == True:
        return (model - data) / cmplxSigma
    else:
        return model


def hanger_params(res, **kwargs):
    """Initialize fitting parameters used by the hanger_fit function.

    Parameters
    ----------
    res : ``scraps.Resonator`` object
        The object you want to calculate parameter guesses for.

{'fraction_non_alphanumeric': 0.09073277679399612, 'fraction_numerical': 0.01770806982038958, 'mean_word_length': 2.4892615475139745, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7538379', 'n_tokens_mistral': 3836, 'n_tokens_neox': 3558, 'n_words': 1373}