Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Dfsu read elements preserve order #690

Merged
merged 29 commits into from
May 24, 2024
Merged
Show file tree
Hide file tree
Changes from 9 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions mikeio/dfsu/_dfsu.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Collection, Sequence, Tuple
from typing import Any, Sequence, Tuple

import numpy as np
import pandas as pd
Expand Down Expand Up @@ -361,7 +361,7 @@ def read(
*,
items: str | int | Sequence[str | int] | None = None,
time: int | str | slice | None = None,
elements: Collection[int] | None = None,
elements: Sequence[int] | None = None,
area: Tuple[float, float, float, float] | None = None,
x: float | None = None,
y: float | None = None,
Expand Down
6 changes: 3 additions & 3 deletions mikeio/dfsu/_layered.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from __future__ import annotations
from pathlib import Path
from typing import Any, Collection, Sequence, Tuple, TYPE_CHECKING
from typing import Any, Sequence, Tuple, TYPE_CHECKING

import numpy as np
from mikecore.DfsuFile import DfsuFile, DfsuFileType
Expand Down Expand Up @@ -178,7 +178,7 @@ def read(
*,
items: str | int | Sequence[str | int] | None = None,
time: int | str | slice | None = None,
elements: Collection[int] | None = None,
elements: Sequence[int] | None = None,
area: Tuple[float, float, float, float] | None = None,
x: float | None = None,
y: float | None = None,
Expand Down Expand Up @@ -246,7 +246,7 @@ def read(
elements = self.geometry.find_index( # type: ignore
x=x, y=y, z=z, area=area, layers=layers
)
if len(elements) == 0:
if len(elements) == 0: # type: ignore
raise ValueError("No elements in selection!")

geometry = (
Expand Down
65 changes: 6 additions & 59 deletions mikeio/spatial/_FM_geometry.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
from functools import cached_property
from pathlib import Path
from typing import (
Collection,
List,
Any,
Literal,
Expand Down Expand Up @@ -944,7 +943,7 @@ def _get_boundary_faces(self) -> np.ndarray:
return all_faces[uf_id[bnd_face_id]]

def isel(
self, idx: Collection[int], keepdims: bool = False, **kwargs: Any
self, idx: Sequence[int], keepdims: bool = False, **kwargs: Any
) -> "GeometryFM2D" | GeometryPoint2D:
"""export a selection of elements to a new geometry

Expand All @@ -953,7 +952,7 @@ def isel(

Parameters
----------
idx : collection(int)
idx : list(int)
collection of element indicies
keepdims : bool, optional
Should the original Geometry type be kept (keepdims=True)
Expand All @@ -970,10 +969,7 @@ def isel(
find_index : find element indicies for points or an area
"""

if self._type == DfsuFileType.DfsuSpectral1D:
return self._nodes_to_geometry(nodes=idx)
jsmariegaard marked this conversation as resolved.
Show resolved Hide resolved
else:
return self.elements_to_geometry(elements=idx, keepdims=keepdims)
return self.elements_to_geometry(elements=idx, keepdims=keepdims)

def find_index(
self,
Expand Down Expand Up @@ -1072,53 +1068,8 @@ def _elements_in_area(
else:
raise ValueError("'area' must be bbox [x0,y0,x1,y1] or polygon")

def _nodes_to_geometry(
self, nodes: Collection[int]
) -> "GeometryFM2D" | GeometryPoint2D:
"""export a selection of nodes to new flexible file geometry

Note: takes only the elements for which all nodes are selected

Parameters
----------
nodes : list(int)
list of node ids

Returns
-------
UnstructuredGeometry
which can be used for further extraction or saved to file
"""
nodes = np.atleast_1d(nodes) # type: ignore
if len(nodes) == 1:
xy = self.node_coordinates[nodes[0], :2]
return GeometryPoint2D(xy[0], xy[1])

elements = []
for j, el_nodes in enumerate(self.element_table):
if np.all(np.isin(el_nodes, nodes)):
elements.append(j)

assert len(elements) > 0, "no elements found"
elements = np.sort(elements) # make sure elements are sorted!

node_ids, elem_tbl = self._get_nodes_and_table_for_elements(elements)
node_coords = self.node_coordinates[node_ids]
codes = self.codes[node_ids]

return GeometryFM2D(
node_coordinates=node_coords,
codes=codes,
node_ids=node_ids,
dfsu_type=self._type,
projection=self.projection_string,
element_table=elem_tbl,
element_ids=self.element_ids[elements],
reindex=True,
)

def elements_to_geometry(
self, elements: int | Collection[int], keepdims: bool = False
self, elements: int | Sequence[int], keepdims: bool = False
) -> "GeometryFM2D" | GeometryPoint2D:
if isinstance(elements, (int, np.integer)):
sel_elements: List[int] = [elements]
Expand All @@ -1129,16 +1080,12 @@ def elements_to_geometry(

return GeometryPoint2D(x=x, y=y, projection=self.projection)

sorted_elements = np.sort(
sel_elements
) # make sure elements are sorted! # TODO is this necessary? If so, should be done in the initialiser

# extract information for selected elements

node_ids, elem_tbl = self._get_nodes_and_table_for_elements(sorted_elements)
node_ids, elem_tbl = self._get_nodes_and_table_for_elements(sel_elements)
node_coords = self.node_coordinates[node_ids]
codes = self.codes[node_ids]
elem_ids = self.element_ids[sorted_elements]
elem_ids = self.element_ids[sel_elements]

return GeometryFM2D(
node_coordinates=node_coords,
Expand Down
22 changes: 9 additions & 13 deletions mikeio/spatial/_FM_geometry_layered.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from __future__ import annotations
from functools import cached_property
from typing import Any, Collection, Iterable, Literal, Sequence, List, Tuple
from typing import Any, Iterable, Literal, Sequence, List, Tuple

import numpy as np
from mikecore.DfsuFile import DfsuFileType
Expand Down Expand Up @@ -69,14 +69,14 @@ def geometry2d(self) -> GeometryFM2D:
return self.to_2d_geometry()

def isel(
self, idx: Collection[int], keepdims: bool = False, **kwargs: Any
self, idx: Sequence[int], keepdims: bool = False, **kwargs: Any
) -> GeometryFM3D | GeometryPoint3D | GeometryFM2D | GeometryFMVerticalColumn:

return self.elements_to_geometry(elements=idx, keepdims=keepdims)

def elements_to_geometry(
self,
elements: int | Collection[int],
elements: int | Sequence[int],
node_layers: Layer = "all",
keepdims: bool = False,
) -> GeometryFM3D | GeometryPoint3D | GeometryFM2D | GeometryFMVerticalColumn:
Expand All @@ -91,21 +91,17 @@ def elements_to_geometry(

return GeometryPoint3D(x=x, y=y, z=z, projection=self.projection)

sorted_elements = np.sort(
sel_elements
) # make sure elements are sorted! # TODO is this necessary?

# create new geometry
new_type = self._type

layers_used = self.layer_ids[sorted_elements]
layers_used = self.layer_ids[sel_elements]
unique_layer_ids = np.unique(layers_used)
n_layers = len(unique_layer_ids)

if n_layers > 1:
bottom: Layer = "bottom"
elem_bot = self.get_layer_elements(layers=bottom)
if np.all(np.in1d(sorted_elements, elem_bot)):
if np.all(np.in1d(sel_elements, elem_bot)):
n_layers = 1

if (
Expand All @@ -119,7 +115,7 @@ def elements_to_geometry(

# extract information for selected elements
if n_layers == 1:
elem2d = self.elem2d_ids[sorted_elements]
elem2d = self.elem2d_ids[sel_elements]
geom2d = self.geometry2d
node_ids, elem_tbl = geom2d._get_nodes_and_table_for_elements(elem2d)
assert len(elem_tbl[0]) <= 4, "Not a 2D element"
Expand All @@ -128,11 +124,11 @@ def elements_to_geometry(
elem_ids = self._element_ids[elem2d]
else:
node_ids, elem_tbl = self._get_nodes_and_table_for_elements(
sorted_elements, node_layers=node_layers
sel_elements, node_layers=node_layers
)
node_coords = self.node_coordinates[node_ids]
codes = self.codes[node_ids]
elem_ids = self._element_ids[sorted_elements]
elem_ids = self._element_ids[sel_elements]

if new_type == DfsuFileType.Dfsu2D:
return GeometryFM2D(
Expand Down Expand Up @@ -183,7 +179,7 @@ def element_coordinates(self) -> np.ndarray:

def _get_nodes_and_table_for_elements(
self,
elements: Collection[int] | np.ndarray,
elements: Sequence[int] | np.ndarray,
node_layers: Layer = "all",
) -> Tuple[Any, Any]:
"""list of nodes and element table for a list of elements
Expand Down
8 changes: 3 additions & 5 deletions mikeio/spatial/_FM_geometry_spectral.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from __future__ import annotations
from typing import Collection, Any, Tuple
from typing import Any, Sequence, Tuple


import numpy as np
Expand Down Expand Up @@ -127,12 +127,12 @@ def directions(self) -> np.ndarray | None:
# TODO reconsider inheritance to avoid overriding method signature
class GeometryFMAreaSpectrum(_GeometryFMSpectrum):
def isel( # type: ignore
self, idx: Collection[int], **kwargs: Any
self, idx: Sequence[int], **kwargs: Any
) -> "GeometryFMPointSpectrum" | "GeometryFMAreaSpectrum":
return self.elements_to_geometry(elements=idx)

def elements_to_geometry( # type: ignore
self, elements: Collection[int], keepdims: bool = False
self, elements: Sequence[int], keepdims: bool = False
) -> "GeometryFMPointSpectrum" | "GeometryFMAreaSpectrum":
"""export a selection of elements to new flexible file geometry
Parameters
Expand All @@ -156,7 +156,6 @@ def elements_to_geometry( # type: ignore
y=coords[1],
)

elements = np.sort(elements) # make sure elements are sorted!
node_ids, elem_tbl = self._get_nodes_and_table_for_elements(elements)
node_coords = self.node_coordinates[node_ids]
codes = self.codes[node_ids]
Expand Down Expand Up @@ -209,7 +208,6 @@ def _nodes_to_geometry(self, nodes):
elements.append(j)

assert len(elements) > 0, "no elements found"
elements = np.sort(elements) # make sure elements are sorted!

node_ids, elem_tbl = self._get_nodes_and_table_for_elements(elements)
node_coords = self.node_coordinates[node_ids]
Expand Down
21 changes: 21 additions & 0 deletions tests/test_dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -696,6 +696,27 @@ def test_da_sel_area_dfsu2d():
assert da1.geometry.n_elements == 14


def test_da_isel_order_is_important_dfsu2d():
filename = "tests/testdata/FakeLake.dfsu"
da = mikeio.read(filename, items=0, time=0)[0]

# select elements sorted
da1 = da.isel(element=[0, 1])
assert da1.values[0] == pytest.approx(-3.2252840995788574)
assert da1.geometry.element_coordinates[0, 0] == pytest.approx(-0.61049269425)

# select elements in arbitrary order
da2 = da.isel(element=[1, 0])
assert da2.values[1] == pytest.approx(-3.2252840995788574)
assert da2.geometry.element_coordinates[1, 0] == pytest.approx(-0.61049269425)

# select same elements multiple times, not sure why, but consistent with NumPy, xarray
da3 = da.isel(element=[1, 0, 1])
assert da3.values[1] == pytest.approx(-3.2252840995788574)
assert da3.geometry.element_coordinates[1, 0] == pytest.approx(-0.61049269425)
assert len(da3.geometry.element_coordinates) == 3


def test_da_sel_area_grid2d():
filename = "tests/testdata/gebco_sound.dfs2"
da = mikeio.read(filename, items=0)[0]
Expand Down
10 changes: 10 additions & 0 deletions tests/test_dfsu.py
Original file line number Diff line number Diff line change
Expand Up @@ -249,6 +249,16 @@ def test_read_area_polygon():
assert subdomain.within(domain)


def test_read_elements():
ds = mikeio.read(filename="tests/testdata/wind_north_sea.dfsu", elements=[0, 10])
assert ds.geometry.element_coordinates[0][0] == pytest.approx(1.4931853081272184)
assert ds.Wind_speed.to_numpy()[0, 0] == pytest.approx(9.530759811401367)

ds2 = mikeio.read(filename="tests/testdata/wind_north_sea.dfsu", elements=[10, 0])
assert ds2.geometry.element_coordinates[1][0] == pytest.approx(1.4931853081272184)
assert ds2.Wind_speed.to_numpy()[0, 1] == pytest.approx(9.530759811401367)


def test_find_index_on_island():
filename = "tests/testdata/FakeLake.dfsu"
dfs = mikeio.open(filename)
Expand Down
29 changes: 29 additions & 0 deletions tests/test_dfsu_layered.py
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,25 @@ def test_read_dfsu3d_column():
assert dscol2._zn.shape == (ds.n_timesteps, 5 * 3)


def test_flip_column_upside_down():
filename = "tests/testdata/oresund_sigma_z.dfsu"
dfs = mikeio.open(filename)

(x, y) = (333934.1, 6158101.5)

ds = dfs.read() # all data in file
dscol = ds.sel(x=x, y=y)
assert dscol.geometry.element_coordinates[0, 2] == pytest.approx(-7.0)
assert dscol.isel(time=-1).Temperature.values[0] == pytest.approx(17.460058)

idx = list(reversed(range(dscol.n_elements)))

dscol_ud = dscol.isel(element=idx)

assert dscol_ud.geometry.element_coordinates[-1, 2] == pytest.approx(-7.0)
assert dscol_ud.isel(time=-1).Temperature.values[-1] == pytest.approx(17.460058)


def test_read_dfsu3d_column_save(tmp_path):
filename = "tests/testdata/oresund_sigma_z.dfsu"
dfs = mikeio.open(filename)
Expand Down Expand Up @@ -596,3 +615,13 @@ def test_read_wildcard_items():
ds = dfs.read(items="Sal*")
assert ds.items[0].name == "Salinity"
assert ds.n_items == 1


def test_read_elements_3d():
ds = mikeio.read("tests/testdata/oresund_sigma_z.dfsu", elements=[0, 10])
assert ds.geometry.element_coordinates[0][0] == pytest.approx(354020.46382194717)
assert ds.Salinity.to_numpy()[0, 0] == pytest.approx(23.18906021118164)

ds2 = mikeio.read("tests/testdata/oresund_sigma_z.dfsu", elements=[10, 0])
assert ds2.geometry.element_coordinates[1][0] == pytest.approx(354020.46382194717)
assert ds2.Salinity.to_numpy()[0, 1] == pytest.approx(23.18906021118164)
10 changes: 10 additions & 0 deletions tests/test_dfsu_spectral.py
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,16 @@ def test_read_area_spectrum_elements(dfsu_area):
ds2 = dfs.read(elements=elems)
assert ds2.shape[1] == len(elems)
assert np.all(ds1[0].to_numpy()[:, elems, ...] == ds2[0].to_numpy())
assert ds2.geometry.element_coordinates[0, 0] == pytest.approx(2.651450863095597)
assert ds2["Energy density"].isel(time=-1).isel(frequency=0).isel(
direction=0
).to_numpy()[0] == pytest.approx(1.770e-12)

ds3 = dfs.read(elements=[4, 3])
assert ds3.geometry.element_coordinates[1, 0] == pytest.approx(2.651450863095597)
assert ds3["Energy density"].isel(time=-1).isel(frequency=0).isel(
direction=0
).to_numpy()[1] == pytest.approx(1.770e-12)


def test_read_area_spectrum_xy(dfsu_area):
Expand Down
Loading
Loading