diff --git a/hopsflow/gaussflow_two.py b/hopsflow/gaussflow_two.py index 55aad3b..d7fb82a 100644 --- a/hopsflow/gaussflow_two.py +++ b/hopsflow/gaussflow_two.py @@ -5,11 +5,14 @@ This is done analytically for a BCF that is a sum of exponentials. """ from dataclasses import dataclass, field -import util -from .util import BCF, expand_t +from . import util +from .util import BCF import numpy as np from numpy.typing import NDArray from numpy.polynomial import Polynomial +from typing import Union, Optional +from itertools import product +from numba import jit @dataclass @@ -219,12 +222,8 @@ class Propagator: self.params = sys self._roots, self._res = calculate_coefficients(sys) - @expand_t - def __call__(self, t: NDArray[np.float64]) -> np.ndarray: - """Get the propagator as array of shape ``(time, 4, 4)``.""" - mat = self._res - residual_matrix = np.array( + self._residual_matrix = np.array( [ [mat[0], mat[1], mat[2], mat[3]], [mat[4], mat[0], mat[5], mat[6]], @@ -233,12 +232,187 @@ class Propagator: ] ) + def propagator(self, t: Union[NDArray[np.float64], float]) -> np.ndarray: + """ + Get the propagator as array of shape ``(time, 4, 4)`` or + ``(4, 4)`` if ``t`` is a float. + """ + t = np.asarray(t) + was_float = False + if t.shape == tuple(): + t = np.array([t]) + was_float = True + + res = np.sum( + self._residual_matrix[None, :, :, :] + * np.exp(self._roots[None, :] * t[:, None])[:, None, None, :], + axis=3, + ).real + + return res[0] if was_float else res + + def __call__(self, *args, **kwargs): + return self.propagator(*args, **kwargs) + + def inv(self, t: float) -> np.ndarray: + """Get the inverse of the propagator matrix at time ``t``.""" + + return np.linalg.inv(self.__call__(t)) + + +class CorrelationMatrix(Propagator): + def __init__( + self, + sys: SystemParams, + initial_corr: np.ndarray, + αs: tuple[Optional[BCF], Optional[BCF]] = (None, None), + ): + super().__init__(sys) + + self.G = self._residual_matrix + self.Ge = -self._roots + + α = [] + αe = [] + + for i, bcf in enumerate(αs): + if bcf: + α.append(bcf.factors) + αe.append(bcf.exponents) + else: + α.append(self.params.G[i]) + αe.append(self.params.W[i]) + + self.αc = np.array(list(map(np.conj, α)), dtype=object) + self.αce = np.array(list(map(np.conj, αe)), dtype=object) + self.α = np.array(α, dtype=object) + self.αe = np.array(αe, dtype=object) + self.initial_corr = initial_corr + + @jit + def __call__( + self, + t: np.ndarray, + s: Optional[np.ndarray] = None, + ) -> NDArray[np.complex128]: + r"""Caluclate the corellation matrix :math:`\langle x_i(t) + x_j(s)\rangle` of the system. + + :param t: The first time point. + :param s: The second timepoint. + :param αs: The (finite temparature) bath correlation + functions. If they're :any:`None` they're chosen to be + the zero temperature BCF. + :param initial_corr: The initial correlation matrix + :math:`\langle x_i[] x_j\rangle`. + """ + + s = s or t + + G = self.G[None, :, :, :] + Ge = self.Ge + αc = self.αc + αce = self.αce + α = self.α + αe = self.αe + initial_corr = self.initial_corr + + if (s > t).any(): + raise ValueError("`s` must be smaller than or equal to `t`") + + Gt = self.propagator(t) + Gs = self.propagator(s) + + result = np.einsum("tik,tjl,kl->tij", Gt, Gs, initial_corr) + # (Gt @ initial_corr[None, :, :]) @ Gs.T + + for l in range(len(α)): + for i, j, m, n, g in product( + range(G.shape[0]), + range(G.shape[0]), + range(Ge.shape[0]), + range(Ge.shape[0]), + range(len(α[l])), + ): + # straight from mathematica + + result[:, i, j] += ( + G[:, i, l, m] + * G[:, j, l, n] + * ( + ( + αc[l, g] + * ( + (-np.exp(s * Ge[n]) + np.exp(s * αce[l, g])) * Ge[m] + - np.exp(s * Ge[n]) * Ge[n] + + np.exp(s * (Ge[m] + Ge[n] + αce[l, g])) + * (Ge[n] - αce[l, g]) + + np.exp(s * αce[l, g]) * αce[l, g] + ) + ) + / ( + np.exp(s * (Ge[n] + αce[l, g])) + * (Ge[n] - αce[l, g]) + * (Ge[m] + αce[l, g]) + ) + + ( + np.exp(-(s * Ge[n]) - t * αe[l, g]) + * α[l, g] + * ( + -( + np.exp(t * Ge[m]) + * (-1 + np.exp(s * (Ge[n] + αe[l, g]))) + * Ge[m] + ) + + ( + np.exp(t * Ge[m]) + - np.exp(t * αe[l, g]) + + np.exp(s * (Ge[m] + Ge[n]) + t * αe[l, g]) + - np.exp(t * Ge[m] + s * (Ge[n] + αe[l, g])) + ) + * Ge[n] + + np.exp(t * αe[l, g]) + * (-1 + np.exp(s * (Ge[m] + Ge[n]))) + * αe[l, g] + ) + ) + / ((-Ge[m] + αe[l, g]) * (Ge[n] + αe[l, g])) + ) + ) / (np.exp(t * Ge[m]) * (Ge[m] + Ge[n])) + + return result + + def system_energy(self, t: np.ndarray) -> float: + corr = np.real(np.diagonal(self.__call__(t), axis1=1, axis2=2)) + + Ω = self.params.Ω + Λ = self.params.Λ + γ = self.params.γ + return ( - np.sum( - residual_matrix * np.exp(self._roots * t)[None, None, :], axis=2 - ).real - * 2 + 1 + / 4 + * ( + Ω * corr[:, 0] # type: ignore + + (Ω + γ) * corr[:, 1] # type: ignore + + Λ * corr[:, 2] # type: ignore + + (Λ + γ) * corr[:, 3] # type: ignore + ) ) - def inv(self, t) -> np.ndarray: - return np.linalg.inv(self.__call__(t)) + +def initial_correlation_pure_osci(n: int, m: int): + r""" + Initial correlation matrix :math:`\langle x_i[] x_j\rangle` for + two HO in product eigenstates with labels ``n`` and ``m``. + """ + corr = np.diag( + np.array([1 + 2 * n, 1 + 2 * n, 1 + 2 * m, 1 + 2 * m], dtype=np.complex128) + ) + + corr[0, 1] = 1j + corr[1, 0] = -1j + corr[2, 3] = 1j + corr[3, 2] = -1j + + return corr