Merge pull request #839 from DHI/derived_variables

Generic derived variables

Author: ecomodeller

docs/examples/Generic.qmd

@@ -169,6 +169,29 @@ ds3 = mikeio.read("gebco_sound_spatial.dfs2")
 ds3.Elevation[0].plot();
 ```
 
+## Derived items
+
+Creating derived items from existing items, e.g. current speed from u and v velocities.
+```{python}
+from mikeio.generic import DerivedItem
+
+fn = "../data/oresundHD_run1.dfsu"
+fn_derived = "oresundHD_speed.dfsu"
+mikeio.generic.transform(
+    fn,
+    fn_derived,
+    [
+        DerivedItem(
+            name="Current speed",
+            type=mikeio.EUMType.Current_Speed,
+            unit=mikeio.EUMUnit.knot,
+            func=lambda x: 1.94 * np.sqrt(x["U velocity"] ** 2 + x["V velocity"] ** 2),
+        )
+    ],
+)
+```
+
+
 ## Time average
 
 ```{python}
@@ -211,6 +234,7 @@ da_q75.plot(title="75th percentile, wind speed", label="m/s")
 
 ```{python}
 import os
+
 os.remove("concat.dfs1")
 os.remove("oresundHD_difference.dfsu")
 os.remove("extracted.dfs1")
@@ -219,5 +243,6 @@ os.remove("gebco_sound_local_datum.dfs2")
 os.remove("gebco_sound_spatial.dfs2")
 os.remove("Avg_NorthSea_HD_and_windspeed.dfsu")
 os.remove(fn_q)
+os.remove("oresundHD_speed.dfsu")
 ```
 

mikeio/dfsu/_dfsu.py

