mirror of
https://github.com/vale981/hopsflow
synced 2025-03-04 08:31:37 -05:00
427 lines
14 KiB
Python
427 lines
14 KiB
Python
import numpy as np
|
||
import itertools
|
||
from dataclasses import dataclass, field
|
||
from typing import Callable, Optional
|
||
import numpy.typing as npt
|
||
from . import util
|
||
from numpy.polynomial import Polynomial
|
||
from .util import BCF, expand_t
|
||
import scipy
|
||
|
||
|
||
@dataclass
|
||
class SystemParams:
|
||
r"""A parameter object to hold information about the physical
|
||
system and the global parameters for the algorithm.
|
||
|
||
:param Ω: the system hamiltonian energy scale, :math:`H_s =
|
||
\Omega (p^2 + q^2)`
|
||
:param η: the coupling strength (in essence a prefactor to the BCF)
|
||
:param α_0: the zero temperature BCF
|
||
|
||
|
||
:attr t_max: the maximum simulation time, will be copied from :attr:`α`
|
||
:attr W: the exponential factors in the BCF expansion
|
||
:attr G: the pre-factors in the BCF expansion
|
||
"""
|
||
|
||
Ω: float
|
||
η: float
|
||
α_0: BCF
|
||
|
||
t_max: float = field(init=False)
|
||
|
||
# BCF coefficients
|
||
W: np.ndarray = field(init=False)
|
||
G: np.ndarray = field(init=False)
|
||
|
||
root_tol: float = 1e-7
|
||
"""
|
||
The relative tolerance of the roots.
|
||
|
||
The roots are being calculated with the numpy polynomial module
|
||
and then refined with newtons method. This parameter is passed as
|
||
``rtol`` to :any:`scipy.optimize.newton`.
|
||
"""
|
||
|
||
def __post_init__(self):
|
||
self.t_max = self.α_0.t_max
|
||
|
||
self.W = self.α_0.exponents
|
||
self.G = self.α_0.factors.copy() * self.η
|
||
|
||
|
||
def construct_polynomials(sys: SystemParams) -> tuple[Polynomial, Polynomial]:
|
||
r"""
|
||
Construct the polynomials required to find the coefficients
|
||
and exponents of the solution.
|
||
|
||
:param sys: a parameter object with the system parameters
|
||
:returns: a list of polynomials
|
||
|
||
- :math:`f_0(z) = \prod_k (z-z_k) (z-z_k^{\ast})`
|
||
- :math:`p(z) = p_1(z) + \sum_n q_n(z)` where
|
||
- :math:`p_1(z) = (z^2 + \Omega^2)\prod_k (z-z_k)(z-z_k^{\ast})`
|
||
- :math:`q_n(z)=\Omega f_n(z) \prod_{k\neq n}(z-z_k) (z-z_k^{\ast})`
|
||
|
||
"""
|
||
# we begin by calculating all the constants we'll need
|
||
φ = -np.angle(sys.G)
|
||
G_abs = np.abs(sys.G)
|
||
γ, δ = sys.W.real, sys.W.imag
|
||
s, c = np.sin(φ), np.cos(φ)
|
||
|
||
# the roots of the denominator of the laplace transform of α
|
||
roots_z = -sys.W
|
||
|
||
# now the polynomials
|
||
f_0 = util.poly_real(
|
||
Polynomial.fromroots(np.concatenate((roots_z, roots_z.conj())))
|
||
)
|
||
p_1 = Polynomial([sys.Ω ** 2, 0, 1]) * f_0
|
||
|
||
q = [
|
||
-G_c
|
||
* sys.Ω
|
||
* Polynomial((δ_c * c_c + γ_c * s_c, s_c))
|
||
* (
|
||
util.poly_real(
|
||
Polynomial.fromroots(
|
||
np.concatenate(
|
||
(
|
||
util.except_element(roots_z, i),
|
||
util.except_element(roots_z, i).conj(),
|
||
)
|
||
)
|
||
)
|
||
)
|
||
if len(roots_z) > 1
|
||
else 1 # corner case
|
||
)
|
||
for G_c, γ_c, δ_c, s_c, c_c, i in zip(G_abs, γ, δ, s, c, range(len(c)))
|
||
]
|
||
|
||
p = p_1 + sum(q)
|
||
|
||
return f_0, p
|
||
|
||
|
||
def calculate_coefficients(sys: SystemParams) -> tuple[np.ndarray, np.ndarray]:
|
||
r"""
|
||
Calculates the coefficients required to construct the
|
||
propagator matrix :math:`G`.
|
||
|
||
:param sys: a parameter object with the system parameters
|
||
:returns: the exponents and the residuals which play a role as pre-factors
|
||
"""
|
||
|
||
f_0, p = construct_polynomials(sys)
|
||
|
||
master_roots = p.roots()
|
||
|
||
if np.unique(master_roots).size != master_roots.size:
|
||
raise RuntimeError(
|
||
"""The roots of the polynomial are not unique.
|
||
You can try to alter the number of terms in the expansion of the BCF."""
|
||
)
|
||
|
||
# improving the accuracy of the roots seems to be essential
|
||
p_prime = p.deriv()
|
||
p_prime_2 = p_prime.deriv()
|
||
for i, root in enumerate(master_roots):
|
||
res = scipy.optimize.newton(
|
||
p,
|
||
root,
|
||
p_prime,
|
||
rtol=sys.root_tol,
|
||
fprime2=p_prime_2,
|
||
maxiter=100000, # we can allow that, there aren't that many roots
|
||
)
|
||
master_roots[i] = res
|
||
|
||
resiquals = f_0(master_roots) / p.deriv()(master_roots)
|
||
return master_roots, resiquals
|
||
|
||
|
||
class Propagator:
|
||
"""The propagator matrix :math:`G` for the system.
|
||
|
||
You can get inidividual matrix elements as functions of time
|
||
through indexing this object.
|
||
|
||
Calling it with a time argument returns the whole matrix as
|
||
array of shape ``(time, 2, 2)``.
|
||
|
||
:param sys: a parameter object with the system parameters
|
||
"""
|
||
|
||
def __init__(self, sys: SystemParams):
|
||
self.params = sys
|
||
self._roots, self._res = calculate_coefficients(sys)
|
||
self._roots_exp = np.expand_dims(self._roots, axis=1)
|
||
self._res_exp = np.expand_dims(self._res, axis=1)
|
||
self._Ω = sys.Ω
|
||
|
||
self._res_times_roots = self._res_exp * self._roots_exp
|
||
self._res_times_roots_squared = self._res_times_roots * self._roots_exp
|
||
self._elements = np.array([[self.el_11, self.el_12], [self.el_21, self.el_22]])
|
||
pass
|
||
|
||
@expand_t
|
||
def el_11(self, t):
|
||
return (self._res_times_roots * np.exp(t * self._roots_exp)).real.sum(axis=0)
|
||
|
||
@expand_t
|
||
def el_12(self, t):
|
||
return self._Ω * (self._res_exp * np.exp(t * self._roots_exp)).real.sum(axis=0)
|
||
|
||
@expand_t
|
||
def el_21(self, t):
|
||
return (self._res_times_roots_squared * np.exp(t * self._roots_exp)).real.sum(
|
||
axis=0
|
||
) / self._Ω
|
||
|
||
def el_22(self, t):
|
||
return self.el_11(t)
|
||
|
||
def __getitem__(self, indices) -> Callable[[npt.ArrayLike], np.ndarray]:
|
||
"""Get an individual matrix elment as function."""
|
||
return self._elements[indices]
|
||
|
||
@expand_t
|
||
def __call__(self, t) -> np.ndarray:
|
||
"""Get the propagator as array of shape ``(time, 2, 2)``."""
|
||
|
||
g_12 = self._res_exp * np.exp(t * self._roots_exp)
|
||
diag = self._roots_exp * g_12
|
||
g_21 = self._roots_exp * diag / self._Ω
|
||
|
||
return (
|
||
np.array([[diag, g_12 * self._Ω], [g_21, diag]])
|
||
.real.sum(axis=2)
|
||
.swapaxes(0, 2)
|
||
.swapaxes(1, 2)
|
||
)
|
||
|
||
def inv(self, t) -> np.ndarray:
|
||
return np.linalg.inv(self.__call__(t))
|
||
|
||
|
||
class Flow:
|
||
r"""
|
||
A collection of methods to calculate the time derivative
|
||
of the bath energy expectation value :math:`\frac{d}{dt}\langle H_B\rangle`
|
||
which can be retrieved as a function of time by calling this object.
|
||
|
||
:param system: the system parameters, see :any:`SystemParams`
|
||
:param α: the (finite temperature) BCF
|
||
:param α_0_dot: the zero temperature BCF time derivative
|
||
:param n: the excitation number of the initial state of the system :math:`|n\rangle`
|
||
|
||
"""
|
||
|
||
def __init__(self, system: SystemParams, α: BCF, α_0_dot: BCF, n: int):
|
||
#: the BCF
|
||
self.α = α
|
||
|
||
#: the zero temperature BCF
|
||
self.α_0 = system.α_0
|
||
|
||
#: the exponential factors in the BCF derivative
|
||
#: expansion :math:`\dot{\alpha}_0=\sum_k P_k e^{-L_k \cdot t}`
|
||
self.L = α_0_dot.exponents
|
||
|
||
#: the pre-factors in the BCF derivative expansion
|
||
self.P = α_0_dot.factors.copy() * system.η
|
||
|
||
self.C, self.B = calculate_coefficients(system)
|
||
|
||
#: the pre-factors of the exponential sum of :math:`A=G_{11}=\sum_k A_k e^{-C_k \cdot t}`
|
||
self.A = self.B * self.C
|
||
|
||
#: the exponents of the exponential sum of :math:`B=G_{12}=\sum_k B_k e^{-C_k \cdot t}`
|
||
#:
|
||
#: - note the minus sign
|
||
self.C = -self.C # mind the extra -
|
||
|
||
#: the pre-factors of the exponential sum of ``B``
|
||
self.B = system.Ω * self.B
|
||
|
||
#: the exponential factors in the BCF
|
||
#: expansion :math:`\alpha=\sum_k G_k e^{-W_k \cdot t}`
|
||
self.W = α.exponents
|
||
|
||
#: the pre-factors factors in the BCF expansion
|
||
self.G = α.factors.copy() * system.η
|
||
|
||
#: the exponential factors in the zero temperature BCF
|
||
#: expansion :math:`\alpha=\sum_k U_k e^{-Q_k \cdot t}`
|
||
self.Q = system.α_0.exponents
|
||
|
||
#: the pre-factors factors in the zero temperature BCF expansion
|
||
self.U = system.α_0.factors.copy() * system.η
|
||
|
||
#: the expectation value :math:`\langle q(0)^2\rangle`
|
||
self.q_s_0 = 1 + 2 * n
|
||
|
||
#: the expectation value :math:`\langle q(0)^2\rangle`
|
||
self.p_s_0 = 1 + 2 * n
|
||
|
||
#: the expectation value :math:`\langle q(0)p(0)\rangle`
|
||
self.qp_0 = 1j
|
||
|
||
#: the propagator matrix :math:`G(t)`, see :any:`Propagator`
|
||
self.prop = Propagator(system)
|
||
|
||
#: the coefficient matrix :math:`\Gamma^1`
|
||
self.Γ1 = (self.B[:, None] * self.P[None, :]) / (
|
||
self.L[None, :] - self.C[:, None]
|
||
)
|
||
|
||
#: the coefficient matrix :math:`\Gamma^2`
|
||
self.Γ2 = (self.B[:, None] * self.G[None, :]) / (
|
||
self.C[:, None] - self.W[None, :]
|
||
)
|
||
|
||
#: the coefficient matrix :math:`\Gamma^3`
|
||
self.Γ3 = (self.B[:, None] * self.G.conj()[None, :]) / (
|
||
self.C[:, None] + self.W.conj()[None, :]
|
||
)
|
||
|
||
#: the coefficient matrix :math:`\Gamma^A`
|
||
self.ΓA = (self.A[:, None] * self.P[None, :]) / (
|
||
self.L[None, :] - self.C[:, None]
|
||
)
|
||
|
||
#: the system parameters
|
||
self.system = system
|
||
|
||
def A_conv(self, t: npt.ArrayLike):
|
||
r"""The integral :math:`\int_0^t A(s)\dot{\alpha}_0(t-s)ds`."""
|
||
result = np.zeros_like(t, dtype="complex128")
|
||
|
||
for (n, m) in itertools.product(range(len(self.A)), range(len(self.P))):
|
||
result += self.ΓA[n, m] * (np.exp(-self.C[n] * t) - np.exp(-self.L[m] * t))
|
||
|
||
return result
|
||
|
||
def B_conv(self, t: npt.ArrayLike):
|
||
r"""The integral :math:`\int_0^t B(s)\dot{\alpha}_0(t-s)ds`."""
|
||
result = np.zeros_like(t, dtype="complex128")
|
||
for (n, m) in itertools.product(range(len(self.B)), range(len(self.P))):
|
||
result += self.Γ1[n, m] * (np.exp(-self.C[n] * t) - np.exp(-self.L[m] * t))
|
||
|
||
return result
|
||
|
||
def flow_nontherm(self, t: npt.ArrayLike):
|
||
r"""The part of the flow that **does not** involve :math:`\alpha`."""
|
||
a, b = self.prop[0, 0](t), self.prop[0, 1](t)
|
||
ac, bc = self.A_conv(t), self.B_conv(t)
|
||
|
||
return (
|
||
-1
|
||
/ 2
|
||
* (
|
||
self.q_s_0 * a * ac
|
||
+ self.p_s_0 * b * bc
|
||
+ self.qp_0 * a * bc
|
||
+ self.qp_0.conjugate() * b * ac
|
||
).imag
|
||
)
|
||
|
||
def flow_therm_zero(self, t: npt.ArrayLike):
|
||
r"""The part of the flow that **does** involve :math:`\alpha`."""
|
||
t = np.asarray(t)
|
||
result = np.zeros_like(t, dtype="float")
|
||
|
||
for (m, k, n, l) in itertools.product(
|
||
range(len(self.B)),
|
||
range(len(self.P)),
|
||
range(len(self.B)),
|
||
range(len(self.G)),
|
||
):
|
||
g_1_2 = (
|
||
self.Γ1[m, k]
|
||
* self.Γ2[n, l]
|
||
* (
|
||
(1 - np.exp(-(self.C[m] + self.W[l]) * t)) / (self.C[m] + self.W[l])
|
||
+ (np.exp(-(self.C[m] + self.C[n]) * t) - 1)
|
||
/ (self.C[m] + self.C[n])
|
||
+ (np.exp(-(self.L[k] + self.W[l]) * t) - 1)
|
||
/ (self.L[k] + self.W[l])
|
||
+ (1 - np.exp(-(self.L[k] + self.C[n]) * t))
|
||
/ (self.L[k] + self.C[n])
|
||
)
|
||
).imag
|
||
|
||
g_1_3 = (
|
||
self.Γ1[m, k]
|
||
* self.Γ3[n, l]
|
||
* (
|
||
(1 - np.exp(-(self.C[m] + self.C[n]) * t)) / (self.C[m] + self.C[n])
|
||
- (1 - np.exp(-(self.L[k] + self.C[n]) * t))
|
||
/ (self.L[k] + self.C[n])
|
||
- (
|
||
np.exp(-(self.C[n] + self.W[l].conj()) * t)
|
||
- np.exp(-(self.C[m] + self.C[n]) * t)
|
||
)
|
||
/ (self.C[m] - self.W[l].conj())
|
||
+ (
|
||
np.exp(-(self.C[n] + self.W[l].conj()) * t)
|
||
- np.exp(-(self.L[k] + self.C[n]) * t)
|
||
)
|
||
/ (self.L[k] - self.W[l].conj())
|
||
)
|
||
).imag
|
||
|
||
result += -1 / 2 * (g_1_2 + g_1_3)
|
||
|
||
return result
|
||
|
||
def flow_therm_nonzero(self, t: npt.ArrayLike):
|
||
r"""The part of the flow that **does** involve :math:`\alpha-\alpha_0`."""
|
||
t = np.asarray(t)
|
||
result = np.zeros_like(t, dtype="float")
|
||
|
||
for (n, m) in itertools.product(
|
||
range(len(self.A)),
|
||
range(len(self.G)),
|
||
):
|
||
result += (
|
||
self.A[n]
|
||
* self.G[m]
|
||
/ (self.C[n] + self.W[m])
|
||
* (1 - np.exp(-(self.C[n] + self.W[m]) * t))
|
||
).real
|
||
|
||
for (n, m) in itertools.product(
|
||
range(len(self.A)),
|
||
range(len(self.U)),
|
||
):
|
||
result -= (
|
||
self.A[n]
|
||
* self.U[m]
|
||
/ (self.C[n] + self.Q[m])
|
||
* (1 - np.exp(-(self.C[n] + self.Q[m]) * t))
|
||
).real
|
||
|
||
return (
|
||
1
|
||
/ 2
|
||
* (
|
||
-self.system.Ω * result
|
||
+ self.prop.el_12(t)
|
||
* (self.α.approx(t) - self.α_0.approx(t))
|
||
* self.system.η
|
||
)
|
||
).real
|
||
|
||
def __call__(self, t: npt.ArrayLike) -> np.ndarray:
|
||
r"""
|
||
The flow. Time derivative of the bath energy
|
||
expectation value :math:`\frac{d}{dt}\langle H_B\rangle`.
|
||
"""
|
||
|
||
return (
|
||
self.flow_nontherm(t) + self.flow_therm_zero(t) + self.flow_therm_nonzero(t)
|
||
)
|