Repository: DHI/mikeio
Pull request #226 (2021-08-25 15:52:45)

Dataset math operations
Allow basic math operations +, - and * on Datasets. Both with scalar and other Datasets.


Patch:
diff --git a/docs/understanding_dataset.rst b/docs/understanding_dataset.rst
index 662a3ea5..0dae2376 100644
--- a/docs/understanding_dataset.rst
+++ b/docs/understanding_dataset.rst
@@ -112,6 +112,15 @@ including different ways of *selecting* data:
 * `average() <api.html#mikeio.Dataset.average>`_ - Compute the weighted average along the specified axis.
 * `aggregate() <api.html#mikeio.Dataset.aggregate>`_ - Aggregate along an axis
 
+*Mathematical operations* +, - and * with numerical values:
+* ds + value
+* ds - value
+* ds * value
+
+and + and - between two Datasets (if number of items and shapes conform):
+* ds1 + ds2
+* ds1 - ds2
+
 Other methods that also return a Dataset:
 
 * `interp_time() <api.html#mikeio.Dataset.interp_time>`_ - Temporal interpolation (`Time interpolation notebook <https://nbviewer.jupyter.org/github/DHI/mikeio/blob/main/notebooks/Time%20interpolation.ipynb>`_)
diff --git a/mikeio/dataset.py b/mikeio/dataset.py
index b1e5df4a..4d484cc3 100644
--- a/mikeio/dataset.py
+++ b/mikeio/dataset.py
@@ -201,6 +201,76 @@ class Dataset(TimeSeries):
 
         raise ValueError(f"indexing with a {type(x)} is not (yet) supported")
 
+    def __radd__(self, other):
+        return self.__add__(other)
+
+    def __add__(self, other):
+        if isinstance(other, self.__class__):
+            return self._add_dataset(other)
+        else:
+            return self._add_value(other)
+
+    def __rsub__(self, other):
+        ds = self.__mul__(-1.0)
+        return other + ds
+
+    def __sub__(self, other):
+        if isinstance(other, self.__class__):
+            return self._add_dataset(other, sign=-1.0)
+        else:
+            return self._add_value(-other)
+
+    def __rmul__(self, other):
+        return self.__mul__(other)
+
+    def __mul__(self, other):
+        if isinstance(other, self.__class__):
+            raise NotImplemented("Multiplication is not implemented for two Datasets")
+        else:
+            return self._multiply_value(other)
+
+    def _add_dataset(self, other, sign=1.0):        
+        self._check_datasets_match(other)
+        try:
+            data = [self[x] + sign*other[y] for x, y in zip(self.items, other.items)]
+        except:
+            raise ValueError("Could not add data in Dataset")
+        time = self.time.copy()
+        items = deepcopy(self.items)
+        return Dataset(data, time, items) 
+
+    def _check_datasets_match(self, other):
+        if self.n_items != other.n_items:
+            raise ValueError(f"Number of items must match ({self.n_items} and {other.n_items})")
+        for j in range(self.n_items):
+            if self.items[j].type != other.items[j].type:
+                raise ValueError(f"Item types must match. Item {j}: {self.items[j].type} != {other.items[j].type}")
+            if self.items[j].unit != other.items[j].unit:
+                raise ValueError(f"Item units must match. Item {j}: {self.items[j].unit} != {other.items[j].unit}")            
+        if not np.all(self.time == other.time):
+            raise ValueError("All timesteps must match")
+        if self.shape != other.shape:
+            raise ValueError("shape must match")    
+
+    def _add_value(self, value):
+        try:
+            data = [value + self[x] for x in self.items]
+        except:
+            raise ValueError(f"{value} could not be added to Dataset")        
+        items = deepcopy(self.items)
+        time = self.time.copy()
+        return Dataset(data, time, items)
+
+
+    def _multiply_value(self, value):
+        try:
+            data = [value * self[x] for x in self.items]
+        except:
+            raise ValueError(f"{value} could not be multiplied to Dataset")        
+        items = deepcopy(self.items)
+        time = self.time.copy()
+        return Dataset(data, time, items)
+
     def copy(self):
         """Returns a copy of this dataset."""
 
