From f8b554fcfe29254c56b7ecb06f495be4a03ebdd7 Mon Sep 17 00:00:00 2001 From: Valentin Boettcher Date: Tue, 29 Nov 2022 18:11:07 -0500 Subject: [PATCH] do not allow adding values with different bath numbers --- hopsflow/util.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/hopsflow/util.py b/hopsflow/util.py index 055df69..8a2531a 100644 --- a/hopsflow/util.py +++ b/hopsflow/util.py @@ -189,10 +189,15 @@ class EnsembleValue: def __add__( self, other: Union["EnsembleValue", float, int, np.ndarray] ) -> EnsembleValue: - if type(self) == type(other): + if isinstance(other, EnsembleValue): if len(self) != len(other): raise RuntimeError("Can only add values of equal length.") + if self.num_baths != other.num_baths: + raise RuntimeError( + "Can only add values pertaining to an equal number of baths." + ) + left = self._value right = other._value