from __future__ import annotations
import gc
import h5py
import numba
import numpy as np
import scipy.linalg as sla
[docs]
class TwoTimeVar:
"""Class to store two-time variables."""
def __init__(self, shape, ntimes=1, dt=0.1, dtype=np.float64, restricted=True):
"""Constructor for TwoTimeVar.
Parameters
----------
shape : tuple[int]
shape of a one-particle operator's matrix.
ntimes : int, optional
number of time points other than t=0; by default 1
dt : float, optional
timestep size, by default 0.1
dtype : optional
scalar datatype, by default np.float64
restricted : bool, optional
unused, by default True
"""
self.shape = shape
self.dtype = dtype
self.ntimes = ntimes
self.dt = dt
self.retcomp = None
self.lescomp = None
self.propagator = None
self.restricted = restricted
[docs]
def alloc(self, ntimes=None):
"""Allocate memory.
Parameters
----------
ntimes : int, optional
number of time points; by default None
"""
if ntimes is not None:
self.ntimes = ntimes
nt = self.ntimes
ntimepair_tri = (nt + 1) * (nt + 2) // 2
self.propagator = np.zeros((nt + 1, *self.shape), dtype=self.dtype)
self.retcomp = np.zeros((ntimepair_tri, *self.shape), dtype=self.dtype)
self.lescomp = np.zeros((ntimepair_tri, *self.shape), dtype=self.dtype)
[docs]
def free(self):
"""Free memory."""
self.ntimes = 0
self.retcomp = None
self.lescomp = None
gc.collect()
[docs]
def save(self, filename, groupname):
"""Save variable to HDF5 file.
Parameters
----------
filename : str
file name.
groupname : str
HDF5 group name.
"""
with h5py.File(filename, "w") as f:
group = f.create_group(groupname)
group.create_dataset("dt", data=(self.dt,))
group.create_dataset("ntimes", data=(self.ntimes,))
group.create_dataset("retcomp", data=self.retcomp)
group.create_dataset("lescomp", data=self.lescomp)
group.create_dataset("propagator", data=self.propagator)
[docs]
@staticmethod
def tri_idx(i, j):
"""Index into a packed lower triangular array.
Parameters
----------
i : int
row index.
j : int
column index
Returns
-------
int
index into the packed lower triangular array.
Raises
------
ValueError
Raised if you try to access an element in the upper triangle.
"""
if i < j:
msg = "Must have i >= j"
raise ValueError(msg)
return i * (i + 1) // 2 + j
[docs]
def ret_at(self, i, j):
return self.retcomp[self.tri_idx(i, j)]
# note that the lesser component is stored as upper triangle
[docs]
def les_at(self, i, j):
return self.lescomp[self.tri_idx(j, i)]
[docs]
def from_time_independent_H(self, H, dm=None, mu=0.0):
r"""Driver to calculate free :math:`G` from time-independent :math:`H`."""
H = np.asarray(H, dtype=self.dtype)
eigs, vecs = sla.eigh(H)
eigs -= mu
retcomp_fast_time_independent_H(eigs, vecs, self.ntimes, self.dt, self.retcomp)
if dm is None:
nocc = np.sum(eigs <= mu)
dm = vecs[:, :nocc] @ vecs[:, :nocc].conj().T
calc_lescomp_free(dm, self.ntimes, self.lescomp, self.retcomp)
[docs]
def from_time_dependent_H(self, Hfunc, dm=None, mu=0.0):
r"""Driver to calculate free :math:`G` from time-dependent :math:`H`."""
self.calc_propagator(Hfunc)
retcomp_fast_from_propagator(self.ntimes, self.propagator, self.retcomp)
if dm is None:
H = Hfunc(0.0)
eigs, vecs = sla.eigh(H)
eigs -= mu
nocc = np.sum(eigs <= mu)
dm = vecs[:, :nocc] @ vecs[:, :nocc].conj().T
calc_lescomp_free(dm, self.ntimes, self.lescomp, self.retcomp)
[docs]
def calc_propagator(self, Hfunc):
r"""Second order Magnus scheme to get :math:`U(t)` from :math:`H(t)`.
The propagator is stored in the `propagator` attribute.
.. math::
U(t) = \mathrm{Texp}\left(-i\int_{t_0}^t H(t')\, dt'\right)
Parameters
----------
Hfunc : callable
Time-dependent Hamiltonian.
"""
dt = self.dt
prop = self.propagator
prop[0][np.diag_indices_from(prop[0])] = 1.0
for i in range(1, self.ntimes + 1):
prop[i] = sla.expm(-1j * dt * Hfunc((i - 1 / 2) * dt)) @ prop[i - 1]
[docs]
@numba.njit(parallel=True)
def retcomp_fast_from_propagator(ntimes, prop, retcomp):
r"""Calculate :math:`G^\mathrm{R}(t_1, t_2)` from propagator.
Handles only the case :math:`t_1 \geq t_2`.
.. math::
G^\mathrm{R}(t_1, t_2) = -i \theta(t_1 - t_2) U(t_1) U^\dagger(t_2)
Parameters
----------
ntimes : int
number of time points.
prop : ndarray
propagator; shape (ntimes + 1, N, N).
retcomp : ndarray
retarded Green's function; shape (ntimepair_tri, N, N).
"""
for ij in numba.prange((ntimes + 1) ** 2):
i = ij // (ntimes + 1)
j = ij % (ntimes + 1)
if i >= j:
retcomp[i * (i + 1) // 2 + j][:] = -1j * prop[i] @ prop[j].conj().T
[docs]
@numba.njit(parallel=True)
def retcomp_fast_time_independent_H(heigs, heigvecs, ntimes, dt, retcomp):
r"""Calculate :math:`G^\mathrm{R}(t_1, t_2)` from time-indep. Hamiltonian.
Handles only the case :math:`t_1 \geq t_2`.
.. math::
G^\mathrm{R}(t_1, t_2) &= -i \theta(t_1 - t_2) U(t_1) U^\dagger(t_2) \\
&= -i \theta(t_1 - t_2) \exp\left(-i H (t_1 - t_2)\right).
Parameters
----------
ntimes : int
number of time points.
prop : ndarray
propagator; shape (ntimes + 1, N, N).
retcomp : ndarray
retarded Green's function; shape (ntimepair_tri, N, N).
"""
for ij in numba.prange((ntimes + 1) ** 2):
i = ij // (ntimes + 1)
j = ij % (ntimes + 1)
if i >= j:
coefs = -1j * dt * (i - j)
eigvals_exp = -1j * np.exp(coefs * heigs)
tmp = heigvecs * eigvals_exp.reshape(1, -1)
retcomp[i * (i + 1) // 2 + j][:] = tmp @ heigvecs.conj().T
return retcomp
[docs]
@numba.njit(parallel=True)
def calc_lescomp_free(dm, ntimes, lescomp, retcomp):
r"""Calculate :math:`G^{<}(t_1, t_2)` from density matrix and :math:`G^\mathrm{R}`.
Handles only the case :math:`t_1 \geq t_2`.
.. math::
G^{<}(t_1, t_2) = G^\mathrm{R}(t_1, t_0) G^{<}(t_0, t_0) G^\mathrm{R}(t_2, t_0)^\dagger
where :math:`G^{<}(t_0, t_0) = -i \rho` is calculated from the density matrix.
Parameters
----------
dm : ndarray
initial density matrix.
ntimes : int
number of time points.
lescomp : ndarray
lesser Green's function; shape (ntimepair_tri, N, N).
retcomp : ndarray
retarded Green's function; shape (ntimepair_tri, N, N).
"""
idm = dm * 1j
for ji in numba.prange((ntimes + 1) ** 2):
j = ji // (ntimes + 1)
i = ji % (ntimes + 1)
if j >= i:
lescomp[j * (j + 1) // 2 + i][:] = (
retcomp[i * (i + 1) // 2] @ idm @ retcomp[j * (j + 1) // 2].conj().T
)
return lescomp
[docs]
def test_free():
G = TwoTimeVar((2, 2), ntimes=1000, dt=1.0, dtype=np.complex128)
G.alloc()
H = np.array([[-1.0, 0.5j], [-0.5j, 1.0]])
G.from_time_independent_H(H)
G2 = TwoTimeVar((2, 2), ntimes=1000, dt=1.0, dtype=np.complex128)
G2.alloc()
def Hfunc(_):
return H
G2.from_time_dependent_H(Hfunc)
import IPython
IPython.embed()
if __name__ == "__main__":
test_free()