uxarray 2026.3.0__tar.gz → 2026.4.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {uxarray-2026.3.0 → uxarray-2026.4.1}/PKG-INFO +1 -1
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/core/dataset.py +20 -2
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/io/utils.py +84 -21
- uxarray-2026.4.1/uxarray/remap/accessor.py +229 -0
- uxarray-2026.4.1/uxarray/remap/yac.py +407 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray.egg-info/PKG-INFO +1 -1
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray.egg-info/SOURCES.txt +1 -0
- uxarray-2026.3.0/uxarray/remap/accessor.py +0 -112
- {uxarray-2026.3.0 → uxarray-2026.4.1}/.codecov.yml +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/.git_archival.txt +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/.gitattributes +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/CITATION.cff +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/CODE_OF_CONDUCT.md +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/CONTRIBUTING.md +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/INSTALLATION.md +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/LICENSE +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/MANIFEST.in +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/README.md +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/build.sh +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/pyproject.toml +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/setup.cfg +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/setup.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/__init__.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/constants.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/conventions/__init__.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/conventions/descriptors.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/conventions/ugrid.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/core/__init__.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/core/accessors.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/core/aggregation.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/core/api.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/core/dataarray.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/core/gradient.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/core/utils.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/core/zonal.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/cross_sections/__init__.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/cross_sections/dataarray_accessor.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/cross_sections/grid_accessor.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/cross_sections/sample.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/formatting_html.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/grid/__init__.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/grid/arcs.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/grid/area.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/grid/bounds.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/grid/connectivity.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/grid/coordinates.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/grid/dual.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/grid/geometry.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/grid/grid.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/grid/integrate.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/grid/intersections.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/grid/neighbors.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/grid/point_in_face.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/grid/slice.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/grid/utils.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/grid/validation.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/io/__init__.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/io/_delaunay.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/io/_esmf.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/io/_exodus.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/io/_fesom2.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/io/_geopandas.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/io/_geos.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/io/_healpix.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/io/_icon.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/io/_mpas.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/io/_scrip.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/io/_structured.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/io/_topology.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/io/_ugrid.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/io/_vertices.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/io/_voronoi.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/plot/__init__.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/plot/accessor.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/plot/constants.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/plot/matplotlib.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/plot/utils.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/remap/__init__.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/remap/bilinear.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/remap/inverse_distance_weighted.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/remap/nearest_neighbor.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/remap/spatial_coords_remap.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/remap/utils.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/subset/__init__.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/subset/dataarray_accessor.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/subset/grid_accessor.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/utils/__init__.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray/utils/computing.py +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray.egg-info/dependency_links.txt +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray.egg-info/requires.txt +0 -0
- {uxarray-2026.3.0 → uxarray-2026.4.1}/uxarray.egg-info/top_level.txt +0 -0
|
@@ -90,6 +90,24 @@ class UxDataset(xr.Dataset):
|
|
|
90
90
|
else:
|
|
91
91
|
self._uxgrid = uxgrid
|
|
92
92
|
|
|
93
|
+
# As of xarray's 2026.4.0, `xr.Dataset(xr.Dataset)` is prohibited;
|
|
94
|
+
# hence this check, i.e. if we get `xr.Dataset` as input, use its `data_vars`
|
|
95
|
+
# as `dict` and handle `coords` and `attrs` properly as well
|
|
96
|
+
if args and isinstance(args[0], xr.Dataset):
|
|
97
|
+
ds = args[0]
|
|
98
|
+
# Replacee only args[0], `ds`, with `ds.data_vars` as `dict`
|
|
99
|
+
args = (dict(ds.data_vars),) + args[1:]
|
|
100
|
+
# coords not passed positionally
|
|
101
|
+
if len(args) < 2:
|
|
102
|
+
kwargs.setdefault(
|
|
103
|
+
"coords", dict(ds.coords)
|
|
104
|
+
) # Set it as kwarg only if not explicitly provided
|
|
105
|
+
# attrs not passed positionally
|
|
106
|
+
if len(args) < 3:
|
|
107
|
+
kwargs.setdefault(
|
|
108
|
+
"attrs", ds.attrs
|
|
109
|
+
) # Set it as kwarg only if not explicitly provided
|
|
110
|
+
|
|
93
111
|
super().__init__(*args, **kwargs)
|
|
94
112
|
|
|
95
113
|
# declare plotting accessor
|
|
@@ -627,9 +645,9 @@ class UxDataset(xr.Dataset):
|
|
|
627
645
|
"""
|
|
628
646
|
if grid_format == "HEALPix":
|
|
629
647
|
ds = self.rename_dims({"n_face": "cell"})
|
|
630
|
-
return xr.Dataset(ds)
|
|
648
|
+
return xr.Dataset(ds.data_vars, coords=ds.coords, attrs=ds.attrs)
|
|
631
649
|
|
|
632
|
-
return xr.Dataset(self)
|
|
650
|
+
return xr.Dataset(self.data_vars, coords=self.coords, attrs=self.attrs)
|
|
633
651
|
|
|
634
652
|
def get_dual(self):
|
|
635
653
|
"""Compute the dual mesh for a dataset, returns a new dataset object.
|
|
@@ -7,51 +7,114 @@ from uxarray.io._mpas import _mpas_to_ugrid_dims
|
|
|
7
7
|
from uxarray.io._ugrid import _is_ugrid, _read_ugrid
|
|
8
8
|
|
|
9
9
|
|
|
10
|
+
def _is_exodus(dataset: xr.Dataset) -> bool:
|
|
11
|
+
"""Check whether a dataset looks like an Exodus mesh."""
|
|
12
|
+
has_packed_coords = "coord" in dataset
|
|
13
|
+
has_split_coords = {"coordx", "coordy"}.issubset(dataset.variables)
|
|
14
|
+
has_connectivity = any(
|
|
15
|
+
name.startswith("connect") for name in dataset.variables
|
|
16
|
+
) or any("num_nod_per_el" in dim for dim in dataset.dims)
|
|
17
|
+
|
|
18
|
+
return has_connectivity and (has_packed_coords or has_split_coords)
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def _is_scrip(dataset: xr.Dataset) -> bool:
|
|
22
|
+
"""Check whether a dataset looks like an unstructured SCRIP grid."""
|
|
23
|
+
required_vars = {
|
|
24
|
+
"grid_center_lon",
|
|
25
|
+
"grid_center_lat",
|
|
26
|
+
"grid_corner_lon",
|
|
27
|
+
"grid_corner_lat",
|
|
28
|
+
}
|
|
29
|
+
unstructured_markers = {"grid_imask", "grid_rank", "grid_area"}
|
|
30
|
+
|
|
31
|
+
return required_vars.issubset(dataset.variables) and any(
|
|
32
|
+
marker in dataset for marker in unstructured_markers
|
|
33
|
+
)
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def _is_mpas(dataset: xr.Dataset) -> bool:
|
|
37
|
+
"""Check whether a dataset looks like an MPAS grid."""
|
|
38
|
+
if "verticesOnCell" not in dataset:
|
|
39
|
+
return False
|
|
40
|
+
|
|
41
|
+
companion_groups = (
|
|
42
|
+
{"nEdgesOnCell"},
|
|
43
|
+
{"latCell", "lonCell"},
|
|
44
|
+
{"latVertex", "lonVertex"},
|
|
45
|
+
{"xCell", "yCell", "zCell"},
|
|
46
|
+
{"xVertex", "yVertex", "zVertex"},
|
|
47
|
+
)
|
|
48
|
+
|
|
49
|
+
return any(group.issubset(dataset.variables) for group in companion_groups)
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def _is_esmf(dataset: xr.Dataset) -> bool:
|
|
53
|
+
"""Check whether a dataset looks like an ESMF mesh."""
|
|
54
|
+
return "maxNodePElement" in dataset.dims and "elementConn" in dataset
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def _is_geos_cs(dataset: xr.Dataset) -> bool:
|
|
58
|
+
"""Check whether a dataset looks like a GEOS cube-sphere grid."""
|
|
59
|
+
required_dims = {"nf", "YCdim", "XCdim"}
|
|
60
|
+
required_vars = {"corner_lons", "corner_lats"}
|
|
61
|
+
|
|
62
|
+
return required_dims.issubset(dataset.sizes) and required_vars.issubset(
|
|
63
|
+
dataset.variables
|
|
64
|
+
)
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
def _is_icon(dataset: xr.Dataset) -> bool:
|
|
68
|
+
"""Check whether a dataset looks like an ICON grid."""
|
|
69
|
+
required_vars = {"vertex_of_cell", "clon", "clat", "vlon", "vlat"}
|
|
70
|
+
return required_vars.issubset(dataset.variables)
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
def _is_fesom2(dataset: xr.Dataset) -> bool:
|
|
74
|
+
"""Check whether a dataset looks like a FESOM2 grid."""
|
|
75
|
+
return "triag_nodes" in dataset
|
|
76
|
+
|
|
77
|
+
|
|
10
78
|
def _parse_grid_type(dataset):
|
|
11
|
-
"""
|
|
12
|
-
UGrid, SCRIP, Exodus, ESMF, and shape file.
|
|
79
|
+
"""Determine the grid type represented by an input dataset.
|
|
13
80
|
|
|
14
81
|
Parameters
|
|
15
82
|
----------
|
|
16
83
|
dataset : Xarray dataset
|
|
17
|
-
|
|
84
|
+
Xarray dataset containing grid topology information.
|
|
18
85
|
|
|
19
86
|
Returns
|
|
20
87
|
-------
|
|
21
|
-
|
|
22
|
-
|
|
88
|
+
tuple[str, str | None, str | None]
|
|
89
|
+
A 3-tuple of ``(mesh_type, lon_name, lat_name)``. ``mesh_type`` is one
|
|
90
|
+
of ``"Exodus"``, ``"Scrip"``, ``"UGRID"``, ``"MPAS"``, ``"ESMF"``,
|
|
91
|
+
``"GEOS-CS"``, ``"ICON"``, ``"FESOM2"``, or ``"Structured"``. The
|
|
92
|
+
longitude and latitude coordinate names are only returned for structured
|
|
93
|
+
grids and are otherwise ``None``.
|
|
23
94
|
|
|
24
95
|
Raises
|
|
25
96
|
------
|
|
26
97
|
RuntimeError
|
|
27
|
-
|
|
28
|
-
ValueError
|
|
29
|
-
If file is not in UGRID format
|
|
98
|
+
If the dataset format cannot be recognized.
|
|
30
99
|
"""
|
|
31
100
|
|
|
32
101
|
_structured, lon_name, lat_name = _is_structured(dataset)
|
|
33
102
|
|
|
34
|
-
if
|
|
35
|
-
# exodus with coord or coordx
|
|
36
|
-
mesh_type = "Exodus"
|
|
37
|
-
elif "coordx" in dataset:
|
|
103
|
+
if _is_exodus(dataset):
|
|
38
104
|
mesh_type = "Exodus"
|
|
39
|
-
elif
|
|
40
|
-
# scrip with grid_center_lon
|
|
105
|
+
elif _is_scrip(dataset):
|
|
41
106
|
mesh_type = "Scrip"
|
|
42
107
|
elif _is_ugrid(dataset):
|
|
43
|
-
# ugrid topology is present
|
|
44
108
|
mesh_type = "UGRID"
|
|
45
|
-
elif
|
|
109
|
+
elif _is_mpas(dataset):
|
|
46
110
|
mesh_type = "MPAS"
|
|
47
|
-
elif
|
|
111
|
+
elif _is_esmf(dataset):
|
|
48
112
|
mesh_type = "ESMF"
|
|
49
|
-
elif
|
|
50
|
-
# expected dimensions for a GEOS cube sphere grid
|
|
113
|
+
elif _is_geos_cs(dataset):
|
|
51
114
|
mesh_type = "GEOS-CS"
|
|
52
|
-
elif
|
|
115
|
+
elif _is_icon(dataset):
|
|
53
116
|
mesh_type = "ICON"
|
|
54
|
-
elif
|
|
117
|
+
elif _is_fesom2(dataset):
|
|
55
118
|
mesh_type = "FESOM2"
|
|
56
119
|
elif _structured:
|
|
57
120
|
mesh_type = "Structured"
|
|
@@ -0,0 +1,229 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from typing import TYPE_CHECKING
|
|
4
|
+
|
|
5
|
+
if TYPE_CHECKING:
|
|
6
|
+
from uxarray.core.dataarray import UxDataArray
|
|
7
|
+
from uxarray.core.dataset import UxDataset
|
|
8
|
+
from uxarray.grid.grid import Grid
|
|
9
|
+
|
|
10
|
+
from uxarray.remap.bilinear import _bilinear
|
|
11
|
+
from uxarray.remap.inverse_distance_weighted import _inverse_distance_weighted_remap
|
|
12
|
+
from uxarray.remap.nearest_neighbor import _nearest_neighbor_remap
|
|
13
|
+
|
|
14
|
+
_VALID_BACKENDS = ("uxarray", "yac")
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
def _validate_backend(backend: str) -> None:
|
|
18
|
+
if backend not in _VALID_BACKENDS:
|
|
19
|
+
raise ValueError(
|
|
20
|
+
f"Invalid backend '{backend}'. Expected one of {_VALID_BACKENDS}."
|
|
21
|
+
)
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
class RemapAccessor:
|
|
25
|
+
"""Expose remapping methods on UxDataArray and UxDataset objects."""
|
|
26
|
+
|
|
27
|
+
def __init__(self, ux_obj: UxDataArray | UxDataset):
|
|
28
|
+
self.ux_obj = ux_obj
|
|
29
|
+
|
|
30
|
+
def __repr__(self) -> str:
|
|
31
|
+
prefix = f"<{type(self.ux_obj).__name__}.remap>\n"
|
|
32
|
+
return (
|
|
33
|
+
prefix
|
|
34
|
+
+ "Supported methods:\n"
|
|
35
|
+
+ " • nearest_neighbor(destination_grid, remap_to='faces')\n"
|
|
36
|
+
+ " • inverse_distance_weighted(destination_grid, remap_to='faces', power=2, k=8)\n"
|
|
37
|
+
)
|
|
38
|
+
|
|
39
|
+
def __call__(
|
|
40
|
+
self,
|
|
41
|
+
*args,
|
|
42
|
+
backend: str = "uxarray",
|
|
43
|
+
yac_method: str | None = None,
|
|
44
|
+
yac_options: dict | None = None,
|
|
45
|
+
**kwargs,
|
|
46
|
+
) -> UxDataArray | UxDataset:
|
|
47
|
+
"""
|
|
48
|
+
Shortcut for nearest-neighbor remapping.
|
|
49
|
+
|
|
50
|
+
Calling `.remap(...)` with no explicit method will invoke
|
|
51
|
+
`nearest_neighbor(...)`.
|
|
52
|
+
|
|
53
|
+
When ``backend="yac"``, this generic entrypoint can also be used to
|
|
54
|
+
select a YAC-specific interpolation method through ``yac_method``.
|
|
55
|
+
"""
|
|
56
|
+
nn_kwargs: dict = {"backend": backend, "yac_options": yac_options}
|
|
57
|
+
if yac_method is not None:
|
|
58
|
+
nn_kwargs["yac_method"] = yac_method
|
|
59
|
+
return self.nearest_neighbor(*args, **nn_kwargs, **kwargs)
|
|
60
|
+
|
|
61
|
+
def nearest_neighbor(
|
|
62
|
+
self,
|
|
63
|
+
destination_grid: Grid,
|
|
64
|
+
remap_to: str = "faces",
|
|
65
|
+
backend: str = "uxarray",
|
|
66
|
+
yac_method: str | None = "nnn",
|
|
67
|
+
yac_options: dict | None = None,
|
|
68
|
+
**kwargs,
|
|
69
|
+
) -> UxDataArray | UxDataset:
|
|
70
|
+
"""
|
|
71
|
+
Perform nearest-neighbor remapping.
|
|
72
|
+
|
|
73
|
+
Each destination point takes the value of its closest source point.
|
|
74
|
+
|
|
75
|
+
Parameters
|
|
76
|
+
----------
|
|
77
|
+
destination_grid : Grid
|
|
78
|
+
The UXarray grid to which data will be interpolated.
|
|
79
|
+
remap_to : {'nodes', 'edges', 'faces'}, default='faces'
|
|
80
|
+
Which grid element receives the remapped values.
|
|
81
|
+
|
|
82
|
+
backend : {'uxarray', 'yac'}, default='uxarray'
|
|
83
|
+
Remapping backend to use. When set to 'yac', requires YAC to be
|
|
84
|
+
available on PYTHONPATH.
|
|
85
|
+
yac_method : {'nnn', 'average', 'conservative'}, optional
|
|
86
|
+
YAC interpolation method. Defaults to 'nnn' when backend='yac'.
|
|
87
|
+
yac_options : dict, optional
|
|
88
|
+
YAC interpolation configuration options.
|
|
89
|
+
|
|
90
|
+
Returns
|
|
91
|
+
-------
|
|
92
|
+
UxDataArray or UxDataset
|
|
93
|
+
A new object with data mapped onto `destination_grid`.
|
|
94
|
+
|
|
95
|
+
Notes
|
|
96
|
+
-----
|
|
97
|
+
When ``backend="yac"``, remapping uses YAC's low-level ``yac.core``
|
|
98
|
+
Python bindings. See the YAC documentation and installation guide:
|
|
99
|
+
|
|
100
|
+
- https://dkrz-sw.gitlab-pages.dkrz.de/yac/
|
|
101
|
+
- https://dkrz-sw.gitlab-pages.dkrz.de/yac/d1/d9f/installing_yac.html
|
|
102
|
+
"""
|
|
103
|
+
|
|
104
|
+
_validate_backend(backend)
|
|
105
|
+
if backend == "yac":
|
|
106
|
+
from uxarray.remap.yac import _yac_remap
|
|
107
|
+
|
|
108
|
+
yac_kwargs = yac_options or {}
|
|
109
|
+
return _yac_remap(
|
|
110
|
+
self.ux_obj, destination_grid, remap_to, yac_method, yac_kwargs
|
|
111
|
+
)
|
|
112
|
+
return _nearest_neighbor_remap(self.ux_obj, destination_grid, remap_to)
|
|
113
|
+
|
|
114
|
+
def inverse_distance_weighted(
|
|
115
|
+
self,
|
|
116
|
+
destination_grid: Grid,
|
|
117
|
+
remap_to: str = "faces",
|
|
118
|
+
power=2,
|
|
119
|
+
k=8,
|
|
120
|
+
backend: str = "uxarray",
|
|
121
|
+
yac_method: str | None = None,
|
|
122
|
+
yac_options: dict | None = None,
|
|
123
|
+
**kwargs,
|
|
124
|
+
) -> UxDataArray | UxDataset:
|
|
125
|
+
"""
|
|
126
|
+
Perform inverse-distance-weighted (IDW) remapping.
|
|
127
|
+
|
|
128
|
+
Each destination point is a weighted average of nearby source points,
|
|
129
|
+
with weights proportional to 1/(distance**power).
|
|
130
|
+
|
|
131
|
+
Parameters
|
|
132
|
+
----------
|
|
133
|
+
destination_grid : Grid
|
|
134
|
+
The UXarray grid to which data will be interpolated.
|
|
135
|
+
remap_to : {'nodes', 'edges', 'faces'}, default='faces'
|
|
136
|
+
Which grid element receives the remapped values.
|
|
137
|
+
power : int, default=2
|
|
138
|
+
Exponent controlling distance decay. Larger values make the
|
|
139
|
+
interpolation more local.
|
|
140
|
+
k : int, default=8
|
|
141
|
+
Number of nearest source points to include in the weighted average.
|
|
142
|
+
|
|
143
|
+
backend : {'uxarray', 'yac'}, default='uxarray'
|
|
144
|
+
Remapping backend to use. When set to 'yac', requires YAC to be
|
|
145
|
+
available on PYTHONPATH.
|
|
146
|
+
yac_method : {'nnn', 'conservative'}, optional
|
|
147
|
+
YAC interpolation method. Required when backend='yac'.
|
|
148
|
+
yac_options : dict, optional
|
|
149
|
+
YAC interpolation configuration options.
|
|
150
|
+
|
|
151
|
+
Returns
|
|
152
|
+
-------
|
|
153
|
+
UxDataArray or UxDataset
|
|
154
|
+
A new object with data mapped onto `destination_grid`.
|
|
155
|
+
|
|
156
|
+
Notes
|
|
157
|
+
-----
|
|
158
|
+
When ``backend="yac"``, this method delegates to YAC's ``average``
|
|
159
|
+
interpolation method through the low-level ``yac.core`` Python
|
|
160
|
+
bindings. See the YAC documentation and installation guide:
|
|
161
|
+
|
|
162
|
+
- https://dkrz-sw.gitlab-pages.dkrz.de/yac/
|
|
163
|
+
- https://dkrz-sw.gitlab-pages.dkrz.de/yac/d1/d9f/installing_yac.html
|
|
164
|
+
"""
|
|
165
|
+
|
|
166
|
+
_validate_backend(backend)
|
|
167
|
+
if backend == "yac":
|
|
168
|
+
raise NotImplementedError(
|
|
169
|
+
"inverse_distance_weighted with backend='yac' is not currently "
|
|
170
|
+
"exposed through the UXarray YAC accessor. "
|
|
171
|
+
"Use backend='uxarray' for IDW, or use the YAC backend through "
|
|
172
|
+
".remap(..., backend='yac', yac_method=..., yac_options=...)."
|
|
173
|
+
)
|
|
174
|
+
return _inverse_distance_weighted_remap(
|
|
175
|
+
self.ux_obj, destination_grid, remap_to, power, k
|
|
176
|
+
)
|
|
177
|
+
|
|
178
|
+
def bilinear(
|
|
179
|
+
self,
|
|
180
|
+
destination_grid: Grid,
|
|
181
|
+
remap_to: str = "faces",
|
|
182
|
+
backend: str = "uxarray",
|
|
183
|
+
yac_method: str | None = "average",
|
|
184
|
+
yac_options: dict | None = None,
|
|
185
|
+
**kwargs,
|
|
186
|
+
) -> UxDataArray | UxDataset:
|
|
187
|
+
"""
|
|
188
|
+
Perform bilinear remapping.
|
|
189
|
+
|
|
190
|
+
Parameters
|
|
191
|
+
---------
|
|
192
|
+
destination_grid : Grid
|
|
193
|
+
Destination Grid for remapping
|
|
194
|
+
remap_to : {'nodes', 'edges', 'faces'}, default='faces'
|
|
195
|
+
Which grid element receives the remapped values.
|
|
196
|
+
|
|
197
|
+
backend : {'uxarray', 'yac'}, default='uxarray'
|
|
198
|
+
Remapping backend to use. When set to 'yac', bilinear remapping is
|
|
199
|
+
routed through YAC's average interpolation.
|
|
200
|
+
yac_method : {'average'}, optional
|
|
201
|
+
YAC interpolation method for the bilinear convenience wrapper.
|
|
202
|
+
Only ``'average'`` is supported here.
|
|
203
|
+
yac_options : dict, optional
|
|
204
|
+
YAC interpolation configuration options for the average method.
|
|
205
|
+
|
|
206
|
+
Returns
|
|
207
|
+
-------
|
|
208
|
+
UxDataArray or UxDataset
|
|
209
|
+
A new object with data mapped onto `destination_grid`.
|
|
210
|
+
"""
|
|
211
|
+
|
|
212
|
+
_validate_backend(backend)
|
|
213
|
+
if backend == "yac":
|
|
214
|
+
from uxarray.remap.yac import _yac_remap
|
|
215
|
+
|
|
216
|
+
if yac_method not in (None, "average"):
|
|
217
|
+
raise ValueError(
|
|
218
|
+
"bilinear with backend='yac' only supports yac_method='average'. "
|
|
219
|
+
"Use .remap(..., backend='yac', yac_method=...) for other YAC methods."
|
|
220
|
+
)
|
|
221
|
+
yac_kwargs = yac_options or {}
|
|
222
|
+
return _yac_remap(
|
|
223
|
+
self.ux_obj,
|
|
224
|
+
destination_grid,
|
|
225
|
+
remap_to,
|
|
226
|
+
yac_method or "average",
|
|
227
|
+
yac_kwargs,
|
|
228
|
+
)
|
|
229
|
+
return _bilinear(self.ux_obj, destination_grid, remap_to)
|
|
@@ -0,0 +1,407 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import importlib
|
|
4
|
+
import importlib.util
|
|
5
|
+
import sys
|
|
6
|
+
from dataclasses import dataclass
|
|
7
|
+
from pathlib import Path
|
|
8
|
+
from types import ModuleType
|
|
9
|
+
from typing import Any
|
|
10
|
+
from uuid import uuid4
|
|
11
|
+
|
|
12
|
+
import numpy as np
|
|
13
|
+
|
|
14
|
+
import uxarray.core.dataarray
|
|
15
|
+
from uxarray.remap.utils import (
|
|
16
|
+
LABEL_TO_COORD,
|
|
17
|
+
_assert_dimension,
|
|
18
|
+
_construct_remapped_ds,
|
|
19
|
+
_get_remap_dims,
|
|
20
|
+
_to_dataset,
|
|
21
|
+
)
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
class YacNotAvailableError(RuntimeError):
|
|
25
|
+
"""Raised when the YAC backend is requested but unavailable."""
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
@dataclass
|
|
29
|
+
class _YacOptions:
|
|
30
|
+
method: str
|
|
31
|
+
kwargs: dict[str, Any]
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def _load_yac_core_from_file() -> ModuleType | None:
|
|
35
|
+
if "yac.core" in sys.modules:
|
|
36
|
+
return sys.modules["yac.core"]
|
|
37
|
+
|
|
38
|
+
for path_entry in sys.path:
|
|
39
|
+
pkg_dir = Path(path_entry) / "yac"
|
|
40
|
+
if not pkg_dir.is_dir():
|
|
41
|
+
continue
|
|
42
|
+
|
|
43
|
+
matches = sorted(pkg_dir.glob("core*.so"))
|
|
44
|
+
if not matches:
|
|
45
|
+
matches = sorted(pkg_dir.glob("core*.pyd"))
|
|
46
|
+
if not matches:
|
|
47
|
+
continue
|
|
48
|
+
|
|
49
|
+
pkg = sys.modules.get("yac")
|
|
50
|
+
if pkg is None:
|
|
51
|
+
pkg = ModuleType("yac")
|
|
52
|
+
sys.modules["yac"] = pkg
|
|
53
|
+
pkg.__path__ = [str(pkg_dir)]
|
|
54
|
+
|
|
55
|
+
spec = importlib.util.spec_from_file_location("yac.core", matches[0])
|
|
56
|
+
if spec is None or spec.loader is None:
|
|
57
|
+
continue
|
|
58
|
+
|
|
59
|
+
module = importlib.util.module_from_spec(spec)
|
|
60
|
+
sys.modules["yac.core"] = module
|
|
61
|
+
spec.loader.exec_module(module)
|
|
62
|
+
setattr(pkg, "core", module)
|
|
63
|
+
return module
|
|
64
|
+
|
|
65
|
+
return None
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
def _import_yac():
|
|
69
|
+
module = _load_yac_core_from_file()
|
|
70
|
+
if module is not None:
|
|
71
|
+
return module
|
|
72
|
+
|
|
73
|
+
try:
|
|
74
|
+
return importlib.import_module("yac.core")
|
|
75
|
+
except Exception as exc: # pragma: no cover - fallback depends on local install
|
|
76
|
+
raise YacNotAvailableError(
|
|
77
|
+
"YAC backend requested but 'yac.core' is not available. "
|
|
78
|
+
"Build YAC with Python bindings and ensure its site-packages and "
|
|
79
|
+
"shared libraries are discoverable."
|
|
80
|
+
) from exc
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
def _normalize_yac_method(yac_method: str | None) -> _YacOptions:
|
|
84
|
+
if not yac_method:
|
|
85
|
+
raise ValueError(
|
|
86
|
+
"backend='yac' requires yac_method to be set to 'nnn', 'average', or 'conservative'."
|
|
87
|
+
)
|
|
88
|
+
method = yac_method.lower()
|
|
89
|
+
if method not in {"nnn", "average", "conservative"}:
|
|
90
|
+
raise ValueError(f"Unsupported YAC method: {yac_method!r}")
|
|
91
|
+
return _YacOptions(method=method, kwargs={})
|
|
92
|
+
|
|
93
|
+
|
|
94
|
+
def _get_location(yac_core, dim: str):
|
|
95
|
+
mapping = {
|
|
96
|
+
"n_face": yac_core.yac_location.YAC_LOC_CELL,
|
|
97
|
+
"n_node": yac_core.yac_location.YAC_LOC_CORNER,
|
|
98
|
+
"n_edge": yac_core.yac_location.YAC_LOC_EDGE,
|
|
99
|
+
}
|
|
100
|
+
try:
|
|
101
|
+
return mapping[dim]
|
|
102
|
+
except KeyError as exc:
|
|
103
|
+
raise ValueError(f"Unsupported remap dimension for YAC: {dim!r}") from exc
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
def _get_lon_lat(grid, dim: str) -> tuple[np.ndarray, np.ndarray]:
|
|
107
|
+
attr_map = {
|
|
108
|
+
"n_face": ("face_lon", "face_lat"),
|
|
109
|
+
"n_node": ("node_lon", "node_lat"),
|
|
110
|
+
"n_edge": ("edge_lon", "edge_lat"),
|
|
111
|
+
}
|
|
112
|
+
try:
|
|
113
|
+
lon_attr, lat_attr = attr_map[dim]
|
|
114
|
+
except KeyError as exc:
|
|
115
|
+
raise ValueError(f"Unsupported remap dimension for YAC: {dim!r}") from exc
|
|
116
|
+
|
|
117
|
+
lon = getattr(grid, lon_attr, None)
|
|
118
|
+
lat = getattr(grid, lat_attr, None)
|
|
119
|
+
if lon is None or lat is None:
|
|
120
|
+
raise ValueError(
|
|
121
|
+
f"Grid does not provide {lon_attr}/{lat_attr} required for YAC remapping."
|
|
122
|
+
)
|
|
123
|
+
return np.deg2rad(np.asarray(lon.values, dtype=np.float64)), np.deg2rad(
|
|
124
|
+
np.asarray(lat.values, dtype=np.float64)
|
|
125
|
+
)
|
|
126
|
+
|
|
127
|
+
|
|
128
|
+
def _coerce_enum(enum_type, value: Any):
|
|
129
|
+
if not isinstance(value, str):
|
|
130
|
+
return value
|
|
131
|
+
|
|
132
|
+
normalized = value.upper()
|
|
133
|
+
for member in enum_type:
|
|
134
|
+
if member.name == normalized or member.name.endswith(f"_{normalized}"):
|
|
135
|
+
return member
|
|
136
|
+
|
|
137
|
+
raise ValueError(f"Unsupported value {value!r} for enum {enum_type.__name__}.")
|
|
138
|
+
|
|
139
|
+
|
|
140
|
+
class _YacRemapper:
|
|
141
|
+
"""Build and reuse YAC interpolation weights for one source dimension.
|
|
142
|
+
|
|
143
|
+
Each instance owns the YAC source/target field registration for a single
|
|
144
|
+
source location type (faces, nodes, or edges) and one requested YAC method.
|
|
145
|
+
The resulting weights can then be applied repeatedly to batches of values
|
|
146
|
+
that share the same source dimension.
|
|
147
|
+
"""
|
|
148
|
+
|
|
149
|
+
def __init__(
|
|
150
|
+
self,
|
|
151
|
+
src_grid,
|
|
152
|
+
tgt_grid,
|
|
153
|
+
src_dim: str,
|
|
154
|
+
tgt_dim: str,
|
|
155
|
+
yac_method: str,
|
|
156
|
+
yac_kwargs: dict[str, Any],
|
|
157
|
+
):
|
|
158
|
+
yac_core = _import_yac()
|
|
159
|
+
self._frac_mask_fallback_value = yac_kwargs.get("frac_mask_fallback_value")
|
|
160
|
+
self._src_location = _get_location(yac_core, src_dim)
|
|
161
|
+
self._tgt_location = _get_location(yac_core, tgt_dim)
|
|
162
|
+
|
|
163
|
+
define_edges = "n_edge" in (src_dim, tgt_dim)
|
|
164
|
+
unique = uuid4().hex
|
|
165
|
+
self._src_grid = yac_core.BasicGrid.from_uxgrid(
|
|
166
|
+
f"uxarray_src_{unique}",
|
|
167
|
+
src_grid,
|
|
168
|
+
def_edges=define_edges,
|
|
169
|
+
)
|
|
170
|
+
self._tgt_grid = yac_core.BasicGrid.from_uxgrid(
|
|
171
|
+
f"uxarray_tgt_{unique}",
|
|
172
|
+
tgt_grid,
|
|
173
|
+
def_edges=define_edges,
|
|
174
|
+
)
|
|
175
|
+
src_lon, src_lat = _get_lon_lat(src_grid, src_dim)
|
|
176
|
+
tgt_lon, tgt_lat = _get_lon_lat(tgt_grid, tgt_dim)
|
|
177
|
+
|
|
178
|
+
self._src_field = yac_core.InterpField(
|
|
179
|
+
self._src_grid.add_coordinates(self._src_location, src_lon, src_lat)
|
|
180
|
+
)
|
|
181
|
+
self._tgt_field = yac_core.InterpField(
|
|
182
|
+
self._tgt_grid.add_coordinates(self._tgt_location, tgt_lon, tgt_lat)
|
|
183
|
+
)
|
|
184
|
+
|
|
185
|
+
stack = yac_core.InterpolationStack()
|
|
186
|
+
if yac_method == "nnn":
|
|
187
|
+
weight_type = _coerce_enum(
|
|
188
|
+
yac_core.yac_interp_nnn_weight_type,
|
|
189
|
+
yac_kwargs.get("reduction_type", yac_kwargs.get("nnn_type")),
|
|
190
|
+
)
|
|
191
|
+
if weight_type is None:
|
|
192
|
+
weight_type = yac_core.yac_interp_nnn_weight_type.YAC_INTERP_NNN_AVG
|
|
193
|
+
stack.add_nnn(
|
|
194
|
+
nnn_type=weight_type,
|
|
195
|
+
n=yac_kwargs.get("n", 1),
|
|
196
|
+
max_search_distance=yac_kwargs.get("max_search_distance", 0.0),
|
|
197
|
+
scale=yac_kwargs.get("scale", 1.0),
|
|
198
|
+
)
|
|
199
|
+
elif yac_method == "average":
|
|
200
|
+
reduction_type = _coerce_enum(
|
|
201
|
+
yac_core.yac_interp_avg_weight_type,
|
|
202
|
+
yac_kwargs.get("reduction_type", yac_kwargs.get("weight_type")),
|
|
203
|
+
)
|
|
204
|
+
if reduction_type is None:
|
|
205
|
+
reduction_type = (
|
|
206
|
+
yac_core.yac_interp_avg_weight_type.YAC_INTERP_AVG_ARITHMETIC
|
|
207
|
+
)
|
|
208
|
+
stack.add_average(
|
|
209
|
+
reduction_type=reduction_type,
|
|
210
|
+
partial_coverage=yac_kwargs.get("partial_coverage", False),
|
|
211
|
+
)
|
|
212
|
+
elif yac_method == "conservative":
|
|
213
|
+
normalisation = _coerce_enum(
|
|
214
|
+
yac_core.yac_interp_method_conserv_normalisation,
|
|
215
|
+
yac_kwargs.get("normalisation"),
|
|
216
|
+
)
|
|
217
|
+
if normalisation is None:
|
|
218
|
+
normalisation = yac_core.yac_interp_method_conserv_normalisation.YAC_INTERP_CONSERV_DESTAREA
|
|
219
|
+
stack.add_conservative(
|
|
220
|
+
order=yac_kwargs.get("order", 1),
|
|
221
|
+
enforced_conserv=yac_kwargs.get("enforced_conserv", False),
|
|
222
|
+
partial_coverage=yac_kwargs.get("partial_coverage", False),
|
|
223
|
+
normalisation=normalisation,
|
|
224
|
+
)
|
|
225
|
+
fixed_value = yac_kwargs.get("fixed_value", 0.0)
|
|
226
|
+
if fixed_value is not None:
|
|
227
|
+
stack.add_fixed(float(fixed_value))
|
|
228
|
+
|
|
229
|
+
self._weights = yac_core.compute_weights(
|
|
230
|
+
stack,
|
|
231
|
+
self._src_field,
|
|
232
|
+
self._tgt_field,
|
|
233
|
+
)
|
|
234
|
+
self._interpolations: dict[int, Any] = {}
|
|
235
|
+
self._src_size = self._src_grid.get_data_size(self._src_location)
|
|
236
|
+
self._tgt_size = self._tgt_grid.get_data_size(self._tgt_location)
|
|
237
|
+
|
|
238
|
+
def apply(
|
|
239
|
+
self, values: np.ndarray, frac_mask: np.ndarray | None = None
|
|
240
|
+
) -> np.ndarray:
|
|
241
|
+
"""Apply the pre-computed interpolation weights to *values*.
|
|
242
|
+
|
|
243
|
+
The interpolation method (NNN or conservative) is determined by
|
|
244
|
+
*yac_method* passed to the constructor and is fixed for the lifetime of
|
|
245
|
+
this remapper instance. This method simply executes the weight
|
|
246
|
+
application; it does not select or alter the interpolation algorithm.
|
|
247
|
+
|
|
248
|
+
Parameters
|
|
249
|
+
----------
|
|
250
|
+
values : np.ndarray
|
|
251
|
+
1-D or 2-D array of source-grid values. The trailing dimension must
|
|
252
|
+
equal the number of source points registered with YAC
|
|
253
|
+
(``self._src_size``). When 2-D, the leading dimension is treated as
|
|
254
|
+
the YAC collection size and is remapped in one batched call.
|
|
255
|
+
frac_mask : np.ndarray, optional
|
|
256
|
+
Optional fractional source mask with the same shape as ``values``.
|
|
257
|
+
When provided, it is forwarded to YAC's interpolation call.
|
|
258
|
+
|
|
259
|
+
Returns
|
|
260
|
+
-------
|
|
261
|
+
np.ndarray
|
|
262
|
+
Array of remapped values on the destination grid with the same
|
|
263
|
+
number of leading collections as the input.
|
|
264
|
+
"""
|
|
265
|
+
values = np.ascontiguousarray(values, dtype=np.float64)
|
|
266
|
+
if values.ndim == 1:
|
|
267
|
+
values = values.reshape(1, -1)
|
|
268
|
+
elif values.ndim != 2:
|
|
269
|
+
raise ValueError(
|
|
270
|
+
f"YAC remap expects a 1-D or 2-D array, got {values.ndim}-D input."
|
|
271
|
+
)
|
|
272
|
+
if values.shape[1] != self._src_size:
|
|
273
|
+
raise ValueError(
|
|
274
|
+
f"YAC remap expects {self._src_size} values, got {values.shape[1]}."
|
|
275
|
+
)
|
|
276
|
+
|
|
277
|
+
if frac_mask is not None:
|
|
278
|
+
frac_mask = np.ascontiguousarray(frac_mask, dtype=np.float64)
|
|
279
|
+
if frac_mask.ndim == 1:
|
|
280
|
+
frac_mask = frac_mask.reshape(1, -1)
|
|
281
|
+
elif frac_mask.ndim != 2:
|
|
282
|
+
raise ValueError(
|
|
283
|
+
"YAC fractional mask expects a 1-D or 2-D array, "
|
|
284
|
+
f"got {frac_mask.ndim}-D input."
|
|
285
|
+
)
|
|
286
|
+
if frac_mask.shape != values.shape:
|
|
287
|
+
raise ValueError(
|
|
288
|
+
"YAC fractional mask must match remap input shape. "
|
|
289
|
+
f"Got mask shape {frac_mask.shape} and value shape {values.shape}."
|
|
290
|
+
)
|
|
291
|
+
|
|
292
|
+
collection_size = values.shape[0]
|
|
293
|
+
interpolation = self._interpolations.get(collection_size)
|
|
294
|
+
if interpolation is None:
|
|
295
|
+
interpolation = self._weights.get_interpolation(
|
|
296
|
+
collection_size=collection_size,
|
|
297
|
+
frac_mask_fallback_value=self._frac_mask_fallback_value,
|
|
298
|
+
)
|
|
299
|
+
self._interpolations[collection_size] = interpolation
|
|
300
|
+
|
|
301
|
+
out = (
|
|
302
|
+
interpolation(values, frac_mask=frac_mask)
|
|
303
|
+
if frac_mask is not None
|
|
304
|
+
else interpolation(values)
|
|
305
|
+
)
|
|
306
|
+
return np.asarray(out, dtype=np.float64)
|
|
307
|
+
|
|
308
|
+
|
|
309
|
+
def _prepare_frac_mask(frac_mask, da_t, src_values, src_dim: str) -> np.ndarray:
|
|
310
|
+
"""Normalize a fractional mask to the flattened shape expected by YAC."""
|
|
311
|
+
if hasattr(frac_mask, "dims"):
|
|
312
|
+
other_dims = [d for d in da_t.dims if d != src_dim]
|
|
313
|
+
frac_mask_values = np.asarray(frac_mask.transpose(*other_dims, src_dim).values)
|
|
314
|
+
else:
|
|
315
|
+
frac_mask_values = np.asarray(frac_mask)
|
|
316
|
+
|
|
317
|
+
if frac_mask_values.shape != src_values.shape:
|
|
318
|
+
raise ValueError(
|
|
319
|
+
"YAC fractional mask must match the remapped source variable shape. "
|
|
320
|
+
f"Got mask shape {frac_mask_values.shape} and source shape {src_values.shape}."
|
|
321
|
+
)
|
|
322
|
+
return frac_mask_values.reshape(-1, frac_mask_values.shape[-1])
|
|
323
|
+
|
|
324
|
+
|
|
325
|
+
def _yac_remap(source, destination_grid, remap_to: str, yac_method: str, yac_kwargs):
|
|
326
|
+
"""Remap a UXarray object through YAC and reconstruct the UXarray result.
|
|
327
|
+
|
|
328
|
+
This is the main integration boundary between the public UXarray remap
|
|
329
|
+
accessor and the lower-level ``yac.core`` bindings. It normalizes the
|
|
330
|
+
requested YAC method, validates method-specific constraints, batches each
|
|
331
|
+
remapped variable by its source dimension, and returns a remapped
|
|
332
|
+
``UxDataArray`` or ``UxDataset`` with UXarray metadata preserved.
|
|
333
|
+
"""
|
|
334
|
+
_assert_dimension(remap_to)
|
|
335
|
+
destination_dim = LABEL_TO_COORD[remap_to]
|
|
336
|
+
options = _normalize_yac_method(yac_method)
|
|
337
|
+
options.kwargs.update(yac_kwargs or {})
|
|
338
|
+
ds, is_da, name = _to_dataset(source)
|
|
339
|
+
dims_to_remap = _get_remap_dims(ds)
|
|
340
|
+
|
|
341
|
+
if options.method == "conservative":
|
|
342
|
+
if destination_dim != "n_face":
|
|
343
|
+
raise ValueError(
|
|
344
|
+
"YAC conservative remapping requires the destination to be "
|
|
345
|
+
"face-centered (remap_to='faces'). "
|
|
346
|
+
f"Got remap_to={remap_to!r} which maps to dimension {destination_dim!r}."
|
|
347
|
+
)
|
|
348
|
+
non_face_src = dims_to_remap - {"n_face"}
|
|
349
|
+
if non_face_src:
|
|
350
|
+
raise ValueError(
|
|
351
|
+
"YAC conservative remapping requires all source data to be "
|
|
352
|
+
f"face-centered (dimension 'n_face'). "
|
|
353
|
+
f"Found non-face source dimension(s): {non_face_src}. "
|
|
354
|
+
"Use yac_method='nnn' for node- or edge-centered data."
|
|
355
|
+
)
|
|
356
|
+
remappers: dict[str, _YacRemapper] = {}
|
|
357
|
+
remapped_vars = {}
|
|
358
|
+
|
|
359
|
+
for src_dim in dims_to_remap:
|
|
360
|
+
remappers[src_dim] = _YacRemapper(
|
|
361
|
+
ds.uxgrid,
|
|
362
|
+
destination_grid,
|
|
363
|
+
src_dim,
|
|
364
|
+
destination_dim,
|
|
365
|
+
options.method,
|
|
366
|
+
options.kwargs,
|
|
367
|
+
)
|
|
368
|
+
|
|
369
|
+
for var_name, da in ds.data_vars.items():
|
|
370
|
+
src_dim = next((d for d in da.dims if d in dims_to_remap), None)
|
|
371
|
+
if src_dim is None:
|
|
372
|
+
remapped_vars[var_name] = da
|
|
373
|
+
continue
|
|
374
|
+
|
|
375
|
+
other_dims = [d for d in da.dims if d != src_dim]
|
|
376
|
+
da_t = da.transpose(*other_dims, src_dim)
|
|
377
|
+
src_values = np.asarray(da_t.values)
|
|
378
|
+
flat_src = src_values.reshape(-1, src_values.shape[-1])
|
|
379
|
+
frac_masks = yac_kwargs.get("frac_masks")
|
|
380
|
+
frac_mask = (
|
|
381
|
+
frac_masks.get(var_name)
|
|
382
|
+
if isinstance(frac_masks, dict) and var_name in frac_masks
|
|
383
|
+
else yac_kwargs.get("frac_mask")
|
|
384
|
+
)
|
|
385
|
+
flat_frac_mask = None
|
|
386
|
+
if frac_mask is not None:
|
|
387
|
+
flat_frac_mask = _prepare_frac_mask(frac_mask, da_t, src_values, src_dim)
|
|
388
|
+
remapper = remappers[src_dim]
|
|
389
|
+
out_flat = remapper.apply(flat_src, frac_mask=flat_frac_mask)
|
|
390
|
+
|
|
391
|
+
out_shape = src_values.shape[:-1] + (remapper._tgt_size,)
|
|
392
|
+
out_values = out_flat.reshape(out_shape)
|
|
393
|
+
coords = {dim: da.coords[dim] for dim in other_dims if dim in da.coords}
|
|
394
|
+
da_out = uxarray.core.dataarray.UxDataArray(
|
|
395
|
+
out_values,
|
|
396
|
+
dims=other_dims + [destination_dim],
|
|
397
|
+
coords=coords,
|
|
398
|
+
name=da.name,
|
|
399
|
+
attrs=da.attrs,
|
|
400
|
+
uxgrid=destination_grid,
|
|
401
|
+
)
|
|
402
|
+
remapped_vars[var_name] = da_out
|
|
403
|
+
|
|
404
|
+
ds_remapped = _construct_remapped_ds(
|
|
405
|
+
source, remapped_vars, destination_grid, remap_to
|
|
406
|
+
)
|
|
407
|
+
return ds_remapped[name] if is_da else ds_remapped
|
|
@@ -80,6 +80,7 @@ uxarray/remap/inverse_distance_weighted.py
|
|
|
80
80
|
uxarray/remap/nearest_neighbor.py
|
|
81
81
|
uxarray/remap/spatial_coords_remap.py
|
|
82
82
|
uxarray/remap/utils.py
|
|
83
|
+
uxarray/remap/yac.py
|
|
83
84
|
uxarray/subset/__init__.py
|
|
84
85
|
uxarray/subset/dataarray_accessor.py
|
|
85
86
|
uxarray/subset/grid_accessor.py
|
|
@@ -1,112 +0,0 @@
|
|
|
1
|
-
from __future__ import annotations
|
|
2
|
-
|
|
3
|
-
from typing import TYPE_CHECKING
|
|
4
|
-
|
|
5
|
-
if TYPE_CHECKING:
|
|
6
|
-
from uxarray.core.dataarray import UxDataArray
|
|
7
|
-
from uxarray.core.dataset import UxDataset
|
|
8
|
-
from uxarray.grid.grid import Grid
|
|
9
|
-
|
|
10
|
-
from uxarray.remap.bilinear import _bilinear
|
|
11
|
-
from uxarray.remap.inverse_distance_weighted import _inverse_distance_weighted_remap
|
|
12
|
-
from uxarray.remap.nearest_neighbor import _nearest_neighbor_remap
|
|
13
|
-
|
|
14
|
-
|
|
15
|
-
class RemapAccessor:
|
|
16
|
-
"""Expose remapping methods on UxDataArray and UxDataset objects."""
|
|
17
|
-
|
|
18
|
-
def __init__(self, ux_obj: UxDataArray | UxDataset):
|
|
19
|
-
self.ux_obj = ux_obj
|
|
20
|
-
|
|
21
|
-
def __repr__(self) -> str:
|
|
22
|
-
prefix = f"<{type(self.ux_obj).__name__}.remap>\n"
|
|
23
|
-
return (
|
|
24
|
-
prefix
|
|
25
|
-
+ "Supported methods:\n"
|
|
26
|
-
+ " • nearest_neighbor(destination_grid, remap_to='faces')\n"
|
|
27
|
-
+ " • inverse_distance_weighted(destination_grid, remap_to='faces', power=2, k=8)\n"
|
|
28
|
-
)
|
|
29
|
-
|
|
30
|
-
def __call__(self, *args, **kwargs) -> UxDataArray | UxDataset:
|
|
31
|
-
"""
|
|
32
|
-
Shortcut for nearest-neighbor remapping.
|
|
33
|
-
|
|
34
|
-
Calling `.remap(...)` with no explicit method will invoke
|
|
35
|
-
`nearest_neighbor(...)`.
|
|
36
|
-
"""
|
|
37
|
-
return self.nearest_neighbor(*args, **kwargs)
|
|
38
|
-
|
|
39
|
-
def nearest_neighbor(
|
|
40
|
-
self, destination_grid: Grid, remap_to: str = "faces", **kwargs
|
|
41
|
-
) -> UxDataArray | UxDataset:
|
|
42
|
-
"""
|
|
43
|
-
Perform nearest-neighbor remapping.
|
|
44
|
-
|
|
45
|
-
Each destination point takes the value of its closest source point.
|
|
46
|
-
|
|
47
|
-
Parameters
|
|
48
|
-
----------
|
|
49
|
-
destination_grid : Grid
|
|
50
|
-
The UXarray grid to which data will be interpolated.
|
|
51
|
-
remap_to : {'nodes', 'edges', 'faces'}, default='faces'
|
|
52
|
-
Which grid element receives the remapped values.
|
|
53
|
-
|
|
54
|
-
Returns
|
|
55
|
-
-------
|
|
56
|
-
UxDataArray or UxDataset
|
|
57
|
-
A new object with data mapped onto `destination_grid`.
|
|
58
|
-
"""
|
|
59
|
-
|
|
60
|
-
return _nearest_neighbor_remap(self.ux_obj, destination_grid, remap_to)
|
|
61
|
-
|
|
62
|
-
def inverse_distance_weighted(
|
|
63
|
-
self, destination_grid: Grid, remap_to: str = "faces", power=2, k=8, **kwargs
|
|
64
|
-
) -> UxDataArray | UxDataset:
|
|
65
|
-
"""
|
|
66
|
-
Perform inverse-distance-weighted (IDW) remapping.
|
|
67
|
-
|
|
68
|
-
Each destination point is a weighted average of nearby source points,
|
|
69
|
-
with weights proportional to 1/(distance**power).
|
|
70
|
-
|
|
71
|
-
Parameters
|
|
72
|
-
----------
|
|
73
|
-
destination_grid : Grid
|
|
74
|
-
The UXarray grid to which data will be interpolated.
|
|
75
|
-
remap_to : {'nodes', 'edges', 'faces'}, default='faces'
|
|
76
|
-
Which grid element receives the remapped values.
|
|
77
|
-
power : int, default=2
|
|
78
|
-
Exponent controlling distance decay. Larger values make the
|
|
79
|
-
interpolation more local.
|
|
80
|
-
k : int, default=8
|
|
81
|
-
Number of nearest source points to include in the weighted average.
|
|
82
|
-
|
|
83
|
-
Returns
|
|
84
|
-
-------
|
|
85
|
-
UxDataArray or UxDataset
|
|
86
|
-
A new object with data mapped onto `destination_grid`.
|
|
87
|
-
"""
|
|
88
|
-
|
|
89
|
-
return _inverse_distance_weighted_remap(
|
|
90
|
-
self.ux_obj, destination_grid, remap_to, power, k
|
|
91
|
-
)
|
|
92
|
-
|
|
93
|
-
def bilinear(
|
|
94
|
-
self, destination_grid: Grid, remap_to: str = "faces", **kwargs
|
|
95
|
-
) -> UxDataArray | UxDataset:
|
|
96
|
-
"""
|
|
97
|
-
Perform bilinear remapping.
|
|
98
|
-
|
|
99
|
-
Parameters
|
|
100
|
-
---------
|
|
101
|
-
destination_grid : Grid
|
|
102
|
-
Destination Grid for remapping
|
|
103
|
-
remap_to : {'nodes', 'edges', 'faces'}, default='faces'
|
|
104
|
-
Which grid element receives the remapped values.
|
|
105
|
-
|
|
106
|
-
Returns
|
|
107
|
-
-------
|
|
108
|
-
UxDataArray or UxDataset
|
|
109
|
-
A new object with data mapped onto `destination_grid`.
|
|
110
|
-
"""
|
|
111
|
-
|
|
112
|
-
return _bilinear(self.ux_obj, destination_grid, remap_to)
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|