The HEOM method solves the dynamics and steady state of a system and its environment, the latter of which is encoded in a set of auxiliary density matrices.
In this example we show the evolution of a single two-level system in contact with a single bosonic environment.
The properties of the system are encoded in Hamiltonian, and a coupling operator which describes how it is coupled to the environment.
The bosonic environment is implicitly assumed to obey a particular Hamiltonian (see paper), the parameters of which are encoded in the spectral density, and subsequently the free-bath correlation functions.
In the example below we show how to model an Ohmic environment with exponential cut-off in two ways:
First we fit the spectral density with a set of underdamped brownian oscillator functions.
Second, we evaluate the correlation functions, and fit those with a certain choice of exponential functions.
In each case we will use the fit parameters to determine the correlation function expansion co-efficients needed to construct a description of the bath (i.e. a BosonicBath
object) to supply to the HEOMSolver
so that we can solve for the system dynamics.
import contextlib
import dataclasses
import time
import numpy as np
from matplotlib import pyplot as plt
from scipy.optimize import curve_fit
import qutip
from qutip import (
basis,
expect,
liouvillian,
sigmax,
sigmaz,
spost,
spre,
)
from qutip.solver.heom import (
HEOMSolver,
BosonicBath,
)
# Import mpmath functions for evaluation of gamma and zeta
# functions in the expression for the correlation:
from mpmath import mp
mp.dps = 15
mp.pretty = True
%matplotlib inline
Let's define some helper functions for plotting results and timing how long operations take:
def coth(x):
""" Vectorized hyperbolic cotangent of x. """
return 1. / np.tanh(x)
def plot_result_expectations(plots, axes=None):
""" Plot the expectation values of operators as functions of time.
Each plot in plots consists of (solver_result,
measurement_operation, color, label).
"""
if axes is None:
fig, axes = plt.subplots(1, 1, sharex=True, figsize=(8, 8))
fig_created = True
else:
fig = None
fig_created = False
# add kw arguments to each plot if missing
plots = [p if len(p) == 5 else p + ({},) for p in plots]
for result, m_op, color, label, kw in plots:
exp = np.real(expect(result.states, m_op))
kw.setdefault("linewidth", 2)
if color == 'rand':
axes.plot(
result.times, exp,
c=np.random.rand(3,), label=label, **kw,
)
else:
axes.plot(result.times, exp, color, label=label, **kw)
if fig_created:
axes.legend(loc=0, fontsize=12)
axes.set_xlabel("t", fontsize=28)
return fig
@contextlib.contextmanager
def timer(label):
""" Simple utility for timing functions:
with timer("name"):
... code to time ...
"""
start = time.time()
yield
end = time.time()
print(f"{label}: {end - start}")
# Solver options:
options = {
"nsteps": 15000,
"store_states": True,
"rtol": 1e-14,
"atol": 1e-14,
"method": "vern9",
"progress_bar": "enhanced",
}
And let us set up the system Hamiltonian, bath and system measurement operators:
# Defining the system Hamiltonian
eps = 0.0 # Energy of the 2-level system.
Del = 0.2 # Tunnelling term
Hsys = 0.5 * eps * sigmaz() + 0.5 * Del * sigmax()
# Initial state of the system.
rho0 = basis(2, 0) * basis(2, 0).dag()
# Define some operators with which we will measure the system
# 1,1 element of density matrix - corresonding to groundstate
P11p = basis(2, 0) * basis(2, 0).dag()
P22p = basis(2, 1) * basis(2, 1).dag()
# 1,2 element of density matrix - corresonding to coherence
P12p = basis(2, 0) * basis(2, 1).dag()
Before we begin fitting, let us examine the analytic expressions for the correlation and spectral density functions and write Python equivalents.
The correlation function is given by (see, e.g., http://www1.itp.tu-berlin.de/brandes/public_html/publications/notes.pdf for a derivation, equation 7.59, but with a factor of $\pi$ moved into the definition of the correlation function):
\begin{align} C(t) =& \: \frac{1}{\pi}\alpha \omega_{c}^{1 - s} \beta^{- (s + 1)} \: \times \\ & \: \Gamma(s + 1) \left[ \zeta \left(s + 1, \frac{1 + \beta \omega_c - i \omega_c t}{\beta \omega_c}\right) + \zeta \left(s + 1, \frac{1 + i \omega_c t}{\beta \omega_c}\right) \right] \end{align}where $\Gamma$ is the Gamma function and
\begin{equation} \zeta(z, u) \equiv \sum_{n=0}^{\infty} \frac{1}{(n + u)^z}, \; u \neq 0, -1, -2, \ldots \end{equation}is the generalized Zeta function. The Ohmic case is given by $s = 1$.
The corresponding spectral density for the Ohmic case is:
\begin{equation} J(\omega) = \omega \alpha e^{- \frac{\omega}{\omega_c}} \end{equation}def ohmic_correlation(t, alpha, wc, beta, s=1):
""" The Ohmic bath correlation function as a function of t
(and the bath parameters).
"""
corr = (
(1 / np.pi) * alpha * wc**(1 - s) * beta**(-(s + 1)) * mp.gamma(s + 1)
)
z1_u = (1 + beta * wc - 1.0j * wc * t) / (beta * wc)
z2_u = (1 + 1.0j * wc * t) / (beta * wc)
# Note: the arguments to zeta should be in as high precision as possible.
# See http://mpmath.org/doc/current/basics.html#providing-correct-input
return np.array([
complex(corr * (mp.zeta(s + 1, u1) + mp.zeta(s + 1, u2)))
for u1, u2 in zip(z1_u, z2_u)
], dtype=np.complex128)
def ohmic_spectral_density(w, alpha, wc):
""" The Ohmic bath spectral density as a function of w
(and the bath parameters).
"""
return w * alpha * np.e**(-w / wc)
def ohmic_power_spectrum(w, alpha, wc, beta):
""" The Ohmic bath power spectrum as a function of w
(and the bath parameters).
"""
return (
w * alpha * np.e**(-abs(w) / wc) *
((1 / (np.e**(w * beta) - 1)) + 1) * 2
)
Finally, let's set the bath parameters we will work with and write down some measurement operators:
# Bath parameters:
@dataclasses.dataclass
class OhmicBathParameters:
""" Ohmic bath parameters. """
Q: object = dataclasses.field(default_factory=sigmaz, repr=False)
alpha: float = 3.25
T: float = 0.5
wc: float = 1.0
s: float = 1
def __post_init__(self):
self.beta = 1 / self.T
def replace(self, **kw):
return dataclasses.replace(self, **kw)
obp = OhmicBathParameters()
And set the cut-off for the HEOM hierarchy:
# HEOM parameters:
# The max_depth defaults to 5 so that the notebook executes more
# quickly. Change it to 11 to wait longer for more accurate results.
max_depth = 5
We begin by fitting the spectral density, using a series of $k$ underdamped harmonic oscillators case with the Meier-Tannor form (J. Chem. Phys. 111, 3365 (1999); https://doi.org/10.1063/1.479669):
\begin{equation} J_{\mathrm approx}(\omega; a, b, c) = \sum_{i=0}^{k-1} \frac{2 a_i b_i w}{((w + c_i)^2 + b_i^2) ((w - c_i)^2 + b_i^2)} \end{equation}where $a, b$ and $c$ are the fit parameters and each is a vector of length $k$.
# Helper functions for packing the paramters a, b and c into a single numpy
# array as required by SciPy's curve_fit:
def pack(a, b, c):
""" Pack parameter lists for fitting. """
return np.concatenate((a, b, c))
def unpack(params):
""" Unpack parameter lists for fitting. """
N = len(params) // 3
a = np.array(params[:N])
b = np.array(params[N:2 * N])
c = np.array(params[2 * N:])
return a, b, c
# The approximate spectral density and a helper for fitting the approximate
# spectral density to values calculated from the analytical formula:
def spectral_density_approx(w, a, b, c):
""" Calculate the fitted value of the function for the given
parameters.
"""
return np.sum(
2 * a[:, None] * np.multiply.outer(b, w) / (
((w + c[:, None])**2 + b[:, None]**2) *
((w - c[:, None])**2 + b[:, None]**2)
),
axis=0,
)
def fit_spectral_density(J, w, alpha, wc, N):
""" Fit the spectral density with N underdamped oscillators. """
sigma = [0.0001] * len(w)
J_max = abs(max(J, key=abs))
guesses = pack([J_max] * N, [wc] * N, [wc] * N)
lower_bounds = pack([-100 * J_max] * N, [0.1 * wc] * N, [0.1 * wc] * N)
upper_bounds = pack([100 * J_max] * N, [100 * wc] * N, [100 * wc] * N)
params, _ = curve_fit(
lambda x, *params: spectral_density_approx(w, *unpack(params)),
w, J,
p0=guesses,
bounds=(lower_bounds, upper_bounds),
sigma=sigma,
maxfev=1000000000,
)
return unpack(params)
With the spectral density approximation $J_{\mathrm approx}(w; a, b, c)$ implemented above, we can now perform the fit and examine the results.
w = np.linspace(0, 25, 20000)
J = ohmic_spectral_density(w, alpha=obp.alpha, wc=obp.wc)
params_k = [
fit_spectral_density(J, w, alpha=obp.alpha, wc=obp.wc, N=i+1)
for i in range(4)
]
Let's plot the fit for each $k$ and examine how it improves with an increasing number of terms:
for k, params in enumerate(params_k):
lam, gamma, w0 = params
y = spectral_density_approx(w, lam, gamma, w0)
print(f"Parameters [k={k}]: lam={lam}; gamma={gamma}; w0={w0}")
plt.plot(w, J, w, y)
plt.show()
Parameters [k=0]: lam=[6.14746382]; gamma=[1.77939431]; w0=[0.1]
Parameters [k=1]: lam=[2.21188807 3.26249622]; gamma=[1.24955687 1.43449451]; w0=[0.1 1.80554797]
Parameters [k=2]: lam=[1.64959679 2.33983518 1.12738626]; gamma=[1.18453049 1.09371098 1.00881732]; w0=[2.69770735 1.37013881 0.13257063]
Parameters [k=3]: lam=[ 7.91592694 0.60083723 -4.40789276 0.01058515]; gamma=[2.29619001 1.00246811 4.29908162 0.3073633 ]; w0=[0.1 0.1 3.98168646 0.1 ]
The fit with four terms looks good. Let's take a closer look at it by plotting the contribution of each term of the fit:
# The parameters for the fit with four terms:
lam, gamma, w0 = params_k[-1]
print(f"Parameters [k={len(params_k) - 1}]: lam={lam}; gamma={gamma}; w0={w0}")
Parameters [k=3]: lam=[ 7.91592694 0.60083723 -4.40789276 0.01058515]; gamma=[2.29619001 1.00246811 4.29908162 0.3073633 ]; w0=[0.1 0.1 3.98168646 0.1 ]
# Plot the components of the fit separately:
def spectral_density_ith_component(w, i, lam, gamma, w0):
""" Return the i'th term of the approximation for the spectral density. """
return (
2 * lam[i] * gamma[i] * w /
(((w + w0[i])**2 + gamma[i]**2) * ((w - w0[i])**2 + gamma[i]**2))
)
def plot_spectral_density_fit_components(J, w, lam, gamma, w0):
""" Plot the individual components of a fit to the spectral density. """
fig, axes = plt.subplots(1, 1, sharex=True, figsize=(8, 8))
axes.plot(w, J, 'r--', linewidth=2, label="original")
for i in range(len(lam)):
axes.plot(
w, spectral_density_ith_component(w, i, lam, gamma, w0),
linewidth=2,
label=f"fit component {i}",
)
axes.set_xlabel(r'$w$', fontsize=28)
axes.set_ylabel(r'J', fontsize=28)
axes.legend()
return fig
plot_spectral_density_fit_components(J, w, lam, gamma, w0);
And let's also compare the power spectrum of the fit and the analytical spectral density:
def plot_power_spectrum(alpha, wc, beta, lam, gamma, w0, save=True):
""" Plot the power spectrum of a fit against the actual power spectrum. """
w = np.linspace(-10, 10, 50000)
s_orig = ohmic_power_spectrum(w, alpha=alpha, wc=wc, beta=beta)
s_fit = (
spectral_density_approx(w, lam, gamma, w0) *
((1 / (np.e**(w * beta) - 1)) + 1) * 2
)
fig, axes = plt.subplots(1, 1, sharex=True, figsize=(8, 8))
axes.plot(w, s_orig, 'r', linewidth=2, label="original")
axes.plot(w, s_fit, 'b', linewidth=2, label="fit")
axes.set_xlabel(r'$\omega$', fontsize=28)
axes.set_ylabel(r'$S(\omega)$', fontsize=28)
axes.legend()
if save:
fig.savefig('powerspectrum.eps')
plot_power_spectrum(obp.alpha, obp.wc, obp.beta, lam, gamma, w0, save=False)
Now that we have a good fit to the spectral density, we can calculate the Matsubara expansion terms for the BosonicBath
from them. At the same time we will calculate the Matsubara terminator for this expansion.
def matsubara_coefficients_from_spectral_fit(lam, gamma, w0, beta, Q, Nk):
""" Calculate the Matsubara co-efficients for a fit to the spectral
density.
"""
# initial 0 value with the correct dimensions:
terminator = 0. * spre(Q)
# the number of matsubara expansion terms to include in the terminator:
terminator_max_k = 1000
ckAR = []
vkAR = []
ckAI = []
vkAI = []
for lamt, Gamma, Om in zip(lam, gamma, w0):
ckAR.extend([
(lamt / (4 * Om)) * coth(beta * (Om + 1.0j * Gamma) / 2),
(lamt / (4 * Om)) * coth(beta * (Om - 1.0j * Gamma) / 2),
])
for k in range(1, Nk + 1):
ek = 2 * np.pi * k / beta
ckAR.append(
(-2 * lamt * 2 * Gamma / beta) * ek /
(
((Om + 1.0j * Gamma)**2 + ek**2) *
((Om - 1.0j * Gamma)**2 + ek**2)
)
)
terminator_factor = 0
for k in range(Nk + 1, terminator_max_k):
ek = 2 * np.pi * k / beta
ck = (
(-2 * lamt * 2 * Gamma / beta) * ek /
(
((Om + 1.0j * Gamma)**2 + ek**2) *
((Om - 1.0j * Gamma)**2 + ek**2)
)
)
terminator_factor += ck / ek
terminator += terminator_factor * (
2 * spre(Q) * spost(Q.dag())
- spre(Q.dag() * Q)
- spost(Q.dag() * Q)
)
vkAR.extend([
-1.0j * Om + Gamma,
1.0j * Om + Gamma,
])
vkAR.extend([
2 * np.pi * k * obp.T + 0.j
for k in range(1, Nk + 1)
])
ckAI.extend([
-0.25 * lamt * 1.0j / Om,
0.25 * lamt * 1.0j / Om,
])
vkAI.extend([
-(-1.0j * Om - Gamma),
-(1.0j * Om - Gamma),
])
return ckAR, vkAR, ckAI, vkAI, terminator
def generate_spectrum_results(obp, params, Nk, max_depth):
""" Run the HEOM with the given bath parameters and
and return the results of the evolution.
"""
lam, gamma, w0 = params
ckAR, vkAR, ckAI, vkAI, terminator = (
matsubara_coefficients_from_spectral_fit(
lam, gamma, w0, beta=obp.beta, Q=obp.Q, Nk=Nk,
)
)
Ltot = liouvillian(Hsys) + terminator
tlist = np.linspace(0, 30 * np.pi / Del, 600)
with timer("RHS construction time"):
bath = BosonicBath(obp.Q, ckAR, vkAR, ckAI, vkAI)
HEOM_spectral_fit = HEOMSolver(
Ltot, bath, max_depth=max_depth, options=options,
)
with timer("ODE solver time"):
results_spectral_fit = (HEOM_spectral_fit.run(rho0, tlist))
return results_spectral_fit
Below we generate results for different convergence parameters (number of terms in the fit, number of matsubara terms, and depth of the hierarchy). For the parameter choices here, we need a relatively large depth of around '11', which can be a little slow.
# Generate results for different number of lorentzians in fit:
results_spectral_fit_pk = [
generate_spectrum_results(obp, params, Nk=1, max_depth=max_depth)
for params in params_k
]
plot_result_expectations([
(
result, P11p, 'rand',
f"P11 (spectral fit) $k_J$={pk + 1}",
)
for pk, result in enumerate(results_spectral_fit_pk)
]);
RHS construction time: 0.0019350051879882812 Total run time: 0.51s ODE solver time: 0.5141124725341797 RHS construction time: 0.00794839859008789 Total run time: 1.04s ODE solver time: 1.045264482498169 RHS construction time: 0.02301025390625 Total run time: 2.47s ODE solver time: 2.468433380126953 RHS construction time: 0.06010723114013672 Total run time: 8.15s ODE solver time: 8.154273509979248
# generate results for different number of Matsubara terms per Lorentzian
# for max number of Lorentzians:
Nk_list = range(2, 4)
results_spectral_fit_nk = [
generate_spectrum_results(obp, params_k[-1], Nk=Nk, max_depth=max_depth)
for Nk in Nk_list
]
plot_result_expectations([
(
result, P11p, 'rand',
f"P11 (spectral fit) K={nk}",
)
for nk, result in zip(Nk_list, results_spectral_fit_nk)
]);
RHS construction time: 0.13563942909240723 Total run time: 20.75s ODE solver time: 20.753215312957764 RHS construction time: 0.1949019432067871 Total run time: 50.97s ODE solver time: 50.97481369972229
# Generate results for different depths:
Nc_list = range(2, max_depth)
results_spectral_fit_nc = [
generate_spectrum_results(obp, params_k[-1], Nk=1, max_depth=Nc)
for Nc in Nc_list
]
plot_result_expectations([
(
result, P11p, 'rand',
f"P11 (spectral fit) $N_C={nc}$",
)
for nc, result in zip(Nc_list, results_spectral_fit_nc)
]);
RHS construction time: 0.0038518905639648438 Total run time: 0.49s ODE solver time: 0.4922187328338623 RHS construction time: 0.008773565292358398 Total run time: 0.99s ODE solver time: 0.9855506420135498 RHS construction time: 0.021680355072021484 Total run time: 2.99s ODE solver time: 2.9881505966186523
We now combine the fitting and correlation function data into one large plot.
def correlation_approx_matsubara(t, ck, vk):
""" Calculate the approximate real or imaginary part of the
correlation function from the matsubara expansion co-efficients.
"""
ck = np.array(ck)
vk = np.array(vk)
return np.sum(ck[:, None] * np.exp(-vk[:, None] * t), axis=0)
def plot_cr_fit_vs_actual(t, ckAR, vkAR, C, axes):
""" Plot the C_R(t) fit. """
yR = correlation_approx_matsubara(t, ckAR, vkAR)
axes.plot(
t, np.real(C),
"r", linewidth=3, label="Original",
)
axes.plot(
t, np.real(yR),
"g", dashes=[3, 3], linewidth=2, label="Reconstructed",
)
axes.legend(loc=0)
axes.set_ylabel(r'$C_R(t)$', fontsize=28)
axes.set_xlabel(r'$t\;\omega_c$', fontsize=28)
axes.locator_params(axis='y', nbins=4)
axes.locator_params(axis='x', nbins=4)
axes.text(0.15, 0.85, "(a)", fontsize=28, transform=axes.transAxes)
def plot_ci_fit_vs_actual(t, ckAI, vkAI, C, axes):
""" Plot the C_I(t) fit. """
yI = correlation_approx_matsubara(t, ckAI, vkAI)
axes.plot(
t, np.imag(C),
"r", linewidth=3, label="Original",
)
axes.plot(
t, np.real(yI),
"g", dashes=[3, 3], linewidth=2, label="Reconstructed",
)
axes.legend(loc=0)
axes.set_ylabel(r'$C_I(t)$', fontsize=28)
axes.set_xlabel(r'$t\;\omega_c$', fontsize=28)
axes.locator_params(axis='y', nbins=4)
axes.locator_params(axis='x', nbins=4)
axes.text(0.80, 0.80, "(b)", fontsize=28, transform=axes.transAxes)
def plot_jw_fit_vs_actual(bath_fit, obp, axes):
""" Plot the J(w) fit. """
[lam, gamma, w0] = bath_fit
[alpha, wc] = [obp.alpha, obp.wc]
w = np.linspace(0, 25, 20000)
J_orig = ohmic_spectral_density(w, alpha=alpha, wc=wc)
J_fit = spectral_density_approx(w, lam, gamma, w0)
axes.plot(
w, J_orig,
"r", linewidth=3, label=r"$J(\omega)$ original",
)
axes.plot(
w, J_fit,
"g", dashes=[3, 3], linewidth=2, label=r"$J(\omega)$ Fit $k_J = 4$",
)
axes.legend(loc=0)
axes.set_ylabel(r'$J(\omega)$', fontsize=28)
axes.set_xlabel(r'$\omega/\omega_c$', fontsize=28)
axes.locator_params(axis='y', nbins=4)
axes.locator_params(axis='x', nbins=4)
axes.text(0.15, 0.85, "(c)", fontsize=28, transform=axes.transAxes)
def plot_sw_fit_vs_actual(bath_fit, obp, axes):
""" Plot the S(w) fit. """
[lam, gamma, w0] = bath_fit
[alpha, wc, beta] = [obp.alpha, obp.wc, obp.beta]
# avoid the pole in the fit around zero:
w = np.concatenate(
[np.linspace(-10, -0.1, 5000),
np.linspace(0.1, 10, 5000)],
)
s_orig = ohmic_power_spectrum(w, alpha=alpha, wc=wc, beta=beta)
s_fit = (
spectral_density_approx(w, lam, gamma, w0) *
((1 / (np.e**(w * beta) - 1)) + 1) * 2
)
axes.plot(w, s_orig, "r", linewidth=3, label="Original")
axes.plot(w, s_fit, "g", dashes=[3, 3], linewidth=2, label="Reconstructed")
axes.legend()
axes.set_ylabel(r'$S(\omega)$', fontsize=28)
axes.set_xlabel(r'$\omega/\omega_c$', fontsize=28)
axes.locator_params(axis='y', nbins=4)
axes.locator_params(axis='x', nbins=4)
axes.text(0.15, 0.85, "(d)", fontsize=28, transform=axes.transAxes)
def plot_matsubara_spectrum_fit_vs_actual(
t, C, matsubara_fit, bath_fit, obp,
):
""" Plot the Matsubara fit of the spectrum . """
fig = plt.figure(figsize=(12, 10))
grid = plt.GridSpec(2, 2, wspace=0.4, hspace=0.3)
[ckAR, vkAR, ckAI, vkAI] = matsubara_fit
plot_cr_fit_vs_actual(
t, ckAR, vkAR, C,
axes=fig.add_subplot(grid[0, 0]),
)
plot_ci_fit_vs_actual(
t, ckAI, vkAI, C,
axes=fig.add_subplot(grid[0, 1]),
)
plot_jw_fit_vs_actual(
bath_fit, obp,
axes=fig.add_subplot(grid[1, 0]),
)
plot_sw_fit_vs_actual(
bath_fit, obp,
axes=fig.add_subplot(grid[1, 1]),
)
return fig
t = np.linspace(0, 15, 100)
C = ohmic_correlation(t, alpha=obp.alpha, wc=obp.wc, beta=obp.beta)
ckAR, vkAR, ckAI, vkAI, terminator = (
matsubara_coefficients_from_spectral_fit(
lam, gamma, w0, beta=obp.beta, Q=obp.Q, Nk=1,
)
)
matsubara_fit = [ckAR, vkAR, ckAI, vkAI]
bath_fit = [lam, gamma, w0]
plot_matsubara_spectrum_fit_vs_actual(
t, C, matsubara_fit,
bath_fit, obp,
);
Having successfully fitted the spectral density and used the result to calculate the Matsubara expansion and terminator for the HEOM bosonic bath, we now proceed to the second case of fitting the correlation function itself instead.
Here we fit the real and imaginary parts seperately, using the following ansatz
$$C_R^F(t) = \sum_{i=1}^{k_R} c_R^ie^{-\gamma_R^i t}\cos(\omega_R^i t)$$$$C_I^F(t) = \sum_{i=1}^{k_I} c_I^ie^{-\gamma_I^i t}\sin(\omega_I^i t)$$# The approximate correlation functions and a helper for fitting
# the approximate correlation function to values calculated from
# the analytical formula:
def correlation_approx_real(t, a, b, c):
""" Calculate the fitted value of the function for the given parameters.
"""
a = np.array(a)
b = np.array(b)
c = np.array(c)
return np.sum(
a[:, None] * np.exp(b[:, None] * t) * np.cos(c[:, None] * t),
axis=0,
)
def correlation_approx_imag(t, a, b, c):
""" Calculate the fitted value of the function for the given parameters.
"""
a = np.array(a)
b = np.array(b)
c = np.array(c)
return np.sum(
a[:, None] * np.exp(b[:, None] * t) * np.sin(c[:, None] * t),
axis=0,
)
def fit_correlation_real(C, t, wc, N):
""" Fit the spectral density with N underdamped oscillators. """
sigma = [0.1] * len(t)
C_max = abs(max(C, key=abs))
guesses = pack([C_max] * N, [-wc] * N, [wc] * N)
lower_bounds = pack([-20 * C_max] * N, [-np.inf] * N, [0.] * N)
upper_bounds = pack([20 * C_max] * N, [0.1] * N, [np.inf] * N)
params, _ = curve_fit(
lambda x, *params: correlation_approx_real(t, *unpack(params)),
t, C,
p0=guesses,
bounds=(lower_bounds, upper_bounds),
sigma=sigma,
maxfev=1000000000,
)
return unpack(params)
def fit_correlation_imag(C, t, wc, N):
""" Fit the spectral density with N underdamped oscillators. """
sigma = [0.0001] * len(t)
C_max = abs(max(C, key=abs))
guesses = pack([-C_max] * N, [-2] * N, [1] * N)
lower_bounds = pack([-5 * C_max] * N, [-100] * N, [0.] * N)
upper_bounds = pack([5 * C_max] * N, [0.01] * N, [100] * N)
params, _ = curve_fit(
lambda x, *params: correlation_approx_imag(t, *unpack(params)),
t, C,
p0=guesses,
bounds=(lower_bounds, upper_bounds),
sigma=sigma,
maxfev=1000000000,
)
return unpack(params)
t = np.linspace(0, 15, 15000)
C = ohmic_correlation(t, alpha=obp.alpha, wc=obp.wc, beta=obp.beta)
params_k_real = [
fit_correlation_real(np.real(C), t, wc=obp.wc, N=i+1)
for i in range(3)
]
params_k_imag = [
fit_correlation_imag(np.imag(C), t, wc=obp.wc, N=i+1)
for i in range(3)
]
for k, params in enumerate(params_k_real):
lam, gamma, w0 = params
y = correlation_approx_real(t, lam, gamma, w0)
print(f"Parameters [k={k}]: lam={lam}; gamma={gamma}; w0={w0}")
plt.plot(t, np.real(C), label="C_R(t) analytic")
plt.plot(t, y, label=f"C_R(t) k={k + 1}")
plt.legend()
plt.show()
Parameters [k=0]: lam=[1.72456692]; gamma=[-1.29280164]; w0=[1.14098482e-12]
Parameters [k=1]: lam=[0.93982319 0.73272363]; gamma=[-1.1929147 -0.63181829]; w0=[1.40619089e+00 9.76208988e-27]
Parameters [k=2]: lam=[0.50981772 0.68385831 0.44182127]; gamma=[-0.4747359 -0.79206162 -1.01686245]; w0=[5.24433582e-06 9.00794383e-01 2.14916543e+00]
for k, params in enumerate(params_k_imag):
lam, gamma, w0 = params
y = correlation_approx_imag(t, lam, gamma, w0)
print(f"Parameters [k={k}]: lam={lam}; gamma={gamma}; w0={w0}")
plt.plot(t, np.imag(C), label="C_I(t) analytic")
plt.plot(t, y, label=f"C_I(t) k={k + 1}")
plt.legend()
plt.show()
Parameters [k=0]: lam=[-3.35965972]; gamma=[-1.52063694]; w0=[0.80990586]
Parameters [k=1]: lam=[-3.35965972 -0.98438382]; gamma=[-1.16604573 -1.45054232]; w0=[0.29085486 1.78000603]
Parameters [k=2]: lam=[-0.85911264 -3.31883451 -0.30749363]; gamma=[-1.09310191 -0.9960073 -1.1862732 ]; w0=[1.3691296 0.16626634 2.69674089]
Now we construct the BosonicBath
co-efficients and frequencies from the fit to the correlation function:
def matsubara_coefficients_from_corr_fit_real(lam, gamma, w0):
""" Return the matsubara coefficients for the imaginary part
of the correlation function.
"""
ckAR = [0.5 * x + 0j for x in lam] # the 0.5 is from the cosine
# extend the list with the complex conjugates:
ckAR.extend(np.conjugate(ckAR))
vkAR = [-x - 1.0j * y for x, y in zip(gamma, w0)]
vkAR.extend([-x + 1.0j * y for x, y in zip(gamma, w0)])
return ckAR, vkAR
def matsubara_coefficients_from_corr_fit_imag(lam, gamma, w0):
""" Return the matsubara coefficients for the imaginary part
of the correlation function.
"""
ckAI = [-0.5j * x for x in lam] # the 0.5 is from the sine
# extend the list with the complex conjugates:
ckAI.extend(np.conjugate(ckAI))
vkAI = [-x - 1.0j * y for x, y in zip(gamma, w0)]
vkAI.extend([-x + 1.0j * y for x, y in zip(gamma, w0)])
return ckAI, vkAI
ckAR, vkAR = matsubara_coefficients_from_corr_fit_real(*params_k_real[-1])
ckAI, vkAI = matsubara_coefficients_from_corr_fit_imag(*params_k_imag[-1])
def corr_spectrum_approx(w, ckAR, vkAR, ckAI, vkAI):
""" Calculates the approximate power spectrum from ck and vk. """
S = np.zeros(len(w), dtype=np.complex128)
for ck, vk in zip(ckAR, vkAR):
S += (
2 * ck * np.real(vk) /
((w - np.imag(vk))**2 + (np.real(vk)**2))
)
for ck, vk in zip(ckAI, vkAI):
S += (
2 * 1.0j * ck * np.real(vk) /
((w - np.imag(vk))**2 + (np.real(vk)**2))
)
return S
def plot_jw_correlation_fit_vs_actual(matsubara_fit, obp, axes):
""" Plot J(w) from the correlation fit. """
[ckAR, vkAR, ckAI, vkAI] = matsubara_fit
[alpha, wc] = [obp.alpha, obp.wc]
w = np.linspace(0.001, 25, 20000)
J_orig = ohmic_spectral_density(w, alpha=alpha, wc=wc)
J_fit = np.real(
corr_spectrum_approx(w, ckAR, vkAR, ckAI, vkAI) /
(((1 / (np.e**(w * obp.beta) - 1)) + 1) * 2)
)
axes.plot(
w, J_orig,
"r", linewidth=3, label=r"$J(\omega)$ original",
)
axes.plot(
w, J_fit,
"g", dashes=[3, 3], linewidth=2, label=r"$J(\omega)$ fit",
)
axes.legend(loc=0)
axes.set_ylabel(r'$J(\omega)$', fontsize=28)
axes.set_xlabel(r'$\omega/\omega_c$', fontsize=28)
axes.locator_params(axis='y', nbins=4)
axes.locator_params(axis='x', nbins=4)
axes.text(3, 1.1, "(c)", fontsize=28)
def plot_sw_correlation_fit_vs_actual(matsubara_fit, obp, axes):
""" Plot S(W) from the correlation fit. """
[ckAR, vkAR, ckAI, vkAI] = matsubara_fit
[alpha, wc, beta] = [obp.alpha, obp.wc, obp.beta]
# avoid the pole in the fit around zero:
w = np.concatenate([
np.linspace(-10, -0.1, 5000),
np.linspace(0.1, 10, 5000),
])
s_orig = ohmic_power_spectrum(w, alpha=alpha, wc=wc, beta=beta)
s_fit = corr_spectrum_approx(w, ckAR, vkAR, ckAI, vkAI)
axes.plot(
w, s_orig,
"r", linewidth=3, label="Original",
)
axes.plot(
w, s_fit,
"g", dashes=[3, 3], linewidth=2, label="Reconstructed",
)
axes.legend()
axes.set_ylabel(r'$S(\omega)$', fontsize=28)
axes.set_xlabel(r'$\omega/\omega_c$', fontsize=28)
axes.locator_params(axis='y', nbins=4)
axes.locator_params(axis='x', nbins=4)
axes.text(0.15, 0.85, "(d)", fontsize=28, transform=axes.transAxes)
def plot_matsubara_correlation_fit_vs_actual(t, C, matsubara_fit, obp):
fig = plt.figure(figsize=(12, 10))
grid = plt.GridSpec(2, 2, wspace=0.4, hspace=0.3)
ckAR, vkAR, ckAI, vkAI = matsubara_fit
plot_cr_fit_vs_actual(
t, ckAR, vkAR, C,
axes=fig.add_subplot(grid[0, 0]),
)
plot_ci_fit_vs_actual(
t, ckAI, vkAI, C,
axes=fig.add_subplot(grid[0, 1]),
)
plot_jw_correlation_fit_vs_actual(
matsubara_fit, obp,
axes=fig.add_subplot(grid[1, 0]),
)
plot_sw_correlation_fit_vs_actual(
matsubara_fit, obp,
axes=fig.add_subplot(grid[1, 1]),
)
t = np.linspace(0, 15, 100)
C = ohmic_correlation(t, alpha=obp.alpha, wc=obp.wc, beta=obp.beta)
matsubara_fit = [ckAR, vkAR, ckAI, vkAI]
plot_matsubara_correlation_fit_vs_actual(
t, C, matsubara_fit, obp,
)
/usr/share/miniconda3/envs/test-environment/lib/python3.10/site-packages/matplotlib/cbook/__init__.py:1298: ComplexWarning: Casting complex values to real discards the imaginary part return np.asarray(x, float)
def generate_corr_results(params_real, params_imag, max_depth):
ckAR, vkAR = matsubara_coefficients_from_corr_fit_real(
*params_real
)
ckAI, vkAI = matsubara_coefficients_from_corr_fit_imag(
*params_imag
)
tlist = np.linspace(0, 30 * np.pi / Del, 600)
with timer("RHS construction time"):
bath = BosonicBath(obp.Q, ckAR, vkAR, ckAI, vkAI)
HEOM_corr_fit = HEOMSolver(
Hsys, bath, max_depth=max_depth, options=options,
)
with timer("ODE solver time"):
results_corr_fit = (HEOM_corr_fit.run(rho0, tlist))
return results_corr_fit
# Generate results for different number of lorentzians in fit:
results_corr_fit_pk = [
print(f"{pk + 1}") or generate_corr_results(
params_real, params_imag, max_depth=max_depth,
)
for pk, (params_real, params_imag)
in enumerate(zip(params_k_real, params_k_imag))
]
1 RHS construction time: 0.0030128955841064453 Total run time: 0.50s ODE solver time: 0.49604034423828125 2 RHS construction time: 0.02112889289855957 Total run time: 1.74s ODE solver time: 1.741285800933838 3 RHS construction time: 0.2393660545349121 Total run time: 20.41s ODE solver time: 20.410775184631348
plot_result_expectations([
(
result, P11p, 'rand',
f"P11 (correlation fit) k_R=k_I={pk + 1}",
)
for pk, result in enumerate(results_corr_fit_pk)
]);
fig, axes = plt.subplots(1, 1, sharex=True, figsize=(12, 7))
plot_result_expectations([
(
results_corr_fit_pk[0], P11p,
'y', "Correlation Function Fit $k_R=k_I=1$",
),
(
results_corr_fit_pk[2], P11p,
'y-.', "Correlation Function Fit $k_R=k_I=3$",
),
(results_spectral_fit_pk[0], P11p, 'b', "Spectral Density Fit $k_J=1$"),
(results_spectral_fit_pk[2], P11p, 'g--', "Spectral Density Fit $k_J=3$"),
(results_spectral_fit_pk[3], P11p, 'r-.', "Spectral Density Fit $k_J=4$"),
], axes=axes)
axes.set_yticks([0.6, 0.8, 1])
axes.set_ylabel(r'$\rho_{11}$', fontsize=30)
axes.set_xlabel(r'$t\;\omega_c$', fontsize=30)
axes.legend(loc=0, fontsize=20);
qutip.about()
QuTiP: Quantum Toolbox in Python ================================ Copyright (c) QuTiP team 2011 and later. Current admin team: Alexander Pitchford, Nathan Shammah, Shahnawaz Ahmed, Neill Lambert, Eric Giguère, Boxi Li, Jake Lishman, Simon Cross and Asier Galicia. Board members: Daniel Burgarth, Robert Johansson, Anton F. Kockum, Franco Nori and Will Zeng. Original developers: R. J. Johansson & P. D. Nation. Previous lead developers: Chris Granade & A. Grimsmo. Currently developed through wide collaboration. See https://github.com/qutip for details. QuTiP Version: 5.1.0.dev0+50fc9b7 Numpy Version: 1.22.4 Scipy Version: 1.13.0 Cython Version: 3.0.10 Matplotlib Version: 3.5.2 Python Version: 3.10.4 Number of CPUs: 4 BLAS Info: Generic INTEL MKL Ext: False Platform Info: Linux (x86_64) Installation path: /usr/share/miniconda3/envs/test-environment/lib/python3.10/site-packages/qutip ================================================================================ Please cite QuTiP in your publication. ================================================================================ For your convenience a bibtex reference can be easily generated using `qutip.cite()`
This section can include some tests to verify that the expected outputs are generated within the notebook. We put this section at the end of the notebook, so it's not interfering with the user experience. Please, define the tests using assert, so that the cell execution fails if a wrong output is generated.
assert np.allclose(
expect(P11p, results_spectral_fit_pk[2].states),
expect(P11p, results_spectral_fit_pk[3].states),
rtol=1e-2,
)