Skip to content

Commit

Permalink
moved plots, added functions so tutorial looks better
Browse files Browse the repository at this point in the history
  • Loading branch information
@gsuarezr
gsuarezr committed Nov 7, 2023
1 parent 2361a6c commit 3a0c493
Show file tree
Hide file tree
Showing 2 changed files with 263 additions and 3 deletions.
66 changes: 63 additions & 3 deletions qutip/solver/heom/bofin_baths.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
from qutip.core import data as _data
from qutip.core.qobj import Qobj
from qutip.core.superoperator import spre, spost
from qutip.visualization import gen_spectral_plots, gen_corr_plots

__all__ = [
"BathExponent",
Expand Down Expand Up @@ -1287,6 +1288,21 @@ def get_fit(
end = time()
self.fit_time = end - start

def correlation_approx_matsubara(self, t, real=True):
""" Calculate the approximate real or imaginary part of the
correlation function from the matsubara expansion co-efficients.
"""
if real:
ck, vk = self.cvars[0], self.cvars[1]
else:
ck, vk = self.cvars[2], self.cvars[3]
ck = np.array(ck)
vk = np.array(vk)
return np.sum(ck[:, None] * np.exp(-vk[:, None] * t), axis=0)

def fit_plots(self, w, J, t, C, w2, S):
gen_spectral_plots(self, w, J, t, C, w2, S)

def _matsubara_spectral_fit(self):
"""
Obtains the bath exponents from the list of fit parameters
Expand All @@ -1306,6 +1322,10 @@ def _matsubara_spectral_fit(self):
ckAI = []
vkAI = []
UD = UnderDampedBath(self.Q, lam[0], gamma[0], w0[0], self.T, self.Nk)
terminator = 0. * spre(self.Q)
# the number of matsubara expansion terms to include in the terminator:
terminator_max_k = 1000

for lamt, Gamma, Om in zip(lam, gamma, w0):
ck_real, vk_real, ck_imag, vk_imag = UD._matsubara_params(
lamt, 2 * Gamma, Om + 0j, self.T, self.Nk
Expand All @@ -1314,12 +1334,29 @@ def _matsubara_spectral_fit(self):
vkAR.append(vk_real)
vkAI.append(vk_imag)
ckAI.append(ck_imag)
terminator_factor = 0
for k in range(self.Nk + 1, terminator_max_k):
ek = 2 * np.pi * k / self.beta
ck = (
(-2 * lamt * 2 * Gamma / self.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(self.Q) * spost(self.Q.dag())
- spre(self.Q.dag() * self.Q)
- spost(self.Q.dag() * self.Q)
)
self.cvars = [
np.array(ckAR).flatten(),
np.array(vkAR).flatten(),
np.array(ckAI).flatten(),
np.array(vkAI).flatten(),
]
self.terminator = terminator
self.Bath_spec = BosonicBath(
self.Q,
np.array(ckAR).flatten(),
Expand Down Expand Up @@ -1392,6 +1429,9 @@ def summary(self):
print(formatted_line1 + formatted_line2)
print("\t " * 6 + time)

def fit_plots(self, w, J, t, C, w2, S, beta):
gen_corr_plots(self, w, J, t, C, w2, S, beta)

def fit_correlation(
self,
t,
Expand Down Expand Up @@ -1431,6 +1471,12 @@ def fit_correlation(
"""
start = time()
rmse1, rmse2 = [8, 8]
flag = False
if (Nr is not None) & (Ni is not None):
flag = True
Nr -= 1
Ni -= 1

if Nr is None:
Nr = 0
if Ni is None:
Expand All @@ -1448,6 +1494,8 @@ def fit_correlation(
upper=upper,
label="correlation_real",
)
if flag == True:
break
while rmse2 > final_rmse:
Ni += 1
rmse2, params_imag = self._fit(
Expand All @@ -1461,6 +1509,8 @@ def fit_correlation(
upper=upper,
label="correlation_imag",
)
if flag == True:
break
self.Nr = Nr
self.Ni = Ni
self.rmse_real = rmse1
Expand Down Expand Up @@ -1489,6 +1539,13 @@ def corr_spectrum_approx(self, w):
)
return S

def corr_spectral_approx(self, w, beta):
J = np.real(
self.corr_spectrum_approx(w) /
(((1 / (np.e**(w * beta) - 1)) + 1) * 2)
)
return J

def _matsubara_coefficients(self):
lam, gamma, w0 = self.params_real
lam2, gamma2, w02 = self.params_imag
Expand All @@ -1508,7 +1565,7 @@ def _matsubara_coefficients(self):


class OhmicBath(BosonicBath):
def __init__(self, T, Q, alpha, wc, s, Nk=4, method="correlation", rmse=1e-6):
def __init__(self, T, Q, alpha, wc, s, Nk=4, method="spectral", rmse=7e-5):
self.alpha = alpha
self.wc = wc
self.s = s
Expand All @@ -1519,14 +1576,17 @@ def __init__(self, T, Q, alpha, wc, s, Nk=4, method="correlation", rmse=1e-6):
self.beta = 1 / T
if method == "correlation":
self.fit = FitCorr(self.Q)
t = np.linspace(0, 50 / self.wc, 1000)
t = np.linspace(0, 15 / self.wc, 1000)
C = self.ohmic_correlation(t, self.s)
self.fit.fit_correlation(t, C, final_rmse=rmse)
self.bath = self.fit.Bath_corr
else:
self.fit = FitSpectral(self.T, self.Q, self.Nk)
w = np.linspace(0, 50 * self.wc, 1000)
w = np.linspace(0, 15 * self.wc, 20000)
J = self.ohmic_spectral_density(w)
self.fit.get_fit(J, w, final_rmse=rmse)
self.bath = self.fit.Bath_spec
self.terminator = self.fit.terminator

def summary(self):
self.fit.summary()
Expand Down
200 changes: 200 additions & 0 deletions qutip/visualization.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1871,3 +1871,203 @@ def plot_schmidt(ket, theme='light', splitting=None,
ax.set_ylabel("first particles")

return fig, output


# PLOTS FOR FITS (THIS BIT NEEDS REFACTORING)
# SPECTRAL:

def gen_spectral_plots(fs, w, J, t, C, w2, S):
def plot_cr_fit_vs_actual(t, C, func, axes):
""" Plot the C_R(t) fit. """
yR = func(t)

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.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, C, func, axes):
""" Plot the C_I(t) fit. """
yI = func(t)

axes.plot(
t, np.imag(C),
"r", linewidth=3,
)
axes.plot(
t, np.real(yI),
"g", dashes=[3, 3], linewidth=2,
)

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(w, J, func, params, axes):
""" Plot the J(w) fit. """
a, b, c = params
J_fit = func(w, a, b, c)

axes.plot(
w, J,
"r", linewidth=3,
)
axes.plot(
w, J_fit,
"g", dashes=[3, 3], linewidth=2,
)

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(func, axes):
""" Plot the S(w) fit. """

# avoid the pole in the fit around zero:
s_fit = func(w2)

axes.plot(w2, S, "r", linewidth=3)
axes.plot(w2, s_fit, "g", dashes=[3, 3], linewidth=2)

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
):
""" Plot the Matsubara fit of the spectrum . """
fig = plt.figure(figsize=(12, 10))
grid = plt.GridSpec(2, 2, wspace=0.4, hspace=0.3)

plot_cr_fit_vs_actual(
t, C, fs.correlation_approx_matsubara,
axes=fig.add_subplot(grid[0, 0]),
)
plot_ci_fit_vs_actual(
t, C, lambda t: fs.correlation_approx_matsubara(t, real=False),
axes=fig.add_subplot(grid[0, 1]),
)
plot_jw_fit_vs_actual(w, J, fs.spectral_density_approx, fs.params_spec,
axes=fig.add_subplot(grid[1, 0]),
)
plot_sw_fit_vs_actual(fs.spec_spectrum_approx,
axes=fig.add_subplot(grid[1, 1]),
)
fig.legend(loc='upper center', ncol=2, fancybox=True, shadow=True)
return fig
return plot_matsubara_spectrum_fit_vs_actual(t, C)


def gen_corr_plots(fc, w, J, t, C, w2, S, beta):
def plot_cr_fit_vs_actual(t, C, axes):
""" Plot the C_R(t) fit. """
a, b, c = fc.params_real
yR = np.real(fc.corr_approx(t, a, b, c))

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.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, C, axes):
""" Plot the C_I(t) fit. """
a, b, c = fc.params_imag
yI = np.imag(fc.corr_approx(t, a, b, c))
axes.plot(
t, np.imag(C),
"r", linewidth=3,
)
axes.plot(
t, np.real(yI),
"g", dashes=[3, 3], linewidth=2,
)

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(w, J, beta, axes):
""" Plot the J(w) fit. """
J_fit = fc.corr_spectral_approx(w, beta)
axes.plot(
w, J,
"r", linewidth=3,
)
axes.plot(
w, J_fit,
"g", dashes=[3, 3], linewidth=2,
)

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(axes):
""" Plot the S(w) fit. """

# avoid the pole in the fit around zero:
s_fit = fc.corr_spectrum_approx(w2)

axes.plot(w2, S, "r", linewidth=3)
axes.plot(w2, s_fit, "g", dashes=[3, 3], linewidth=2)

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
):
""" Plot the Matsubara fit of the spectrum . """
fig = plt.figure(figsize=(12, 10))
grid = plt.GridSpec(2, 2, wspace=0.4, hspace=0.3)

plot_cr_fit_vs_actual(
t, C,
axes=fig.add_subplot(grid[0, 0]),
)
plot_ci_fit_vs_actual(
t, C,
axes=fig.add_subplot(grid[0, 1]),
)
plot_jw_fit_vs_actual(w, J, beta,
axes=fig.add_subplot(grid[1, 0]),
)
plot_sw_fit_vs_actual(axes=fig.add_subplot(grid[1, 1]),)
fig.legend(loc='upper center', ncol=2, fancybox=True, shadow=True)
return fig
return plot_matsubara_spectrum_fit_vs_actual(t, C)

0 comments on commit 3a0c493

Please sign in to comment.