Merge pull request #792 from DHI/math

Support multiplication and division of datasets.

Author: ecomodeller

docs/user-guide/dataarray.qmd

@@ -170,6 +170,12 @@ da_ratio = da1 / da2
 da_ratio.plot(title="", label="Ratio", vmin=0.8, vmax=1.2, levels=9, cmap="coolwarm")
 ```
 
+## Unit handling
+
+Multiplication and divison of two physical quantities would normally change the unit of the result, but in the case of DataArrays, the type and unit of the result will be the ones of the first operand.
+
+
+
 Other methods that also return a DataArray:
 
 * [`interp_like`](`mikeio.DataArray.interp_like`) - Spatio (temporal) interpolation (see example [Dfsu interpolation](../examples/dfsu/spatial_interpolation.qmd)

docs/user-guide/dataset.qmd

@@ -166,11 +166,13 @@ including different ways of *selecting* data:
 * ds + value
 * ds - value
 * ds * value
-
-and + and - between two Datasets (if number of items and shapes conform):
+* ds / value
+and between two Datasets (if number of items and shapes conform):
 
 * ds1 + ds2
 * ds1 - ds2
+* ds1 * ds2
+* ds1 / ds2
 
 Other methods that also return a Dataset:
 

mikeio/dataset/_dataarray.py

@@ -1781,33 +1781,31 @@ def __radd__(self, other: "DataArray" | float) -> "DataArray":
         return self.__add__(other)
 
     def __add__(self, other: "DataArray" | float) -> "DataArray":
-        return self._apply_math_operation(other, np.add, txt="+")
+        return self._apply_math_operation(other, np.add)
 
     def __rsub__(self, other: "DataArray" | float) -> "DataArray":
         return other + self.__neg__()
 
     def __sub__(self, other: "DataArray" | float) -> "DataArray":
-        return self._apply_math_operation(other, np.subtract, txt="-")
+        return self._apply_math_operation(other, np.subtract)
 
     def __rmul__(self, other: "DataArray" | float) -> "DataArray":
         return self.__mul__(other)
 
     def __mul__(self, other: "DataArray" | float) -> "DataArray":
-        return self._apply_math_operation(
-            other, np.multiply, txt="x"
-        )  # x in place of *
+        return self._apply_math_operation(other, np.multiply)
 
     def __pow__(self, other: float) -> "DataArray":
-        return self._apply_math_operation(other, np.power, txt="**")
+        return self._apply_math_operation(other, np.power)
 
     def __truediv__(self, other: "DataArray" | float) -> "DataArray":
-        return self._apply_math_operation(other, np.divide, txt="/")
+        return self._apply_math_operation(other, np.divide)
 
     def __floordiv__(self, other: "DataArray" | float) -> "DataArray":
-        return self._apply_math_operation(other, np.floor_divide, txt="//")
+        return self._apply_math_operation(other, np.floor_divide)
 
     def __mod__(self, other: float) -> "DataArray":
-        return self._apply_math_operation(other, np.mod, txt="%")
+        return self._apply_math_operation(other, np.mod)
 
     def __neg__(self) -> "DataArray":
         return self._apply_unary_math_operation(np.negative)
@@ -1830,7 +1828,9 @@ def _apply_unary_math_operation(self, func: Callable) -> "DataArray":
         return new_da
 
     def _apply_math_operation(
-        self, other: "DataArray" | float, func: Callable, *, txt: str
+        self,
+        other: "DataArray" | float,
+        func: Callable,
     ) -> "DataArray":
         """Apply a binary math operation with a scalar, an array or another DataArray."""
         try:
@@ -1839,39 +1839,11 @@ def _apply_math_operation(
         except TypeError:
             raise TypeError("Math operation could not be applied to DataArray")
 
-        # TODO: check if geometry etc match if other is DataArray?
-
         new_da = self.copy()  # TODO: alternatively: create new dataset (will validate)
         new_da.values = data
 
-        if not self._keep_EUM_after_math_operation(other, func):
-            other_name = other.name if hasattr(other, "name") else "array"
-            new_da.item = ItemInfo(
-                f"{self.name} {txt} {other_name}", itemtype=EUMType.Undefined
-            )
-
         return new_da
 
-    def _keep_EUM_after_math_operation(
-        self, other: "DataArray" | float, func: Callable
-    ) -> bool:
-        """Does the math operation falsify the EUM?"""
-        if hasattr(other, "shape") and hasattr(other, "ndim"):
-            # other is array-like, so maybe we cannot keep EUM
-            if func == np.subtract or func == np.sum:
-                # +/-: we may want to keep EUM
-                if isinstance(other, DataArray):
-                    if self.type == other.type and self.unit == other.unit:
-                        return True
-                    else:
-                        return False
-                else:
-                    return True  # assume okay, since no EUM
-            return False
-
-        # other is likely scalar, okay to keep EUM
-        return True
-
     # ============= Logical indexing ===========
 
     def __lt__(self, other) -> "DataArray":  # type: ignore

mikeio/dataset/_dataset.py

@@ -1300,7 +1300,7 @@ def _concat_time(
         copy: bool = True,
         keep: Literal["last", "first"] = "last",
     ) -> "Dataset":
-        self._check_all_items_match(other)
+        self._check_n_items(other)
         # assuming time is always first dimension we can skip / keep it by bool
         start_dim = int("time" in self.dims)
         if not np.all(
@@ -1348,25 +1348,19 @@ def _concat_time(
             newdata, time=newtime, items=ds.items, geometry=ds.geometry, zn=zn
         )
 
-    def _check_all_items_match(self, other: "Dataset") -> None:
+    def _check_n_items(self, other: "Dataset") -> None:
         if self.n_items != other.n_items:
             raise ValueError(
                 f"Number of items must match ({self.n_items} and {other.n_items})"
             )
 
-        for this, that in zip(self.items, other.items):
-            if this.name != that.name:
-                raise ValueError(
-                    f"Item names must match. Item: {this.name} != {that.name}"
-                )
-            if this.type != that.type:
-                raise ValueError(
-                    f"Item types must match. Item: {this.type} != {that.type}"
-                )
-            if this.unit != that.unit:
-                raise ValueError(
-                    f"Item units must match. Item: {this.unit} != {that.unit}"
-                )
+    def _check_datasets_match(self, other: "Dataset") -> None:
+        self._check_n_items(other)
+
+        if not np.all(self.time == other.time):
+            raise ValueError("All timesteps must match")
+        if self.shape != other.shape:
+            raise ValueError("shape must match")
 
     # ============ aggregate =============
 
@@ -1765,83 +1759,63 @@ def __radd__(self, other: "Dataset" | float) -> "Dataset":
 
     def __add__(self, other: "Dataset" | float) -> "Dataset":
         if isinstance(other, self.__class__):
-            return self._add_dataset(other)
+            return self._binary_op(other, operator="+")
         else:
-            # float-like
-            return self._add_value(other)  # type: ignore
+            return self._scalar_op(other, operator="+")  # type: ignore
 
     def __rsub__(self, other: "Dataset" | float) -> "Dataset":
-        ds = self.__mul__(-1.0)
-        return other + ds
+        ds = self._scalar_op(-1.0, operator="*")
+        return ds._scalar_op(other, operator="+")  # type: ignore
 
     def __sub__(self, other: "Dataset" | float) -> "Dataset":
         if isinstance(other, self.__class__):
-            return self._add_dataset(other, sign=-1.0)
+            return self._binary_op(other, operator="-")
         else:
-            return self._add_value(-other)  # type: ignore
+            return self._scalar_op(-other, operator="+")  # type: ignore
 
     def __rmul__(self, other: "Dataset" | float) -> "Dataset":
         return self.__mul__(other)
 
     def __mul__(self, other: "Dataset" | float) -> "Dataset":
         if isinstance(other, self.__class__):
-            raise ValueError("Multiplication is not possible for two Datasets")
+            return self._binary_op(other, operator="*")
         else:
-            return self._multiply_value(other)  # type: ignore
-
-    def _add_dataset(self, other: "Dataset", sign: float = 1.0) -> "Dataset":
-        self._check_datasets_match(other)
-        try:
-            data = [
-                self[x].to_numpy() + sign * other[y].to_numpy()
-                for x, y in zip(self.items, other.items)
-            ]
-        except TypeError:
-            raise TypeError("Could not add data in Dataset")
-        newds = self.copy()
-        for new, old in zip(newds, data):
-            new.values = old
-        return newds
-
-    def _check_datasets_match(self, other: "Dataset") -> None:
-        self._check_all_items_match(other)
-
-        if not np.all(self.time == other.time):
-            raise ValueError("All timesteps must match")
-        if self.shape != other.shape:
-            raise ValueError("shape must match")
+            return self._scalar_op(other, operator="*")  # type: ignore
 
-    def _add_value(self, value: float) -> "Dataset":
-        try:
-            data = [value + self[x].to_numpy() for x in self.items]
-        except TypeError:
-            raise TypeError(f"{value} could not be added to Dataset")
-        items = deepcopy(self.items)
-        time = self.time.copy()
-        return Dataset(
-            data,
-            time=time,
-            items=items,
-            geometry=self.geometry,
-            zn=self._zn,
-            validate=False,
-        )
+    def __truediv__(self, other: "Dataset" | float) -> "Dataset":
+        if isinstance(other, self.__class__):
+            return self._binary_op(other, operator="/")
+        else:
+            return self._scalar_op(other, operator="/")  # type: ignore
 
-    def _multiply_value(self, value: float) -> "Dataset":
-        try:
-            data = [value * self[x].to_numpy() for x in self.items]
-        except TypeError:
-            raise TypeError(f"{value} could not be multiplied to Dataset")
-        items = deepcopy(self.items)
-        time = self.time.copy()
-        return Dataset(
-            data,
-            time=time,
-            items=items,
-            geometry=self.geometry,
-            zn=self._zn,
-            validate=False,
-        )
+    def _binary_op(self, other: "Dataset", operator: str) -> "Dataset":
+        self._check_datasets_match(other)
+        match operator:
+            case "+":
+                data = [x + y for x, y in zip(self, other)]
+            case "-":
+                data = [x - y for x, y in zip(self, other)]
+            case "*":
+                data = [x * y for x, y in zip(self, other)]
+            case "/":
+                data = [x / y for x, y in zip(self, other)]
+            case _:
+                raise ValueError(f"Unsupported operator: {operator}")
+        return Dataset(data)
+
+    def _scalar_op(self, value: float, operator: str) -> "Dataset":
+        match operator:
+            case "+":
+                data = [x + value for x in self]
+            case "-":
+                data = [x - value for x in self]
+            case "*":
+                data = [x * value for x in self]
+            case "/":
+                data = [x / value for x in self]
+            case _:
+                raise ValueError(f"Unsupported operator: {operator}")
+        return Dataset(data)
 
     # ===============================================
 

tests/test_dataset.py

@@ -1182,10 +1182,14 @@ def test_add_dataset(ds1, ds2):
 
     ds2b = ds2.copy()
     ds2b[0].item = ItemInfo(EUMType.Wind_Velocity)
-    with pytest.raises(ValueError):
-        # item type does not match
-        ds1 + ds2b
-
+    # item type does not match, but we don't care about the item type, item is defined by the first dataset
+    ds3 = ds2b + ds1
+    assert ds3.items[0].type == EUMType.Wind_Velocity
+    assert ds3.items[0].name == ds2b.items[0].name
+
+    ds4 = ds1 + ds2b
+    assert ds4.items[0].type == EUMType.Undefined
+    assert ds4.items[0].name == ds1.items[0].name
     ds2c = ds2.copy()
     tt = ds2c.time.to_numpy()
     tt[-1] = tt[-1] + np.timedelta64(1, "s")
@@ -1201,6 +1205,45 @@ def test_sub_dataset(ds1, ds2):
     assert np.all(ds3[1].to_numpy() == 1.8)
 
 
+def test_multiply_dataset(ds1, ds2):
+    dsa = mikeio.Dataset(
+        {
+            "Foo": mikeio.DataArray(
+                [1, 2, 3], item=mikeio.ItemInfo("Foo", EUMType.Water_Level)
+            )
+        }
+    )
+    dsb = mikeio.Dataset({"Foo": mikeio.DataArray([4, 5, 6])})
+    dsr = dsa * dsb
+    assert np.all(dsr.Foo.to_numpy() == np.array([4, 10, 18]))
+    assert dsr.Foo.type == EUMType.Water_Level
+
+
+def test_multiply_number_of_items_datasets_must_match():
+    dsa = mikeio.Dataset(
+        {"Foo": mikeio.DataArray([1, 2, 3]), "Bar": mikeio.DataArray([1, 2, 3])}
+    )
+    dsb = mikeio.Dataset({"Bar": mikeio.DataArray([4, 5, 6])})
+    with pytest.raises(ValueError, match="Number of items"):
+        dsa * dsb
+
+
+def test_divide_dataset(ds1, ds2):
+    ds_nom = mikeio.Dataset({"Foo": mikeio.DataArray([1, 2, 3])})
+    ds_denom = mikeio.Dataset({"Foo": mikeio.DataArray([4, 5, 6])})
+    ds3 = ds_nom / ds_denom
+    assert np.all(ds3[0].to_numpy() == np.array([0.25, 0.4, 0.5]))
+
+
+def test_divide_number_of_items_datasets_must_match():
+    dsa = mikeio.Dataset(
+        {"Foo": mikeio.DataArray([1, 2, 3]), "Bar": mikeio.DataArray([1, 2, 3])}
+    )
+    dsb = mikeio.Dataset({"Bar": mikeio.DataArray([4, 5, 6])})
+    with pytest.raises(ValueError, match="Number of items"):
+        dsa / dsb
+
+
 def test_non_equidistant():
     nt = 4
     d = np.random.uniform(size=nt)