Bosonic fitting

doc/changes/2260.feature

@@ -0,0 +1 @@
+Added support for arbitrary spectral densities in HEOM by fitting underdamped modes, added an Ohmic Bath class

qutip/solver/heom/__init__.py

@@ -31,6 +31,9 @@
     "HSolverDL",
     "HierarchyADOs",
     "HierarchyADOsState",
+    "SpectralFitter",
+    "CorrelationFitter",
+    "OhmicBath",
 ]
 
 from .bofin_baths import (
@@ -53,3 +56,9 @@
     HierarchyADOs,
     HierarchyADOsState,
 )
+
+from .bofin_fit import (
+    CorrelationFitter,
+    SpectralFitter,
+    OhmicBath,
+)

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 scipy.integrate import quad_vec
 
 __all__ = [
     "BathExponent",
@@ -243,8 +244,9 @@ def _check_coup_op(self, Q):
 
     def __init__(
             self, Q, ck_real, vk_real, ck_imag, vk_imag, combine=True,
-            tag=None,
+            tag=None, T=None
     ):
+        self.T = T
         self._check_cks_and_vks(ck_real, vk_real, ck_imag, vk_imag)
         self._check_coup_op(Q)
 
@@ -337,6 +339,85 @@ def combine(cls, exponents, rtol=1e-5, atol=1e-7):
 
         return new_exponents
 
+    def spectral_density(self, w):
+        raise NotImplemented
+
+    def correlation_function(
+            self, t, epsabs=1e-6, epsrel=1e-6, method='gk15'):
+        """ Calculates the correlation function
+            by numerically computing the integral
+        Parameters
+        ----------
+        t : np.array or float
+            the time at which to evaluare the correlation function
+        epsabs : float
+            Absolute error tolerance
+        epsrel : float
+            Relative error tolerance
+        """
+        def c_i(w, t): return (1/np.pi)*self.spectral_density(w)*(
+            (1/np.tanh(w/(2*self.T)))*np.cos(w*t) - 1j*np.sin(w*t))
+
+        def c(t): return quad_vec(
+            lambda x: c_i(x, t),
+            0,
+            np.Inf,
+            epsabs=epsabs,
+            epsrel=epsrel,
+            quadrature=method,
+        )[0]
+        return c(t)
+
+    def bose_einstein(self, w):
+        try:
+            return (1 / (np.e**(w / self.T) - 1))
+        except ZeroDivisionError:
+            return 0
+
+    def power_spectrum(self, w):
+        """Calculates the power spectrum from the spectral density
+        """
+        S = self.spectral_density(
+            w)*((self.bose_einstein(w) + 1) * 2)
+        return S
+
+    def correlation_function_approx(self, t):
+        """Computes the correlation function from
+         the exponents"""
+        corr = 0+0j
+        for exp in self.exponents:
+            if (
+                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)
+        return corr
+
+    def power_spectrum_approx(self, w):
+        S = 0+0j
+        for exp in self.exponents:
+            if exp.ck is None:
+                exp.ck = 0
+            if exp.ck2 is None:
+                exp.ck2 = 0
+            if exp.type == BathExponent.types['I']:
+                S += 2*np.real((1j*exp.ck-exp.ck2)/(exp.vk - 1j*w))
+            else:
+                S += 2*np.real((exp.ck+1j*exp.ck2)/(exp.vk - 1j*w))
+
+        return S
+
+    def spectral_density_approx(self, w):
+        J = np.real(
+            self.power_spectrum_approx(w) /
+            ((self.bose_einstein(w) + 1) * 2)
+        )
+        return J
+
 
 class DrudeLorentzBath(BosonicBath):
     """
@@ -373,16 +454,18 @@ class DrudeLorentzBath(BosonicBath):
     def __init__(
         self, Q, lam, gamma, T, Nk, combine=True, tag=None,
     ):
+        self.lam = lam
+        self.gamma = gamma
+        self.T = T
         ck_real, vk_real, ck_imag, vk_imag = self._matsubara_params(
             lam=lam,
             gamma=gamma,
             T=T,
             Nk=Nk,
         )
 
-        super().__init__(
-            Q, ck_real, vk_real, ck_imag, vk_imag, combine=combine, tag=tag,
-        )
+        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,
@@ -427,6 +510,9 @@ def _matsubara_params(self, lam, gamma, T, Nk):
         vk_imag = [gamma]
 
         return ck_real, vk_real, ck_imag, vk_imag
+
+    def spectral_density(self, w):
+        return 2*self.lam*self.gamma*w/(self.gamma**2 + w**2)
 
 
 class DrudeLorentzPadeBath(BosonicBath):
@@ -437,7 +523,7 @@ class DrudeLorentzPadeBath(BosonicBath):
     A Padé approximant is a sum-over-poles expansion (
     see https://en.wikipedia.org/wiki/Pad%C3%A9_approximant).
 
-    The application of the Padé method to spectrum decompoisitions is described
+    The application of the Padé method to spectrum decompositions is described
     in "Padé spectrum decompositions of quantum distribution functions and
     optimal hierarchical equations of motion construction for quantum open
     systems" [1].
@@ -479,6 +565,8 @@ class DrudeLorentzPadeBath(BosonicBath):
     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)
 
         ck_real = [np.real(eta) for eta in eta_p]
@@ -488,9 +576,8 @@ 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,
-        )
+        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,
@@ -565,8 +652,8 @@ def _delta(self, i, j):
 
     def _calc_eps(self, Nk):
         alpha = np.diag([
-                1. / np.sqrt((2 * k + 5) * (2 * k + 3))
-                for k in range(2 * Nk - 1)
+            1. / np.sqrt((2 * k + 5) * (2 * k + 3))
+            for k in range(2 * Nk - 1)
         ], k=1)
         alpha += alpha.transpose()
         evals = eigvalsh(alpha)
@@ -575,14 +662,17 @@ def _calc_eps(self, Nk):
 
     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)
+            1. / 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]]
         return chi
 
+    def spectral_density(self, w):
+        return 2*self.lam*self.gamma*w/(self.gamma**2 + w**2)
+
 
 class _DrudeLorentzTerminator:
     """ A class for calculating the terminator of a Drude-Lorentz bath
@@ -662,12 +752,16 @@ def __init__(
             T=T,
             Nk=Nk,
         )
+        self.lam = lam
+        self.gamma = gamma
+        self.w0 = w0
+        self.T = T
 
-        super().__init__(
-            Q, ck_real, vk_real, ck_imag, vk_imag, combine=combine, tag=tag,
-        )
+        super().__init__(Q, ck_real, vk_real, ck_imag,
+                         vk_imag, combine=combine, tag=tag, T=T)
 
-    def _matsubara_params(self, lam, gamma, w0, T, Nk):
+    @staticmethod
+    def _matsubara_params(lam, gamma, w0, T, Nk):
         """ Calculate the Matsubara coefficents and frequencies. """
         beta = 1/T
         Om = np.sqrt(w0**2 - (gamma/2)**2)
@@ -704,6 +798,10 @@ def _matsubara_params(self, lam, gamma, w0, T, Nk):
 
         return ck_real, vk_real, ck_imag, vk_imag
 
+    def spectral_density(self, w):
+        return self.lam**2 * self.gamma * w / ((w**2 - self.w0**2)**2
+                                               + (self.gamma*w)**2)
+
 
 class FermionicBath(Bath):
     """
@@ -962,8 +1060,8 @@ def _delta(self, i, j):
 
     def _calc_eps(self, Nk):
         alpha = np.diag([
-                1. / np.sqrt((2 * k + 3) * (2 * k + 1))
-                for k in range(2 * Nk - 1)
+            1. / np.sqrt((2 * k + 3) * (2 * k + 1))
+            for k in range(2 * Nk - 1)
         ], k=1)
         alpha += alpha.transpose()
 
@@ -973,8 +1071,8 @@ def _calc_eps(self, Nk):
 
     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)
+            1. / 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)