@@ -298,7 +298,7 @@ def read(
         *,
         items: str | int | Sequence[str | int] | None = None,
         time: int | str | slice | Sequence[int] | None = None,
-        elements: Sequence[int] | np.ndarray | None = None,
+        elements: int | Sequence[int] | np.ndarray | None = None,
         area: tuple[float, float, float, float]
         | Sequence[tuple[float, float]]
         | None = None,

mikeio/generic.py

@@ -10,7 +10,7 @@
 from datetime import datetime, timedelta
 from shutil import copyfile
 from collections.abc import Iterable, Sequence
-from typing import Callable, Union
+from typing import Callable, Mapping, Union
 import warnings
 
 
@@ -30,9 +30,10 @@
 from mikecore.eum import eumQuantity
 from tqdm import tqdm, trange
 
+
 from . import __dfs_version__
 from .dfs._dfs import _get_item_info, _valid_item_numbers
-from .eum import ItemInfo
+from .eum import ItemInfo, EUMType, EUMUnit
 import mikeio
 
 
@@ -94,7 +95,7 @@ def _clone(
     outfilename: str | pathlib.Path,
     start_time: datetime | None = None,
     timestep: float | None = None,
-    items: Sequence[int | DfsDynamicItemInfo] | None = None,
+    items: Sequence[int | DfsDynamicItemInfo | ItemInfo] | None = None,
     datatype: int | None = None,
 ) -> DfsFile:
     source = DfsFileFactory.DfsGenericOpen(str(infilename))
@@ -1011,3 +1012,127 @@ def change_datatype(
 
     dfs_out.Close()
     dfs_in.Close()
+
+
+class DerivedItem:
+    """Item derived from other items.
+
+    Parameters
+    ----------
+    item: ItemInfo
+        ItemInfo object for the derived item
+    func: Callable[[Mapping[str, np.ndarray]], np.ndarray] | None
+        Function to compute the derived item from a mapping of item names to data arrays.
+        If None, the item data will be returned directly from the mapping using the item's name.
+        Default is None.
+
+    Example
+    -------
+    ```{python}
+    import numpy as np
+    import mikeio
+    from mikeio.generic import DerivedItem
+
+    item = DerivedItem(
+            item=ItemInfo("Current Speed", mikeio.EUMType.Current_Speed),
+            func=lambda x: np.sqrt(x["U velocity"] ** 2 + x["V velocity"] ** 2),
+        )
+
+    """
+
+    def __init__(
+        self,
+        name: str,
+        type: EUMType | None = None,
+        unit: EUMUnit | None = None,
+        func: Callable[[Mapping[str, np.ndarray]], np.ndarray] | None = None,
+    ) -> None:
+        """Create a DerivedItem.
+
+        Parameters
+        ----------
+        name: str
+            Name of the derived item.
+        type: EUMType
+            EUMType of the derived item.
+        unit: EUMUnit | None, optional
+            EUMUnit of the derived item, pass None to use the default unit for the type.
+            Default is None.
+        func: Callable[[Mapping[str, np.ndarray]], np.ndarray] | None, optional
+            Function to compute the derived item from a mapping of item names to data arrays.
+
+        """
+        self.item = ItemInfo(name, type, unit)
+        self.func = func
+
+
+def transform(
+    infilename: str | pathlib.Path,
+    outfilename: str | pathlib.Path,
+    vars: Sequence[DerivedItem],
+    keep_existing_items: bool = True,
+) -> None:
+    """Transform a dfs file by applying functions to items.
+
+    Parameters
+    ----------
+    infilename: str | pathlib.Path
+        full path to the input file
+    outfilename: str | pathlib.Path
+        full path to the output file
+    vars: Sequence[DerivedItem]
+        List of derived items to compute.
+    keep_existing_items: bool, optional
+        If True, existing items in the input file will be kept in the output file.
+        If False, only the derived items will be written to the output file.
+        Default is True.
+
+    """
+    dfs_i = DfsFileFactory.DfsGenericOpen(str(infilename))
+
+    item_numbers = _valid_item_numbers(dfs_i.ItemInfo)
+    n_items = len(item_numbers)
+
+    items = [v.item for v in vars]
+    funcs = {v.item.name: v.func for v in vars}
+
+    if keep_existing_items:
+        existing_items = [
+            ItemInfo.from_mikecore_dynamic_item_info(
+                dfs_i.ItemInfo[i],
+            )
+            for i in item_numbers
+        ]
+        items = existing_items + items
+
+    dfs = _clone(
+        str(infilename),
+        str(outfilename),
+        items=items,
+    )
+
+    n_time_steps = dfs_i.FileInfo.TimeAxis.NumberOfTimeSteps
+
+    for timestep in range(n_time_steps):
+        data = {}
+        for item in range(n_items):
+            name = dfs_i.ItemInfo[item].Name
+            data[name] = dfs_i.ReadItemTimeStep(item_numbers[item] + 1, timestep).Data
+
+        for item in items:
+            func = funcs.get(item.name, None)
+            if func is None:
+                darray = data[item.name]
+            else:
+                try:
+                    darray = func(data)
+                except KeyError as e:
+                    missing_key = e.args[0]
+                    keys = ", ".join(data.keys())
+                    raise KeyError(
+                        f"Item '{missing_key}' is not available in the file. Available items: {keys}"
+                    )
+            dfs.WriteItemTimeStepNext(0.0, darray)
+
+    dfs_i.Close()
+    dfs.Close()

tests/test_generic.py

@@ -12,6 +12,8 @@
     fill_corrupt,
     add,
     change_datatype,
+    transform,
+    DerivedItem,
 )
 import pytest
 from mikecore.DfsFileFactory import DfsFileFactory
@@ -672,3 +674,70 @@ def test_change_datatype_dfs0(tmp_path: Path) -> None:
     org = mikeio.read(infilename).to_numpy()
     new = mikeio.read(outfilename).to_numpy()
     assert np.allclose(org, new, rtol=1e-08, atol=1e-10, equal_nan=True)
+
+
+def test_transform_variables(tmp_path: Path) -> None:
+    infilename = "tests/testdata/oresundHD_run1.dfsu"
+    outfilename = tmp_path / "need_for_speed.dfsu"
+
+    items = [
+        DerivedItem(
+            name="Current Speed",
+            type=mikeio.EUMType.Current_Speed,
+            func=lambda x: np.sqrt(x["U velocity"] ** 2 + x["V velocity"] ** 2),
+        )
+    ]
+
+    transform(infilename, outfilename, vars=items, keep_existing_items=False)
+    dfs = mikeio.Dfsu2DH(outfilename)
+    assert dfs.items[0].type == mikeio.EUMType.Current_Speed
+    assert len(dfs.items) == 1
+
+    dfs1 = mikeio.Dfsu2DH(infilename)
+    sel_items = [
+        DerivedItem(name=item.name, type=item.type, unit=item.unit)
+        for item in dfs1.items
+        if item.name != "Surface elevation"
+    ]
+    sel_items.extend(items)
+
+    outfilename2 = tmp_path / "existing_and_speed.dfsu"
+
+    transform(infilename, outfilename2, vars=sel_items, keep_existing_items=False)
+    dfs2 = mikeio.Dfsu2DH(outfilename2)
+    assert dfs2.items[0].name == "Total water depth"  # existing item
+    assert dfs2.items[1].name == "U velocity"  # existing item
+    assert dfs2.items[2].name == "V velocity"  # existing item
+    assert dfs2.items[3].name == "Current Speed"  # derived item
+
+
+def test_transform_can_include_existing_items(tmp_path: Path) -> None:
+    infilename = "tests/testdata/oresundHD_run1.dfsu"
+    outfilename = tmp_path / "need_for_speed.dfsu"
+
+    items = [
+        DerivedItem(
+            name="Current Speed",
+            type=mikeio.EUMType.Current_Speed,
+            func=lambda x: np.sqrt(x["U velocity"] ** 2 + x["V velocity"] ** 2),
+        )
+    ]
+
+    transform(infilename, outfilename, vars=items, keep_existing_items=True)
+    dfs = mikeio.Dfsu2DH(outfilename)
+    assert dfs.items[-1].type == mikeio.EUMType.Current_Speed
+    assert len(dfs.items) == 5
+    ds = dfs.read(time=-1, elements=0)
+    assert ds["U velocity"].values == pytest.approx(0.0292403083)
+    assert ds["V velocity"].values == pytest.approx(0.127983957)
+    assert ds["Current Speed"].values == pytest.approx(0.13128172)
+
+
+def test_transform_func_with_missing_item_reports_existing_items() -> None:
+    infilename = "tests/testdata/oresundHD_run1.dfsu"
+    outfilename = "notgonnahappen.dfsu"
+
+    items = [DerivedItem(name="Not relevant", func=lambda x: x["not in the file"])]
+    with pytest.raises(KeyError) as excinfo:
+        transform(infilename, outfilename, items)
+    assert "U velocity" in str(excinfo.value)