Merge pull request #702 from DHI/neq_dfsu

Non-equidistant time axis for Dfsu

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(

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 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,21 +98,26 @@ 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:
                 assert data._zn is not 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 @@ def __init__(self, filename: str | Path) -> None:
         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 @@ def items(self) -> list[ItemInfo]:
         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 @@ def timestep(self) -> float:
     @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 @@ def read(
 
         shape: Tuple[int, ...]
 
+        t_rel = np.zeros(len(time_steps))
+
         n_steps = len(time_steps)
         shape = (
             (n_elems,)
@@ -447,12 +477,10 @@ def read(
             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 @@ def read(
                     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 @@ def read(
                 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 @@ def read(
             dt=self.timestep,
         )
 
-
     def append(self, ds: Dataset, validate: bool = True) -> None:
         """
         Append data to an existing dfsu file
@@ -524,15 +552,14 @@ def append(self, ds: Dataset, validate: bool = True) -> None:
         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,
         area: Tuple[float, float, float, float] | None,
         x: float | None,
         y: float | None,
     ) -> np.ndarray | None:
-
         """Parse geometry selection
 
         Parameters

mikeio/dfsu/_layered.py

@@ -46,8 +46,10 @@ def __init__(self, filename: str | Path) -> None:
         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 @@ def items(self) -> list[ItemInfo]:
     @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 @@ def timestep(self) -> float:
     @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 @@ def read(
         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 @@ def read(
         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 @@ def read(
                     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 @@ def read(
                 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 @@ def append(self, ds: Dataset, validate: bool = True) -> None:
         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):