diff --git a/qutip/solver/heom/bofin_baths.py b/qutip/solver/heom/bofin_baths.py index a61ea1e662..faf6a0db43 100644 --- a/qutip/solver/heom/bofin_baths.py +++ b/qutip/solver/heom/bofin_baths.py @@ -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", @@ -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 @@ -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 @@ -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(), @@ -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, @@ -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: @@ -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( @@ -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 @@ -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 @@ -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 @@ -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() diff --git a/qutip/visualization.py b/qutip/visualization.py index 36f08cad65..6a14d04ef6 100644 --- a/qutip/visualization.py +++ b/qutip/visualization.py @@ -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)