Merge pull request #666 from DHI/default-dims

Delegate default dims to geometry

mikeio/dataset/_dataarray.py

@@ -11,7 +11,15 @@
     Mapping,
     MutableMapping,
 )
-from typing import Any, Union, Literal, TYPE_CHECKING, overload, Callable, Tuple
+from typing import (
+    Any,
+    Union,
+    Literal,
+    TYPE_CHECKING,
+    overload,
+    Callable,
+    Tuple,
+)
 
 
 import numpy as np
@@ -151,13 +159,15 @@ def __init__(
         *,
         time: pd.DatetimeIndex | str | None = None,
         item: ItemInfo | None = None,
-        geometry: GeometryType = GeometryUndefined(),
+        geometry: GeometryType | None = None,
         zn: np.ndarray | None = None,
         dims: Sequence[str] | None = None,
     ) -> None:
         # TODO: add optional validation validate=True
         self._values = self._parse_data(data)
         self.time: pd.DatetimeIndex = self._parse_time(time)
+
+        geometry = GeometryUndefined() if geometry is None else geometry
         self.dims = self._parse_dims(dims, geometry)
 
         self._check_time_data_length(self.time)
@@ -197,48 +207,22 @@ def _parse_dims(
     def _guess_dims(
         ndim: int, shape: Tuple[int, ...], n_timesteps: int, geometry: GeometryType
     ) -> Tuple[str, ...]:
-        # TODO delete default dims to geometry
-
         # This is not very robust, but is probably a reasonable guess
         time_is_first = (n_timesteps > 1) or (shape[0] == 1 and n_timesteps == 1)
         dims = ["time"] if time_is_first else []
         ndim_no_time = ndim if (len(dims) == 0) else ndim - 1
 
-        if isinstance(geometry, GeometryFMPointSpectrum):
-            if ndim_no_time == 1:
-                dims.append("frequency")
-            if ndim_no_time == 2:
-                dims.append("direction")
-                dims.append("frequency")
-        elif isinstance(geometry, GeometryFM3D):
-            if ndim_no_time > 0:
-                dims.append("element")
-        elif isinstance(geometry, GeometryFM2D):
-            if geometry._type == DfsuFileType.DfsuSpectral1D:
-                if ndim_no_time > 0:
-                    dims.append("node")
-            else:
-                if ndim_no_time > 0:
-                    dims.append("element")
-            if geometry.is_spectral:
-                if ndim_no_time == 2:
-                    dims.append("frequency")
-                elif ndim_no_time == 3:
-                    dims.append("direction")
-                    dims.append("frequency")
-        elif isinstance(geometry, Grid1D):
-            dims.append("x")
-        elif isinstance(geometry, Grid2D):
-            dims.append("y")
-            dims.append("x")
+        if isinstance(geometry, GeometryUndefined):
+            DIMS_MAPPING = {
+                0: [],
+                1: ["x"],
+                2: ["y", "x"],
+                3: ["z", "y", "x"],
+            }
+            spdims = DIMS_MAPPING[ndim_no_time]
         else:
-            # gridded
-            if ndim_no_time > 2:
-                dims.append("z")
-            if ndim_no_time > 1:
-                dims.append("y")
-            if ndim_no_time > 0:
-                dims.append("x")
+            spdims = geometry.default_dims
+        dims.extend(spdims)  # type: ignore
         return tuple(dims)
 
     def _check_time_data_length(self, time: Sized) -> None:

mikeio/spatial/_FM_geometry.py

@@ -60,6 +60,13 @@ def __init__(
         self.x = x
         self.y = y
 
+    @property
+    def default_dims(self) -> Tuple[str, ...]:
+        if self.directions is None:
+            return ("frequency",)
+        else:
+            return ("direction", "frequency")
+
     @property
     def is_layered(self) -> bool:
         return False
@@ -357,6 +364,10 @@ def _reindex(self) -> None:
         self._node_ids = new_node_ids
         self._element_ids = new_element_ids
 
+    @property
+    def default_dims(self) -> Tuple[str, ...]:
+        return ("element",)
+
     @property
     def is_spectral(self) -> bool:
         return False

mikeio/spatial/_geometry.py

@@ -1,6 +1,7 @@
 from abc import ABC, abstractmethod
 
 from collections import namedtuple
+from typing import Tuple
 
 from mikecore.Projections import MapProjection
 
@@ -40,18 +41,35 @@ def is_local_coordinates(self) -> bool:
     def ndim(self) -> int:
         pass
 
+    @property
+    @abstractmethod
+    def default_dims(self) -> Tuple[str, ...]:
+        pass
+
 
-class GeometryUndefined:
+class GeometryUndefined(_Geometry):
     def __repr__(self):
         return "GeometryUndefined()"
 
+    @property
+    def ndim(self) -> int:
+        raise NotImplementedError()
+
+    @property
+    def default_dims(self) -> Tuple[str, ...]:
+        raise NotImplementedError()
+
 
 class GeometryPoint2D(_Geometry):
-    def __init__(self, x: float, y: float, projection:str = "LONG/LAT"):
+    def __init__(self, x: float, y: float, projection: str = "LONG/LAT"):
         super().__init__(projection)
         self.x = x
         self.y = y
 
+    @property
+    def default_dims(self) -> Tuple[str, ...]:
+        return ()
+
     def __repr__(self) -> str:
         return f"GeometryPoint2D(x={self.x}, y={self.y})"
 
@@ -75,7 +93,7 @@ def to_shapely(self):
 
 
 class GeometryPoint3D(_Geometry):
-    def __init__(self, x: float, y: float, z: float, projection:str = "LONG/LAT"):
+    def __init__(self, x: float, y: float, z: float, projection: str = "LONG/LAT"):
         super().__init__(projection)
 
         self.x = x
@@ -85,6 +103,10 @@ def __init__(self, x: float, y: float, z: float, projection:str = "LONG/LAT"):
     def __repr__(self):
         return f"GeometryPoint3D(x={self.x}, y={self.y}, z={self.z})"
 
+    @property
+    def default_dims(self) -> Tuple[str, ...]:
+        return ()
+
     @property
     def wkt(self) -> str:
         return f"POINT Z ({self.x} {self.y} {self.z})"

mikeio/spatial/_grid_geometry.py

@@ -142,6 +142,10 @@ def __init__(
     def ndim(self) -> int:
         return 1
 
+    @property
+    def default_dims(self) -> Tuple[str, ...]:
+        return ("x",)
+
     def __repr__(self) -> str:
         out = ["<mikeio.Grid1D>", _print_axis_txt("x", self.x, self.dx)]
         return "\n".join(out)
@@ -491,6 +495,10 @@ def __init__(
 
         self.plot = _Grid2DPlotter(self)
 
+    @property
+    def default_dims(self) -> Tuple[str, ...]:
+        return ("y", "x")
+
     @property
     def ndim(self) -> int:
         return 2
@@ -1119,6 +1127,10 @@ def __init__(
         self._origin = origin
         self._orientation = orientation
 
+    @property
+    def default_dims(self) -> Tuple[str, ...]:
+        return ("z", "y", "x")
+
     @property
     def ndim(self) -> int:
         return 3