From c07cbdd1ccf8b239d1d3addc51831db4c2369b27 Mon Sep 17 00:00:00 2001 From: Huite Bootsma Date: Wed, 24 Jul 2024 14:06:12 +0200 Subject: [PATCH] Add _create_data_array methods on Ugrid, add from_data on UgridDataArray. Replaced @abstractproperty, @abstractstaticmethod, add self to these abstract methods. Slight robustness fix on earcut tests. --- docs/api.rst | 5 +++ docs/changelog.rst | 4 ++ pyproject.toml | 1 - tests/test_ugrid1d.py | 22 ++++++++++ tests/test_ugrid2d.py | 29 ++++++++++++- tests/test_ugrid_dataset.py | 5 +++ xugrid/core/accessorbase.py | 48 ++++++++++++---------- xugrid/core/wrap.py | 23 +++++++++++ xugrid/ugrid/ugrid1d.py | 33 +++++++++++++++ xugrid/ugrid/ugrid2d.py | 34 ++++++++++++++++ xugrid/ugrid/ugridbase.py | 81 ++++++++++++++++++++++++++----------- 11 files changed, 238 insertions(+), 47 deletions(-) diff --git a/docs/api.rst b/docs/api.rst index eee8a0d37..a06b58ec7 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -44,6 +44,7 @@ UgridDataArray UgridDataArray UgridDataArray.ugrid UgridDataArray.from_structured + UgridDataArray.from_data UgridDataset ------------ @@ -232,6 +233,8 @@ UGRID1D Topology Ugrid1d.from_geodataframe Ugrid1d.from_shapely Ugrid1d.to_shapely + + Ugrid1d.create_data_array Ugrid1d.plot @@ -338,6 +341,8 @@ UGRID2D Topology Ugrid2d.to_shapely Ugrid2d.earcut_triangulate_polygons + Ugrid2d.create_data_array + Ugrid2d.plot UGRID Roles Accessor diff --git a/docs/changelog.rst b/docs/changelog.rst index 4911726b3..1ed5f31af 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -27,6 +27,10 @@ Added - Included ``edge_node_connectivity`` in :meth:`xugrid.Ugrid2d.from_meshkernel`, so the ordering of edges is consistent with ``meshkernel``. +- Added :meth:`xugrid.Ugrid1d.create_data_array`, + :meth:`xugrid.Ugrid2d.create_data_array`, and + :meth:`xugrid.UgridDataArray.from_data` to more easily instantiate a + UgridDataArray from a grid topology and an array of values. Changed ~~~~~~~ diff --git a/pyproject.toml b/pyproject.toml index 54e22d1f2..b7409abf6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,7 +66,6 @@ profile = "black" exclude_lines = [ "pragma: no cover", "@abc.abstractmethod", - "@abc.abstractproperty", ] [tool.pixi.project] diff --git a/tests/test_ugrid1d.py b/tests/test_ugrid1d.py index 534accf9d..fddfe25a6 100644 --- a/tests/test_ugrid1d.py +++ b/tests/test_ugrid1d.py @@ -455,3 +455,25 @@ def test_equals(): grid_copy.attrs["attr"] = "something_else" # Dataset.identical is called so returns False assert not grid.equals(grid_copy) + + +def test_ugrid1d_create_data_array(): + grid = grid1d() + + uda = grid.create_data_array(np.zeros(grid.n_node), facet="node") + assert isinstance(uda, xugrid.UgridDataArray) + + uda = grid.create_data_array(np.zeros(grid.n_edge), facet="edge") + assert isinstance(uda, xugrid.UgridDataArray) + + # Error on facet + with pytest.raises(ValueError, match="Invalid facet"): + grid.create_data_array([1, 2, 3], facet="face") + + # Error on on dimensions + with pytest.raises(ValueError, match="Can only create DataArrays from 1D arrays"): + grid.create_data_array([[1, 2, 3]], facet="node") + + # Error on size + with pytest.raises(ValueError, match="Conflicting sizes"): + grid.create_data_array([1, 2, 3, 4], facet="node") diff --git a/tests/test_ugrid2d.py b/tests/test_ugrid2d.py index e6f4ce7dd..2efbe6837 100644 --- a/tests/test_ugrid2d.py +++ b/tests/test_ugrid2d.py @@ -1464,11 +1464,36 @@ def test_earcut_triangulate_polygons(): grid = xugrid.Ugrid2d.earcut_triangulate_polygons(polygons=gdf) assert isinstance(grid, xugrid.Ugrid2d) - assert grid.n_face == 7 + assert np.allclose(polygon.area, grid.area.sum()) grid, index = xugrid.Ugrid2d.earcut_triangulate_polygons( polygons=gdf, return_index=True ) assert isinstance(grid, xugrid.Ugrid2d) assert isinstance(index, np.ndarray) - assert np.array_equal(index, np.zeros(7, dtype=int)) + assert (index == 0).all() + + +def test_ugrid2d_create_data_array(): + grid = grid2d() + + uda = grid.create_data_array(np.zeros(grid.n_node), facet="node") + assert isinstance(uda, xugrid.UgridDataArray) + + uda = grid.create_data_array(np.zeros(grid.n_edge), facet="edge") + assert isinstance(uda, xugrid.UgridDataArray) + + uda = grid.create_data_array(np.zeros(grid.n_face), facet="face") + assert isinstance(uda, xugrid.UgridDataArray) + + # Error on facet + with pytest.raises(ValueError, match="Invalid facet"): + grid.create_data_array([1, 2, 3, 4], facet="volume") + + # Error on on dimensions + with pytest.raises(ValueError, match="Can only create DataArrays from 1D arrays"): + grid.create_data_array([[1, 2, 3, 4]], facet="face") + + # Error on size + with pytest.raises(ValueError, match="Conflicting sizes"): + grid.create_data_array([1, 2, 3, 4, 5], facet="face") diff --git a/tests/test_ugrid_dataset.py b/tests/test_ugrid_dataset.py index fabddb9d4..bee0a06d4 100644 --- a/tests/test_ugrid_dataset.py +++ b/tests/test_ugrid_dataset.py @@ -123,6 +123,11 @@ def test_init(self): assert isinstance(self.uda.ugrid.grid, xugrid.Ugrid2d) assert self.uda.grid.face_dimension in self.uda.coords + def test_from_data(self): + grid = self.uda.ugrid.grid + uda = xugrid.UgridDataArray.from_data(np.zeros(grid.n_node), grid, facet="node") + assert isinstance(uda, xugrid.UgridDataArray) + def test_reinit_error(self): # Should not be able to initialize using a UgridDataArray. with pytest.raises(TypeError, match="obj must be xarray.DataArray"): diff --git a/xugrid/core/accessorbase.py b/xugrid/core/accessorbase.py index 160ae1901..bf4e87aba 100644 --- a/xugrid/core/accessorbase.py +++ b/xugrid/core/accessorbase.py @@ -9,63 +9,69 @@ class AbstractUgridAccessor(abc.ABC): @abc.abstractmethod - def to_dataset(): + def to_dataset(self): pass @abc.abstractmethod - def assign_node_coords(): + def assign_node_coords(self): pass @abc.abstractmethod - def set_node_coords(): + def set_node_coords(self): pass - @abc.abstractproperty - def crs(): + @property + @abc.abstractmethod + def crs(self): pass @abc.abstractmethod - def set_crs(): + def set_crs(self): pass @abc.abstractmethod - def to_crs(): + def to_crs(self): pass @abc.abstractmethod - def sel(): + def sel(self): pass @abc.abstractmethod - def sel_points(): + def sel_points(self): pass @abc.abstractmethod - def intersect_line(): + def intersect_line(self): pass @abc.abstractmethod - def intersect_linestring(): + def intersect_linestring(self): pass - @abc.abstractproperty - def bounds(): + @property + @abc.abstractmethod + def bounds(self): pass - @abc.abstractproperty - def total_bounds(): + @property + @abc.abstractmethod + def total_bounds(self): pass - @abc.abstractproperty - def name(): + @property + @abc.abstractmethod + def name(self): pass - @abc.abstractproperty - def names(): + @property + @abc.abstractmethod + def names(self): pass - @abc.abstractproperty - def topology(): + @property + @abc.abstractmethod + def topology(self): pass @staticmethod diff --git a/xugrid/core/wrap.py b/xugrid/core/wrap.py index 5a96e87cc..36610552f 100644 --- a/xugrid/core/wrap.py +++ b/xugrid/core/wrap.py @@ -12,6 +12,7 @@ from typing import List, Sequence, Union import xarray as xr +from numpy.typing import ArrayLike from pandas import RangeIndex import xugrid @@ -285,6 +286,28 @@ def from_structured( ) return UgridDataArray(face_da, grid) + @staticmethod + def from_data(data: ArrayLike, grid: UgridType, facet: str) -> UgridDataArray: + """ + Create a UgridDataArray from a grid and a 1D array of values. + + Parameters + ---------- + data: array like + Values for this array. Must be a ``numpy.ndarray`` or castable to + it. + grid: Ugrid1d, Ugrid2d + facet: str + With which facet to associate the data. Options for Ugrid1d are, + ``"node"`` or ``"edge"``. Options for Ugrid2d are ``"node"``, + ``"edge"``, or ``"face"``. + + Returns + ------- + uda: UgridDataArray + """ + return grid.create_data_array(data=data, facet=facet) + class UgridDataset(DatasetForwardMixin): def __init__( diff --git a/xugrid/ugrid/ugrid1d.py b/xugrid/ugrid/ugrid1d.py index 93cb82d6b..56ca7e225 100644 --- a/xugrid/ugrid/ugrid1d.py +++ b/xugrid/ugrid/ugrid1d.py @@ -4,7 +4,9 @@ import numpy as np import pandas as pd import xarray as xr +from numpy.typing import ArrayLike +import xugrid from xugrid import conversion from xugrid.constants import ( BoolArray, @@ -709,3 +711,34 @@ def reindex_like( ), } return obj.isel(indexers, missing_dims="ignore") + + def create_data_array(self, data: ArrayLike, facet: str) -> "xugrid.UgridDataArray": + """ + Create a UgridDataArray from this grid and a 1D array of values. + + Parameters + ---------- + data: array like + Values for this array. Must be a ``numpy.ndarray`` or castable to + it. + grid: Ugrid1d, Ugrid2d + facet: str + With which facet to associate the data. Options for Ugrid1d are, + ``"node"`` or ``"edge"``. Options for Ugrid2d are ``"node"``, + ``"edge"``, or ``"face"``. + + Returns + ------- + uda: UgridDataArray + """ + match facet: + case "node": + dimension = self.node_dimension + case "edge": + dimension = self.edge_dimension + case _: + raise ValueError( + f"Invalid facet: {facet}. Must be one of: node, edge face." + ) + + return self._create_data_array(data, dimension) diff --git a/xugrid/ugrid/ugrid2d.py b/xugrid/ugrid/ugrid2d.py index 74b111074..3c0262330 100644 --- a/xugrid/ugrid/ugrid2d.py +++ b/xugrid/ugrid/ugrid2d.py @@ -7,6 +7,7 @@ import pandas as pd import xarray as xr from numba_celltree import CellTree2d +from numpy.typing import ArrayLike from scipy.sparse import coo_matrix, csr_matrix from scipy.sparse.csgraph import reverse_cuthill_mckee @@ -2243,3 +2244,36 @@ def _bbox_area(bounds): collection = shapely.polygonize(shapely.linestrings(edges)) polygon = max(collection.geoms, key=lambda x: _bbox_area(x.bounds)) return polygon + + def create_data_array(self, data: ArrayLike, facet: str) -> "xugrid.UgridDataArray": + """ + Create a UgridDataArray from this grid and a 1D array of values. + + Parameters + ---------- + data: array like + Values for this array. Must be a ``numpy.ndarray`` or castable to + it. + grid: Ugrid1d, Ugrid2d + facet: str + With which facet to associate the data. Options for Ugrid1d are, + ``"node"`` or ``"edge"``. Options for Ugrid2d are ``"node"``, + ``"edge"``, or ``"face"``. + + Returns + ------- + uda: UgridDataArray + """ + match facet: + case "node": + dimension = self.node_dimension + case "edge": + dimension = self.edge_dimension + case "face": + dimension = self.face_dimension + case _: + raise ValueError( + f"Invalid facet: {facet}. Must be one of: node, edge face." + ) + + return self._create_data_array(data, dimension) diff --git a/xugrid/ugrid/ugridbase.py b/xugrid/ugrid/ugridbase.py index 423051c19..7ff56728e 100644 --- a/xugrid/ugrid/ugridbase.py +++ b/xugrid/ugrid/ugridbase.py @@ -1,11 +1,12 @@ import abc import copy from itertools import chain -from typing import Tuple, Type, Union +from typing import Tuple, Type, Union, cast import numpy as np import pandas as pd import xarray as xr +from numpy.typing import ArrayLike from scipy.sparse import csr_matrix from xugrid.constants import BoolArray, FloatArray, IntArray @@ -85,74 +86,108 @@ def align(obj, grids, old_indexes): class AbstractUgrid(abc.ABC): - @abc.abstractproperty - def topology_dimension(): + @property + @abc.abstractmethod + def topology_dimension(self): pass - @abc.abstractproperty - def core_dimension(): + @property + @abc.abstractmethod + def core_dimension(self): pass - @abc.abstractproperty - def dimensions(): + @property + @abc.abstractmethod + def dimensions(self): pass - @abc.abstractproperty - def mesh(): + @property + @abc.abstractmethod + def mesh(self): pass - @abc.abstractproperty - def meshkernel(): + @property + @abc.abstractmethod + def meshkernel(self): pass - @abc.abstractstaticmethod - def from_dataset(): + @staticmethod + @abc.abstractmethod + def from_dataset(self): pass @abc.abstractmethod - def to_dataset() -> xr.Dataset: + def to_dataset(self) -> xr.Dataset: pass @abc.abstractmethod - def topology_subset(): + def topology_subset(self): pass @abc.abstractmethod - def clip_box(): + def clip_box(self): pass @abc.abstractmethod - def sel_points(): + def sel_points(self): pass @abc.abstractmethod - def intersect_line(): + def intersect_line(self): pass @abc.abstractmethod - def intersect_linestring(): + def intersect_linestring(self): pass @abc.abstractmethod - def sel(): + def sel(self): pass @abc.abstractmethod - def _clear_geometry_properties(): + def _clear_geometry_properties(self): pass - @abc.abstractstaticmethod + @staticmethod + @abc.abstractmethod def merge_partitions(): pass @abc.abstractmethod - def reindex_like(): + def reindex_like(self): pass @abc.abstractmethod def connectivity_matrix(self, dim: str, xy_weights: bool) -> csr_matrix: pass + @abc.abstractmethod + def create_data_array(self, data: ArrayLike, facet: str): + pass + + def _create_data_array(self, data: ArrayLike, dimension: str): + from xugrid import UgridDataArray + + data = np.array(data) + if data.ndim != 1: + raise ValueError( + "Can only create DataArrays from 1D arrays. " + f"Data has {data.ndim} dimensions." + ) + len_data = len(data) + len_grid = self.dimensions[dimension] + if len_data != len_grid: + raise ValueError( + f"Conflicting sizes for dimension {dimension}: length " + f"{len_data} on the data, but length {len_grid} on the grid." + ) + + da = xr.DataArray(data=data, dims=(dimension,)) + + # TODO: is there a better way to do this to satisfy mypy? + grid = cast(UgridType, self) + return UgridDataArray(da, grid) + def _initialize_indexes_attrs(self, name, dataset, indexes, attrs): defaults = conventions.default_topology_attrs(name, self.topology_dimension)