From dd86c5fae18a28f4e7b16294f2fb3f9641aeed58 Mon Sep 17 00:00:00 2001 From: Valentin Boettcher Date: Mon, 11 Jul 2022 14:32:50 +0200 Subject: [PATCH] implement all the power stuff --- hiro_models/model_base.py | 148 +++++++++++++++++++++++++++++++++++--- 1 file changed, 137 insertions(+), 11 deletions(-) diff --git a/hiro_models/model_base.py b/hiro_models/model_base.py index 7223b2e..9c58a7a 100644 --- a/hiro_models/model_base.py +++ b/hiro_models/model_base.py @@ -1,13 +1,11 @@ """A base class for model HOPS configs.""" from dataclasses import dataclass, field -from typing import Any, Optional - +from typing import Any, Optional, Union, ClassVar from hops.util.dynamic_matrix import DynamicMatrix from .utility import JSONEncoder, object_hook import numpy as np from numpy.typing import NDArray -from typing import Union import json import copy import hashlib @@ -18,6 +16,7 @@ import hopsflow from hopsflow.util import EnsembleValue import hashlib import hops.core.hierarchy_parameters as params +from collections.abc import Callable @dataclass @@ -51,14 +50,25 @@ class Model(ABC): # Utility # ########################################################################### - def to_dict(self): - """Returns a dictionary representation of the model configuration.""" + def to_dict(self, extra_fields: "dict[str, Callable[[Model], str]]" = {}): + """Returns a dictionary representation of the model + configuration. - return {key: self.__dict__[key] for key in self.__dict__ if key[0] != "_"} | { - "__version__": self.__version__, - "__base_version__": self.__base_version__, - "__model__": self.__class__.__name__, - } + :param extra_fields: A dictionary whose keys will be added to + the final dict and whose values are callables that take + the model instance as an argument and return the value + that will be assigned the addyd key. + """ + + return ( + {key: self.__dict__[key] for key in self.__dict__ if key[0] != "_"} + | { + "__version__": self.__version__, + "__base_version__": self.__base_version__, + "__model__": self.__class__.__name__, + } + | {key: value(self) for key, value in extra_fields.items()} + ) def to_hashable_dict(self): """ @@ -273,6 +283,42 @@ class Model(ABC): operator = self.system return self.system_expectation(data, operator, real=True, **kwargs) + def system_power(self, data: HIData, **kwargs) -> Optional[EnsembleValue]: + """Calculates the power based on the time dependency of the + system hamiltonian from ``data``. + + The ``kwargs`` are passed on to + :any:`hopsflow.util.ensemble_mean`. + + :returns: See :any:`hopsflow.util.ensemble_mean`. Returns + :any:`None` if the system is static. + """ + + operator = self.system.derivative() + + if (abs(operator(self.t)).sum() == 0).all(): + return None + + return self.system_expectation(data, operator, real=True, **kwargs) + + def energy_change_from_system_power( + self, data: HIData, **kwargs + ) -> Optional[EnsembleValue]: + """Calculates the integrated system power from the hierarchy + data ``data``. + + The ``kwargs`` are passed on to :any:`system_power`. + + :returns: See :any:`system_power`. Returns :any:`None` if the + system is static. + """ + + power = self.system_power(data, **kwargs) + if power is not None: + return power.integrate(self.t) + + return None + def bath_energy_flow(self, data: HIData, **kwargs) -> EnsembleValue: """Calculates the bath energy flow from the hierarchy data ``data``. @@ -317,6 +363,54 @@ class Model(ABC): **kwargs, ) + def interaction_power(self, data: HIData, **kwargs) -> EnsembleValue: + """Calculates interaction power from the hierarchy data + ``data``. + + The ``kwargs`` are passed on to + :any:`hopsflow.util.interaction_energy_ensemble`. + + :returns: See :any:`hopsflow.util.interaction_energy_ensemble`. + """ + + N, kwargs = _get_N_kwargs(kwargs, data) + + return hopsflow.hopsflow.interaction_energy_ensemble( + data.valid_sample_iterator(data.stoc_traj), # type: ignore + data.valid_sample_iterator(data.aux_states), # type: ignore + self.hopsflow_system, + (data.valid_sample_iterator(data.rng_seed), self.hopsflow_therm(data.time[:])), # type: ignore + N=N, + save=f"interaction_power_{self.hexhash}", + power=True, + **kwargs, + ) + + def energy_change_from_interaction_power( + self, data: HIData, **kwargs + ) -> EnsembleValue: + """Calculates the integrated interaction power from the hierarchy data + ``data``. + + The ``kwargs`` are passed on to + :any:`hopsflow.util.energy_change_from_interaction_power`. + + :returns: See :any:`hopsflow.util.energy_change_from_interaction_power`. + """ + + N, kwargs = _get_N_kwargs(kwargs, data) + + return hopsflow.hopsflow.energy_change_from_interaction_power( + np.array(data.time), + data.valid_sample_iterator(data.stoc_traj), # type: ignore + data.valid_sample_iterator(data.aux_states), # type: ignore + self.hopsflow_system, + (data.valid_sample_iterator(data.rng_seed), self.hopsflow_therm(data.time[:])), # type: ignore + N=N, + save=f"interaction_power_{self.hexhash}", # under the hood the power is used + **kwargs, + ) + def bath_energy(self, data: HIData, **kwargs) -> EnsembleValue: """Calculates bath energy by integrating the bath energy flow calculated from the ``data``. @@ -359,7 +453,8 @@ class Model(ABC): return total - (system + bath) def total_energy(self, data: HIData, **kwargs) -> EnsembleValue: - """Calculates the total from the trajectories in ``data``. + """Calculates the total energy from the trajectories in + ``data`` using energy bilance. The ``kwargs`` are passed on to :any:`bath_energy`, :any:`system_energy` and :any:`interaction_energy`. @@ -375,6 +470,37 @@ class Model(ABC): return total + def total_power(self, data: HIData, **kwargs) -> EnsembleValue: + """Calculates the total power from the trajectories in + ``data``. + + The ``kwargs`` are passed on to :any:`system_power` and + :any:`interaction_power`. + + :returns: The total power. + """ + + power = self.interaction_power(data, **kwargs).sum_baths() + system_power = self.system_power(data, **kwargs) + + if system_power is not None: + power = power + system_power + + return power + + def total_energy_from_power(self, data: HIData, **kwargs) -> EnsembleValue: + """Calculates the total energy from the trajectories in + ``data`` using the integrated power. + + The ``kwargs`` are passed on to :any:`total_power`. + + :returns: The total energy. + """ + + return self.total_power(data, **kwargs).integrate(self.t) + + std_extra_fields: ClassVar = {"BCF scaling": lambda model: model.bcf_scale} + def _get_N_kwargs(kwargs: dict, data: HIData) -> tuple[int, dict]: N = kwargs.get("N", data.samples)