@@ -282,7 +352,6 @@ class Dataset(TimeSeries):
     def aggregate(self, axis=1, func=np.nanmean):
         """Aggregate along an axis
 
-
         Parameters
         ----------
         axis: int, optional


Test patch:
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
index a9f643ea..a114e281 100644
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -11,6 +11,31 @@ from mikeio.eum import EUMType, ItemInfo, EUMUnit
 def _get_time(nt):
     return list(rrule(freq=SECONDLY, count=nt, dtstart=datetime(2000, 1, 1)))
 
+@pytest.fixture
+def ds1():
+    nt = 10
+    ne = 7
+
+    d1 = np.zeros([nt, ne]) + 0.1
+    d2 = np.zeros([nt, ne]) + 0.2
+    data = [d1, d2]
+
+    time = _get_time(nt)
+    items = [ItemInfo("Foo"), ItemInfo("Bar")]
+    return Dataset(data, time, items)
+
+@pytest.fixture
+def ds2():
+    nt = 10
+    ne = 7
+
+    d1 = np.zeros([nt, ne]) + 1.0
+    d2 = np.zeros([nt, ne]) + 2.0
+    data = [d1, d2]
+
+    time = _get_time(nt)
+    items = [ItemInfo("Foo"), ItemInfo("Bar")]
+    return Dataset(data, time, items)
 
 def test_get_names():
 
@@ -878,6 +903,58 @@ def test_init():
     assert ds.n_elements == n_elements
     assert ds.items[0].name == "Foo"
 
+def test_add_scalar(ds1):
+    ds2 = ds1 + 10.0
+    assert np.all(ds2[0] - ds1[0] == 10.0)
+
+    ds3 = 10.0 + ds1
+    assert np.all(ds3[0] == ds2[0])
+    assert np.all(ds3[1] == ds2[1])       
+    
+def test_sub_scalar(ds1):
+    ds2 = ds1 - 10.0
+    assert np.all(ds1[0] - ds2[0] == 10.0)
+
+    ds3 = 10.0 - ds1
+    assert np.all(ds3[0] == 9.9)
+    assert np.all(ds3[1] == 9.8) 
+
+def test_mul_scalar(ds1):
+    ds2 = ds1 * 2.0
+    assert np.all(ds2[0]*0.5 == ds1[0]) 
+
+    ds3 = 2.0 * ds1
+    assert np.all(ds3[0] == ds2[0])
+    assert np.all(ds3[1] == ds2[1])    
+
+
+def test_add_dataset(ds1, ds2):
+    ds3 = ds1 + ds2
+    assert np.all(ds3[0] == 1.1)
+    assert np.all(ds3[1] == 2.2)
+
+    ds4 = ds2 + ds1
+    assert np.all(ds3[0] == ds4[0])
+    assert np.all(ds3[1] == ds4[1])    
+
+    ds2b = ds2.copy()
+    ds2b.items[0] = ItemInfo(EUMType.Wind_Velocity)
+    with pytest.raises(ValueError):
+        # item type does not match
+        ds1 + ds2b
+
+    ds2c = ds2.copy()
+    tt = ds2c.time.to_numpy()
+    tt[-1] = tt[-1] + np.timedelta64(1, 's')
+    ds2c.time = pd.DatetimeIndex(tt)
+    with pytest.raises(ValueError):
+        # time does not match
+        ds1 + ds2c
+
+def test_sub_dataset(ds1, ds2):
+    ds3 = ds2 - ds1
+    assert np.all(ds3[0] == 0.9)
+    assert np.all(ds3[1] == 1.8)
 
 def test_non_equidistant():
     nt = 4