Source code for quantify_core.analysis.fitting_models

# Repository: https://gitlab.com/quantify-os/quantify-core
# Licensed according to the LICENCE file on the main branch
"""Models and fit functions to be used with the lmfit fitting framework."""
from __future__ import annotations

from typing import Tuple

import lmfit
import numpy as np
from numpy.typing import NDArray


[docs] def get_model_common_doc() -> str: """ Returns a common docstring to be used with custom fitting :class:`~lmfit.model.Model` s. .. admonition:: Usage example for a custom fitting model :class: dropdown, tip See the usage example at the end of the :class:`~ResonatorModel` source-code: .. literalinclude:: ../../../quantify_core/analysis/fitting_models.py :pyobject: ResonatorModel """ return ( lmfit.models.COMMON_INIT_DOC.replace( ":class:`Model`", ":class:`~lmfit.model.Model`" ) .replace("\n ", "\n") .replace(" : str", " : :obj:`str`") .replace("['x']", ":code:`['x']`") .replace(", optional", "") .replace(" optional", "") .replace("{'raise', 'propagate', 'omit'}", "") )
[docs] def get_guess_common_doc() -> str: """ Returns a common docstring to be used for the :meth:`~lmfit.model.Model.guess` method of custom fitting :class:`~lmfit.model.Model` s. .. admonition:: Usage example for a custom fitting model :class: dropdown, tip See the usage example at the end of the :class:`~ResonatorModel` source-code: .. literalinclude:: ../../../quantify_core/analysis/fitting_models.py :pyobject: ResonatorModel """ return ( lmfit.models.COMMON_GUESS_DOC.replace( " : Parameters", " : :class:`~lmfit.parameter.Parameters`" ) .replace("\n ", "\n") .replace(" optional", "") .replace(" : array_like", " : :class:`~numpy.ndarray`") )
[docs] def mk_seealso( function_name: str, role: str = "func", prefix: str = "\n\n", module_location: str = ".", ) -> str: """ Returns a sphinx `seealso` pointing to a function. Intended to be used for building custom fitting model docstrings. .. admonition:: Usage example for a custom fitting model :class: dropdown, tip See the usage example at the end of the :class:`~ResonatorModel` source-code: .. literalinclude:: ../../../quantify_core/analysis/fitting_models.py :pyobject: ResonatorModel Parameters ---------- function_name name of the function to point to role a sphinx role, e.g. :code:`"func"` prefix string preceding the `seealso` module_location can be used to indicate a function outside this module, e.g., :code:`my_module.submodule` which contains the function. Returns ------- : resulting string """ return f"{prefix}.. seealso:: :{role}:`~{module_location}{function_name}`\n"
# pylint: disable=too-many-arguments # pylint: disable=invalid-name
[docs] def hanger_func_complex_SI( f: float, fr: float, Ql: float, Qe: float, A: float, theta: float, phi_v: float, phi_0: float, alpha: float = 1, ) -> complex: r""" This is the complex function for a hanger (lambda/4 resonator). Parameters ---------- f: frequency fr: resonance frequency A: background transmission amplitude Ql: loaded quality factor of the resonator Qe: magnitude of extrinsic quality factor :code:`Qe = |Q_extrinsic|` theta: phase of extrinsic quality factor (in rad) phi_v: phase to account for propagation delay to sample phi_0: phase to account for propagation delay from sample alpha: slope of signal around the resonance Returns ------- : complex valued transmission See eq. S4 from Bruno et al. (2015) `ArXiv:1502.04082 <https://arxiv.org/abs/1502.04082>`_. .. math:: S_{21} = A \left(1+\alpha \frac{f-f_r}{f_r} \right) \left(1- \frac{\frac{Q_l}{|Q_e|}e^{i\theta} }{1+2iQ_l \frac{f-f_r}{f_r}} \right) e^{i (\phi_v f + \phi_0)} The loaded and extrinsic quality factors are related to the internal and coupled Q according to: .. math:: \frac{1}{Q_l} = \frac{1}{Q_c}+\frac{1}{Q_i} and .. math:: \frac{1}{Q_c} = \mathrm{Re}\left(\frac{1}{|Q_e|e^{-i\theta}}\right) """ slope_corr = 1 + alpha * (f - fr) / fr hanger_contribution = 1 - Ql / Qe * np.exp(1j * theta) / ( 1 + 2.0j * Ql * (f - fr) / fr ) propagation_delay_corr = np.exp(1j * (phi_v * f + phi_0)) S21 = A * slope_corr * hanger_contribution * propagation_delay_corr return S21
[docs] def cos_func( x: float, frequency: float, amplitude: float, offset: float, phase: float = 0, ) -> float: r""" An oscillating cosine function: :math:`y = \mathrm{amplitude} \times \cos(2 \pi \times \mathrm{frequency} \times x + \mathrm{phase}) + \mathrm{offset}` Parameters ---------- x: The independent variable (time, for example) frequency: A generalized frequency (in units of inverse x) amplitude: Amplitude of the oscillation offset: Output signal vertical offset phase: Phase offset / rad Returns ------- : Output signal magnitude """ # pylint: disable=line-too-long return amplitude * np.cos(2 * np.pi * frequency * x + phase) + offset
[docs] def lorentzian_func( x: float, x0: float, width: float, a: float, c: float, ) -> float: r""" A Lorentzian function. .. math:: y = \frac{a*\mathrm{width}}{\pi(\mathrm{width}^2 + (x - x_0)^2)} + c Parameters ---------- x: independent variable x0: horizontal offset width: Lorenztian linewidth a: amplitude c: vertical offset Returns ------- : Lorentzian function """ return a * width / (np.pi * ((x - x0) ** 2) + width**2) + c
[docs] def exp_decay_func( t: float, tau: float, amplitude: float, offset: float, n_factor: float, ) -> float: r""" This is a general exponential decay function: :math:`y = \mathrm{amplitude} \times \exp\left(-(t/\tau)^\mathrm{n\_factor}\right) + \mathrm{offset}` Parameters ---------- t: time tau: decay time amplitude: amplitude of the exponential decay offset: asymptote of the exponential decay, the value at t=infinity n_factor: exponential decay factor Returns ------- : Output of exponential function as a float """ # pylint: disable=line-too-long return amplitude * np.exp(-((t / tau) ** n_factor)) + offset
[docs] def exp_damp_osc_func( t: float, tau: float, n_factor: float, frequency: float, phase: float, amplitude: float, offset: float, ): r""" A sinusoidal oscillation with an exponentially decaying envelope function: :math:`y = \mathrm{amplitude} \times \exp\left(-(t/\tau)^\mathrm{n\_factor}\right)(\cos(2\pi\mathrm{frequency}\times t + \mathrm{phase}) + \mathrm{oscillation_offset}) + \mathrm{exponential_offset}` Parameters ---------- t: time tau: decay time n_factor: exponential decay factor frequency: frequency of the oscillation phase: phase of the oscillation amplitude: initial amplitude of the oscillation oscillation_offset: vertical offset of cosine oscillation relative to exponential asymptote exponential_offset: offset of exponential asymptote Returns ------- : Output of decaying cosine function as a float """ # pylint: disable=line-too-long oscillation = amplitude * (np.cos(2 * np.pi * frequency * t + phase)) exp_decay = np.exp(-((t / tau) ** n_factor)) osc_decay = oscillation * exp_decay + offset return osc_decay
[docs] class LorentzianModel(lmfit.model.Model): """ Model for data which follows a Lorentzian function. Uses the function :func:`~lorentzian_func` as the defining equation. """ # pylint: disable=empty-docstring # pylint: disable=abstract-method # pylint: disable=too-few-public-methods def __init__(self, *args, **kwargs) -> None: # pass in the defining equation so the user doesn't have to later. super().__init__(lorentzian_func, *args, **kwargs) self.set_param_hint("x0", vary=True) self.set_param_hint("a", vary=True) self.set_param_hint("c", vary=True) self.set_param_hint("width", vary=True) # pylint: disable=missing-function-docstring
[docs] def guess(self, data: NDArray, **kws) -> lmfit.parameter.Parameters: """Guess some initial values for the model based on the data.""" x = kws.get("x") if x is None: return None # type: ignore # Guess that the resonance is where the function takes its maximal # value x0_guess = x[np.argmax(data)] self.set_param_hint("x0", value=x0_guess) # assume the user isn't trying to fit just a small part of a resonance curve. xmin = x.min() xmax = x.max() width_max = xmax - xmin delta_x = np.diff(x) # assume f is sorted min_delta_x = delta_x[delta_x > 0].min() # assume data actually samples the resonance reasonably width_min = min_delta_x width_guess = np.sqrt(width_min * width_max) # geometric mean, why not? self.set_param_hint("width", value=width_guess) # The guess for the vertical offset is the mean absolute value of the data c_guess = np.mean(data) self.set_param_hint("c", value=c_guess) # Calculate A_guess from difference between the peak and the backround level a_guess = np.pi * width_guess * (np.max(data) - c_guess) self.set_param_hint("a", value=a_guess) params = self.make_params() return lmfit.models.update_param_vals(params, self.prefix, **kws)
# This class is used a literal include in the docs so the pylint options are here # pylint: disable=empty-docstring # pylint: disable=abstract-method # pylint: disable=too-few-public-methods
[docs] class CosineModel(lmfit.model.Model): """ Exemplary lmfit model with a guess for a cosine. .. note:: The :mod:`lmfit.models` module provides several fitting models that might fit your needs out of the box. """
[docs] def __init__(self, *args, **kwargs): # pass in the model's equation super().__init__(cos_func, *args, **kwargs) # configure constraints that are independent from the data to be fitted self.set_param_hint("frequency", min=0, vary=True) # enforce positive frequency self.set_param_hint("amplitude", min=0, vary=True) # enforce positive amplitude self.set_param_hint("offset", vary=True) self.set_param_hint( "phase", vary=True, min=-np.pi, max=np.pi ) # enforce phase range
# pylint: disable=missing-function-docstring
[docs] def guess(self, data, x, **kws) -> lmfit.parameter.Parameters: """ guess parameters based on the data Parameters ---------- data: np.ndarray Data to fit to x: np.ndarray Independet variable """ self.set_param_hint("offset", value=np.average(data)) self.set_param_hint("amplitude", value=(np.max(data) - np.min(data)) / 2) # Guess frequency and phase using Fourier Transform freq_guess, phase_guess = fft_freq_phase_guess(data, x) phase_wrap = (phase_guess + np.pi) % (2 * np.pi) - np.pi self.set_param_hint("frequency", value=freq_guess) self.set_param_hint("phase", value=phase_wrap) params = self.make_params() return lmfit.models.update_param_vals(params, self.prefix, **kws)
# Same design patter is used in lmfit.models to inherit common docstrings. # We adjust these common docstrings to our docs build pipeline __init__.__doc__ = get_model_common_doc() + mk_seealso("cos_func") guess.__doc__ = get_guess_common_doc()
[docs] class ResonatorModel(lmfit.model.Model): """ Resonator model Implementation and design patterns inspired by the `complex resonator model example <https://lmfit.github.io/lmfit-py/examples/example_complex_resonator_model.html>`_ (`lmfit` documentation). """ # pylint: disable=line-too-long # pylint: disable=empty-docstring # pylint: disable=abstract-method
[docs] def __init__(self, *args, **kwargs): # pass in the defining equation so the user doesn't have to later. super().__init__(hanger_func_complex_SI, *args, **kwargs) self.set_param_hint("Ql", min=0) # Enforce Q is positive self.set_param_hint("Qe", min=0) # Enforce Q is positive # Internal and coupled quality factor can be derived from fitted params self.set_param_hint("Qi", expr="1./(1./Ql-1./Qe*cos(theta))", vary=False) self.set_param_hint("Qc", expr="Qe/cos(theta)", vary=False)
# pylint: disable=too-many-locals # pylint: disable=missing-function-docstring
[docs] def guess(self, data, **kws) -> lmfit.parameter.Parameters: f = kws.get("f") if f is None: return None argmin_s21 = np.abs(data).argmin() fmin = f.min() fmax = f.max() # guess that the resonance is the lowest point fr_guess = f[argmin_s21] # assume the user isn't trying to fit just a small part of a resonance curve. Q_min = 0.1 * (fr_guess / (fmax - fmin)) delta_f = np.diff(f) # assume f is sorted min_delta_f = delta_f[delta_f > 0].min() Q_max = ( fr_guess / min_delta_f ) # assume data actually samples the resonance reasonably Q_guess = np.sqrt(Q_min * Q_max) # geometric mean, why not? (phi_0_guess, phi_v_guess) = resonator_phase_guess( data, f ) # Come up with a guess for phase velocity self.set_param_hint("fr", value=fr_guess, min=fmin, max=fmax) self.set_param_hint("Ql", value=Q_guess * 1.01, min=Q_min, max=Q_max) self.set_param_hint("Qe", value=Q_guess * 0.99, min=0) self.set_param_hint("A", value=np.mean(abs(data)), min=0) # The parameters below need a proper guess. self.set_param_hint("theta", value=0, min=-np.pi / 2, max=np.pi / 2) self.set_param_hint("phi_0", value=phi_0_guess) self.set_param_hint("phi_v", value=phi_v_guess) self.set_param_hint("alpha", value=0, min=-1, max=1) params = self.make_params() return lmfit.models.update_param_vals(params, self.prefix, **kws)
# Same design patter is used in lmfit.models __init__.__doc__ = get_model_common_doc() + mk_seealso("hanger_func_complex_SI") guess.__doc__ = get_guess_common_doc()
[docs] class ExpDecayModel(lmfit.model.Model): """ Model for an exponential decay, such as a qubit T1 measurement. """ # pylint: disable=empty-docstring # pylint: disable=abstract-method # pylint: disable=too-few-public-methods
[docs] def __init__(self, *args, **kwargs): # pass in the defining equation so the user doesn't have to later. super().__init__(exp_decay_func, *args, **kwargs) self.set_param_hint("tau", min=0) # Enforce T1 is positive self.set_param_hint("amplitude", vary=True) self.set_param_hint("offset", vary=True) self.set_param_hint("n_factor", expr="1", vary=False)
# pylint: disable=missing-function-docstring
[docs] def guess(self, data, **kws) -> lmfit.parameter.Parameters: delay = kws.get("delay") if delay is None: return None # To guess the upper amplitude and offset, # use the first and last values of the data self.set_param_hint("offset", value=data[-1]) self.set_param_hint("amplitude", value=data[0] - data[-1]) # The guess for tau is somewhere in the middle of the time range tau = np.median(delay) self.set_param_hint("tau", value=tau, min=0) params = self.make_params() return lmfit.models.update_param_vals(params, self.prefix, **kws)
# Same design patter is used in lmfit.models __init__.__doc__ = get_model_common_doc() + mk_seealso("exp_decay_func") guess.__doc__ = get_guess_common_doc()
[docs] class RabiModel(lmfit.model.Model): r""" Model for a Rabi oscillation as a function of the microwave drive amplitude. Phase of oscillation is fixed at :math:`\pi` in order to ensure that the oscillation is at a minimum when the drive amplitude is 0. """ # pylint: disable=empty-docstring # pylint: disable=abstract-method # pylint: disable=too-few-public-methods
[docs] def __init__(self, *args, **kwargs): # pass in the defining equation so the user doesn't have to later. super().__init__(cos_func, *args, **kwargs) # Enforce oscillation frequency is positive self.set_param_hint("frequency", min=0) # Fix the phase at pi so that the ouput is at a minimum when x=0 self.set_param_hint("phase", expr="3.141592653589793", vary=False) # Pi-pulse amplitude can be derived from the oscillation frequency self.set_param_hint("amp180", expr="1/(2*frequency)", vary=False)
# pylint: disable=missing-function-docstring
[docs] def guess(self, data, **kws) -> lmfit.parameter.Parameters: drive_amp = kws.get("drive_amp") if drive_amp is None: return None amp_guess = abs(max(data) - min(data)) / 2 # amp is positive by convention offs_guess = np.mean(data) (freq_guess, _) = fft_freq_phase_guess(data, drive_amp) self.set_param_hint("frequency", value=freq_guess, min=0) self.set_param_hint("amplitude", value=amp_guess, min=0) self.set_param_hint("offset", value=offs_guess) params = self.make_params() return lmfit.models.update_param_vals(params, self.prefix, **kws)
# Same design patter is used in lmfit.models __init__.__doc__ = get_model_common_doc() + mk_seealso("cos_func") guess.__doc__ = get_guess_common_doc()
[docs] class DecayOscillationModel(lmfit.model.Model): r""" Model for a decaying oscillation which decays to a point with 0 offset from the centre of the of the oscillation (as in a Ramsey experiment, for example). """ # pylint: disable=empty-docstring # pylint: disable=abstract-method # pylint: disable=too-few-public-methods
[docs] def __init__(self, *args, **kwargs): # pass in the defining equation so the user doesn't have to later. super().__init__(exp_damp_osc_func, *args, **kwargs) # Enforce oscillation frequency is positive self.set_param_hint("frequency", min=0) # Enforce amplitude is positive self.set_param_hint("amplitude", min=0) # Enforce decay time is positive self.set_param_hint("tau", min=0) # Fix the n_factor at 1 self.set_param_hint("n_factor", expr="1", vary=False)
# pylint: disable=missing-function-docstring
[docs] def guess(self, data, **kws) -> lmfit.parameter.Parameters: t = kws.get("t") if t is None: raise ValueError( 'Time variable "t" must be specified in order to guess parameters' ) return None amp_guess = abs(max(data) - min(data)) / 2 # amp is positive by convention exp_offs_guess = np.mean(data) tau_guess = 2 / 3 * np.max(t) (freq_guess, phase_guess) = fft_freq_phase_guess(data, t) self.set_param_hint("frequency", value=freq_guess, min=0) self.set_param_hint("amplitude", value=amp_guess, min=0) self.set_param_hint("offset", value=exp_offs_guess) self.set_param_hint("phase", value=phase_guess) self.set_param_hint("tau", value=tau_guess, min=0) params = self.make_params() return lmfit.models.update_param_vals(params, self.prefix, **kws)
# Same design patter is used in lmfit.models __init__.__doc__ = get_model_common_doc() + mk_seealso("exp_damp_osc_func") guess.__doc__ = get_guess_common_doc()
[docs] def resonator_phase_guess(s21: np.ndarray, freq: np.ndarray) -> Tuple[float, float]: """ Guesses the phase velocity in resonator spectroscopy, based on the median of all the differences between consecutive phases. Parameters ---------- s21: Resonator S21 data freq: Frequency of the spectroscopy pulse Returns ------- phi_0: Guess for the phase offset phi_v: Guess for the phase velocity """ phase = np.angle(s21) med_diff = np.median(np.diff(phase)) freq_step = np.median(np.diff(freq)) phi_v = med_diff / freq_step phi_0 = phase[0] - phi_v * freq[0] return phi_0, phi_v
[docs] def fft_freq_phase_guess(data: np.ndarray, t: np.ndarray) -> Tuple[float, float]: """ Guess for a cosine fit using FFT, only works for evenly spaced points. Parameters ---------- data: Input data to FFT t: Independent variable (e.g. time) Returns ------- freq_guess: Guess for the frequency of the cosine function ph_guess: Guess for the phase of the cosine function """ # Only first half of array is used, because the second half contains the # negative frequecy components, and we want a positive frequency. power = np.fft.fft(data)[: len(data) // 2] freq = np.fft.fftfreq(len(data), t[1] - t[0])[: len(power)] power[0] = 0 # Removes DC component from fourier transform # Use absolute value of complex valued spectrum abs_power = np.abs(power) freq_guess = abs(freq[abs_power == max(abs_power)][0]) # the condition data == max(data) can have several solutions # (for example when discretization is visible) # to prevent errors we pick the first solution ph_guess = 2 * np.pi - (2 * np.pi * t[data == max(data)] * freq_guess)[0] return freq_guess, ph_guess