put some more functionality into the base class

model_base.py

@@ -1,13 +1,25 @@
 """A base class for model HOPS configs."""
 
+from typing import Optional
 from utility import JSONEncoder, object_hook
 import numpy as np
+from numpy.typing import NDArray
 import json
 import copy
 import hashlib
+from abc import ABC, abstractmethod
+import qutip as qt
+from hops.core.hierarchy_data import HIData
+import hopsflow
+from hopsflow.util import EnsembleReturn
+import hashlib
 
 
-class Model:
+class MyABC(ABC):
+    pass
+
+
+class Model(ABC):
     """
     A base class with some data management functionality.
     """
@@ -29,11 +41,9 @@ class Model:
     def to_json(self):
         """Returns a json representation of the model configuration."""
 
-        return json.dumps(
+        return JSONEncoder.dumps(
             {key: self.__dict__[key] for key in self.__dict__ if key[0] != "_"}
-            | {"__version__": self.__version__},
+            | {"__version__": self.__version__}
-            cls=JSONEncoder,
-            ensure_ascii=False,
         )
 
     def __hash__(self):
@@ -77,3 +87,131 @@ class Model:
         """Return a deep copy of the model."""
 
         return copy.deepcopy(self)
+
+    @property
+    @abstractmethod
+    def system(self) -> qt.Qobj:
+        """The system hamiltonian."""
+
+        pass
+
+    @property
+    @abstractmethod
+    def coupling_operators(self) -> list[np.ndarray]:
+        """The bath coupling operators :math:`L`."""
+        pass
+
+    @abstractmethod
+    def bcf_coefficients(
+        self, n: Optional[int] = None
+    ) -> tuple[list[NDArray[np.complex128]], list[NDArray[np.complex128]]]:
+        """
+        The normalized zero temperature BCF fit coefficients
+        :math:`[G_i], [W_i]` with ``n`` terms.
+        """
+
+        pass
+
+    @property
+    @abstractmethod
+    def thermal_processes(self) -> list[Optional[hopsflow.hopsflow.StocProc]]:
+        """
+        The thermal noise stochastic processes for each bath.
+        :any:`None` means zero temperature.
+        """
+
+        pass
+
+    @property
+    @abstractmethod
+    def bcf_scales(self) -> list[float]:
+        """The scaling factors for the bath correlation functions."""
+
+        pass
+
+    @property
+    def hopsflow_system(self) -> hopsflow.hopsflow.SystemParams:
+        """The :any:`hopsflow` system config for the system."""
+
+        g, w = self.bcf_coefficients()
+
+        return hopsflow.hopsflow.SystemParams(
+            L=self.coupling_operators,
+            G=[g_i * scale for g_i, scale in zip(g, self.bcf_scales)],
+            W=w,
+            fock_hops=True,
+        )
+
+    def hopsflow_therm(
+        self, τ: NDArray[np.float64]
+    ) -> Optional[hopsflow.hopsflow.ThermalParams]:
+        """The :any:`hopsflow` thermal config for the system."""
+
+        processes = self.thermal_processes
+        scales = self.bcf_scales
+
+        for process, scale in zip(processes, scales):
+            if process:
+                process.set_scale(scale)
+                process.calc_deriv = True
+
+        return hopsflow.hopsflow.ThermalParams(processes, τ)
+
+    ###########################################################################
+    #                            Derived Quantities                           #
+    ###########################################################################
+
+    def system_expectation(
+        self, data: HIData, operator: qt.Qobj, **kwargs
+    ) -> EnsembleReturn:
+        """Calculates the expectation value of ``operator`` from the
+        hierarchy data ``data``.
+
+        The ``kwargs`` are passed on to
+        :any:`hopsflow.util.ensemble_mean`.
+
+        :returns: See :any:`hopsflow.util.ensemble_mean`.
+        """
+
+        operator_hash = JSONEncoder.hexhash(operator).encode("utf-8")
+
+        return hopsflow.util.operator_expectation_ensemble(
+            data.stoc_traj,  # type: ignore
+            operator.full(),
+            kwargs.get("N", data.samples),
+            nonlinear=True,  # always nonlinear
+            save=f"{operator_hash}_{self.__hash__()}",
+        )
+
+    def system_energy(self, data: HIData, **kwargs) -> EnsembleReturn:
+        """Calculates the system energy from the hierarchy data
+        ``data``.
+
+        The ``kwargs`` are passed on to
+        :any:`hopsflow.util.ensemble_mean`.
+
+        :returns: See :any:`hopsflow.util.ensemble_mean`.
+        """
+
+        operator = self.system.full()
+        return self.system_expectation(data, operator, **kwargs)
+
+    def bath_energy_flow(self, data: HIData, **kwargs) -> EnsembleReturn:
+        """Calculates the bath energy flow from the hierarchy data
+        ``data``.
+
+        The ``kwargs`` are passed on to
+        :any:`hopsflow.util.ensemble_mean`.
+
+        :returns: See :any:`hopsflow.util.ensemble_mean`.
+        """
+
+        return hopsflow.hopsflow.heat_flow_ensemble(
+            data.stoc_traj,  # type: ignore
+            data.aux_states,  # type: ignore
+            self.hopsflow_system,
+            kwargs.get("N", data.samples),
+            (data.rng_seed, self.hopsflow_therm),  # type: ignore
+            save=f"flow_{self.__hash__()}",
+            **kwargs,
+        )

two_qubit_model.py

@@ -29,6 +29,7 @@ r"""
 
 import copy
 import dataclasses
+import hopsflow
 from dataclasses import dataclass, field
 import functools
 import itertools
@@ -251,6 +252,12 @@ class TwoQubitModel(Model):
 
         return qt.tensor(qt.identity(2), coupling_op)
 
+    @property
+    def coupling_operators(self) -> list[np.ndarray]:
+        """The bath coupling operators :math:`L`."""
+
+        return [self.bath_coupling(i).full() for i in (0, 1)]
+
     def bcf_scale(self, i: int) -> float:
         """
         The BCF scaling factor of the ``i``th bath.
@@ -258,6 +265,12 @@ class TwoQubitModel(Model):
 
         return float(self.δ[i]) ** 2
 
+    @property
+    def bcf_scales(self) -> list[float]:
+        """The scaling factors for the bath correlation functions."""
+
+        return [self.bcf_scale(i) for i in (1, 2)]
+
     def η(self, i: int):
         """The BCF pre-factor :math:`η` of the ``i``th bath."""
         ω_c = float(self.ω_c[i])
@@ -332,14 +345,22 @@ class TwoQubitModel(Model):
         )
 
     def bcf_coefficients(
-        self, i: int, n: Optional[int] = None
+        self, n: Optional[int] = None
-    ) -> tuple[NDArray[np.complex128], NDArray[np.complex128]]:
+    ) -> tuple[list[NDArray[np.complex128]], list[NDArray[np.complex128]]]:
         """
         The normalizedzero temperature BCF fit coefficients
         :math:`G_i,W_i` the ``i``th bath with ``n`` terms.
         """
