update hopsflow for two baths

This commit is contained in:
Valentin Boettcher 2022-03-04 16:30:42 +01:00
parent 7a2231f532
commit df67a6b95e
2 changed files with 182 additions and 146 deletions

View file

@ -22,13 +22,12 @@ class SystemParams:
"""A parameter object to hold information about the physical """A parameter object to hold information about the physical
system and global HOPS parameters. system and global HOPS parameters.
:param L: the coupling operator as system matrix :param L: the coupling operators as system matrices
:param G: the coupling factors in the exponential expansion of the BCF :param G: the coupling factors in the exponential expansion of the BCF
:param W: the exponents 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 bcf_scale: BCF scale factors
:param nonlinear: whether the trajectory was obtained through the nonlinear HOPS :param nonlinear: whether the trajectory was obtained through the nonlinear HOPS
:param g: individual scaling factors for the hierarchy states :param fock_hops: whether the now fock hops hierarchy normalization is used
:attr dim: the dimensionality of the system
""" """
__slots__ = [ __slots__ = [
@ -36,54 +35,54 @@ class SystemParams:
"G", "G",
"W", "W",
"bcf_scale", "bcf_scale",
"g",
"nonlinear", "nonlinear",
"coupling_flow_prefactors",
"dim", "dim",
] ]
def __init__( def __init__(
self, self,
L: np.ndarray, L: list[np.ndarray],
G: np.ndarray, G: list[np.ndarray],
W: np.ndarray, W: list[np.ndarray],
bcf_scale: float = 1, bcf_scale: Optional[list[float]] = None,
nonlinear: bool = False, nonlinear: bool = False,
g: Optional[np.ndarray] = None, fock_hops: bool = True,
): ):
"""Class initializer. Computes the most useful attributes """Class initializer. Computes the most useful attributes
ahead of time.""" ahead of time."""
self.L = L self.L = L
self.G = G self.G = [g * scale for g, scale in zip(G, bcf_scale)] if bcf_scale else G
self.W = W self.W = W
self.bcf_scale = bcf_scale
self.g = g self.coupling_flow_prefactors = [
W * (-np.sqrt(G) if fock_hops else 1j * G) for G, W in zip(self.G, self.W)
]
self.nonlinear = nonlinear self.nonlinear = nonlinear
self.dim = L.shape[0]
#: the dimensionality of the system
self.dim = L[0].shape[0]
class HOPSRun: class HOPSRun:
"""A parameter object to hold data commonly used by the energy """A parameter object to hold data commonly used by the energy
flow calculations. flow calculations.
:param ψ_0: the HOPS trajectory ``(time, dim-state)`` :param ψ_0: The HOPS trajectory ``(time, dim-state)``.
:param ψ_1: the first HOPS hierarchy states ``(time, hierarchy-width * dim-state)`` :param ψ_1: The first HOPS hierarchy states ``(time,
- will be reshaped to ``(time, hierarchy-width, dim_state)`` 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 It will be reshaped to a list containing arrays of the
:attr ψ_coup: ψ_0 with the coupling operator applied shape``(time, hierarchy-width, dim_state)`` (one for each
:attr hierarch_width: the width of the hierarchy (number of exponential terms) bath).
:attr t_steps: the number of timesteps
:attr nonlinear: whether the trajectory was obtained through the nonlinear HOPS
""" """
__slots__ = [ __slots__ = [
"ψ_0", "ψ_0",
"ψ_1", "ψ_1_n",
"norm_squared", "norm_squared",
"ψ_coup", "ψ_coups",
"hierarchy_width",
"t_steps", "t_steps",
"nonlinear", "nonlinear",
] ]
@ -94,17 +93,31 @@ class HOPSRun:
self.ψ_0 = ψ_0 self.ψ_0 = ψ_0
self.nonlinear = params.nonlinear self.nonlinear = params.nonlinear
self.norm_squared = np.sum(self.ψ_0.conj() * self.ψ_0, axis=1).real self.norm_squared = (
self.ψ_coup = util.apply_operator(self.ψ_0, params.L) np.sum(self.ψ_0.conj() * self.ψ_0, axis=1).real
if params.nonlinear
else None
)
"""
The squared norm of the state, may be ``None`` for linear HOPS.
"""
self.ψ_coups = [util.apply_operator(self.ψ_0, L) for L in params.L]
"""
ψ_0 with the coupling operator applied for each bath.
"""
self.t_steps = self.ψ_0.shape[0] self.t_steps = self.ψ_0.shape[0]
self.hierarchy_width = ψ_1.shape[1] // params.dim """The number of timesteps."""
ψ_1_rs = ψ_1.reshape(self.t_steps, self.hierarchy_width, params.dim) hierarchy_width = ψ_1.shape[1] // params.dim
ψ_1_rs = ψ_1.reshape(self.t_steps, hierarchy_width, params.dim)
if params.g: offsets = np.cumsum([0] + [len(g) for g in params.G])
ψ_1_rs /= params.g[None, :, None] self.ψ_1_n = [
ψ_1_rs[:, offsets[i] : offsets[i + 1], :] for i in range(len(offsets) - 1)
self.ψ_1 = ψ_1_rs ]
"""The first hierarchy states for each bath."""
def normalize_maybe(self, array: np.ndarray): def normalize_maybe(self, array: np.ndarray):
"""Normalize the ``array`` if nonlinear HOPS is being used.""" """Normalize the ``array`` if nonlinear HOPS is being used."""
@ -118,12 +131,14 @@ class ThermalParams:
"""Aparameter object to hold information abouth the thermal """Aparameter object to hold information abouth the thermal
stochastic process. stochastic process.
:param ξ: the thermal stochastic process :param ξ: the thermal stochastic processes
:param τ: the time points of the trajectory :param τ: the time points of the trajectory
:param rand_skip: how many random numbers to skip when
parametrizing the stochastic process
:param num_deriv: whether to calculate the derivative of the :param num_deriv: whether to calculate the derivative of the
process numerically or use it directly from the process numerically or use it directly from the
:class:`StocProc`. The latter alternative is strongly preferred :class:`StocProc`. The latter alternative is strongly
if available. preferred if available.
:param dx: step size for the calculation of the derivative of ξ, :param dx: step size for the calculation of the derivative of ξ,
only relevant if numerical differentiation is used only relevant if numerical differentiation is used
:param order: order the calculation of the derivative of ξ, must :param order: order the calculation of the derivative of ξ, must
@ -131,25 +146,18 @@ class ThermalParams:
numerical differentiation is used numerical differentiation is used
""" """
__slots__ = ["ξ", "τ", "dx", "order", "num_deriv", "rand_skip"] __slots__ = ["ξs", "τ", "dx", "order", "num_deriv", "rand_skip"]
def __init__( def __init__(
self, self,
ξ: StocProc, ξs: list[Optional[StocProc]],
τ: np.ndarray, τ: np.ndarray,
rand_skip: int = 0, rand_skip: int = 0,
num_deriv: bool = False, num_deriv: bool = False,
dx: float = 1e-3, dx: float = 1e-3,
order: int = 3, order: int = 3,
): ):
"""Class initializer. Computes the most useful attributes self.ξs = ξs
ahead of time.
The process ξ is intialized with ξ_coeff and its derivative is
being calculated.
"""
self.ξ = ξ
self.τ = τ self.τ = τ
self.dx = dx self.dx = dx
self.order = order self.order = order
@ -167,40 +175,52 @@ class ThermalRunParams:
:param seed: the seed used to generate the random numbers for the :param seed: the seed used to generate the random numbers for the
process realization process realization
:attr ξ_dot: the process derivative evaluated at τ :attr :attr ξ_dots: the process derivatives evaluated at τ :attr
:attr ξ_values: the process evaluated at τ :attr ξ_values: the processes evaluated at τ
:attr ξ_coeff: the coefficients of the realization of the thermal :attr ξ_coeff: the coefficients of the realization of the thermal
stochastic process stochastic processes
""" """
__slots__ = ["ξ_coeff", "ξ_dot", "ξ_values"] __slots__ = ["ξ_coeff", "ξ_dots", "ξ_values"]
def __init__(self, therm_params: ThermalParams, seed: int): def __init__(self, therm_params: ThermalParams, seed: int):
"""Class initializer. Computes the most useful attributes """Class initializer. Computes the most useful attributes
ahead of time. ahead of time.
The process ξ is intialized with ξ_coeff and its derivative is The processes ξs are intialized with ξ_coeff and their
being calculated. derivative is being calculated.
""" """
np.random.seed(seed) np.random.seed(seed)
np.random.rand(therm_params.rand_skip) np.random.rand(therm_params.rand_skip)
self.ξ_coeff = util.uni_to_gauss(np.random.rand(therm_params.ξ.get_num_y() * 2)) # type: ignore
therm_params.ξ.new_process(self.ξ_coeff)
self.ξ_values = therm_params.ξ(therm_params.τ) self.ξ_coeff = []
self.ξ_dot: np.ndarray = ( self.ξ_values = []
self.ξ_dots = []
for ξ in therm_params.ξs:
if ξ:
coeff = util.uni_to_gauss(np.random.rand((ξ.get_num_y() or 0) * 2))
self.ξ_coeff.append(coeff)
ξ.new_process(coeff)
self.ξ_values.append(ξ(therm_params.τ))
self.ξ_dots.append(
scipy.misc.derivative( scipy.misc.derivative(
therm_params.ξ, ξ,
therm_params.τ, therm_params.τ,
dx=therm_params.dx, dx=therm_params.dx,
order=therm_params.order, order=therm_params.order,
) )
if therm_params.num_deriv if therm_params.num_deriv
else therm_params.ξ.dot(therm_params.τ) else ξ.dot(therm_params.τ)
) )
else:
self.ξ_coeff.append(0)
self.ξ_values.append(0)
self.ξ_dots.append(None)
############################################################################### ###############################################################################
@ -212,8 +232,8 @@ def flow_trajectory_coupling(
run: HOPSRun, run: HOPSRun,
params: SystemParams, params: SystemParams,
) -> np.ndarray: ) -> np.ndarray:
r"""Calculates the $\langle L^\dagger \dot{B} + c.c.$ part of the r"""Calculates the :math:`\langle L^\dagger \dot{B} + c.c.` part of the
energy flow for a single trajectory. energy flow for a single trajectory and for each bath.
:param run: a parameter object for the current trajectory, see :any:`HOPSRun` :param run: a parameter object for the current trajectory, see :any:`HOPSRun`
:param therma_run: a parameter object for the current thermal process, see :any:`HOPSRun` :param therma_run: a parameter object for the current thermal process, see :any:`HOPSRun`
@ -223,18 +243,24 @@ def flow_trajectory_coupling(
""" """
# here we apply the prefactors to each hierarchy state # here we apply the prefactors to each hierarchy state
ψ_hops = util.mulitply_hierarchy( ψ_hopses = [
(1j * params.W * params.G * params.bcf_scale), run.ψ_1 util.mulitply_hierarchy(f, ψ_1)
) for f, ψ_1 in zip(params.coupling_flow_prefactors, run.ψ_1_n)
]
flow = 2 * util.dot_with_hierarchy(run.ψ_coup.conj(), ψ_hops).real flows = np.array(
[
2 * util.dot_with_hierarchy(ψ_coup.conj(), ψ_hops).real
for ψ_coup, ψ_hops in zip(run.ψ_coups, ψ_hopses)
]
) # fromiter crashed...
# optionally normalize # optionally normalize
return run.normalize_maybe(flow) return run.normalize_maybe(flows)
def flow_trajectory_therm(run: HOPSRun, therm_run: ThermalRunParams) -> np.ndarray: def flow_trajectory_therm(run: HOPSRun, therm_run: ThermalRunParams) -> np.ndarray:
r"""Calculates the $\langle L^\dagger \dot{\xi} + c.c.$ part of the r"""Calculates the :math:`\langle L^\dagger \dot{\xi} + c.c.` part of the
energy flow for a single trajectory. energy flow for a single trajectory.
:param run: a parameter object, see :any:`HOPSRun` :param run: a parameter object, see :any:`HOPSRun`
@ -242,18 +268,25 @@ def flow_trajectory_therm(run: HOPSRun, therm_run: ThermalRunParams) -> np.ndarr
:returns: the value of the flow for each time step :returns: the value of the flow for each time step
""" """
flow = ( flows = np.array(
[
2 2
* ( * (
run.normalize_maybe(np.sum(run.ψ_coup.conj() * run.ψ_0, axis=1)) run.normalize_maybe(np.sum(ψ_coup.conj() * run.ψ_0, axis=1)) * ξ_dot
* therm_run.ξ_dot # note: this is already scaled by stocproc if ξ_dot is not None
else np.zeros(
ψ_coup.shape[0]
) # note: this is already scaled by stocproc
).real ).real
for ψ_coup, ξ_dot in zip(run.ψ_coups, therm_run.ξ_dots)
]
) )
return flow
return flows
def flow_trajectory( def flow_trajectory(
run: HOPSRun, params: SystemParams, therm_run: Optional[ThermalRunParams] run: HOPSRun, params: SystemParams, therm_run: Optional[ThermalRunParams] = None
) -> np.ndarray: ) -> np.ndarray:
r"""Calculates the total energy flow for a trajectory. r"""Calculates the total energy flow for a trajectory.
@ -270,36 +303,36 @@ def flow_trajectory(
return flow return flow
def interaction_energy_coupling(run: HOPSRun, params: SystemParams) -> np.ndarray: # def interaction_energy_coupling(run: HOPSRun, params: SystemParams) -> np.ndarray:
"""Calculates the coupling part of the interaction energy # """Calculates the coupling part of the interaction energy
expectation value for a trajectory. # expectation value for a trajectory.
:param run: a parameter object, see :any:`HOPSRun` # :param run: a parameter object, see :any:`HOPSRun`
:param params: a parameter object for the system, see # :param params: a parameter object for the system, see
:any:`SystemParams` # :any:`SystemParams`
:returns: the value of the coupling interaction energy for each time step # :returns: the value of the coupling interaction energy for each time step
""" # """
ψ_hops = util.mulitply_hierarchy(-1j * params.G * params.bcf_scale, run.ψ_1) # ψ_hops = util.mulitply_hierarchy(-1j * params.G * params.bcf_scale, run.ψ_1)
energy = util.dot_with_hierarchy(run.ψ_coup.conj(), ψ_hops).real * 2 # energy = util.dot_with_hierarchy(run.ψ_coup.conj(), ψ_hops).real * 2
return run.normalize_maybe(energy) # return run.normalize_maybe(energy)
def interaction_energy_therm(run: HOPSRun, therm_run: ThermalRunParams) -> np.ndarray: # def interaction_energy_therm(run: HOPSRun, therm_run: ThermalRunParams) -> np.ndarray:
r"""Calculates the thermal part of the interaction energy. # r"""Calculates the thermal part of the interaction energy.
:param run: a parameter object, see :any:`HOPSRun` # :param run: a parameter object, see :any:`HOPSRun`
:param therm_run: a parameter object, see :any:`ThermalParams` # :param therm_run: a parameter object, see :any:`ThermalParams`
:returns: the value of the thermal interaction for each time step # :returns: the value of the thermal interaction for each time step
""" # """
energy = ( # energy = (
run.normalize_maybe((run.ψ_coup.conj() * run.ψ_0).sum(axis=1)) # run.normalize_maybe((run.ψ_coup.conj() * run.ψ_0).sum(axis=1))
* therm_run.ξ_values # * therm_run.ξ_values
).real * 2 # ).real * 2
return energy # return energy
############################################################################### ###############################################################################
@ -367,56 +400,56 @@ def heat_flow_ensemble(
) )
def _interaction_energy_ensemble_body( # def _interaction_energy_ensemble_body(
ψs: Union[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray, int]], # ψs: Union[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray, int]],
params: SystemParams, # params: SystemParams,
thermal: ThermalParams, # thermal: ThermalParams,
) -> np.ndarray: # ) -> np.ndarray:
ψ_0, ψ_1 = ψs[0:2] # ψ_0, ψ_1 = ψs[0:2]
ys = None # ys = None
if len(ψs) == 3: # if len(ψs) == 3:
ys = ψs[-1] # ys = ψs[-1]
run = HOPSRun(ψ_0, ψ_1, params) # run = HOPSRun(ψ_0, ψ_1, params)
energy = interaction_energy_coupling(run, params) # energy = interaction_energy_coupling(run, params)
if isinstance(ys, int): # if isinstance(ys, int):
therm_run = ThermalRunParams(thermal, ys) # therm_run = ThermalRunParams(thermal, ys)
energy += interaction_energy_therm(run, therm_run) # energy += interaction_energy_therm(run, therm_run)
return energy # return energy
def interaction_energy_ensemble( # def interaction_energy_ensemble(
ψ_0s: Iterator[np.ndarray], # ψ_0s: Iterator[np.ndarray],
ψ_1s: Iterator[np.ndarray], # ψ_1s: Iterator[np.ndarray],
params: SystemParams, # params: SystemParams,
N: Optional[int], # N: Optional[int],
therm_args: Optional[Tuple[Iterator[int], ThermalParams]] = None, # therm_args: Optional[Tuple[Iterator[int], ThermalParams]] = None,
**kwargs, # **kwargs,
) -> util.EnsembleReturn: # ) -> util.EnsembleReturn:
"""Calculates the heat flow for an ensemble of trajectories. # """Calculates the heat flow for an ensemble of trajectories.
:param ψ_0s: array of trajectories ``(N, time-steps, dim-state)`` # :param ψ_0s: array of trajectories ``(N, time-steps, dim-state)``
:param ψ_0s: array of the first HOPS hierarchy states ``(N, time, # :param ψ_0s: array of the first HOPS hierarchy states ``(N, time,
hierarchy-width * dim-state)`` # hierarchy-width * dim-state)``
:param params: a parameter object for the system, see # :param params: a parameter object for the system, see
:any:`SystemParams` # :any:`SystemParams`
:param therm_args: the realization parameters and the parameter # :param therm_args: the realization parameters and the parameter
object, see :any:`ThermalParams` # object, see :any:`ThermalParams`
The ``**kwargs`` are passed to :any:`hopsflow.utility.ensemble_mean`. # The ``**kwargs`` are passed to :any:`hopsflow.utility.ensemble_mean`.
:returns: the value of the flow for each time step # :returns: the value of the flow for each time step
""" # """
return util.ensemble_mean( # return util.ensemble_mean(
iter(zip(ψ_0s, ψ_1s, therm_args[0])) if therm_args else iter(zip(ψ_0s, ψ_1s)), # iter(zip(ψ_0s, ψ_1s, therm_args[0])) if therm_args else iter(zip(ψ_0s, ψ_1s)),
_interaction_energy_ensemble_body, # _interaction_energy_ensemble_body,
N, # N,
(params, therm_args[1] if therm_args else None), # (params, therm_args[1] if therm_args else None),
**kwargs, # **kwargs,
) # )
# def _shift_expectation_body( # def _shift_expectation_body(

View file

@ -344,6 +344,9 @@ def ensemble_mean(
), ),
cls=JSONEncoder, cls=JSONEncoder,
ensure_ascii=False, ensure_ascii=False,
default=lambda obj: obj.__dict__
if hasattr(obj, "__dict__")
else "<not serializable>",
).encode("utf-8") ).encode("utf-8")
if save: if save:
@ -353,7 +356,7 @@ def ensemble_mean(
) )
if not overwrite_cache and path.exists(): if not overwrite_cache and path.exists():
logging.info(f"Loading cache from: {path}") logging.warning(f"Loading cache from: {path}")
return np.load(str(path), allow_pickle=True) return np.load(str(path), allow_pickle=True)
if N == 1: if N == 1: