Merge branch 'bath_pr_refactor2' of https://github.com/gsuarezr/qutip…

…_gsoc_app into bath_pr_refactor2

qutip/core/environment.py

@@ -33,7 +33,8 @@
 except ModuleNotFoundError:
     _mpmath_available = False
 
-from ..utilities import (n_thermal, iterated_fit)
+from ..utilities import (n_thermal, iterated_fit, aaa,
+                         matrix_pencil, prony, esprit)
 from .superoperator import spre, spost
 from .qobj import Qobj
 
@@ -767,6 +768,133 @@ def approx_by_sd_fit(
             ckAR, vkAR, ckAI, vkAI, combine=combine, T=self.T, tag=tag)
         return approx_env, fit_info
 
+    def approx_by_aaa(
+        self,
+        wlist: ArrayLike,
+        tol: float = 1e-13,
+        N_max: int = 10,
+        combine: bool = True,
+        tag: Any = None,
+    ) -> tuple[ExponentialBosonicEnvironment, dict[str, Any]]:
+
+        _, pol, res, _, _ = aaa(self.power_spectrum, wlist,
+                                tol=tol,
+                                max_iter=N_max*2)
+        mask = np.imag(pol) < 0
+
+        new_pols, new_res = pol[mask], res[mask]
+
+        # Create complex conjugates for both vk and ck
+        vk = 1j*new_pols
+        ck = -1j*new_res
+        ckAR = []
+        vkAR = []
+        ckAI = []
+        vkAI = []
+
+        for term in range(len(vk)):
+            a, b, c, d = np.real(ck[term]), -np.real(vk[term]), - \
+                np.imag(vk[term]), np.imag(ck[term])
+            ckAR.extend([(a + 1j * d) / 2, (a - 1j * d) / 2])
+            vkAR.extend([-b - 1j * c, -b + 1j * c])
+            ckAI.extend([-1j * (a + 1j * d) / 2, 1j * (a - 1j * d) / 2])
+
+        cls = ExponentialBosonicEnvironment(
+            ck_real=ckAR, vk_real=vkAR, ck_imag=ckAI,
+            vk_imag=vkAR, T=self.T, combine=combine, tag=tag)
+        return cls
+
+    def approx_by_mp(
+        self,
+        tlist: ArrayLike,
+        Nr: int = 3,
+        Ni: int = 3,
+        combine: bool = True,
+        tag: Any = None,
+    ) -> tuple[ExponentialBosonicEnvironment, dict[str, Any]]:
+        # amp, phases = matrix_pencil(self.correlation_function(tlist), N)
+        # ckAR = amp.real
+        # ckAI = amp.imag
+        # vk = -((len(tlist)-1)/tlist[-1]) * \
+        #     (np.log(np.abs(phases))+1j*np.angle(phases))
+        amp, phases = matrix_pencil(self.correlation_function(tlist).real, Nr)
+        amp2, phases2 = matrix_pencil(self.correlation_function(tlist).imag, Ni)
+        ckAR = amp
+        ckAI = amp2
+        vkAR = -((len(tlist)-1)/tlist[-1]) * \
+            (np.log(np.abs(phases))+1j*np.angle(phases))
+        vkAI = -((len(tlist)-1)/tlist[-1]) * \
+            (np.log(np.abs(phases2))+1j*np.angle(phases2))
+
+        cls = ExponentialBosonicEnvironment(
+            ck_real=ckAR, vk_real=vkAR, ck_imag=ckAI,
+            vk_imag=vkAI, T=self.T, combine=combine, tag=tag)
+        return cls, (amp, phases)
+
+    def approx_by_prony(
+        self,
+        tlist: ArrayLike,
+        Nr: int = 3,
+        Ni: int = 3,
+        combine: bool = True,
+        tag: Any = None,
+    ) -> tuple[ExponentialBosonicEnvironment, dict[str, Any]]:
+        # amp, phases = prony(self.correlation_function(tlist), N)
+        # ckAR = amp.real
+        # ckAI = amp.imag
+        # vk = -((len(tlist)-1)/tlist[-1]) * \
+        #     (np.log(np.abs(phases))+1j*np.angle(phases))
+        # TODO: Why doesn't the heom construction work when one fits the
+        # complex signal rather than the real and imaginary parts separately?
+        # Probably a dumb reason but I couldn't figure it out
+        # I was passing exponents incorrectly, they should be passed the same
+        # way as they are passed to AAA, however Neill and Paul may like
+        # having these separate so ask before modification (I typically find
+        # better fits with less exponent's when fitting the complex signal)
+        amp, phases = prony(self.correlation_function(tlist).real, Nr)
+        amp2, phases2 = prony(self.correlation_function(tlist).imag, Ni)
+        ckAR = amp
+        ckAI = amp2
+        vkAR = -((len(tlist)-1)/tlist[-1]) * \
+            (np.log(np.abs(phases))+1j*np.angle(phases))
+        vkAI = -((len(tlist)-1)/tlist[-1]) * \
+            (np.log(np.abs(phases2))+1j*np.angle(phases2))
+
+        cls = ExponentialBosonicEnvironment(
+            ck_real=ckAR, vk_real=vkAR, vk_imag=vkAI,
+            ck_imag=ckAI, combine=combine, tag=tag)
+        return cls, (amp, phases)
+
+    def approx_by_esprit(
+        self,
+        tlist: ArrayLike,
+        Nr: int = 3,
+        Ni: int = 3,
+        combine: bool = True,
+        tag: Any = None,
+    ) -> tuple[ExponentialBosonicEnvironment, dict[str, Any]]:
+        # amp, phases = esprit(self.correlation_function(tlist), N)
+        # ckAR = amp.real
+        # ckAI = amp.imag
+        # vk = -((len(tlist)-1)/tlist[-1]) * \
+        #     (np.log(np.abs(phases))+1j*np.angle(phases))
+        # cls = ExponentialBosonicEnvironment(
+        #     ck_real=ckAR, vk_real=vk, vk_imag=vk,
+        #     ck_imag=ckAI, combine=combine, tag=tag)
+        amp, phases = esprit(self.correlation_function(tlist).real, Nr)
+        amp2, phases2 = esprit(self.correlation_function(tlist).imag, Ni)
+        ckAR = amp
+        ckAI = amp2
+        vkAR = -((len(tlist)-1)/tlist[-1]) * \
+            (np.log(np.abs(phases))+1j*np.angle(phases))
+        vkAI = -((len(tlist)-1)/tlist[-1]) * \
+            (np.log(np.abs(phases2))+1j*np.angle(phases2))
+
+        cls = ExponentialBosonicEnvironment(
+            ck_real=ckAR, vk_real=vkAR, vk_imag=vkAI,
+            ck_imag=ckAI, combine=combine, tag=tag)
+        return cls, (amp, phases)
+
 
 class _BosonicEnvironment_fromCF(BosonicEnvironment):
     def __init__(self, C, tlist, tMax, T, tag, args):