-        n = n or self.bcf_terms[i]
+        g, w = [], []
-        return self.bcf(i).exponential_coefficients(n)
+
+        for i in 0, 1:
+            n = n or self.bcf_terms[i]
+            g_n, w_n = self.bcf(i).exponential_coefficients(n)
+
+            g.append(g_n)
+            w.append(w_n)
+
+        return g, w
 
     @staticmethod
     def basis(n_1: int = 1, n_2: int = 1) -> qt.Qobj:
@@ -382,21 +403,19 @@ class TwoQubitModel(Model):
             negative_frequencies=False,
         )
 
+    @property
+    def thermal_processes(self) -> list[Optional[sp.StocProc]]:
+        """
+        The thermal noise stochastic processes for each bath.
+        :any:`None` means zero temperature.
+        """
+
+        return [self.thermal_process(i) for i in (0, 1)]
+
     ###########################################################################
     #                                 Utility                                 #
     ###########################################################################
 
-    def effective_coupling(self, i: int) -> float:
-        """
-        The effective coupling strength of bath ``i``.  The maximum pre-factor times the bcf scale
-        divided by the minimal damping of the bcf expansion.
-        """
-
-        g, w = self.bcf_coefficients(i)
-
-        term = np.argmax(np.abs(g))
-        return self.bcf_scale(i) * np.real(g[term] / w.real[term])
-
     @property
     def hops_config(self):
         """
@@ -404,14 +423,13 @@ class TwoQubitModel(Model):
         for this system.
         """
 
-        g_1, w_1 = self.bcf_coefficients(0)
+        g, w = self.bcf_coefficients()
-        g_2, w_2 = self.bcf_coefficients(1)
 
         system = params.SysP(
             H_sys=self.system.full(),
             L=[self.bath_coupling(0).full(), self.bath_coupling(1).full()],
-            g=[g_1, g_2],
+            g=g,
-            w=[w_1, w_2],
+            w=w,
             bcf_scale=[self.bcf_scale(0), self.bcf_scale(1)],
             T=self.T,
             description=self.description,

utility.py

@@ -6,6 +6,7 @@ import json
 import numpy as np
 import qutip as qt
 from typing import Any
+import hashlib
 
 
 @beartype
@@ -54,6 +55,37 @@ class JSONEncoder(json.JSONEncoder):
     def _(self, obj: StocProcTolerances):
         return {"type": "StocProcTolerances", "value": dataclasses.asdict(obj)}
 
+    @classmethod
+    def dumps(cls, data, **kwargs) -> str:
+        """Like :any:`json.dumps`, just for this encoder.
+
+        The ``kwargs`` are passed on to :any:`json.dumps`.
+        """
+
+        return json.dumps(
+            data,
+            **kwargs,
+            cls=cls,
+            ensure_ascii=False,
+        )
+
+    @classmethod
+    def hash(cls, data, **kwargs):
+        """
+        Like :any:`dumps`, only that the result is being piped into
+        :any:`hashlib.sha256`. A ``sha256`` hash is being returned.
+        """
+
+        return hashlib.sha256(cls.dumps(data, **kwargs).encode("utf-8"))
+
+    @classmethod
+    def hexhash(cls, data, **kwargs) -> str:
+        """
+        Like :any:`hash`, only that a hexdigest is being returned.
+        """
+
+        return cls.hash(data, **kwargs).hexdigest()
+
 
 def object_hook(dct: dict[str, Any]):
     """A custom decoder for the types introduced in :any:`JSONEncoder`."""