Repository: DHI/mikeio
Pull request #702 (2024-06-12 14:52:37)

Non-equidistant time axis for Dfsu


Patch:
diff --git a/mikeio/dfs/_dfs.py b/mikeio/dfs/_dfs.py
index cf3b4f9f..28e6dc83 100644
--- a/mikeio/dfs/_dfs.py
+++ b/mikeio/dfs/_dfs.py
@@ -46,8 +46,9 @@ def _read_item_time_step(
     it: int,
     error_bad_data: bool = True,
     fill_bad_data_value: float = np.nan,
-) -> Tuple[DfsFile, np.ndarray]:
+) -> Tuple[DfsFile, np.ndarray, float]:
     itemdata = dfs.ReadItemTimeStep(item_numbers[item] + 1, it)
+    t = itemdata.Time
     if itemdata is not None:
         d = itemdata.Data
         d[d == deletevalue] = np.nan
@@ -60,7 +61,7 @@ def _read_item_time_step(
             d[:] = fill_bad_data_value
             dfs.Close()
             dfs = DfsFileFactory.DfsGenericOpen(filename)
-    return dfs, d
+    return dfs, d, t
 
 
 def _fuzzy_item_search(
diff --git a/mikeio/dfsu/_dfsu.py b/mikeio/dfsu/_dfsu.py
index e8830227..8a0fcf65 100644
--- a/mikeio/dfsu/_dfsu.py
+++ b/mikeio/dfsu/_dfsu.py
@@ -1,11 +1,13 @@
 from __future__ import annotations
 from dataclasses import dataclass
+from datetime import datetime
 from pathlib import Path
 
 from typing import Any, Literal, Sequence, Tuple
 
 import numpy as np
 import pandas as pd
+from mikecore.DfsFile import TimeAxisType
 from mikecore.DfsFactory import DfsFactory
 from mikecore.DfsuBuilder import DfsuBuilder
 from mikecore.DfsuFile import DfsuFile, DfsuFileType
@@ -29,7 +31,7 @@ from ..spatial import (
 from ..spatial import Grid2D
 from .._track import _extract_track
 from ._common import get_elements_from_source, get_nodes_from_source
-from ..eum import ItemInfo
+from ..eum import ItemInfo, TimeStepUnit
 
 
 def write_dfsu(filename: str | Path, data: Dataset) -> None:
@@ -44,10 +46,6 @@ def write_dfsu(filename: str | Path, data: Dataset) -> None:
     """
     filename = str(filename)
 
-    if not data.is_equidistant:
-        raise ValueError("Non-equidistant time axis is not supported.")
-
-    dt = data.timestep
     geometry = data.geometry
     dfsu_filetype = DfsuFileType.Dfsu2D
 
@@ -70,7 +68,16 @@ def write_dfsu(filename: str | Path, data: Dataset) -> None:
     factory = DfsFactory()
     proj = factory.CreateProjection(geometry.projection_string)
     builder.SetProjection(proj)
-    builder.SetTimeInfo(data.time[0], dt)
+
+    if data.is_equidistant:
+        temporal_axis = factory.CreateTemporalEqCalendarAxis(
+            TimeStepUnit.SECOND, data.time[0], 0, data.timestep
+        )
+    else:
+        temporal_axis = factory.CreateTemporalNonEqCalendarAxis(
+            TimeStepUnit.SECOND, data.time[0]
+        )
+    builder.SetTemporalAxis(temporal_axis)
     builder.SetZUnit(eumUnit.eumUmeter)
 
     if dfsu_filetype != DfsuFileType.Dfsu2D:
@@ -91,6 +98,11 @@ def write_dfsu_data(dfs: DfsuFile, ds: Dataset, is_layered: bool) -> None:
     n_time_steps = len(ds.time)
     data = ds
 
+    if data.is_equidistant:
+        t_rel = np.zeros(data.n_timesteps)
+    else:
+        t_rel = (data.time - data.time[0]).total_seconds()
+
     for i in range(n_time_steps):
         if is_layered:
             if "time" in data.dims:
@@ -98,14 +110,14 @@ def write_dfsu_data(dfs: DfsuFile, ds: Dataset, is_layered: bool) -> None:
                 zn = data._zn[i]
             else:
                 zn = data._zn
-            dfs.WriteItemTimeStepNext(0, zn.astype(np.float32))
+            dfs.WriteItemTimeStepNext(t_rel[i], zn.astype(np.float32))
         for da in data:
             if "time" in data.dims:
                 d = da.to_numpy()[i, :]
             else:
                 d = da.to_numpy()
             d[np.isnan(d)] = data.deletevalue
-            dfs.WriteItemTimeStepNext(0, d.astype(np.float32))
+            dfs.WriteItemTimeStepNext(t_rel[i], d.astype(np.float32))
     dfs.Close()
 
 
@@ -124,8 +136,10 @@ def _validate_elements_and_geometry_sel(elements: Any, **kwargs: Any) -> None:
 class _DfsuInfo:
     filename: str
     type: DfsuFileType
-    time: pd.DatetimeIndex
+    start_time: datetime
+    equidistant: bool
     timestep: float
+    n_timesteps: int
     items: list[ItemInfo]
     deletevalue: float
 
@@ -138,21 +152,19 @@ def _get_dfsu_info(filename: str | Path) -> _DfsuInfo:
     dfs = DfsuFile.Open(filename)
     type = DfsuFileType(dfs.DfsuFileType)
     deletevalue = dfs.DeleteValueFloat
-    freq = pd.Timedelta(seconds=dfs.TimeStepInSeconds)
-    time = pd.date_range(
-        start=dfs.StartDateTime,
-        periods=dfs.NumberOfTimeSteps,
-        freq=freq,
-    )
+
     timestep = dfs.TimeStepInSeconds
     items = _get_item_info(dfs.ItemInfo)
+    equidistant = dfs.FileInfo.TimeAxis.TimeAxisType == TimeAxisType.CalendarEquidistant
     dfs.Close()
     return _DfsuInfo(
         filename=filename,
         type=type,
-        time=time,
         timestep=timestep,
+        equidistant=equidistant,
+        n_timesteps=dfs.NumberOfTimeSteps,
         items=items,
+        start_time=dfs.FileInfo.TimeAxis.StartDateTime,
         deletevalue=deletevalue,
     )
 
@@ -272,8 +284,10 @@ class Dfsu2DH:
         self._filename = info.filename
         self._type = info.type
         self._deletevalue = info.deletevalue
-        self._time = info.time
+        self._equidistant = info.equidistant
+        self._start_time = info.start_time
         self._timestep = info.timestep
+        self._n_timesteps = info.n_timesteps
         self._items = info.items
         self._geometry = self._read_geometry(self._filename)
 
@@ -317,14 +331,14 @@ class Dfsu2DH:
         return self._items
 
     @property
-    def start_time(self) -> pd.Timestamp:
+    def start_time(self) -> datetime:
         """File start time"""
-        return self._time[0]
+        return self._start_time
 
     @property
     def n_timesteps(self) -> int:
         """Number of time steps"""
-        return len(self._time)
+        return self._n_timesteps
 
     @property
     def timestep(self) -> float:
@@ -334,11 +348,25 @@ class Dfsu2DH:
     @property
     def end_time(self) -> pd.Timestamp:
         """File end time"""
-        return self._time[-1]
+        if self._equidistant:
+            return self.time[-1]
+        else:
+            # read the last timestep
+            ds = self.read(items=0, time=-1)
+            return ds.time[-1]
 
     @property
     def time(self) -> pd.DatetimeIndex:
-        return self._time
+        if self._equidistant:
+            return pd.date_range(
+                start=self.start_time,
+                periods=self.n_timesteps,
+                freq=f"{self.timestep}S",
+            )
+        else:
+            raise NotImplementedError(
+                "Non-equidistant time axis. Read the data to get time."
+            )
 
     @staticmethod
     def _read_geometry(filename: str) -> GeometryFM2D:
@@ -436,6 +464,8 @@ class Dfsu2DH:
 
         shape: Tuple[int, ...]
 
+        t_rel = np.zeros(len(time_steps))
+
         n_steps = len(time_steps)
         shape = (
             (n_elems,)
@@ -447,12 +477,10 @@ class Dfsu2DH:
             data: np.ndarray = np.ndarray(shape=shape, dtype=dtype)
             data_list.append(data)
 
-        time = self.time
-
         for i in trange(n_steps, disable=not self.show_progress):
             it = time_steps[i]
             for item in range(n_items):
-                dfs, d = _read_item_time_step(
+                dfs, d, t = _read_item_time_step(
                     dfs=dfs,
                     filename=self._filename,
                     time=time,
@@ -464,6 +492,7 @@ class Dfsu2DH:
                     error_bad_data=error_bad_data,
                     fill_bad_data_value=fill_bad_data_value,
                 )
+                t_rel[i] = t
 
                 if elements is not None:
                     d = d[elements]
@@ -473,7 +502,7 @@ class Dfsu2DH:
                 else:
                     data_list[item][i] = d
 
-        time = self.time[time_steps]
+        time = pd.to_datetime(t_rel, unit="s", origin=self.start_time)
 
         dfs.Close()
 
@@ -499,7 +528,6 @@ class Dfsu2DH:
             dt=self.timestep,
         )
 
-
     def append(self, ds: Dataset, validate: bool = True) -> None:
         """
         Append data to an existing dfsu file
@@ -524,7 +552,7 @@ class Dfsu2DH:
         dfs = DfsFileFactory.DfsuFileOpenAppend(str(self._filename), parameters=None)
         write_dfsu_data(dfs=dfs, ds=ds, is_layered=False)
         info = _get_dfsu_info(self._filename)
-        self._time = info.time
+        self._n_timesteps = info.n_timesteps
 
     def _parse_geometry_sel(
         self,
@@ -532,7 +560,6 @@ class Dfsu2DH:
         x: float | None,
         y: float | None,
     ) -> np.ndarray | None:
-
         """Parse geometry selection
 
         Parameters
diff --git a/mikeio/dfsu/_layered.py b/mikeio/dfsu/_layered.py
index 9e0e46a7..604a987f 100644
--- a/mikeio/dfsu/_layered.py
+++ b/mikeio/dfsu/_layered.py
@@ -46,8 +46,10 @@ class DfsuLayered:
         self._filename = info.filename
         self._type = info.type
         self._deletevalue = info.deletevalue
-        self._time = info.time
+        self._equidistant = info.equidistant
+        self._start_time = info.start_time
         self._timestep = info.timestep
+        self._n_timesteps = info.n_timesteps
         self._geometry = self._read_geometry(self._filename)
         # 3d files have a zn item
         self._items = self._read_items(self._filename)
@@ -98,12 +100,12 @@ class DfsuLayered:
     @property
     def start_time(self) -> pd.Timestamp:
         """File start time"""
-        return self._time[0]
+        return self._start_time
 
     @property
     def n_timesteps(self) -> int:
         """Number of time steps"""
-        return len(self._time)
+        return self._n_timesteps
 
     @property
     def timestep(self) -> float:
@@ -113,11 +115,25 @@ class DfsuLayered:
     @property
     def end_time(self) -> pd.Timestamp:
         """File end time"""
-        return self._time[-1]
+        if self._equidistant:
+            return self.time[-1]
+        else:
+            # read the last timestep
+            ds = self.read(items=0, time=-1)
+            return ds.time[-1]
 
     @property
     def time(self) -> pd.DatetimeIndex:
-        return self._time
+        if self._equidistant:
+            return pd.date_range(
+                start=self.start_time,
+                periods=self.n_timesteps,
+                freq=f"{self.timestep}S",
+            )
+        else:
+            raise NotImplementedError(
+                "Non-equidistant time axis. Read the data to get time."
+            )
 
     @property
     def geometry(self) -> GeometryFM3D | GeometryFMVerticalProfile:
@@ -285,7 +301,7 @@ class DfsuLayered:
         deletevalue = self.deletevalue
 
         data_list = []
-
+        t_seconds = np.zeros(len(time_steps))
         n_steps = len(time_steps)
         item0_is_node_based = False
         for item in range(n_items):
@@ -299,12 +315,10 @@ class DfsuLayered:
         if single_time_selected and not keepdims:
             data = data[0]
 
-        time = self.time
-
         for i in trange(n_steps, disable=not self.show_progress):
             it = time_steps[i]
             for item in range(n_items):
-                dfs, d = _read_item_time_step(
+                dfs, d, t = _read_item_time_step(
                     dfs=dfs,
                     filename=self._filename,
                     time=time,
@@ -316,6 +330,7 @@ class DfsuLayered:
                     error_bad_data=error_bad_data,
                     fill_bad_data_value=fill_bad_data_value,
                 )
+                t_seconds[i] = t
 
                 if elements is not None:
                     if item == 0 and item0_is_node_based:
@@ -328,7 +343,7 @@ class DfsuLayered:
                 else:
                     data_list[item][i] = d
 
-        time = self.time[time_steps]
+        time = pd.to_datetime(t_seconds, unit="s", origin=self.start_time)
 
         dfs.Close()
 
@@ -389,7 +404,7 @@ class DfsuLayered:
         dfs = DfsFileFactory.DfsuFileOpenAppend(str(self._filename), parameters=None)
         write_dfsu_data(dfs=dfs, ds=ds, is_layered=ds.geometry.is_layered)
         info = _get_dfsu_info(self._filename)
-        self._time = info.time
+        self._n_timesteps = info.n_timesteps
 
 
 class Dfsu2DV(DfsuLayered):
diff --git a/mikeio/dfsu/_spectral.py b/mikeio/dfsu/_spectral.py
index 128e7db3..aeb5d3c6 100644
--- a/mikeio/dfsu/_spectral.py
+++ b/mikeio/dfsu/_spectral.py
@@ -32,8 +32,10 @@ class DfsuSpectral:
         self._filename = info.filename
         self._type = info.type
         self._deletevalue = info.deletevalue
-        self._time = info.time
+        self._equidistant = info.equidistant
+        self._start_time = info.start_time
         self._timestep = info.timestep
+        self._n_timesteps = info.n_timesteps
         self._items = info.items
         self._geometry = self._read_geometry(self._filename)
 
@@ -90,12 +92,12 @@ class DfsuSpectral:
     @property
     def start_time(self) -> pd.Timestamp:
         """File start time"""
-        return self._time[0]
+        return self._start_time
 
     @property
     def n_timesteps(self) -> int:
         """Number of time steps"""
-        return len(self._time)
+        return self._n_timesteps
 
     @property
     def timestep(self) -> float:
@@ -105,11 +107,25 @@ class DfsuSpectral:
     @property
     def end_time(self) -> pd.Timestamp:
         """File end time"""
-        return self._time[-1]
+        if self._equidistant:
+            return self.time[-1]
+        else:
+            # read the last timestep
+            ds = self.read(items=0, time=-1)
+            return ds.time[-1]
 
     @property
     def time(self) -> pd.DatetimeIndex:
-        return self._time
+        if self._equidistant:
+            return pd.date_range(
+                start=self.start_time,
+                periods=self.n_timesteps,
+                freq=f"{self.timestep}S",
+            )
+        else:
+            raise NotImplementedError(
+                "Non-equidistant time axis. Read the data to get time."
+            )
 
     @staticmethod
     def _read_geometry(
diff --git a/mikeio/pfs/_pfssection.py b/mikeio/pfs/_pfssection.py
index f95be0cc..f3b4cbbb 100644
--- a/mikeio/pfs/_pfssection.py
+++ b/mikeio/pfs/_pfssection.py
@@ -275,7 +275,7 @@ class PfsSection(SimpleNamespace, MutableMapping[str, Any]):
     def _param_match(parampat: Any, v: Any, case: bool) -> Any:
         if parampat is None:
             return False
-        if type(v) != type(parampat):
+        if type(v) is not type(parampat):
             return False
         if isinstance(v, str):
             vv = str(v) if case else str(v).lower()


Test patch:
diff --git a/tests/test_dfsu.py b/tests/test_dfsu.py
index 4d04a81a..807abd6f 100644
--- a/tests/test_dfsu.py
+++ b/tests/test_dfsu.py
@@ -587,13 +587,28 @@ def test_write_from_dfsu_2_time_steps(tmp_path):
     assert dfs.end_time != newdfs.end_time
 
 
-def test_write_non_equidistant_is_not_possible(tmp_path):
+def test_write_non_equidistant_is_possible(tmp_path):
     sourcefilename = "tests/testdata/HD2D.dfsu"
     fp = tmp_path / "simple.dfsu"
     ds = mikeio.read(sourcefilename, time=[0, 1, 3])
+    assert not ds.is_equidistant
 
-    with pytest.raises(ValueError):
-        ds.to_dfs(fp)
+    ds.to_dfs(fp)
+
+    ds2 = mikeio.read(fp)
+
+    assert all(ds.time == ds2.time)
+
+    dfs = mikeio.open(fp)
+
+    dfs = mikeio.open(fp)
+
+    # it is not possible to get all time without reading the entire file
+    with pytest.raises(NotImplementedError):
+        dfs.time
+
+    # but getting the end time is not that expensive
+    assert dfs.end_time == ds.time[-1]
 
 
 def test_temporal_resample_by_reading_selected_timesteps(tmp_path):
@@ -960,18 +975,6 @@ def test_interp_like_fm_dataset():
     assert isinstance(dsi, Dataset)
     assert isinstance(dsi.geometry, GeometryFM2D)
 
-
-def test_writing_non_equdistant_dfsu_is_not_possible(tmp_path):
-    ds = mikeio.read("tests/testdata/wind_north_sea.dfsu")
-    dss = ds.isel(time=[0, 2, 3])
-    assert not dss.is_equidistant
-
-    # TODO which exception should be raised, when trying to write non-equidistant dfsu? This operation is not supported
-    with pytest.raises(ValueError, match="equidistant"):
-        fp = tmp_path / "not_gonna_work.dfsu"
-        dss.to_dfs(fp)
-
-
 def test_append_dfsu_2d(tmp_path):
     ds = mikeio.read("tests/testdata/consistency/oresundHD.dfsu", time=[0, 1])
     ds2 = mikeio.read("tests/testdata/consistency/oresundHD.dfsu", time=[2, 3])
diff --git a/tests/test_dfsu_layered.py b/tests/test_dfsu_layered.py
index d21dc957..fb9f2eed 100644
--- a/tests/test_dfsu_layered.py
+++ b/tests/test_dfsu_layered.py
@@ -1,4 +1,5 @@
 import numpy as np
+import pandas as pd
 import pytest
 import mikeio
 from mikeio import Mesh
@@ -644,3 +645,30 @@ def test_read_elements_3d():
     assert ds2.geometry.element_coordinates[1][0] == pytest.approx(354020.46382194717)
     assert ds2.Salinity.to_numpy()[0, 1] == pytest.approx(23.18906021118164)
 
+
+def test_write_3d_non_equidistant(tmp_path):
+    sourcefilename = "tests/testdata/basin_3d.dfsu"
+    fp = tmp_path / "simple.dfsu"
+    ds = mikeio.read(sourcefilename)
+
+    # manipulate time
+    ds.time = pd.DatetimeIndex(["2000-01-01", "2000-01-02", "2000-01-10"])
+
+    assert not ds.is_equidistant
+
+    ds.to_dfs(fp)
+
+    ds2 = mikeio.read(fp)
+
+    assert all(ds.time == ds2.time)
+    assert not ds2.is_equidistant
+
+    dfs = mikeio.open(fp)
+
+    # it is not possible to get all time without reading the entire file
+    with pytest.raises(NotImplementedError):
+        dfs.time
+
+    # but getting the end time is not that expensive
+    assert dfs.end_time == pd.Timestamp("2000-01-10")
+