@@ -2236,6 +2364,49 @@ def from_spectral_density(cls, **kwargs) -> FermionicEnvironment:
                                   "currently not implemented.")
 
 
+class _FermionicEnvironment_fromSD(BosonicEnvironment):
+    def __init__(self, J, wlist, wMax, T, tag, args):
+        super().__init__(T, tag)
+        self._sd = _real_interpolation(J, wlist, 'spectral density', args)
+        if wlist is not None:
+            self.wMax = max(np.abs(wlist[0]), np.abs(wlist[-1]))
+        else:
+            self.wMax = wMax
+
+    def correlation_function_plus(self, t, *, eps=1e-10):
+        if self.T is None:
+            raise ValueError('The temperature must be specified for this '
+                             'operation.')
+        if self.wMax is None:
+            raise ValueError('The support of the spectral density (wMax) '
+                             'must be specified for this operation.')
+        return self._cf_from_ps_plus(t, self.wMax, eps=eps)
+
+    def correlation_function_minus(self, t, *, eps=1e-10):
+        if self.T is None:
+            raise ValueError('The temperature must be specified for this '
+                             'operation.')
+        if self.wMax is None:
+            raise ValueError('The support of the spectral density (wMax) '
+                             'must be specified for this operation.')
+        return self._cf_from_ps_minus(t, self.wMax, eps=eps)
+
+    def spectral_density(self, w):
+        w = np.asarray(w, dtype=float)
+
+        result = np.zeros_like(w)
+        positive_mask = (w > 0)
+        result[positive_mask] = self._sd(w[positive_mask])
+
+        return result.item() if w.ndim == 0 else result
+
+    def power_spectrum_plus(self, w, *, eps=1e-10):
+        return self._ps_plus_from_sd(w, eps)
+
+    def power_spectrum_minus(self, w, *, eps=1e-10):
+        return self._ps_minus_from_sd(w, eps)
+
+
 class LorentzianEnvironment(FermionicEnvironment):
     r"""
     Describes a Lorentzian fermionic environment with the spectral density