Source code for pykbe.dyson

from __future__ import annotations

import numpy as np
import scipy.linalg as sla

from pykbe import gregory
from pykbe.misc import accum_mat_into_tens3, accum_mat_into_tens4, unpack_dim3_into_dim2
from pykbe.misc import parallel_set as pset




class DysonWork:
    def __init__(self, shape, integ_ord=3, dtype=np.complex128):
        self.shape = shape
        self.integ_ord = integ_ord
        self.dtype = dtype

        self.gtemp = np.zeros((integ_ord, *shape), dtype=dtype)
        self.buf1 = np.zeros((integ_ord, *shape), dtype=dtype)
        self.buf2 = np.zeros((integ_ord, integ_ord, *shape), dtype=dtype)



    def set_integ_ord(self, integ_ord):
        if integ_ord != self.integ_ord:
            self.integ_ord = integ_ord
            self.gtemp = np.zeros((integ_ord, *self.shape), dtype=self.dtype)
            self.buf1 = np.zeros((integ_ord, *self.shape), dtype=self.dtype)
            self.buf2 = np.zeros((integ_ord, integ_ord, *self.dtype), dtype=self.dtype)




    def dyson_start_ret(self, G, mu, ham, Sigma, dt):
        k = self.integ_ord
        sz = np.prod(G.shape)
        N = G.shape[0]
        Dk = gregory.poly_differentiation_weights(k)
        for i in range(k + 1):
            pset(G.ret_at(i, i), -1.0j)

        for j in range(k):
            # main loop
            dim = k - j
            M = np.zeros((dim, N, dim, N), dtype=self.dtype)
            Q = np.zeros((N, dim, N), dtype=self.dtype)

            for n in range(j + 1, k + 1):
                p = n - (j + 1)

                accum_mat_into_tens4(
                    ham(n), M, n - (j + 1), n - (j + 1), alpha=-1.0, conj=False, beta=mu
                )

                for l in range(j + 1):
                    accum_mat_into_tens3(
                        G.ret_at(j, l), Q, p, alpha=+1j / dt * Dk[n, l], conj=True
                    )

                for l in range(j + 1, k):
                    q = l - (j + 1)

                    ijnl = gregory.poly_integration_weights(j)[n, l]
                    wt = -dt * ijnl

                    add_const = -1j / dt * Dk[n, l]

                    if n >= l:
                        stemp = Sigma.ret_at(n, l)
                        conjsigma = False
                    else:
                        stemp = Sigma.ret_at(l, n)
                        conjsigma = True
                        wt *= -1
                    accum_mat_into_tens4(
                        stemp, M, q, p, alpha=wt, conj=conjsigma, beta=add_const
                    )
            m_reshape = M.reshape((dim * N, dim * N))
            q_reshape = Q.reshape((sz, dim * sz))
            X = sla.solve(
                m_reshape.T,
                q_reshape.T,
                overwrite_a=True,
                overwrite_b=True,
                check_finite=False,
            ).T
            for n in range(j + 1, k + 1):
                p = n - (j + 1)
                unpack_dim3_into_dim2(G.ret_at(n, j), X, p)