# 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
