diff --git a/python/hopsflow/hopsflow.py b/python/hopsflow/hopsflow.py new file mode 100644 index 0000000..51b2fbd --- /dev/null +++ b/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