From fa3c8f752e4bf5382469ace9903f72224008adb3 Mon Sep 17 00:00:00 2001 From: Gerardo Date: Thu, 25 Jan 2024 00:26:42 +0100 Subject: [PATCH] Moved Nk of the SpectralFitter class from initialization to get_fit, this makes more sense as it allows one to compare different number of exponents withouthaving to initialize the class multiple times --- qutip/solver/heom/bofin_baths.py | 435 +++++++++++++++++++------------ qutip/solver/heom/bofin_fit.py | 296 +++++++++++++-------- 2 files changed, 455 insertions(+), 276 deletions(-) diff --git a/qutip/solver/heom/bofin_baths.py b/qutip/solver/heom/bofin_baths.py index 345db33ddc..b020159f2d 100644 --- a/qutip/solver/heom/bofin_baths.py +++ b/qutip/solver/heom/bofin_baths.py @@ -33,7 +33,7 @@ def _isequal(Q1, Q2, tol): - """ Return true if Q1 and Q2 are equal to within the given tolerance. """ + """Return true if Q1 and Q2 are equal to within the given tolerance.""" return _data.iszero(_data.sub(Q1.data, Q2.data), tol=tol) @@ -98,6 +98,7 @@ class BathExponent: All of the parameters are also available as attributes. """ + types = enum.Enum("ExponentType", ["R", "I", "RI", "+", "-"]) def _check_ck2(self, type, ck2): @@ -114,9 +115,7 @@ def _check_ck2(self, type, ck2): def _check_sigma_bar_k_offset(self, type, offset): if type in (self.types["+"], self.types["-"]): if offset is None: - raise ValueError( - "+ and - bath exponents require sigma_bar_k_offset" - ) + raise ValueError("+ and - bath exponents require sigma_bar_k_offset") else: if offset is not None: raise ValueError( @@ -128,8 +127,15 @@ def _type_is_fermionic(self, type): return type in (self.types["+"], self.types["-"]) def __init__( - self, type, dim, Q, ck, vk, ck2=None, - sigma_bar_k_offset=None, tag=None, + self, + type, + dim, + Q, + ck, + vk, + ck2=None, + sigma_bar_k_offset=None, + tag=None, ): if not isinstance(type, self.types): type = self.types[type] @@ -250,8 +256,7 @@ def _check_coup_op(self, Q): raise ValueError("The coupling operator Q must be a Qobj.") def __init__( - self, Q, ck_real, vk_real, ck_imag, vk_imag, combine=True, - tag=None, T=None + self, Q, ck_real, vk_real, ck_imag, vk_imag, combine=True, tag=None, T=None ): self.T = T self._check_cks_and_vks(ck_real, vk_real, ck_imag, vk_imag) @@ -310,9 +315,8 @@ def combine(cls, exponents, rtol=1e-5, atol=1e-7): e1 = remaining.pop(0) group = [e1] for e2 in remaining[:]: - if ( - np.isclose(e1.vk, e2.vk, rtol=rtol, atol=atol) and - _isequal(e1.Q, e2.Q, tol=atol) + if np.isclose(e1.vk, e2.vk, rtol=rtol, atol=atol) and _isequal( + e1.Q, e2.Q, tol=atol ): group.append(e2) remaining.remove(e2) @@ -326,23 +330,35 @@ def combine(cls, exponents, rtol=1e-5, atol=1e-7): ): # the group is either type I or R ck = sum(exp.ck for exp in combine) - new_exponents.append(BathExponent( - exp1.type, None, exp1.Q, ck, exp1.vk, tag=exp1.tag, - )) + new_exponents.append( + BathExponent( + exp1.type, + None, + exp1.Q, + ck, + exp1.vk, + tag=exp1.tag, + ) + ) else: # the group includes both type I and R exponents - ck_R = ( - sum(exp.ck for exp in combine if exp.type == exp.types.R) + - sum(exp.ck for exp in combine if exp.type == exp.types.RI) + ck_R = sum(exp.ck for exp in combine if exp.type == exp.types.R) + sum( + exp.ck for exp in combine if exp.type == exp.types.RI ) - ck_I = ( - sum(exp.ck for exp in combine if exp.type == exp.types.I) + - sum(exp.ck2 for exp in combine if exp.type == exp.types.RI) + ck_I = sum(exp.ck for exp in combine if exp.type == exp.types.I) + sum( + exp.ck2 for exp in combine if exp.type == exp.types.RI + ) + new_exponents.append( + BathExponent( + "RI", + None, + exp1.Q, + ck_R, + exp1.vk, + ck2=ck_I, + tag=exp1.tag, + ) ) - new_exponents.append(BathExponent( - "RI", None, exp1.Q, ck_R, exp1.vk, ck2=ck_I, - tag=exp1.tag, - )) return new_exponents @@ -357,15 +373,16 @@ def spectral_density(self, w): """ raise NotImplementedError( - "The spectral density of this bath has not ben specified") + "The spectral density of this bath has not ben specified" + ) - def correlation_function( - self, t, **kwargs): + def correlation_function(self, t, **kwargs): """ Calculates the correlation function by numerically computing the integral (see equation 6 in 10.1103/PhysRevResearch.5.013181) - where 2n+1 was replaced by cotangent + where 2n+1 was replaced by cotangent for odd functions + Ask if it affects on other places->symmetry consideration Parameters ---------- t : np.array or float @@ -376,16 +393,31 @@ def correlation_function( ---------- The correlation function as an array or float at time t """ - def integrand(w, t): return (1/np.pi)*self.spectral_density(w)*( - (2*self._bose_einstein(w)+1)*np.cos(w*t) - 1j*np.sin(w*t)) - - result = quad_vec( - lambda w: integrand(w, t), - 0, - np.Inf, - **kwargs - ) - return result[0] + if self.spectral_density(1) == -self.spectral_density(-1): + + def integrand(w, t): + return ( + (1 / np.pi) + * self.spectral_density(w) + * ( + (2 * self._bose_einstein(w) + 1) * np.cos(w * t) + - 1j * np.sin(w * t) + ) + ) + + result = quad_vec(lambda w: integrand(w, t), 0, np.Inf, **kwargs) + return result[0] + else: + + def integrand(w, t): + return ( + (1 / np.pi) + * self.spectral_density(w) + * ((self._bose_einstein(w) + 1) * np.exp(1j * w * t)) + ) + + result = quad_vec(lambda w: integrand(w, t), -np.Inf, np.Inf, **kwargs) + return result[0] def _bose_einstein(self, w): """ @@ -404,7 +436,8 @@ def _bose_einstein(self, w): if self.T is None: raise NotImplementedError( - "Bath temperature must be specified for this operation") + "Bath temperature must be specified for this operation" + ) if self.T == 0: return np.zeros_like(w) @@ -413,7 +446,7 @@ def _bose_einstein(self, w): # should normally be zero w = np.array(w, dtype=float) w[w == 0.0] += 1e-6 - return (1 / (np.exp(w / self.T) - 1)) + return 1 / (np.exp(w / self.T) - 1) def power_spectrum(self, w): """ @@ -429,8 +462,7 @@ def power_spectrum(self, w): ---------- The power spectrum of the mode with energy w """ - S = self.spectral_density( - w)*((self._bose_einstein(w) + 1) * 2) + S = self.spectral_density(w) * ((self._bose_einstein(w) + 1) * 2) return S def correlation_function_approx(self, t): @@ -448,17 +480,17 @@ def correlation_function_approx(self, t): ---------- The correlation function of the bath at time t """ - corr = 0+0j + corr = 0 + 0j for exp in self.exponents: if ( - exp.type == BathExponent.types['R'] or - exp.type == BathExponent.types['RI'] + exp.type == BathExponent.types["R"] + or exp.type == BathExponent.types["RI"] ): corr += exp.ck * np.exp(-exp.vk * t) - if exp.type == BathExponent.types['I']: - corr += 1j*exp.ck * np.exp(-exp.vk * t) - if exp.type == BathExponent.types['RI']: - corr += 1j*exp.ck2 * np.exp(-exp.vk * t) + if exp.type == BathExponent.types["I"]: + corr += 1j * exp.ck * np.exp(-exp.vk * t) + if exp.type == BathExponent.types["RI"]: + corr += 1j * exp.ck2 * np.exp(-exp.vk * t) return corr def power_spectrum_approx(self, w): @@ -475,7 +507,7 @@ def power_spectrum_approx(self, w): ---------- The power spectrum of the mode with energy w """ - S = 0+0j + S = 0 + 0j for exp in self.exponents: ck = exp.ck ck2 = exp.ck2 @@ -483,10 +515,10 @@ def power_spectrum_approx(self, w): ck = 0 if ck2 is None: ck2 = 0 - if exp.type == BathExponent.types['I']: - S += 2*np.real((1j*ck-ck2)/(exp.vk - 1j*w)) + if exp.type == BathExponent.types["I"]: + S += 2 * np.real((1j * ck - ck2) / (exp.vk - 1j * w)) else: - S += 2*np.real((ck+1j*ck2)/(exp.vk - 1j*w)) + S += 2 * np.real((ck + 1j * ck2) / (exp.vk - 1j * w)) return S @@ -504,10 +536,7 @@ def spectral_density_approx(self, w): ---------- The spectral density of the mode with energy w """ - J = np.real( - self.power_spectrum_approx(w) / - ((self._bose_einstein(w) + 1) * 2) - ) + J = np.real(self.power_spectrum_approx(w) / ((self._bose_einstein(w) + 1) * 2)) return J @@ -545,7 +574,14 @@ class DrudeLorentzBath(BosonicBath): """ def __init__( - self, Q, lam, gamma, T, Nk, combine=True, tag=None, + self, + Q, + lam, + gamma, + T, + Nk, + combine=True, + tag=None, ): self.lam = lam self.gamma = gamma @@ -556,11 +592,15 @@ def __init__( Nk=Nk, ) - super().__init__(Q, ck_real, vk_real, ck_imag, - vk_imag, combine=combine, tag=tag, T=T) + super().__init__( + Q, ck_real, vk_real, ck_imag, vk_imag, combine=combine, tag=tag, T=T + ) self._dl_terminator = _DrudeLorentzTerminator( - Q=Q, lam=lam, gamma=gamma, T=T, + Q=Q, + lam=lam, + gamma=gamma, + T=T, ) def terminator(self): @@ -589,13 +629,23 @@ def terminator(self): def _matsubara_params(self, lam, gamma, T, Nk): # should this take only self now? - """ Calculate the Matsubara coefficients and frequencies. """ + """Calculate the Matsubara coefficients and frequencies.""" ck_real = [lam * gamma / np.tan(gamma / (2 * T))] - ck_real.extend([ - (8 * lam * gamma * T * np.pi * k * T / - ((2 * np.pi * k * T)**2 - gamma**2)) - for k in range(1, Nk + 1) - ]) + ck_real.extend( + [ + ( + 8 + * lam + * gamma + * T + * np.pi + * k + * T + / ((2 * np.pi * k * T) ** 2 - gamma**2) + ) + for k in range(1, Nk + 1) + ] + ) vk_real = [gamma] vk_real.extend([2 * np.pi * k * T for k in range(1, Nk + 1)]) @@ -618,7 +668,7 @@ def spectral_density(self, w): ---------- The spectral density of the mode with energy w """ - return 2*self.lam*self.gamma*w/(self.gamma**2 + w**2) + return 2 * self.lam * self.gamma * w / (self.gamma**2 + w**2) def correlation_function(self, t, **kwargs): """ @@ -634,17 +684,20 @@ def correlation_function(self, t, **kwargs): kwargs : This may be the Number of exponents to use defaults to 1000 and it should be denoted by Nk """ - if 'Nk' in kwargs.keys(): - Nk = kwargs['Nk'] + if "Nk" in kwargs.keys(): + Nk = kwargs["Nk"] else: Nk = 1000 ck_real, vk_real, ck_imag, vk_imag = self._matsubara_params( - self.lam, self.gamma, self.T, Nk=Nk) + self.lam, self.gamma, self.T, Nk=Nk + ) def C(c, v): - return np.sum([c[i] * np.exp(-np.array(v[i] * t)) - for i in range(len(c))], axis=0) - return C(ck_real, vk_real)+1j*C(ck_imag, vk_imag) + return np.sum( + [c[i] * np.exp(-np.array(v[i] * t)) for i in range(len(c))], axis=0 + ) + + return C(ck_real, vk_real) + 1j * C(ck_imag, vk_imag) class DrudeLorentzPadeBath(BosonicBath): @@ -695,9 +748,7 @@ class DrudeLorentzPadeBath(BosonicBath): bath an exponent is from. """ - def __init__( - self, Q, lam, gamma, T, Nk, combine=True, tag=None - ): + def __init__(self, Q, lam, gamma, T, Nk, combine=True, tag=None): self.lam = lam self.gamma = gamma eta_p, gamma_p = self._corr(lam=lam, gamma=gamma, T=T, Nk=Nk) @@ -709,11 +760,15 @@ def __init__( ck_imag = [np.imag(eta_p[0])] vk_imag = [gamma_p[0]] - super().__init__(Q, ck_real, vk_real, ck_imag, - vk_imag, combine=combine, tag=tag, T=T) + super().__init__( + Q, ck_real, vk_real, ck_imag, vk_imag, combine=combine, tag=tag, T=T + ) self._dl_terminator = _DrudeLorentzTerminator( - Q=Q, lam=lam, gamma=gamma, T=T, + Q=Q, + lam=lam, + gamma=gamma, + T=T, ) def terminator(self): @@ -741,7 +796,7 @@ def terminator(self): return delta, L def _corr(self, lam, gamma, T, Nk): - beta = 1. / T + beta = 1.0 / T kappa, epsilon = self._kappa_epsilon(Nk) eta_p = [lam * gamma * (self._cot(gamma * beta / 2.0) - 1.0j)] @@ -749,15 +804,19 @@ def _corr(self, lam, gamma, T, Nk): for ll in range(1, Nk + 1): eta_p.append( - (kappa[ll] / beta) * 4 * lam * gamma * (epsilon[ll] / beta) - / ((epsilon[ll]**2 / beta**2) - gamma**2) + (kappa[ll] / beta) + * 4 + * lam + * gamma + * (epsilon[ll] / beta) + / ((epsilon[ll] ** 2 / beta**2) - gamma**2) ) gamma_p.append(epsilon[ll] / beta) return eta_p, gamma_p def _cot(self, x): - return 1. / np.tan(x) + return 1.0 / np.tan(x) def _kappa_epsilon(self, Nk): eps = self._calc_eps(Nk) @@ -768,12 +827,11 @@ def _kappa_epsilon(self, Nk): for j in range(Nk): term = prefactor for k in range(Nk - 1): - term *= ( - (chi[k]**2 - eps[j]**2) / - (eps[k]**2 - eps[j]**2 + self._delta(j, k)) + term *= (chi[k] ** 2 - eps[j] ** 2) / ( + eps[k] ** 2 - eps[j] ** 2 + self._delta(j, k) ) for k in [Nk - 1]: - term /= (eps[k]**2 - eps[j]**2 + self._delta(j, k)) + term /= eps[k] ** 2 - eps[j] ** 2 + self._delta(j, k) kappa.append(term) epsilon = [0] + eps @@ -784,23 +842,21 @@ def _delta(self, i, j): return 1.0 if i == j else 0.0 def _calc_eps(self, Nk): - alpha = np.diag([ - 1. / np.sqrt((2 * k + 5) * (2 * k + 3)) - for k in range(2 * Nk - 1) - ], k=1) + alpha = np.diag( + [1.0 / np.sqrt((2 * k + 5) * (2 * k + 3)) for k in range(2 * Nk - 1)], k=1 + ) alpha += alpha.transpose() evals = eigvalsh(alpha) - eps = [-2. / val for val in evals[0: Nk]] + eps = [-2.0 / val for val in evals[0:Nk]] return eps def _calc_chi(self, Nk): - alpha_p = np.diag([ - 1. / np.sqrt((2 * k + 7) * (2 * k + 5)) - for k in range(2 * Nk - 2) - ], k=1) + alpha_p = np.diag( + [1.0 / np.sqrt((2 * k + 7) * (2 * k + 5)) for k in range(2 * Nk - 2)], k=1 + ) alpha_p += alpha_p.transpose() evals = eigvalsh(alpha_p) - chi = [-2. / val for val in evals[0: Nk - 1]] + chi = [-2.0 / val for val in evals[0 : Nk - 1]] return chi def spectral_density(self, w): @@ -817,7 +873,7 @@ def spectral_density(self, w): ---------- The spectral density of the mode with energy w """ - return 2*self.lam*self.gamma*w/(self.gamma**2 + w**2) + return 2 * self.lam * self.gamma * w / (self.gamma**2 + w**2) def correlation_function(self, t, **kwargs): """ @@ -833,22 +889,21 @@ def correlation_function(self, t, **kwargs): kwargs : This may be the Number of exponents to use, defaults to 1000 and it should be denoted by Nk """ - if 'Nk' in kwargs.keys(): - Nk = kwargs['Nk'] + if "Nk" in kwargs.keys(): + Nk = kwargs["Nk"] else: Nk = 1000 - eta_p, gamma_p = self._corr( - lam=self.lam, gamma=self.gamma, T=self.T, Nk=Nk) + eta_p, gamma_p = self._corr(lam=self.lam, gamma=self.gamma, T=self.T, Nk=Nk) t = np.array(t) - C = np.sum([ - eta_p[i] * np.exp(-gamma_p[i] * t) - for i in range(len(gamma_p))], axis=0) + C = np.sum( + [eta_p[i] * np.exp(-gamma_p[i] * t) for i in range(len(gamma_p))], axis=0 + ) return C class _DrudeLorentzTerminator: - """ A class for calculating the terminator of a Drude-Lorentz bath - expansion. + """A class for calculating the terminator of a Drude-Lorentz bath + expansion. """ def __init__(self, Q, lam, gamma, T): @@ -858,7 +913,7 @@ def __init__(self, Q, lam, gamma, T): self.T = T def terminator(self, exponents): - """ Calculate the terminator for a Drude-Lorentz bath. """ + """Calculate the terminator for a Drude-Lorentz bath.""" Q = self.Q lam = self.lam gamma = self.gamma @@ -874,7 +929,7 @@ def terminator(self, exponents): else: delta -= 1j * exp.ck / exp.vk - op = -2*spre(Q)*spost(Q.dag()) + spre(Q.dag()*Q) + spost(Q.dag()*Q) + op = -2 * spre(Q) * spost(Q.dag()) + spre(Q.dag() * Q) + spost(Q.dag() * Q) L_bnd = -delta * op return delta, L_bnd @@ -917,7 +972,15 @@ class UnderDampedBath(BosonicBath): """ def __init__( - self, Q, lam, gamma, w0, T, Nk, combine=True, tag=None, + self, + Q, + lam, + gamma, + w0, + T, + Nk, + combine=True, + tag=None, ): ck_real, vk_real, ck_imag, vk_imag = self._matsubara_params( lam=lam, @@ -930,37 +993,36 @@ def __init__( self.gamma = gamma self.w0 = w0 - super().__init__(Q, ck_real, vk_real, ck_imag, - vk_imag, combine=combine, tag=tag, T=T) + super().__init__( + Q, ck_real, vk_real, ck_imag, vk_imag, combine=combine, tag=tag, T=T + ) @staticmethod def _matsubara_params(lam, gamma, w0, T, Nk): - """ Calculate the Matsubara coefficients and frequencies. """ - beta = 1/T - Om = np.sqrt(w0**2 - (gamma/2)**2) - Gamma = gamma/2. - - ck_real = ([ - (lam**2 / (4 * Om)) - * (1 / np.tanh(beta * (Om + 1.0j * Gamma) / 2)), - (lam**2 / (4*Om)) - * (1 / np.tanh(beta * (Om - 1.0j * Gamma) / 2)), - ]) - - ck_real.extend([ - (-2 * lam**2 * gamma / beta) * (2 * np.pi * k / beta) - / ( - ((Om + 1.0j * Gamma)**2 + (2 * np.pi * k/beta)**2) - * ((Om - 1.0j * Gamma)**2 + (2 * np.pi * k / beta)**2) - ) - for k in range(1, Nk + 1) - ]) + """Calculate the Matsubara coefficients and frequencies.""" + beta = 1 / T + Om = np.sqrt(w0**2 - (gamma / 2) ** 2) + Gamma = gamma / 2.0 + + ck_real = [ + (lam**2 / (4 * Om)) * (1 / np.tanh(beta * (Om + 1.0j * Gamma) / 2)), + (lam**2 / (4 * Om)) * (1 / np.tanh(beta * (Om - 1.0j * Gamma) / 2)), + ] + + ck_real.extend( + [ + (-2 * lam**2 * gamma / beta) + * (2 * np.pi * k / beta) + / ( + ((Om + 1.0j * Gamma) ** 2 + (2 * np.pi * k / beta) ** 2) + * ((Om - 1.0j * Gamma) ** 2 + (2 * np.pi * k / beta) ** 2) + ) + for k in range(1, Nk + 1) + ] + ) vk_real = [-1.0j * Om + Gamma, 1.0j * Om + Gamma] - vk_real.extend([ - 2 * np.pi * k * T - for k in range(1, Nk + 1) - ]) + vk_real.extend([2 * np.pi * k * T for k in range(1, Nk + 1)]) ck_imag = [ 1.0j * lam**2 / (4 * Om), @@ -985,8 +1047,12 @@ def spectral_density(self, w): ---------- The spectral density of the mode with energy w """ - return self.lam**2 * self.gamma * w / ((w**2 - self.w0**2)**2 - + (self.gamma*w)**2) + return ( + self.lam**2 + * self.gamma + * w + / ((w**2 - self.w0**2) ** 2 + (self.gamma * w) ** 2) + ) class FermionicBath(Bath): @@ -1065,12 +1131,28 @@ def __init__(self, Q, ck_plus, vk_plus, ck_minus, vk_minus, tag=None): exponents = [] for ckp, vkp, ckm, vkm in zip(ck_plus, vk_plus, ck_minus, vk_minus): - exponents.append(BathExponent( - "+", 2, Q, ckp, vkp, sigma_bar_k_offset=1, tag=tag, - )) - exponents.append(BathExponent( - "-", 2, Q, ckm, vkm, sigma_bar_k_offset=-1, tag=tag, - )) + exponents.append( + BathExponent( + "+", + 2, + Q, + ckp, + vkp, + sigma_bar_k_offset=1, + tag=tag, + ) + ) + exponents.append( + BathExponent( + "-", + 2, + Q, + ckm, + vkm, + sigma_bar_k_offset=-1, + tag=tag, + ) + ) super().__init__(exponents) @@ -1118,13 +1200,18 @@ def __init__(self, Q, gamma, w, mu, T, Nk, tag=None): ck_minus, vk_minus = self._corr(gamma, w, mu, T, Nk, sigma=-1.0) super().__init__( - Q, ck_plus, vk_plus, ck_minus, vk_minus, tag=tag, + Q, + ck_plus, + vk_plus, + ck_minus, + vk_minus, + tag=tag, ) def _corr(self, gamma, w, mu, T, Nk, sigma): - beta = 1. / T - kappa = [0.] - kappa.extend([1. for _ in range(1, Nk + 1)]) + beta = 1.0 / T + kappa = [0.0] + kappa.extend([1.0 for _ in range(1, Nk + 1)]) epsilon = [0] epsilon.extend([(2 * ll - 1) * np.pi for ll in range(1, Nk + 1)]) @@ -1136,8 +1223,11 @@ def f(x): for ll in range(1, Nk + 1): eta_list.append( - -1.0j * (kappa[ll] / beta) * gamma * w**2 / - (-(epsilon[ll]**2 / beta**2) + w**2) + -1.0j + * (kappa[ll] / beta) + * gamma + * w**2 + / (-(epsilon[ll] ** 2 / beta**2) + w**2) ) gamma_list.append(epsilon[ll] / beta - sigma * 1.0j * mu) @@ -1197,17 +1287,22 @@ def __init__(self, Q, gamma, w, mu, T, Nk, tag=None): ck_minus, vk_minus = self._corr(gamma, w, mu, T, Nk, sigma=-1.0) super().__init__( - Q, ck_plus, vk_plus, ck_minus, vk_minus, tag=tag, + Q, + ck_plus, + vk_plus, + ck_minus, + vk_minus, + tag=tag, ) def _corr(self, gamma, w, mu, T, Nk, sigma): - beta = 1. / T + beta = 1.0 / T kappa, epsilon = self._kappa_epsilon(Nk) def f_approx(x): f = 0.5 for ll in range(1, Nk + 1): - f = f - 2 * kappa[ll] * x / (x**2 + epsilon[ll]**2) + f = f - 2 * kappa[ll] * x / (x**2 + epsilon[ll] ** 2) return f eta_list = [0.5 * gamma * w * f_approx(1.0j * beta * w)] @@ -1215,8 +1310,11 @@ def f_approx(x): for ll in range(1, Nk + 1): eta_list.append( - -1.0j * (kappa[ll] / beta) * gamma * w**2 - / (-(epsilon[ll]**2 / beta**2) + w**2) + -1.0j + * (kappa[ll] / beta) + * gamma + * w**2 + / (-(epsilon[ll] ** 2 / beta**2) + w**2) ) gamma_list.append(epsilon[ll] / beta - sigma * 1.0j * mu) @@ -1231,12 +1329,11 @@ def _kappa_epsilon(self, Nk): for j in range(Nk): term = prefactor for k in range(Nk - 1): - term *= ( - (chi[k]**2 - eps[j]**2) / - (eps[k]**2 - eps[j]**2 + self._delta(j, k)) + term *= (chi[k] ** 2 - eps[j] ** 2) / ( + eps[k] ** 2 - eps[j] ** 2 + self._delta(j, k) ) for k in [Nk - 1]: - term /= (eps[k]**2 - eps[j]**2 + self._delta(j, k)) + term /= eps[k] ** 2 - eps[j] ** 2 + self._delta(j, k) kappa.append(term) epsilon = [0] + eps @@ -1247,22 +1344,20 @@ def _delta(self, i, j): return 1.0 if i == j else 0.0 def _calc_eps(self, Nk): - alpha = np.diag([ - 1. / np.sqrt((2 * k + 3) * (2 * k + 1)) - for k in range(2 * Nk - 1) - ], k=1) + alpha = np.diag( + [1.0 / np.sqrt((2 * k + 3) * (2 * k + 1)) for k in range(2 * Nk - 1)], k=1 + ) alpha += alpha.transpose() evals = eigvalsh(alpha) - eps = [-2. / val for val in evals[0: Nk]] + eps = [-2.0 / val for val in evals[0:Nk]] return eps def _calc_chi(self, Nk): - alpha_p = np.diag([ - 1. / np.sqrt((2 * k + 5) * (2 * k + 3)) - for k in range(2 * Nk - 2) - ], k=1) + alpha_p = np.diag( + [1.0 / np.sqrt((2 * k + 5) * (2 * k + 3)) for k in range(2 * Nk - 2)], k=1 + ) alpha_p += alpha_p.transpose() evals = eigvalsh(alpha_p) - chi = [-2. / val for val in evals[0: Nk - 1]] + chi = [-2.0 / val for val in evals[0 : Nk - 1]] return chi diff --git a/qutip/solver/heom/bofin_fit.py b/qutip/solver/heom/bofin_fit.py index fd6477897c..60da6e0ed4 100644 --- a/qutip/solver/heom/bofin_fit.py +++ b/qutip/solver/heom/bofin_fit.py @@ -10,8 +10,10 @@ import numpy as np from time import time + try: from mpmath import mp + _mpmath_available = True except ModuleNotFoundError: _mpmath_available = False @@ -33,9 +35,6 @@ class SpectralFitter: T : float Bath temperature. - Nk : int - Number of exponential terms used to approximate the bath correlation - functions. w : np.array The range on which to perform the fit @@ -43,10 +42,9 @@ class SpectralFitter: The spectral density to be fitted as an array or function """ - def __init__(self, T, Q, w, J, Nk=14): + def __init__(self, T, Q, w, J): self.Q = Q self.T = T - self.Nk = Nk self.set_function(w, J) @@ -92,13 +90,13 @@ def _spectral_density_approx(self, w, a, b, c): * a[i] * b[i] * w - / (((w + c[i]) ** 2 + b[i] ** 2) - * ((w - c[i]) ** 2 + b[i] ** 2)) + / (((w + c[i]) ** 2 + b[i] ** 2) * ((w - c[i]) ** 2 + b[i] ** 2)) ) return tot def get_fit( self, + Nk=5, N=None, final_rmse=5e-6, lower=None, @@ -121,6 +119,9 @@ def get_fit( N : optional,int Number of underdamped oscillators to use, if set to False it is found automatically. + Nk: optional,int + Number of exponential terms used to approximate the bath correlation + functions. Defaults to 5 final_rmse : float Desired normalized root mean squared error . lower : list @@ -141,22 +142,34 @@ def get_fit( start = time() rmse, params = _run_fit( - self._spectral_density_approx, self._J_array, self._w, final_rmse, - N=N, sigma=sigma, label="Spectral Density", guesses=guesses, - lower=lower, upper=upper) + self._spectral_density_approx, + self._J_array, + self._w, + final_rmse, + N=N, + sigma=sigma, + label="Spectral Density", + guesses=guesses, + lower=lower, + upper=upper, + ) spec_n = len(params[0]) end = time() fit_time = end - start - bath = self._generate_bath(params) + bath = self._generate_bath(params, Nk) bath.spectral_density = self._J_fun - summary = gen_summary( - fit_time, rmse, N, "The Spectral Density", *params) + summary = gen_summary(fit_time, rmse, N, "The Spectral Density", *params) self.fitInfo = { - "fit_time": fit_time, "rmse": rmse, "N": spec_n, "params": params, - "Nk": self.Nk, "summary": summary} + "fit_time": fit_time, + "rmse": rmse, + "N": spec_n, + "params": params, + "Nk": Nk, + "summary": summary, + } return bath, self.fitInfo - def _generate_bath(self, params): + def _generate_bath(self, params, Nk): """ Obtains the bath exponents from the list of fit parameters some transformations are done, to reverse the ones in the UnderDampedBath @@ -179,9 +192,7 @@ def _generate_bath(self, params): for i in range(len(w0)) ] ) - lam = np.sqrt( - lam + 0j - ) + lam = np.sqrt(lam + 0j) # both w0, and lam modifications are needed to input the # right value of the fit into the Underdamped bath ckAR = [] @@ -191,7 +202,8 @@ def _generate_bath(self, params): for lamt, Gamma, Om in zip(lam, gamma, w0): coeffs = UnderDampedBath._matsubara_params( - lamt, 2 * Gamma, Om + 0j, self.T, self.Nk) + lamt, 2 * Gamma, Om + 0j, self.T, Nk + ) ckAR.extend(coeffs[0]) vkAR.extend(coeffs[1]) ckAI.extend(coeffs[2]) @@ -236,10 +248,9 @@ def set_function(self, t, C): else: self._t = t self._C_array = C - self._C_fun_r = InterpolatedUnivariateSpline( - t, np.real(C)) + self._C_fun_r = InterpolatedUnivariateSpline(t, np.real(C)) self._C_fun_i = InterpolatedUnivariateSpline(t, np.imag(C)) - self._C_fun = lambda t: self._C_fun_r(t)+1j*self._C_fun_i(t) + self._C_fun = lambda t: self._C_fun_r(t) + 1j * self._C_fun_i(t) def _corr_approx(self, t, a, b, c): """This is the form of the correlation function to be used for fitting @@ -309,19 +320,33 @@ def get_fit( # Fit real part start_real = time() rmse_real, params_real = _run_fit( - funcx=lambda * args: np.real(self._corr_approx(*args)), + funcx=lambda *args: np.real(self._corr_approx(*args)), funcy=np.real(self._C_array), - x=self._t, final_rmse=final_rmse, label="correlation_real", - sigma=sigma, N=Nr, guesses=guesses, lower=lower, upper=upper) + x=self._t, + final_rmse=final_rmse, + label="correlation_real", + sigma=sigma, + N=Nr, + guesses=guesses, + lower=lower, + upper=upper, + ) end_real = time() # Fit imaginary part start_imag = time() rmse_imag, params_imag = _run_fit( lambda *args: np.imag(self._corr_approx(*args)), - np.imag(self._C_array), self._t, final_rmse, - label="correlation_imag", sigma=sigma, N=Ni, - guesses=guesses, lower=lower, upper=upper) + np.imag(self._C_array), + self._t, + final_rmse, + label="correlation_imag", + sigma=sigma, + N=Ni, + guesses=guesses, + lower=lower, + upper=upper, + ) end_imag = time() # Calculate Fit Times @@ -330,21 +355,33 @@ def get_fit( # Generate summary full_summary = two_column_summary( - params_real, params_imag, fit_time_real, fit_time_imag, Nr, Ni, - rmse_imag, rmse_real) - - self.fitInfo = {"Nr": len(params_real[0]), "Ni": len(params_imag[0]), - "fit_time_real": fit_time_real, - "fit_time_imag": fit_time_imag, - "rmse_real": rmse_real, "rmse_imag": rmse_imag, - "params_real": params_real, - "params_imag": params_imag, "summary": full_summary} + params_real, + params_imag, + fit_time_real, + fit_time_imag, + Nr, + Ni, + rmse_imag, + rmse_real, + ) + + self.fitInfo = { + "Nr": len(params_real[0]), + "Ni": len(params_imag[0]), + "fit_time_real": fit_time_real, + "fit_time_imag": fit_time_imag, + "rmse_real": rmse_real, + "rmse_imag": rmse_imag, + "params_real": params_real, + "params_imag": params_imag, + "summary": full_summary, + } bath = self._generate_bath(params_real, params_imag) bath.correlation_function = self._C_fun return bath, self.fitInfo def _generate_bath(self, params_real, params_imag): - """ Calculate the Matsubara coefficients and frequencies for the + """Calculate the Matsubara coefficients and frequencies for the fitted underdamped oscillators and generate the corresponding bosonic bath @@ -374,8 +411,7 @@ def _generate_bath(self, params_real, params_imag): vkAI = [-x - 1.0j * y for x, y in zip(gamma2, w02)] vkAI.extend([-x + 1.0j * y for x, y in zip(gamma2, w02)]) - return BosonicBath( - self.Q, ckAR, vkAR, ckAI, vkAI, T=self.T) + return BosonicBath(self.Q, ckAR, vkAR, ckAI, vkAI, T=self.T) class OhmicBath: @@ -404,9 +440,7 @@ def __init__(self, T, Q, alpha, wc, s): self.Q = Q self.T = T if _mpmath_available is False: - print( - "The mpmath module is needed for the description" - " of Ohmic baths") + print("The mpmath module is needed for the description" " of Ohmic baths") def spectral_density(self, w): """Calculates the Ohmic spectral density @@ -438,27 +472,41 @@ def correlation_function(self, t): (1 / np.pi) * self.alpha * self.wc ** (1 - self.s) - * (1/self.T) ** (-(self.s + 1)) + * (1 / self.T) ** (-(self.s + 1)) * mp.gamma(self.s + 1) ) - z1_u = (1 + (1/self.T) * self.wc - 1.0j * - self.wc * t) / ((1/self.T) * self.wc) - z2_u = (1 + 1.0j * self.wc * t) / ((1/self.T) * self.wc) + z1_u = (1 + (1 / self.T) * self.wc - 1.0j * self.wc * t) / ( + (1 / self.T) * self.wc + ) + z2_u = (1 + 1.0j * self.wc * t) / ((1 / self.T) * self.wc) return np.array( [ - complex( - corr * (mp.zeta(self.s + 1, u1) + - mp.zeta(self.s + 1, u2))) - for u1, u2 in zip(z1_u, z2_u)], - dtype=np.complex128,) + complex(corr * (mp.zeta(self.s + 1, u1) + mp.zeta(self.s + 1, u2))) + for u1, u2 in zip(z1_u, z2_u) + ], + dtype=np.complex128, + ) else: - corr = (1 / np.pi)*self.alpha*self.wc**(self.s+1) * \ - mp.gamma(self.s+1)*(1+1j*self.wc*t)**(-(self.s+1)) + corr = ( + (1 / np.pi) + * self.alpha + * self.wc ** (self.s + 1) + * mp.gamma(self.s + 1) + * (1 + 1j * self.wc * t) ** (-(self.s + 1)) + ) return np.array(corr, dtype=np.complex128) def make_correlation_fit( - self, x, rmse=1e-5, lower=None, upper=None, - sigma=None, guesses=None, Nr=None, Ni=None): + self, + x, + rmse=1e-5, + lower=None, + upper=None, + sigma=None, + guesses=None, + Nr=None, + Ni=None, + ): """ Provides a fit to the spectral density or corelation function with N underdamped oscillators baths, This function gets the @@ -492,14 +540,28 @@ def make_correlation_fit( """ C = self.correlation_function(x) fc = CorrelationFitter(self.Q, self.T, x, C) - bath, fitInfo = fc.get_fit(final_rmse=rmse, - lower=lower, upper=upper, - sigma=sigma, guesses=guesses, - Nr=Nr, Ni=Ni) + bath, fitInfo = fc.get_fit( + final_rmse=rmse, + lower=lower, + upper=upper, + sigma=sigma, + guesses=guesses, + Nr=Nr, + Ni=Ni, + ) return bath, fitInfo - def make_spectral_fit(self, x, rmse=1e-5, lower=None, upper=None, - sigma=None, guesses=None, N=None, Nk=5): + def make_spectral_fit( + self, + x, + rmse=1e-5, + lower=None, + upper=None, + sigma=None, + guesses=None, + N=None, + Nk=5, + ): """ Provides a fit to the spectral density or corelation function with N underdamped oscillators baths, This function gets the @@ -533,11 +595,12 @@ def make_spectral_fit(self, x, rmse=1e-5, lower=None, upper=None, """ J = self.spectral_density(x) fs = SpectralFitter(T=self.T, Q=self.Q, w=x, J=J, Nk=Nk) - bath, fitInfo = fs.get_fit(N=N, final_rmse=rmse, lower=lower, - upper=upper, - sigma=sigma, guesses=guesses) + bath, fitInfo = fs.get_fit( + N=N, final_rmse=rmse, lower=lower, upper=upper, sigma=sigma, guesses=guesses + ) return bath, fitInfo + # Utility functions @@ -550,15 +613,12 @@ def unpack(params): """Unpack parameter lists for fitting.""" N = len(params) // 3 a = params[:N] - b = params[N: 2 * N] - c = params[2 * N:] + b = params[N : 2 * N] + c = params[2 * N :] return a, b, c -def _leastsq( - func, y, x, guesses=None, - lower=None, upper=None, sigma=None -): +def _leastsq(func, y, x, guesses=None, lower=None, upper=None, sigma=None): """ Performs nonlinear least squares to fit the function func to x and y @@ -631,14 +691,20 @@ def _rmse(func, x, y, lam, gamma, w0): to zero the better the fit. """ yhat = func(x, lam, gamma, w0) - rmse = np.sqrt(np.mean((yhat - y) ** 2) / len(y)) / \ - (np.max(y) - np.min(y)) + rmse = np.sqrt(np.mean((yhat - y) ** 2) / len(y)) / (np.max(y) - np.min(y)) return rmse def _fit( - func, C, t, N=4, label="correlation_real", - guesses=None, lower=None, upper=None, sigma=None + func, + C, + t, + N=4, + label="correlation_real", + guesses=None, + lower=None, + upper=None, + sigma=None, ): """ Performs a fit the function func to t and C, with N number of @@ -672,12 +738,12 @@ def _fit( """ try: if None in [guesses, lower, upper, sigma]: - raise Exception( - "No parameters for the fit provided, using default ones" - ) + raise Exception("No parameters for the fit provided, using default ones") except Exception: sigma = 1e-4 C_max = abs(max(C, key=abs)) + if C_max == 0: ## When The correlation does not have imaginary or real part + C_max = 1e-12 # so no error on bounds wc = t[np.argmax(C)] if label == "correlation_real": guesses = pack([C_max] * N, [-wc] * N, [wc] * N) @@ -690,10 +756,8 @@ def _fit( else: guesses = pack([C_max] * N, [wc] * N, [wc] * N) - lower = pack([-100 * C_max] * N, - [0.1 * wc] * N, [0.1 * wc] * N) - upper = pack([100 * C_max] * N, - [100 * wc] * N, [100 * wc] * N) + lower = pack([-100 * C_max] * N, [0.1 * wc] * N, [0.1 * wc] * N) + upper = pack([100 * C_max] * N, [100 * wc] * N, [100 * wc] * N) lam, gamma, w0 = _leastsq( func, @@ -704,13 +768,27 @@ def _fit( lower=lower, upper=upper, ) - rmse = _rmse(func, t, C, lam=lam, gamma=gamma, w0=w0) - params = [lam, gamma, w0] + if C_max == 0: + rmse = 0 + params = [0, 0, 0] + else: + rmse = _rmse(func, t, C, lam=lam, gamma=gamma, w0=w0) + params = [lam, gamma, w0] return rmse, params -def _run_fit(funcx, funcy, x, final_rmse, label=None, N=None, - sigma=None, guesses=None, lower=None, upper=None): +def _run_fit( + funcx, + funcy, + x, + final_rmse, + label=None, + N=None, + sigma=None, + guesses=None, + lower=None, + upper=None, +): """ It iteratively tries to fit the fucx to funcy on the interval x, if the Ns are provided the fit is done with N modes, if they are @@ -756,7 +834,7 @@ def _run_fit(funcx, funcy, x, final_rmse, label=None, N=None, flag = False else: flag = True - N = N-1 + N = N - 1 rmse1 = 8 while rmse1 > final_rmse: N += 1 @@ -777,23 +855,21 @@ def _run_fit(funcx, funcy, x, final_rmse, label=None, N=None, def gen_summary(time, rmse, N, label, lam, gamma, w0): - summary = f"Result of fitting {label} "\ - f"with {N} terms: \n \n {'Parameters': <10}|"\ + summary = ( + f"Result of fitting {label} " + f"with {N} terms: \n \n {'Parameters': <10}|" f"{'lam': ^10}|{'gamma': ^10}|{'w0': >5} \n " + ) for k in range(len(gamma)): - summary += ( - f"{k+1: <10}|{lam[k]: ^10.2e}|{gamma[k]:^10.2e}|" - f"{w0[k]:>5.2e}\n " - ) - summary += f"\nA normalized RMSE of {rmse: .2e}"\ - f" was obtained for the {label}\n" + summary += f"{k+1: <10}|{lam[k]: ^10.2e}|{gamma[k]:^10.2e}|" f"{w0[k]:>5.2e}\n " + summary += f"\nA normalized RMSE of {rmse: .2e}" f" was obtained for the {label}\n" summary += f" The current fit took {time: 2f} seconds" return summary def two_column_summary( - params_real, params_imag, fit_time_real, fit_time_imag, Nr, Ni, - rmse_imag, rmse_real): + params_real, params_imag, fit_time_real, fit_time_imag, Nr, Ni, rmse_imag, rmse_real +): lam, gamma, w0 = params_real lam2, gamma2, w02 = params_imag @@ -802,24 +878,32 @@ def two_column_summary( fit_time_real, rmse_real, Nr, - "The Real Part Of \n the Correlation Function", lam, gamma, - w0) + "The Real Part Of \n the Correlation Function", + lam, + gamma, + w0, + ) summary_imag = gen_summary( fit_time_imag, rmse_imag, Ni, - "The Imaginary Part \n Of the Correlation Function", lam2, - gamma2, w02) + "The Imaginary Part \n Of the Correlation Function", + lam2, + gamma2, + w02, + ) full_summary = "Fit correlation class instance: \n \n" lines_real = summary_real.splitlines() lines_imag = summary_imag.splitlines() max_lines = max(len(lines_real), len(lines_imag)) # Fill the shorter string with blank lines - lines_real = lines_real[:-1] + (max_lines - len(lines_real) - ) * [""] + [lines_real[-1]] - lines_imag = lines_imag[:-1] + (max_lines - len(lines_imag) - ) * [""] + [lines_imag[-1]] + lines_real = ( + lines_real[:-1] + (max_lines - len(lines_real)) * [""] + [lines_real[-1]] + ) + lines_imag = ( + lines_imag[:-1] + (max_lines - len(lines_imag)) * [""] + [lines_imag[-1]] + ) # Find the maximum line length in each column max_length1 = max(len(line) for line in lines_real) max_length2 = max(len(line) for line in lines_imag)