From 62e83e1ad5b5e043090682cbb6623ecd9a06f4d7 Mon Sep 17 00:00:00 2001 From: Huite Bootsma Date: Fri, 9 Jun 2023 16:13:32 +0200 Subject: [PATCH 1/9] WIP --- xugrid/regrid/regridder.py | 7 ++++--- xugrid/regrid/structured.py | 12 +++++++++++- xugrid/regrid/unstructured.py | 6 ++++++ 3 files changed, 21 insertions(+), 4 deletions(-) diff --git a/xugrid/regrid/regridder.py b/xugrid/regrid/regridder.py index cf134e19a..92ba86713 100644 --- a/xugrid/regrid/regridder.py +++ b/xugrid/regrid/regridder.py @@ -222,11 +222,12 @@ def to_dataset(self) -> xr.Dataset: """ Store the computed weights and target in a dataset for re-use. """ - ds = xr.Dataset( + weights_ds = xr.Dataset( {f"__regrid_{k}": v for k, v in zip(self._weights._fields, self._weights)} ) - ugrid_ds = self._target.ugrid_topology.to_dataset() - return xr.merge((ds, ugrid_ds)) + source_ds = self._source.to_dataset("source") + target_ds = self._target.to_dataset("target") + return xr.merge((weights_ds, source_ds, target_ds)) @staticmethod def _csr_from_dataset(dataset: xr.Dataset) -> WeightMatrixCSR: diff --git a/xugrid/regrid/structured.py b/xugrid/regrid/structured.py index 0cac81606..639cc8457 100644 --- a/xugrid/regrid/structured.py +++ b/xugrid/regrid/structured.py @@ -351,8 +351,11 @@ def linear_weights(self, other: "StructuredGrid1d"): neighbour, ) return self.sorted_output(source_index, target_index, weights) + - + def to_dataset(self): + pass + class StructuredGrid2d(StructuredGrid1d): """ e.g. (x,y) -> (x,y) @@ -489,6 +492,13 @@ def linear_weights(self, other): weights_x, ) + def to_dataset(self, name: str): + ds_x = self.xbounds.to_dataset() + ds_y = self.ybounds.to_dataset() + ds = xr.merge([ds_x, ds_y]) + ds[name] = xr.DataArray(-1, attrs={"type": "StructuredGrid2d"}) + return ds + class StructuredGrid3d: """ diff --git a/xugrid/regrid/unstructured.py b/xugrid/regrid/unstructured.py index 2cc0db89a..052aba1ce 100644 --- a/xugrid/regrid/unstructured.py +++ b/xugrid/regrid/unstructured.py @@ -1,4 +1,5 @@ import numpy as np +import xarray as xr import xugrid as xu from xugrid.constants import FloatDType @@ -138,3 +139,8 @@ def barycentric(self, other): order = np.argsort(target_index) return source_index[order], target_index[order], weights[order] + + def to_dataset(self, name: str): + ds = self.ugrid_topology.to_dataset() + ds[name] = xr.DataArray(-1, attrs={"type": "UnstructuredGrid2d"}) + return ds From 2c25757b87e4f355980e0433da737148eda59684 Mon Sep 17 00:00:00 2001 From: Huite Bootsma Date: Fri, 9 Jun 2023 16:13:32 +0200 Subject: [PATCH 2/9] WIP --- xugrid/regrid/regridder.py | 7 ++++--- xugrid/regrid/structured.py | 12 +++++++++++- xugrid/regrid/unstructured.py | 6 ++++++ 3 files changed, 21 insertions(+), 4 deletions(-) diff --git a/xugrid/regrid/regridder.py b/xugrid/regrid/regridder.py index cf134e19a..92ba86713 100644 --- a/xugrid/regrid/regridder.py +++ b/xugrid/regrid/regridder.py @@ -222,11 +222,12 @@ def to_dataset(self) -> xr.Dataset: """ Store the computed weights and target in a dataset for re-use. """ - ds = xr.Dataset( + weights_ds = xr.Dataset( {f"__regrid_{k}": v for k, v in zip(self._weights._fields, self._weights)} ) - ugrid_ds = self._target.ugrid_topology.to_dataset() - return xr.merge((ds, ugrid_ds)) + source_ds = self._source.to_dataset("source") + target_ds = self._target.to_dataset("target") + return xr.merge((weights_ds, source_ds, target_ds)) @staticmethod def _csr_from_dataset(dataset: xr.Dataset) -> WeightMatrixCSR: diff --git a/xugrid/regrid/structured.py b/xugrid/regrid/structured.py index 0cac81606..639cc8457 100644 --- a/xugrid/regrid/structured.py +++ b/xugrid/regrid/structured.py @@ -351,8 +351,11 @@ def linear_weights(self, other: "StructuredGrid1d"): neighbour, ) return self.sorted_output(source_index, target_index, weights) + - + def to_dataset(self): + pass + class StructuredGrid2d(StructuredGrid1d): """ e.g. (x,y) -> (x,y) @@ -489,6 +492,13 @@ def linear_weights(self, other): weights_x, ) + def to_dataset(self, name: str): + ds_x = self.xbounds.to_dataset() + ds_y = self.ybounds.to_dataset() + ds = xr.merge([ds_x, ds_y]) + ds[name] = xr.DataArray(-1, attrs={"type": "StructuredGrid2d"}) + return ds + class StructuredGrid3d: """ diff --git a/xugrid/regrid/unstructured.py b/xugrid/regrid/unstructured.py index 2cc0db89a..052aba1ce 100644 --- a/xugrid/regrid/unstructured.py +++ b/xugrid/regrid/unstructured.py @@ -1,4 +1,5 @@ import numpy as np +import xarray as xr import xugrid as xu from xugrid.constants import FloatDType @@ -138,3 +139,8 @@ def barycentric(self, other): order = np.argsort(target_index) return source_index[order], target_index[order], weights[order] + + def to_dataset(self, name: str): + ds = self.ugrid_topology.to_dataset() + ds[name] = xr.DataArray(-1, attrs={"type": "UnstructuredGrid2d"}) + return ds From a59df71015e20bfbd72cfbdc3269647e4ceca42b Mon Sep 17 00:00:00 2001 From: Hendrik Kok Date: Fri, 16 Jun 2023 11:40:56 +0200 Subject: [PATCH 3/9] add source and target grid topology to weight and dataset output . + Use this info when initializing regridder from weight and dataset. + small fix to rename() in Ugrid1d + 2d --- tests/test_regrid/test_regridder.py | 2 +- tests/test_ugrid1d.py | 1 + tests/test_ugrid2d.py | 1 + xugrid/regrid/regridder.py | 49 ++++++++++++++++++----------- xugrid/regrid/structured.py | 34 ++++++++++++-------- xugrid/regrid/unstructured.py | 4 +-- xugrid/ugrid/ugridbase.py | 1 + 7 files changed, 57 insertions(+), 35 deletions(-) diff --git a/tests/test_regrid/test_regridder.py b/tests/test_regrid/test_regridder.py index 1f9ed132e..6b06fb8d1 100644 --- a/tests/test_regrid/test_regridder.py +++ b/tests/test_regrid/test_regridder.py @@ -172,4 +172,4 @@ def test_regridder_from_dataset(cls, disk, quads_1): dataset = regridder.to_dataset() new_regridder = cls.from_dataset(dataset) new_result = new_regridder.regrid(disk) - assert new_result.equals(result) + assert np.array_equal(new_result.values, result.values, equal_nan=True) diff --git a/tests/test_ugrid1d.py b/tests/test_ugrid1d.py index 87a431ddb..2153be724 100644 --- a/tests/test_ugrid1d.py +++ b/tests/test_ugrid1d.py @@ -347,6 +347,7 @@ def test_ugrid1d_rename(): "node_x": "__renamed_node_x", "node_y": "__renamed_node_y", } + assert renamed.name == "__renamed" def test_ugrid1d_rename_with_dataset(): diff --git a/tests/test_ugrid2d.py b/tests/test_ugrid2d.py index 079dbf809..a94fa2b8d 100644 --- a/tests/test_ugrid2d.py +++ b/tests/test_ugrid2d.py @@ -984,6 +984,7 @@ def test_ugrid2d_rename(): "node_x": "__renamed_node_x", "node_y": "__renamed_node_y", } + assert renamed.name == "__renamed" def test_ugrid2d_rename_with_dataset(): diff --git a/xugrid/regrid/regridder.py b/xugrid/regrid/regridder.py index 92ba86713..abc35a83c 100644 --- a/xugrid/regrid/regridder.py +++ b/xugrid/regrid/regridder.py @@ -54,11 +54,13 @@ def _regrid(source: FloatArray, A: WeightMatrixCSR, size: int): return numba.njit(_regrid, parallel=True, cache=True) -def setup_grid(obj): +def setup_grid(obj, **kwargs): if isinstance(obj, (xu.Ugrid2d, xu.UgridDataArray, xu.UgridDataset)): return UnstructuredGrid2d(obj) elif isinstance(obj, (xr.DataArray, xr.Dataset)): - return StructuredGrid2d(obj, name_y="y", name_x="x") + return StructuredGrid2d( + obj, name_y=kwargs.get("name_y", "y"), name_x=kwargs.get("name_x", "x") + ) else: raise TypeError() @@ -120,10 +122,7 @@ def _setup_regrid(self, func) -> Callable: return def _regrid_array(self, source): - if hasattr(self, "_source"): - source_grid = self._source - else: - source_grid = source + source_grid = self._source first_dims_shape = source.shape[: -source_grid.ndim] # The regridding can be mapped over additional dimensions (e.g. for every time slice). @@ -225,8 +224,8 @@ def to_dataset(self) -> xr.Dataset: weights_ds = xr.Dataset( {f"__regrid_{k}": v for k, v in zip(self._weights._fields, self._weights)} ) - source_ds = self._source.to_dataset("source") - target_ds = self._target.to_dataset("target") + source_ds = self._source.to_dataset("__source") + target_ds = self._target.to_dataset("__target") return xr.merge((weights_ds, source_ds, target_ds)) @staticmethod @@ -257,10 +256,19 @@ def _weights_from_dataset( """ @classmethod - def from_weights(cls, weights, target: "xugrid.Ugrid2d"): + def from_weights( + cls, weights, target: Union["xugrid.Ugrid2d", xr.DataArray, xr.Dataset] + ): instance = cls.__new__(cls) - instance._weights = weights - instance._target = UnstructuredGrid2d(target) + instance._weights = cls._weights_from_dataset(weights) + instance._target = setup_grid(target) + unstructured = weights["__source_type"].attrs["type"] == "UnstructuredGrid2d" + if unstructured: + instance._source = setup_grid(xu.Ugrid2d.from_dataset(weights, "__source")) + else: + instance._source = setup_grid( + weights, name_x="__source_x", name_y="__source_y" + ) return instance @classmethod @@ -269,9 +277,12 @@ def from_dataset(cls, dataset: xr.Dataset): Reconstruct the regridder from a dataset with source, target indices and weights. """ - target = xu.Ugrid2d.from_dataset(dataset) - weights = cls._weights_from_dataset(dataset) - return cls.from_weights(weights, target) + unstructured = dataset["__target_type"].attrs["type"] == "UnstructuredGrid2d" + if unstructured: + target = xu.Ugrid2d.from_dataset(dataset, "__target") + + # weights = cls._weights_from_dataset(dataset) + return cls.from_weights(dataset, target) class CentroidLocatorRegridder(BaseRegridder): @@ -308,7 +319,7 @@ def _regrid(source: FloatArray, A: WeightMatrixCOO, size: int): @property def weights(self): - return self._weights + return self.to_dataset() @weights.setter def weights(self, weights: WeightMatrixCOO, target: "xugrid.Ugrid2d"): @@ -335,7 +346,7 @@ def _compute_weights(self, source, target, relative: bool) -> None: @property def weights(self): - return self._weights + return self.to_dataset() @weights.setter def weights(self, weights: WeightMatrixCSR): @@ -399,8 +410,8 @@ def _compute_weights(self, source, target) -> None: @classmethod def from_weights( cls, - weights: WeightMatrixCSR, - target: "xugrid.Ugrid2d", + weights: xr.Dataset, + target: Union["xugrid.Ugrid2d", xr.DataArray, xr.Dataset], method: Union[str, Callable] = "mean", ): instance = super().from_weights(weights, target) @@ -498,7 +509,7 @@ def _compute_weights(self, source, target): @property def weights(self): - return self._weights + return self.to_dataset() @weights.setter def weights(self, weights: WeightMatrixCSR): diff --git a/xugrid/regrid/structured.py b/xugrid/regrid/structured.py index 639cc8457..9ea54fcc1 100644 --- a/xugrid/regrid/structured.py +++ b/xugrid/regrid/structured.py @@ -29,8 +29,7 @@ class StructuredGrid1d: """ def __init__(self, obj: Union[xr.DataArray, xr.Dataset], name: str): - bounds_name_left = f"{name}bounds_left" # e.g. xbounds - bounds_name_right = f"{name}bounds_right" # e.g. xbounds + bounds_name = f"{name}bounds" # e.g. xbounds size_name = f"d{name}" # e.g. dx index = obj.indexes[name] @@ -46,10 +45,8 @@ def __init__(self, obj: Union[xr.DataArray, xr.Dataset], name: str): else: raise ValueError(f"{name} is not monotonic for array {obj.name}") - if bounds_name_left in obj.coords: - start = obj[bounds_name_left].values - end = obj[bounds_name_right].values - bounds = np.column_stack((start, end)) + if bounds_name in obj.coords: + bounds = obj[bounds_name].values else: if size_name in obj.coords: # works for scalar size and array size @@ -351,11 +348,22 @@ def linear_weights(self, other: "StructuredGrid1d"): neighbour, ) return self.sorted_output(source_index, target_index, weights) - - def to_dataset(self): - pass - + def to_dataset(self, name: str): + export_name = name + "_" + self.name + return xr.DataArray( + name=name, + data=np.nan, + dims=[export_name, export_name + "nbounds"], + coords={ + export_name: self.midpoints, + export_name + + "bounds": ([export_name, export_name + "nbounds"], self.bounds), + export_name + "nbounds": np.arange(2), + }, + ) + + class StructuredGrid2d(StructuredGrid1d): """ e.g. (x,y) -> (x,y) @@ -493,10 +501,10 @@ def linear_weights(self, other): ) def to_dataset(self, name: str): - ds_x = self.xbounds.to_dataset() - ds_y = self.ybounds.to_dataset() + ds_x = self.xbounds.to_dataset(name) + ds_y = self.ybounds.to_dataset(name) ds = xr.merge([ds_x, ds_y]) - ds[name] = xr.DataArray(-1, attrs={"type": "StructuredGrid2d"}) + ds[name + "_type"] = xr.DataArray(-1, attrs={"type": "StructuredGrid2d"}) return ds diff --git a/xugrid/regrid/unstructured.py b/xugrid/regrid/unstructured.py index 052aba1ce..abc9b023f 100644 --- a/xugrid/regrid/unstructured.py +++ b/xugrid/regrid/unstructured.py @@ -141,6 +141,6 @@ def barycentric(self, other): return source_index[order], target_index[order], weights[order] def to_dataset(self, name: str): - ds = self.ugrid_topology.to_dataset() - ds[name] = xr.DataArray(-1, attrs={"type": "UnstructuredGrid2d"}) + ds = self.ugrid_topology.rename(name).to_dataset() + ds[name + "_type"] = xr.DataArray(-1, attrs={"type": "UnstructuredGrid2d"}) return ds diff --git a/xugrid/ugrid/ugridbase.py b/xugrid/ugrid/ugridbase.py index df48bdb03..065b205d1 100644 --- a/xugrid/ugrid/ugridbase.py +++ b/xugrid/ugrid/ugridbase.py @@ -192,6 +192,7 @@ def rename(self, name: str): name_dict[name_key] = name_value new = self.copy() + new.name = name new._attrs = new_attrs new._indexes = {k: name_dict[v] for k, v in new._indexes.items()} if new._dataset is not None: From 4060a4a5ced055f155a812bb7d58c785b689c623 Mon Sep 17 00:00:00 2001 From: Hendrik Kok Date: Fri, 16 Jun 2023 11:40:56 +0200 Subject: [PATCH 4/9] add source and target grid topology to weight and dataset output . + Use this info when initializing regridder from weight and dataset. + small fix to rename() in Ugrid1d + 2d --- tests/test_regrid/test_regridder.py | 2 +- tests/test_ugrid1d.py | 1 + tests/test_ugrid2d.py | 1 + xugrid/regrid/regridder.py | 49 ++++++++++++++++++----------- xugrid/regrid/structured.py | 34 ++++++++++++-------- xugrid/regrid/unstructured.py | 4 +-- xugrid/ugrid/ugridbase.py | 1 + 7 files changed, 57 insertions(+), 35 deletions(-) diff --git a/tests/test_regrid/test_regridder.py b/tests/test_regrid/test_regridder.py index 1f9ed132e..6b06fb8d1 100644 --- a/tests/test_regrid/test_regridder.py +++ b/tests/test_regrid/test_regridder.py @@ -172,4 +172,4 @@ def test_regridder_from_dataset(cls, disk, quads_1): dataset = regridder.to_dataset() new_regridder = cls.from_dataset(dataset) new_result = new_regridder.regrid(disk) - assert new_result.equals(result) + assert np.array_equal(new_result.values, result.values, equal_nan=True) diff --git a/tests/test_ugrid1d.py b/tests/test_ugrid1d.py index 87a431ddb..2153be724 100644 --- a/tests/test_ugrid1d.py +++ b/tests/test_ugrid1d.py @@ -347,6 +347,7 @@ def test_ugrid1d_rename(): "node_x": "__renamed_node_x", "node_y": "__renamed_node_y", } + assert renamed.name == "__renamed" def test_ugrid1d_rename_with_dataset(): diff --git a/tests/test_ugrid2d.py b/tests/test_ugrid2d.py index 079dbf809..a94fa2b8d 100644 --- a/tests/test_ugrid2d.py +++ b/tests/test_ugrid2d.py @@ -984,6 +984,7 @@ def test_ugrid2d_rename(): "node_x": "__renamed_node_x", "node_y": "__renamed_node_y", } + assert renamed.name == "__renamed" def test_ugrid2d_rename_with_dataset(): diff --git a/xugrid/regrid/regridder.py b/xugrid/regrid/regridder.py index 92ba86713..abc35a83c 100644 --- a/xugrid/regrid/regridder.py +++ b/xugrid/regrid/regridder.py @@ -54,11 +54,13 @@ def _regrid(source: FloatArray, A: WeightMatrixCSR, size: int): return numba.njit(_regrid, parallel=True, cache=True) -def setup_grid(obj): +def setup_grid(obj, **kwargs): if isinstance(obj, (xu.Ugrid2d, xu.UgridDataArray, xu.UgridDataset)): return UnstructuredGrid2d(obj) elif isinstance(obj, (xr.DataArray, xr.Dataset)): - return StructuredGrid2d(obj, name_y="y", name_x="x") + return StructuredGrid2d( + obj, name_y=kwargs.get("name_y", "y"), name_x=kwargs.get("name_x", "x") + ) else: raise TypeError() @@ -120,10 +122,7 @@ def _setup_regrid(self, func) -> Callable: return def _regrid_array(self, source): - if hasattr(self, "_source"): - source_grid = self._source - else: - source_grid = source + source_grid = self._source first_dims_shape = source.shape[: -source_grid.ndim] # The regridding can be mapped over additional dimensions (e.g. for every time slice). @@ -225,8 +224,8 @@ def to_dataset(self) -> xr.Dataset: weights_ds = xr.Dataset( {f"__regrid_{k}": v for k, v in zip(self._weights._fields, self._weights)} ) - source_ds = self._source.to_dataset("source") - target_ds = self._target.to_dataset("target") + source_ds = self._source.to_dataset("__source") + target_ds = self._target.to_dataset("__target") return xr.merge((weights_ds, source_ds, target_ds)) @staticmethod @@ -257,10 +256,19 @@ def _weights_from_dataset( """ @classmethod - def from_weights(cls, weights, target: "xugrid.Ugrid2d"): + def from_weights( + cls, weights, target: Union["xugrid.Ugrid2d", xr.DataArray, xr.Dataset] + ): instance = cls.__new__(cls) - instance._weights = weights - instance._target = UnstructuredGrid2d(target) + instance._weights = cls._weights_from_dataset(weights) + instance._target = setup_grid(target) + unstructured = weights["__source_type"].attrs["type"] == "UnstructuredGrid2d" + if unstructured: + instance._source = setup_grid(xu.Ugrid2d.from_dataset(weights, "__source")) + else: + instance._source = setup_grid( + weights, name_x="__source_x", name_y="__source_y" + ) return instance @classmethod @@ -269,9 +277,12 @@ def from_dataset(cls, dataset: xr.Dataset): Reconstruct the regridder from a dataset with source, target indices and weights. """ - target = xu.Ugrid2d.from_dataset(dataset) - weights = cls._weights_from_dataset(dataset) - return cls.from_weights(weights, target) + unstructured = dataset["__target_type"].attrs["type"] == "UnstructuredGrid2d" + if unstructured: + target = xu.Ugrid2d.from_dataset(dataset, "__target") + + # weights = cls._weights_from_dataset(dataset) + return cls.from_weights(dataset, target) class CentroidLocatorRegridder(BaseRegridder): @@ -308,7 +319,7 @@ def _regrid(source: FloatArray, A: WeightMatrixCOO, size: int): @property def weights(self): - return self._weights + return self.to_dataset() @weights.setter def weights(self, weights: WeightMatrixCOO, target: "xugrid.Ugrid2d"): @@ -335,7 +346,7 @@ def _compute_weights(self, source, target, relative: bool) -> None: @property def weights(self): - return self._weights + return self.to_dataset() @weights.setter def weights(self, weights: WeightMatrixCSR): @@ -399,8 +410,8 @@ def _compute_weights(self, source, target) -> None: @classmethod def from_weights( cls, - weights: WeightMatrixCSR, - target: "xugrid.Ugrid2d", + weights: xr.Dataset, + target: Union["xugrid.Ugrid2d", xr.DataArray, xr.Dataset], method: Union[str, Callable] = "mean", ): instance = super().from_weights(weights, target) @@ -498,7 +509,7 @@ def _compute_weights(self, source, target): @property def weights(self): - return self._weights + return self.to_dataset() @weights.setter def weights(self, weights: WeightMatrixCSR): diff --git a/xugrid/regrid/structured.py b/xugrid/regrid/structured.py index 639cc8457..9ea54fcc1 100644 --- a/xugrid/regrid/structured.py +++ b/xugrid/regrid/structured.py @@ -29,8 +29,7 @@ class StructuredGrid1d: """ def __init__(self, obj: Union[xr.DataArray, xr.Dataset], name: str): - bounds_name_left = f"{name}bounds_left" # e.g. xbounds - bounds_name_right = f"{name}bounds_right" # e.g. xbounds + bounds_name = f"{name}bounds" # e.g. xbounds size_name = f"d{name}" # e.g. dx index = obj.indexes[name] @@ -46,10 +45,8 @@ def __init__(self, obj: Union[xr.DataArray, xr.Dataset], name: str): else: raise ValueError(f"{name} is not monotonic for array {obj.name}") - if bounds_name_left in obj.coords: - start = obj[bounds_name_left].values - end = obj[bounds_name_right].values - bounds = np.column_stack((start, end)) + if bounds_name in obj.coords: + bounds = obj[bounds_name].values else: if size_name in obj.coords: # works for scalar size and array size @@ -351,11 +348,22 @@ def linear_weights(self, other: "StructuredGrid1d"): neighbour, ) return self.sorted_output(source_index, target_index, weights) - - def to_dataset(self): - pass - + def to_dataset(self, name: str): + export_name = name + "_" + self.name + return xr.DataArray( + name=name, + data=np.nan, + dims=[export_name, export_name + "nbounds"], + coords={ + export_name: self.midpoints, + export_name + + "bounds": ([export_name, export_name + "nbounds"], self.bounds), + export_name + "nbounds": np.arange(2), + }, + ) + + class StructuredGrid2d(StructuredGrid1d): """ e.g. (x,y) -> (x,y) @@ -493,10 +501,10 @@ def linear_weights(self, other): ) def to_dataset(self, name: str): - ds_x = self.xbounds.to_dataset() - ds_y = self.ybounds.to_dataset() + ds_x = self.xbounds.to_dataset(name) + ds_y = self.ybounds.to_dataset(name) ds = xr.merge([ds_x, ds_y]) - ds[name] = xr.DataArray(-1, attrs={"type": "StructuredGrid2d"}) + ds[name + "_type"] = xr.DataArray(-1, attrs={"type": "StructuredGrid2d"}) return ds diff --git a/xugrid/regrid/unstructured.py b/xugrid/regrid/unstructured.py index 052aba1ce..abc9b023f 100644 --- a/xugrid/regrid/unstructured.py +++ b/xugrid/regrid/unstructured.py @@ -141,6 +141,6 @@ def barycentric(self, other): return source_index[order], target_index[order], weights[order] def to_dataset(self, name: str): - ds = self.ugrid_topology.to_dataset() - ds[name] = xr.DataArray(-1, attrs={"type": "UnstructuredGrid2d"}) + ds = self.ugrid_topology.rename(name).to_dataset() + ds[name + "_type"] = xr.DataArray(-1, attrs={"type": "UnstructuredGrid2d"}) return ds diff --git a/xugrid/ugrid/ugridbase.py b/xugrid/ugrid/ugridbase.py index df48bdb03..065b205d1 100644 --- a/xugrid/ugrid/ugridbase.py +++ b/xugrid/ugrid/ugridbase.py @@ -192,6 +192,7 @@ def rename(self, name: str): name_dict[name_key] = name_value new = self.copy() + new.name = name new._attrs = new_attrs new._indexes = {k: name_dict[v] for k, v in new._indexes.items()} if new._dataset is not None: From 940689fae6514650786e056bfecee42a3cc38c50 Mon Sep 17 00:00:00 2001 From: Hendrik Kok Date: Fri, 16 Jun 2023 13:18:19 +0200 Subject: [PATCH 5/9] new grid definition in fixtures for failing structured 1d-test (non-equidistance) --- tests/fixtures/fixture_regridder.py | 27 ++++++++++++++++----------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/tests/fixtures/fixture_regridder.py b/tests/fixtures/fixture_regridder.py index 7366153d4..9f9ec620d 100644 --- a/tests/fixtures/fixture_regridder.py +++ b/tests/fixtures/fixture_regridder.py @@ -133,17 +133,22 @@ def grid_data_d(): @pytest.fixture(scope="function") def grid_data_e(): return xr.DataArray( - data=np.arange(12).reshape((4, 3)), - dims=["y", "x"], - coords={ - "y": np.array([175, 125, 75, 25]), - "x": np.array([30, 67.5, 105]), - "dx": 25, - "dy": -50.0, - "xbounds_left": ("x", np.array([17.5, 42.5, 92.5])), - "xbounds_right": ("x", np.array([42.5, 92.5, 117.5])), - }, - ) + data=np.zeros((4, 3, 2)), + dims=["y", "x",'nbounds'], + coords={ + "y": np.array([175, 125, 75, 25]), + "x": np.array([30, 67.5, 105]), + "dx": 25, + "dy": -50.0, + "xbounds": ( + ["x", "nbounds"], + np.column_stack( + (np.array([17.5, 42.5, 92.5]), np.array([42.5, 92.5, 117.5])) + ), + ), + "nbounds": np.arange(2), + }, +) @pytest.fixture(scope="function") From b0617c378c3fda46a305c8da7aa2c48a19279c95 Mon Sep 17 00:00:00 2001 From: Hendrik Kok Date: Fri, 16 Jun 2023 13:20:36 +0200 Subject: [PATCH 6/9] linter --- tests/fixtures/fixture_regridder.py | 30 ++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/tests/fixtures/fixture_regridder.py b/tests/fixtures/fixture_regridder.py index 9f9ec620d..a79c77ed2 100644 --- a/tests/fixtures/fixture_regridder.py +++ b/tests/fixtures/fixture_regridder.py @@ -133,22 +133,22 @@ def grid_data_d(): @pytest.fixture(scope="function") def grid_data_e(): return xr.DataArray( - data=np.zeros((4, 3, 2)), - dims=["y", "x",'nbounds'], - coords={ - "y": np.array([175, 125, 75, 25]), - "x": np.array([30, 67.5, 105]), - "dx": 25, - "dy": -50.0, - "xbounds": ( - ["x", "nbounds"], - np.column_stack( - (np.array([17.5, 42.5, 92.5]), np.array([42.5, 92.5, 117.5])) + data=np.zeros((4, 3, 2)), + dims=["y", "x", "nbounds"], + coords={ + "y": np.array([175, 125, 75, 25]), + "x": np.array([30, 67.5, 105]), + "dx": 25, + "dy": -50.0, + "xbounds": ( + ["x", "nbounds"], + np.column_stack( + (np.array([17.5, 42.5, 92.5]), np.array([42.5, 92.5, 117.5])) + ), ), - ), - "nbounds": np.arange(2), - }, -) + "nbounds": np.arange(2), + }, + ) @pytest.fixture(scope="function") From 2d410698839913668c1a4af12ea33a255ec1a003 Mon Sep 17 00:00:00 2001 From: Huite Bootsma Date: Fri, 23 Jun 2023 08:19:36 +0200 Subject: [PATCH 7/9] typo --- tests/test_regrid/test_regridder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_regrid/test_regridder.py b/tests/test_regrid/test_regridder.py index 6b06fb8d1..391d2873e 100644 --- a/tests/test_regrid/test_regridder.py +++ b/tests/test_regrid/test_regridder.py @@ -103,7 +103,7 @@ def test_overlap_regridder(disk, quads_1): assert broadcasted.shape == (5, 100) -def test_lineair_interpolator_structured( +def test_linear_interpolator_structured( grid_data_a, grid_data_a_layered, grid_data_b, expected_results_linear ): regridder = BarycentricInterpolator(source=grid_data_a, target=grid_data_b) From fd7574dbaaa231be570080046ae4cf4886ce2b7d Mon Sep 17 00:00:00 2001 From: Huite Bootsma Date: Fri, 23 Jun 2023 08:21:07 +0200 Subject: [PATCH 8/9] Add type hints, docstring should be numpy --- xugrid/regrid/structured.py | 324 ++++++++++++++++++++---------------- 1 file changed, 180 insertions(+), 144 deletions(-) diff --git a/xugrid/regrid/structured.py b/xugrid/regrid/structured.py index 9ea54fcc1..bd48d87e9 100644 --- a/xugrid/regrid/structured.py +++ b/xugrid/regrid/structured.py @@ -6,11 +6,12 @@ While the unstructured logic would work for structured data as well, it is much less efficient than utilizing the structure of the coordinates. """ -from typing import Union +from typing import Any, Tuple, Union import numpy as np import xarray as xr +from xugrid.constants import FloatArray, IntArray from xugrid.regrid.overlap_1d import overlap_1d, overlap_1d_nd from xugrid.regrid.unstructured import UnstructuredGrid2d from xugrid.regrid.utils import broadcast @@ -76,40 +77,46 @@ def __init__(self, obj: Union[xr.DataArray, xr.Dataset], name: str): self.grid = obj @property - def ndim(self): + def ndim(self) -> int: return 1 @property - def dims(self): + def dims(self) -> Tuple[str]: return (self.name,) @property - def size(self): + def size(self) -> int: return len(self.bounds) @property - def length(self): + def length(self) -> FloatArray: return abs(np.diff(self.bounds, axis=1)) - def flip_if_needed(self, index): + def flip_if_needed(self, index: IntArray) -> IntArray: if self.flipped: return self.size - index - 1 else: return index - def valid_nodes_within_bounds(self, other): + def valid_nodes_within_bounds( + self, other: "StructuredGrid1d" + ) -> Tuple[IntArray, IntArray]: """ Retruns nodes when midpoints are within bounding box of overlaying grid. In cases that midpoints (and bounding boxes) are flipped, computed indexes are fliped as well. - Args: - self (StructuredGrid1d): source grid - other (StructuredGrid1d): target grid + Parameters + ---------- + other: StructuredGrid1d + The target grid topology - Returns: - valid_self_index (np.array): valid source indexes - valid_other_index (np.array): valid target indexes + Returns + ------- + valid_self_index: np.array + valid source indexes + valid_other_index: np.array + valid target indexes """ start = np.searchsorted(self.bounds[:, 0], other.midpoints, side=self.side) end = np.searchsorted(self.bounds[:, 1], other.midpoints, side=self.side) @@ -124,19 +131,26 @@ def valid_nodes_within_bounds(self, other): valid_other_index = other.flip_if_needed(valid_other_index) return valid_self_index, valid_other_index - def valid_nodes_within_bounds_and_extend(self, other): + def valid_nodes_within_bounds_and_extend( + self, other: "StructuredGrid1d" + ) -> Tuple[IntArray, IntArray]: """ - returns all valid nodes for linear interpolation. In addition to valid_nodes_within_bounds() - is checked if target midpoints are not outside outer source boundary midpoints. In that case - there is no interpolation possible. + Returns all valid nodes for linear interpolation. In addition to + valid_nodes_within_bounds() is checked if target midpoints are not + outside outer source boundary midpoints. In that case there is no + interpolation possible. - Args: - self (StructuredGrid1d): source grid - other (StructuredGrid1d): target grid + Parameters + ---------- + other: StructuredGrid1d + The target grid. - Returns: - valid_self_index (np.array): valid source indexes - valid_other_index (np.array): valid target indexes + Returns + ------- + valid_self_index: np.array + valid source indexes + valid_other_index: np.array + valid target indexes """ source_index, target_index = self.valid_nodes_within_bounds(other) valid = (other.midpoints[target_index] > self.midpoints[0]) & ( @@ -144,19 +158,26 @@ def valid_nodes_within_bounds_and_extend(self, other): ) return source_index[valid], target_index[valid] - def overlap_1d_structured(self, other): + def overlap_1d_structured( + self, other: "StructuredGrid1d" + ) -> Tuple[IntArray, IntArray, FloatArray]: """ Returns source and target nodes and overlapping length. It utilises overlap_1d() and does an aditional flip in cases of reversed midpoints - Args: - self (StructuredGrid1d): source grid - other (StructuredGrid1d): target grid + Parameters + ---------- + other: StructuredGrid1d + The target grid. - Returns: - valid_self_index (np.array): valid source indexes - valid_other_index (np.array): valid target indexes - weights (np.array): length of overlap + Returns + ------- + valid_self_index: np.array + valid source indexes + valid_other_index: np.array + valid target indexes + weights: np.array + length of overlap """ source_index, target_index, weights = overlap_1d(self.bounds, other.bounds) source_index = self.flip_if_needed(source_index) @@ -165,116 +186,124 @@ def overlap_1d_structured(self, other): def centroids_to_linear_sets( self, - other, source_index: np.array, target_index: np.array, weights: np.array, neighbour: np.array, - ): + ) -> Tuple[IntArray, IntArray, FloatArray]: """ - Returns for every target node an pair of connected source nodes based on - centroids connection inputs - - Args: - self (StructuredGrid1d): source grid - other (StructuredGrid1d): target grid - source_index (np.array): source index (centroids) - target_index (np.array): target index (centroids) - weights (np.array): weights (centroids) + Returns for every target node an pair of connected source nodes based + on centroids connection inputs. - Returns:biliniar + Parameters + ---------- + source_index: np.array + target_index: np.array + weights: np.array - source_index (np.array): source index (linear) - target_index (np.array): target index (linear) - weights (np.array): weights (linear) + Returns + ------- + valid_self_index: np.array + valid_other_index: np.array + weights: np.array + lineair interpolation weights. """ - # if source_index is flipped, source-index is decreasing and neighbour need to be flipped + # if source_index is flipped, source-index is decreasing and neighbour + # need to be flipped if self.flipped: neighbour = -neighbour source_index = np.column_stack((source_index, source_index + neighbour)).ravel() target_index = np.repeat(target_index, 2) weights = np.column_stack((weights, 1.0 - weights)).ravel() - # correct for possibility of out of bound due to column-stack source_index + 1 and -1 + # correct for possibility of out of bound due to column-stack + # source_index + 1 and -1 valid = np.logical_and(source_index <= self.size - 1, source_index >= 0) return source_index[valid], target_index[valid], weights[valid] - def get_midpoint_index(self, array_index): + def maybe_reverse_index(self, index: IntArray) -> IntArray: """ - Returns midpoint array indexes for given array_index. - - Args: - array_index (np.array): array_index + Flips the index if needed for descending coordinates. - Returns:biliniar - midpoint_index (np.array): midpoint_index + Parameters + ---------- + index: np.ndarray + Returns + ------- + checked_index: np.ndarray """ if self.flipped: - return self.size - array_index - 1 + return self.size - index - 1 else: - return array_index + return index - def compute_distance_to_centroids(self, other, source_index, target_index): + def compute_distance_to_centroids( + self, other: "StructuredGrid1d", source_index: IntArray, target_index: IntArray + ) -> Tuple[FloatArray, IntArray]: """ - computes linear weights bases on centroid indexes. + Computes linear weights bases on centroid indexes. - Args: - self (StructuredGrid1d): source grid - other (StructuredGrid1d): target grid - other (_type_): _description_ - source_index (np.array): source index (centroids) - target_index (np.array): target index (centroids) + Parameters + ---------- + other: StructuredGrid1d + source_index: np.array + target_index: np.array - Raises: - ValueError: for not enought midpoints + Raises + ------ + ValueError + When the coordinate contains only a single point. - Returns: - weights (np.array): weights + Returns + ------- + weights: np.array + neighbor: np.narray """ - - source_midpoint_index = self.get_midpoint_index(source_index) - target_midpoints_index = other.get_midpoint_index(target_index) - neighbour = np.ones(target_midpoints_index.size, dtype=int) + source_midpoint_index = self.maybe_reverse_index(source_index) + target_midpoints_index = other.maybe_reverse_index(target_index) + neighbor = np.ones(target_midpoints_index.size, dtype=int) # cases where midpoint target < midpoint source condition = ( other.midpoints[target_midpoints_index] < self.midpoints[source_midpoint_index] ) - neighbour[condition] = -neighbour[condition] + neighbor[condition] = -neighbor[condition] - if not self.midpoints.size > 2: + if self.midpoints.size < 2: raise ValueError( - "source index must larger than 2. Cannot interpolate with one point" + f"Coordinate {self.name} has size: {self.midpoints.size}. " + "At least two points are required for interpolation." ) weights = ( other.midpoints[target_midpoints_index] - self.midpoints[source_midpoint_index] ) / ( - self.midpoints[source_midpoint_index + neighbour] + self.midpoints[source_midpoint_index + neighbor] - self.midpoints[source_midpoint_index] ) weights[weights < 0.0] = 0.0 weights[weights > 1.0] = 1.0 - - return weights, neighbour + return weights, neighbor def sorted_output( - self, source_index: np.array, target_index: np.array, weights: np.array - ): + self, source_index: IntArray, target_index: IntArray, weights: FloatArray + ) -> Tuple[IntArray, IntArray, FloatArray]: """ - Returns sorted input based on target index. The regridder needs input sorted on - row index of WeightMatrixCOO (target index) - - Args: - source_index (np.array): source index - target_index (np.array): target index - weights (np.array): weights - - Returns: - source_index (np.array): source index - target_index (np.array): target index - weights (np.array): weights + Returns sorted input based on target index. The regridder needs input + sorted on row index of WeightMatrixCOO (target index). + + Parameters + ---------- + source_index: np.array + target_index: np.array + weights: np.array + + Returns + ------- + source_index: np.array + target_index: np.array + weights: np.array """ sorter_target = np.argsort(target_index) return ( @@ -283,65 +312,72 @@ def sorted_output( weights[sorter_target], ) - def overlap(self, other: "StructuredGrid1d", relative: bool): + def overlap( + self, other: "StructuredGrid1d", relative: bool + ) -> Tuple[IntArray, IntArray, FloatArray]: """ Returns source and target indexes and overlapping length - Args: - self (StructuredGrid1d): source grid - other (StructuredGrid1d): target grid - relative (bool): overlapping length target relative to source length + Parameters + ---------- + other: StructuredGrid1d + relative: bool + Overlapping length target relative to source length - Returns: - source_index (np.array): source indexes - target_index (np.array): target indexes - weights (np.array): overlapping length + Returns + ------- + source_index: np.array + target_index: np.array + weights: np.array + Overlapping length """ source_index, target_index, weights = self.overlap_1d_structured(other) if relative: weights /= self.length()[source_index] return self.sorted_output(source_index, target_index, weights) - def locate_centroids(self, other: "StructuredGrid1d"): + def locate_centroids( + self, other: "StructuredGrid1d" + ) -> Tuple[IntArray, IntArray, FloatArray]: """ - Returns source and target indexes based on nearest neighbor of centroids + Returns source and target indexes based on nearest neighbor of + centroids. - Args: - self (StructuredGrid1d): source grid - other (StructuredGrid1d): target grid + Parameters + ---------- + other: StructuredGrid1d - Returns: - source_index (np.array): source indexes - target_index (np.array): target indexes - weights (np.array): array of ones + Returns + ------- + source_index: np.array + target_index: np.array + weights: np.array """ source_index, target_index = self.valid_nodes_within_bounds(other) weights = np.ones(source_index.size, dtype=float) return self.sorted_output(source_index, target_index, weights) - def linear_weights(self, other: "StructuredGrid1d"): + def linear_weights( + self, other: "StructuredGrid1d" + ) -> Tuple[IntArray, IntArray, FloatArray]: """ - Returns linear source and target indexes and corresponding weights + Returns linear source and target indexes and corresponding weights. - Args: - self (StructuredGrid1d): source grid - other (StructuredGrid1d): target grid + Parameters + ---------- + other: StructuredGrid1d - Raises: - ValueError: when number of nodes is to small to compute linear weights - - Returns: - source_index (np.array): overlaying self indexes - target_index (np.array): overlaying other indexes - weights (np.array): array linear weights + Returns + ------- + source_index: np.array + target_index: np.array + weights: np.array """ - source_index, target_index = self.valid_nodes_within_bounds_and_extend(other) weights, neighbour = self.compute_distance_to_centroids( other, source_index, target_index ) source_index, target_index, weights = self.centroids_to_linear_sets( - other, source_index, target_index, weights, @@ -349,7 +385,7 @@ def linear_weights(self, other: "StructuredGrid1d"): ) return self.sorted_output(source_index, target_index, weights) - def to_dataset(self, name: str): + def to_dataset(self, name: str) -> xr.DataArray: export_name = name + "_" + self.name return xr.DataArray( name=name, @@ -379,26 +415,26 @@ def __init__( self.ybounds = StructuredGrid1d(obj, name_y) @property - def ndim(self): + def ndim(self) -> int: return 2 @property - def dims(self): + def dims(self) -> Tuple[str, str]: return self.ybounds.dims + self.xbounds.dims # ("y", "x") @property - def size(self): + def size(self) -> int: return self.ybounds.size * self.xbounds.size @property - def shape(self): + def shape(self) -> Tuple[int, int]: return (self.ybounds.size, self.xbounds.size) @property - def area(self): + def area(self) -> FloatArray: return np.multiply.outer(self.ybounds.length, self.xbounds.length) - def convert_to(self, matched_type): + def convert_to(self, matched_type: Any) -> Any: if isinstance(self, matched_type): return self elif isinstance(self, UnstructuredGrid2d): @@ -410,14 +446,14 @@ def convert_to(self, matched_type): def broadcast_sorted( self, - other, - source_index_y: np.array, - source_index_x: np.array, - target_index_y: np.array, - target_index_x: np.array, - weights_y: np.array, - weights_x: np.array, - ): + other: Any, + source_index_y: IntArray, + source_index_x: IntArray, + target_index_y: IntArray, + target_index_x: IntArray, + weights_y: FloatArray, + weights_x: FloatArray, + ) -> Tuple[IntArray, IntArray, FloatArray]: source_index, target_index, weights = broadcast( self.shape, other.shape, @@ -428,7 +464,7 @@ def broadcast_sorted( sorter = np.argsort(target_index) return source_index[sorter], target_index[sorter], weights[sorter] - def overlap(self, other, relative: bool): + def overlap(self, other, relative: bool) -> Tuple[IntArray, IntArray, FloatArray]: """ Returns ------- @@ -452,7 +488,7 @@ def overlap(self, other, relative: bool): weights_x, ) - def locate_centroids(self, other): + def locate_centroids(self, other) -> Tuple[IntArray, IntArray, FloatArray]: """ Returns ------- @@ -476,7 +512,7 @@ def locate_centroids(self, other): weights_x, ) - def linear_weights(self, other): + def linear_weights(self, other) -> Tuple[IntArray, IntArray, FloatArray]: """ Returns ------- @@ -500,7 +536,7 @@ def linear_weights(self, other): weights_x, ) - def to_dataset(self, name: str): + def to_dataset(self, name: str) -> xr.Dataset: ds_x = self.xbounds.to_dataset(name) ds_y = self.ybounds.to_dataset(name) ds = xr.merge([ds_x, ds_y]) From c792381f22a3fa5e31e41c894db511115d224503 Mon Sep 17 00:00:00 2001 From: Huite Bootsma Date: Fri, 23 Jun 2023 08:24:29 +0200 Subject: [PATCH 9/9] Reflow comment, black --- xugrid/regrid/regridder.py | 9 +++++---- xugrid/regrid/structured.py | 2 +- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/xugrid/regrid/regridder.py b/xugrid/regrid/regridder.py index abc35a83c..8ae06a22c 100644 --- a/xugrid/regrid/regridder.py +++ b/xugrid/regrid/regridder.py @@ -125,10 +125,11 @@ def _regrid_array(self, source): source_grid = self._source first_dims_shape = source.shape[: -source_grid.ndim] - # The regridding can be mapped over additional dimensions (e.g. for every time slice). - # This is the `extra_index` iteration in _regrid(). - # But it should work consistently even if no additional present: in that case we create - # a 1-sized additional dimension in front, so the `extra_index` iteration always applies. + # The regridding can be mapped over additional dimensions, e.g. for + # every time slice. This is the `extra_index` iteration in _regrid(). + # But it should work consistently even if no additional present: in + # that case we create a 1-sized additional dimension in front, so the + # `extra_index` iteration always applies. if source.ndim == source_grid.ndim: source = source[np.newaxis] diff --git a/xugrid/regrid/structured.py b/xugrid/regrid/structured.py index bd48d87e9..c56d52533 100644 --- a/xugrid/regrid/structured.py +++ b/xugrid/regrid/structured.py @@ -273,7 +273,7 @@ def compute_distance_to_centroids( if self.midpoints.size < 2: raise ValueError( f"Coordinate {self.name} has size: {self.midpoints.size}. " - "At least two points are required for interpolation." + "At least two points are required for interpolation." ) weights = ( other.midpoints[target_midpoints_index]