begin implementation of proper module

python/hopsflow/hopsflow.py

@@ -0,0 +1,78 @@
+"""Calculate the reservoir energy change from HOPS data."""
+
+import qutip
+import numpy as np
+import scipy.misc
+
+###############################################################################
+#                              Single Trajectory                              #
+###############################################################################
+
+
+def flow_trajectory_coupling(ψ_0, ψ_1, L, G, W, bcf_scale=1, g=None, norm=None):
+    """Calculates the $\langle L^\dagger \dot{B} + c.c.$ part of the
+    energy flow for a single trajectory.
+
+    :param ψ_0: the trajectory ``(time, dim-state)``
+    :param ψ_1: the first HOPS hierarchy states ``(time, hierarchy-width * dim-state)``
+    :param L: the coupling operator as system matrix
+    :param G: the coupling factors in the exponential expansion of the BCF
+    :param W: the exponents in the exponential expansion of the BCF
+    :param bcf_scale: BCF scale factor
+    :param g: individual scaling factors for the hierarchy states
+    :param norm: the norm of the state, may be ``None`` for linear HOPS
+
+    :returns: the value of the flow for each time step
+    """
+
+    t_steps, dim = ψ_0.shape
+    hierarchy_width = ψ_1.shape[1] / dim
+
+    ψ_coup = np.array((L @ ψ_0.T).T)
+
+    # here we apply the prefactors to each hierarchy state
+    ψ_hops = (1j * W * G * bcf_scale)[None, :, None] * ψ_1.reshape(
+        t_steps, hierarchy_width, dim
+    )
+
+    # if there are per-term scaling factors
+    if len(g) == hierarchy_width:
+        ψ_hops /= g[None, :, None]
+
+    flow = 2 * np.sum(ψ_coup.conj()[:, None, :] * ψ_hops, axis=(1, 2)).real
+
+    # optionally normalize
+    if norm:
+        flow /= norm
+
+    return flow
+
+
+def flow_trajectory_therm(ψ_0, L, ξ, ξ_coeff, τ, norm=None, dx=1e-3, order=3):
+    """Calculates the $\langle L^\dagger \dot{\xi} + c.c.$ part of the
+    energy flow for a single trajectory.
+
+    :param ψ_0: the trajectory ``(time, dim-state)``
+    :param L: the coupling operator as system matrix
+    :param ξ: the thermal stochastic process
+    :param ξ_coeff: the coefficients of the realization of the
+                    thermal stochastic process
+    :param τ: the time points of the trajectory
+    :param norm: the norm of the state, may be ``None`` for linear HOPS
+    :param dx: step size for the calculation of the derivative of ξ
+    :param order: order the calculation of the derivative of ξ, must be odd,
+                  see :py:`scipy.misc.derivative`
+
+    :returns: the value of the flow for each time step
+    """
+
+    ξ.new_process(ξ_coeff)
+    ξ_dot = scipy.misc.derivative(ξ, τ, dx=dx, order=order)
+    ψ_coup = np.array((L @ ψ_0.T).T) # TODO: prevent calculating this twice
+
+    flow = 2 * (np.sum(ψ_coup.conj() * ψ_0)) * ξ_dot
+
+    if norm:
+        flow /= norm
+
+    return flow