master-thesis/python/energy_flow_proper/hopsflow/hopsflow.py
2021-11-11 19:51:34 +01:00

354 lines
12 KiB
Python

"""Calculate the reservoir energy change from HOPS data.
The functions in this module mostly take parameter objects which hold
relevant information and compute commonly used values ahead of time.
The parameter objects are passed separately but may depend on each other.
"""
import numpy as np
import scipy.misc
import util as util
from typing import Optional, Tuple
from stocproc import StocProc
###############################################################################
# Interface/Parameter Object#
###############################################################################
class SystemParams:
"""A parameter object to hold information about the physical
system and global HOPS parameters.
: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 nonlinear: whether the trajectory was obtained through the nonlinear HOPS
:param g: individual scaling factors for the hierarchy states
:attr dim: the dimensionality of the system
"""
__slots__ = [
"L",
"G",
"W",
"bcf_scale",
"g",
"nonlinear",
"dim",
]
def __init__(
self,
L: np.ndarray,
G: np.ndarray,
W: np.ndarray,
bcf_scale: float = 1,
nonlinear: bool = False,
g: Optional[np.ndarray] = None,
):
"""Class initializer. Computes the most useful attributes
ahead of time."""
self.L = L
self.G = G
self.W = W
self.bcf_scale = bcf_scale
self.g = g
self.nonlinear = nonlinear
self.dim = L.shape[0]
class HOPSRun:
"""A parameter object to hold data commonly used by the energy
flow calculations.
:param ψ_0: the HOPS trajectory ``(time, dim-state)``
:param ψ_1: the first HOPS hierarchy states ``(time, hierarchy-width * dim-state)``
- will be reshaped to ``(time, hierarchy-width, dim_state)``
- will automatically be rescaled if scaling factors ``g`` are given
:attr norm_squared: the squared norm of the state, may be ``None`` for linear HOPS
:attr ψ_coup: ψ_0 with the coupling operator applied
:attr hierarch_width: the width of the hierarchy (number of exponential terms)
:attr t_steps: the number of timesteps
:attr nonlinear: whether the trajectory was obtained through the nonlinear HOPS
"""
__slots__ = [
"ψ_0",
"ψ_1",
"norm_squared",
"ψ_coup",
"hierarchy_width",
"t_steps",
"nonlinear",
]
def __init__(self, ψ_0: np.ndarray, ψ_1: np.ndarray, params: SystemParams):
"""Class initializer. Computes the most useful attributes
ahead of time."""
self.ψ_0 = ψ_0
self.nonlinear = params.nonlinear
self.norm_squared = np.sum(self.ψ_0.conj() * self.ψ_0, axis=1).real
self.ψ_coup = util.apply_operator(self.ψ_0, params.L)
self.t_steps = self.ψ_0.shape[0]
self.hierarchy_width = ψ_1.shape[1] // params.dim
ψ_1_rs = ψ_1.reshape(self.t_steps, self.hierarchy_width, params.dim)
if params.g:
ψ_1_rs /= params.g[None, :, None]
self.ψ_1 = ψ_1_rs
def normalize_maybe(self, array: np.ndarray):
"""Normalize the ``array`` if nonlinear HOPS is being used."""
if self.nonlinear:
return array / self.norm_squared
return array
class ThermalParams:
"""Aparameter object to hold information abouth the thermal
stochastic process.
:param ξ: the thermal stochastic process
:param τ: the time points of the trajectory
:param num_deriv: whether to calculate the derivative of the
process numerically or use it directly from the
:py:`StocProc`. The latter alternative is strongly preferred
if available.
:param dx: step size for the calculation of the derivative of ξ,
only relevant if numerical differentiation is used
:param order: order the calculation of the derivative of ξ, must
be odd, see :py:`scipy.misc.derivative`, only relevant if
numerical differentiation is used
"""
__slots__ = ["ξ", "τ", "dx", "order", "num_deriv"]
def __init__(
self,
ξ: StocProc,
τ: np.ndarray,
num_deriv: bool = False,
dx: float = 1e-3,
order: int = 3,
):
"""Class initializer. Computes the most useful attributes
ahead of time.
The process ξ is intialized with ξ_coeff and its derivative is
being calculated.
"""
self.ξ = ξ
self.τ = τ
self.dx = 1e-3
self.order = order
self.num_deriv = num_deriv
class ThermalRunParams:
"""A parameter object to hold information abouth the thermal
stochastic process.
:param therm_params: information abouth the thermal stochastic
process
:param ξ_coeff: the coefficients of the realization of the thermal
stochastic process
:attr ξ_dot: the process derivative evaluated at τ :attr
ξ_values: the process evaluated at τ
"""
__slots__ = ["ξ_coeff", "ξ_dot", "ξ_values"]
def __init__(self, therm_params: ThermalParams, ξ_coeff: np.ndarray):
"""Class initializer. Computes the most useful attributes
ahead of time.
The process ξ is intialized with ξ_coeff and its derivative is
being calculated.
"""
self.ξ_coeff = ξ_coeff
therm_params.ξ.new_process(self.ξ_coeff)
self.ξ_values = therm_params.ξ(therm_params.τ)
self.ξ_dot: np.ndarray = (
scipy.misc.derivative(
therm_params.ξ,
therm_params.τ,
dx=therm_params.dx,
order=therm_params.order,
)
if therm_params.num_deriv
else therm_params.ξ.dot(therm_params.τ)
)
###############################################################################
# Single Trajectory #
###############################################################################
def flow_trajectory_coupling(
run: HOPSRun,
params: SystemParams,
) -> np.ndarray:
r"""Calculates the $\langle L^\dagger \dot{B} + c.c.$ part of the
energy flow for a single trajectory.
:param run: a parameter object for the current trajectory, see :py:`HOPSRun`
:param therma_run: a parameter object for the current thermal process, see :py:`HOPSRun`
:param params: a parameter object for the system, see :py:`SystemParams`
:returns: the value of the flow for each time step
"""
# here we apply the prefactors to each hierarchy state
ψ_hops = util.mulitply_hierarchy(
(1j * params.W * params.G * params.bcf_scale), run.ψ_1
)
flow = 2 * util.dot_with_hierarchy(run.ψ_coup.conj(), ψ_hops).real
# optionally normalize
return run.normalize_maybe(flow)
def flow_trajectory_therm(run: HOPSRun, therm_run: ThermalRunParams) -> np.ndarray:
r"""Calculates the $\langle L^\dagger \dot{\xi} + c.c.$ part of the
energy flow for a single trajectory.
:param run: a parameter object, see :py:`HOPSRun`
:param therm_run: a parameter object, see :py:`ThermalRunParams`
:returns: the value of the flow for each time step
"""
flow = 2 * ((np.sum(run.ψ_coup.conj() * run.ψ_0, axis=1)) * therm_run.ξ_dot).real
return run.normalize_maybe(flow)
def flow_trajectory(
run: HOPSRun, params: SystemParams, therm_run: Optional[ThermalRunParams]
) -> np.ndarray:
r"""Calculates the total energy flow for a trajectory.
:param run: a parameter object, see :py:`HOPSRun`
:param therm: a parameter object, see :py:`ThermalParams`
:param params: a parameter object for the system, see :py:`SystemParams`
:returns: the value of the flow for each time step
"""
flow = flow_trajectory_coupling(run, params)
if therm_run:
flow += flow_trajectory_therm(run, therm_run)
return flow
def interaction_energy_coupling(run: HOPSRun, params: SystemParams) -> np.ndarray:
"""Calculates the coupling part of the interaction energy
expectation value for a trajectory.
:param run: a parameter object, see :py:`HOPSRun`
:param params: a parameter object for the system, see
:py:`SystemParams`
:returns: the value of the coupling interaction energy for each time step
"""
ψ_hops = util.mulitply_hierarchy(-1j * params.G * params.bcf_scale, run.ψ_1)
energy = util.dot_with_hierarchy(run.ψ_coup.conj(), ψ_hops).real * 2
return run.normalize_maybe(energy)
def interaction_energy_therm(run: HOPSRun, therm_run: ThermalRunParams) -> np.ndarray:
r"""Calculates the thermal part of the interaction energy.
:param run: a parameter object, see :py:`HOPSRun`
:param therm_run: a parameter object, see :py:`ThermalParams`
:returns: the value of the thermal interaction for each time step
"""
energy = ((run.ψ_coup.conj() * run.ψ_0).sum(axis=1) * therm_run.ξ_dot).real * 2
return energy
###############################################################################
# Ensemble #
###############################################################################
def heat_flow_ensemble(
ψ_0s: np.ndarray,
ψ_1s: np.ndarray,
params: SystemParams,
therm_args: Optional[Tuple[np.ndarray, ThermalParams]] = None,
) -> np.ndarray:
"""Calculates the heat flow for an ensemble of trajectories.
:param ψ_0s: array of trajectories ``(N, time-steps, dim-state)``
:param ψ_0s: array of the first HOPS hierarchy states ``(N, time,
hierarchy-width * dim-state)``
:param params: a parameter object for the system, see
:py:`SystemParams`
:param therm_args: the realization parameters and the parameter
object, see :py:`ThermalParams`
:returns: the value of the flow for each time step
"""
N = ψ_0s.shape[0]
flow = np.zeros(ψ_0s.shape[1])
for i in range(0, N):
run = HOPSRun(ψ_0s[i], ψ_1s[i], params)
flow += flow_trajectory_coupling(run, params)
if therm_args:
therm_run = ThermalRunParams(therm_args[1], therm_args[0][i])
flow += flow_trajectory_therm(run, therm_run)
return flow / N
def interaction_energy_ensemble(
ψ_0s: np.ndarray,
ψ_1s: np.ndarray,
params: SystemParams,
therm_args: Optional[Tuple[np.ndarray, ThermalParams]] = None,
) -> np.ndarray:
"""Calculates the heat flow for an ensemble of trajectories.
:param ψ_0s: array of trajectories ``(N, time-steps, dim-state)``
:param ψ_0s: array of the first HOPS hierarchy states ``(N, time,
hierarchy-width * dim-state)``
:param params: a parameter object for the system, see
:py:`SystemParams`
:param therm_args: the realization parameters and the parameter
object, see :py:`ThermalParams`
:returns: the value of the flow for each time step
"""
N = ψ_0s.shape[0]
energy = np.zeros(ψ_0s.shape[1])
for i in range(0, N):
run = HOPSRun(ψ_0s[i], ψ_1s[i], params)
energy += interaction_energy_coupling(run, params)
if therm_args:
therm_run = ThermalRunParams(therm_args[1], therm_args[0][i])
energy += interaction_energy_therm(run, therm_run)
return energy / N