master-thesis/python/graveyard/billohops/hops.py

"""A pedadogical exercise implementation of HOPS."""
import numpy as np
import numpy.typing as npt
import typing as t
import pdb
import scipy
import scipy.misc


class Hops:
    def __init__(self, η, H_s, L, W, k_max, seed=None, solver_args=dict()):
        """Implements the integration of hops with
        exactly one exponential term.

        :param η: random process from `stocproc`
        :param H_s: the system hamiltonian
        :param L: the interaction operator
        :param W: the exponent of the BCF
        :param k_max: the depth of the hirarchy
        :returns: the step function for the integration of hops
        """
        self.η = η
        self.H_s = H_s
        self.L = L
        self.W = W
        self.k_max = k_max
        self.dim = self.H_s.shape[0]
        self.step = self._make_hops_step()
        self.seed = seed
        self._solver_args = solver_args

    def _make_hops_step(self):
        H_si = -1j * self.H_s
        B = -np.conj(self.L).T
        K = np.diag(np.arange(0, self.k_max + 1))

        def step(t, ψ):
            ψ = ψ.reshape((self.dim, self.k_max + 1), order="F")

            # 1. Apply system H and tho constant contributions
            ψ_1 = H_si @ ψ + self.W * (ψ @ K) + self.L @ (ψ * self.η(t))

            # 2. Now the shifted orders, we set the truncator to zero
            zeros = np.zeros((1, self.dim)).T
            ψ_ext = np.hstack((zeros, ψ, zeros))

            ψ_2 = B @ ψ_ext[:, 0:-2]
            ψ_3 = self.L @ (ψ_ext[:, 2:] @ K)

            res = np.array((ψ_1 + ψ_2 + ψ_3)).reshape(
                ((self.k_max + 1) * self.dim,), order="F"
            )

            return res

        return step

    def integrate_hops_trajectory(self, ψ_0, τ_max, seed=None):
        seed = self.seed if seed is None else seed
        ψ_0_ext = np.concatenate((ψ_0, np.zeros(self.k_max * self.dim))) + 0j
        self.η.new_process(seed=seed)
        return scipy.integrate.solve_ivp(
            self.step,
            (0, τ_max),
            ψ_0_ext,
            vectorized=False,
            dense_output=True,
            **self._solver_args
        )

    def integrate_hops_ensemble(self, ψ_0, τ, N):
        τ_max = np.max(τ)
        ρ = np.zeros((len(τ), self.dim, self.dim), dtype="complex128")
        j = np.zeros(len(τ))

        for i in range(0, N):
            traj = self.integrate_hops_trajectory(
                ψ_0, τ_max, seed=(None if self.seed is None else self.seed + i)
            ).sol(τ)

            ψ = traj[0:2, :].T
            ρ += ψ[:, :, np.newaxis] * ψ.conj()[:, np.newaxis, :]
            j_c = self.energy_tranfser_for_trajectory(np.copy(traj), τ)
            j += j_c

        return τ, ρ / N, j

    def energy_tranfser_for_trajectory(self, ψ, τ):
        """Calculate the energy transfer for one trajectory ``ψ`` which
        includes at least the first HOPS hirarchy."""
        ψ_0 = np.array(self.L @ ψ[0:2, :])
        ψ_1 = ψ[2:4, :]

        return 2 * np.real(-1j * self.W * np.sum(ψ_0.conj().T * ψ_1.T, axis=1))