Source code for py4mulas.opera_kernels

import numpy as np

from .operators import KspaceOpera
from ._common import invert

__alla__ = ["Kubo"]



class Kubo:
    def __init__(self, alpha: KspaceOpera, beta: KspaceOpera, reduce: bool = False):
        self.alpha = alpha
        self.beta = beta
        self.reduce = reduce



    def __call__(
        self, k_args: list[np.ndarray], en: np.ndarray, psi: np.ndarray, eta: float
    ) -> np.ndarray:
        alpha = self.alpha(k_args, en, psi)
        opera_kernel = self.beta(k_args, en, psi)
        opera_kernel *= np.swapaxes(alpha, 1, 2)
        if not self.reduce:
            return opera_kernel
        delta_e = en[:, None, :] - en[:, :, None]
        denom = delta_e * (delta_e + 1j * eta)
        opera_kernel *= 1j * invert(denom)
        # reduce into a vector for symmetrized Kubo formula
        reduced = opera_kernel.sum(axis=2) - opera_kernel.sum(axis=1)
        return reduced