diff --git a/test/core/test_topological_agg.py b/test/core/test_topological_agg.py index 94ce1f58d..632fbff8d 100644 --- a/test/core/test_topological_agg.py +++ b/test/core/test_topological_agg.py @@ -1,5 +1,6 @@ import uxarray as ux +import numpy as np import pytest @@ -32,3 +33,31 @@ def test_node_to_edge_aggs(gridpath): grid_reduction = getattr(uxds['areaTriangle'], agg_func)(destination='edge') assert 'n_edge' in grid_reduction.dims + + +def test_node_to_face_numpy_dask_allclose(gridpath): + # the numpy (eager) and dask (chunked) branches must agree + pytest.importorskip("dask") # dask-backed branch requires dask + uxds = ux.open_dataset(gridpath("mpas", "QU", "oQU480.231010.nc"), gridpath("mpas", "QU", "oQU480.231010.nc")) + uxda = uxds['areaTriangle'] + + for agg_func in AGGS: + numpy_result = getattr(uxda, agg_func)(destination='face') + dask_result = getattr(uxda.chunk(), agg_func)(destination='face') + + assert numpy_result.dims == dask_result.dims + assert np.allclose(numpy_result.values, dask_result.values, equal_nan=True) + + +def test_node_to_edge_numpy_dask_allclose(gridpath): + # the numpy (eager) and dask (chunked) branches must agree + pytest.importorskip("dask") # dask-backed branch requires dask + uxds = ux.open_dataset(gridpath("mpas", "QU", "oQU480.231010.nc"), gridpath("mpas", "QU", "oQU480.231010.nc")) + uxda = uxds['areaTriangle'] + + for agg_func in AGGS: + numpy_result = getattr(uxda, agg_func)(destination='edge') + dask_result = getattr(uxda.chunk(), agg_func)(destination='edge') + + assert numpy_result.dims == dask_result.dims + assert np.allclose(numpy_result.values, dask_result.values, equal_nan=True) diff --git a/uxarray/core/aggregation.py b/uxarray/core/aggregation.py index c9d3dd6a7..8d78d0208 100644 --- a/uxarray/core/aggregation.py +++ b/uxarray/core/aggregation.py @@ -83,30 +83,35 @@ def _node_to_face_aggregation(uxda, aggregation, aggregation_func_kwargs): aggregated_var = _apply_node_to_face_aggregation_numpy( uxda, NUMPY_AGGREGATIONS[aggregation], aggregation_func_kwargs ) + return uxarray.core.dataarray.UxDataArray( + uxgrid=uxda.uxgrid, + data=aggregated_var, + dims=uxda.dims, + name=uxda.name, + ).rename({"n_node": "n_face"}) elif isinstance(uxda.data, da.Array): - # apply aggregation on dask array, TODO: - aggregated_var = _apply_node_to_face_aggregation_numpy( + # apply aggregation lazily on a dask array + return _apply_node_to_face_aggregation_dask( uxda, NUMPY_AGGREGATIONS[aggregation], aggregation_func_kwargs ) else: raise ValueError - return uxarray.core.dataarray.UxDataArray( - uxgrid=uxda.uxgrid, - data=aggregated_var, - dims=uxda.dims, - name=uxda.name, - ).rename({"n_node": "n_face"}) - -def _apply_node_to_face_aggregation_numpy( - uxda, aggregation_func, aggregation_func_kwargs +def _node_to_face_kernel( + data, + face_node_conn, + n_nodes_per_face, + n_face, + aggregation_func, + aggregation_func_kwargs, ): - """Applies a Node to Face Topological aggregation on a Numpy array.""" - data = uxda.values - face_node_conn = uxda.uxgrid.face_node_connectivity.values - n_nodes_per_face = uxda.uxgrid.n_nodes_per_face.values + """Node-to-face topological aggregation on a single numpy block. + ``data`` has the node dimension as its last axis, shape ``(..., n_node)``; + the result has the face dimension as its last axis, shape ``(..., n_face)``. + Shared verbatim by the numpy path and the (blockwise) dask path. + """ ( change_ind, n_nodes_per_face_sorted_ind, @@ -114,7 +119,7 @@ def _apply_node_to_face_aggregation_numpy( size_counts, ) = get_face_node_partitions(n_nodes_per_face) - result = np.empty(shape=(data.shape[:-1]) + (uxda.uxgrid.n_face,)) + result = np.empty(shape=(data.shape[:-1]) + (n_face,)) for e, start, end in zip(element_sizes, change_ind[:-1], change_ind[1:]): face_inds = n_nodes_per_face_sorted_ind[start:end] @@ -131,9 +136,55 @@ def _apply_node_to_face_aggregation_numpy( return result -def _apply_node_to_face_aggregation_dask(*args, **kwargs): - """Applies a Node to Face Topological aggregation on a Dask array.""" - pass +def _apply_node_to_face_aggregation_numpy( + uxda, aggregation_func, aggregation_func_kwargs +): + """Applies a Node to Face Topological aggregation on a Numpy array.""" + return _node_to_face_kernel( + uxda.values, + uxda.uxgrid.face_node_connectivity.values, + uxda.uxgrid.n_nodes_per_face.values, + uxda.uxgrid.n_face, + aggregation_func, + aggregation_func_kwargs, + ) + + +def _apply_node_to_face_aggregation_dask( + uxda, aggregation_func, aggregation_func_kwargs +): + """Applies a Node to Face Topological aggregation on a Dask array, lazily""" + import xarray as xr + + uxgrid = uxda.uxgrid + n_face = uxgrid.n_face + + # n_node must live in a single chunk: any node can feed any face. The grid + # metadata are core-dim inputs, so their core dims must be single-chunk too. + da_in = uxda.chunk({"n_node": -1}) + face_node_conn = uxgrid.face_node_connectivity.chunk( + {"n_face": -1, "n_max_face_nodes": -1} + ) + n_nodes_per_face = uxgrid.n_nodes_per_face.chunk({"n_face": -1}) + + result = xr.apply_ufunc( + _node_to_face_kernel, + da_in, + face_node_conn, + n_nodes_per_face, + input_core_dims=[["n_node"], ["n_face", "n_max_face_nodes"], ["n_face"]], + output_core_dims=[["n_face"]], + dask="parallelized", + output_dtypes=[np.float64], + dask_gufunc_kwargs={"output_sizes": {"n_face": n_face}}, + kwargs={ + "n_face": n_face, + "aggregation_func": aggregation_func, + "aggregation_func_kwargs": aggregation_func_kwargs, + }, + ) + + return uxarray.core.dataarray.UxDataArray(result, uxgrid=uxgrid, name=uxda.name) def _node_to_edge_aggregation(uxda, aggregation, aggregation_func_kwargs): @@ -146,39 +197,96 @@ def _node_to_edge_aggregation(uxda, aggregation, aggregation_func_kwargs): f"{uxda.uxgrid.n_face}." ) + aggregation_func = NUMPY_AGGREGATIONS[aggregation] + if isinstance(uxda.data, np.ndarray): # apply aggregation using numpy aggregation_var = _apply_node_to_edge_aggregation_numpy( - uxda, NUMPY_AGGREGATIONS[aggregation], aggregation_func_kwargs + uxda, aggregation_func, aggregation_func_kwargs ) + return uxarray.core.dataarray.UxDataArray( + uxgrid=uxda.uxgrid, + data=aggregation_var, + dims=uxda.dims, + name=uxda.name, + ).rename({"n_node": "n_edge"}) elif isinstance(uxda.data, da.Array): - # apply aggregation on dask array, TODO: - aggregation_var = _apply_node_to_edge_aggregation_numpy( - uxda, NUMPY_AGGREGATIONS[aggregation], aggregation_func_kwargs + # apply aggregation lazily on a dask array + return _apply_node_to_edge_aggregation_dask( + uxda, aggregation_func, aggregation_func_kwargs ) else: raise ValueError - return uxarray.core.dataarray.UxDataArray( - uxgrid=uxda.uxgrid, - data=aggregation_var, - dims=uxda.dims, - name=uxda.name, - ).rename({"n_node": "n_edge"}) + +def _node_to_edge_kernel( + data, edge_node_conn, aggregation_func, aggregation_func_kwargs +): + """Node-to-edge topological aggregation on a single numpy block. + + ``data`` has the node dimension as its last axis, shape ``(..., n_node)``; + the result has the edge dimension as its last axis, shape ``(..., n_edge)``. + Each edge reduces over its two endpoint nodes. Shared verbatim by the numpy + path and the (blockwise) dask path. + """ + return aggregation_func( + data[..., edge_node_conn], axis=-1, **aggregation_func_kwargs + ) def _apply_node_to_edge_aggregation_numpy( uxda, aggregation_func, aggregation_func_kwargs ): """Applies a Node to Edge topological aggregation on a numpy array.""" - data = uxda.values - edge_node_conn = uxda.uxgrid.edge_node_connectivity.values - result = aggregation_func( - data[..., edge_node_conn], axis=-1, **aggregation_func_kwargs + return _node_to_edge_kernel( + uxda.values, + uxda.uxgrid.edge_node_connectivity.values, + aggregation_func, + aggregation_func_kwargs, ) - return result -def _apply_node_to_edge_aggregation_dask(*args, **kwargs): - """Applies a Node to Edge topological aggregation on a dask array.""" - pass +def _apply_node_to_edge_aggregation_dask( + uxda, aggregation_func, aggregation_func_kwargs +): + """Applies a Node to Edge topological aggregation on a Dask array, lazily. + + Mirrors the node-to-face dask path: the data field and the edge-node + connectivity enter ``apply_ufunc`` as dask inputs (nothing computed + eagerly), ``n_node`` is a single (uncrunked) core dimension, and the leading + dimensions are processed blockwise. Each edge reduces over its two endpoint + nodes. + """ + import xarray as xr + + uxgrid = uxda.uxgrid + n_edge = uxgrid.n_edge + + # n_node must live in a single chunk: any node can feed any edge. The grid + # metadata is a core-dim input, so its core dims must be single-chunk too. + da_in = uxda.chunk({"n_node": -1}) + edge_node_conn = uxgrid.edge_node_connectivity.chunk({"n_edge": -1, "two": -1}) + + # match the numpy path's dtype (float for numeric aggs, bool for all/any) + out_dtype = aggregation_func( + np.empty((1, edge_node_conn.shape[-1]), dtype=uxda.dtype), + axis=-1, + **aggregation_func_kwargs, + ).dtype + + result = xr.apply_ufunc( + _node_to_edge_kernel, + da_in, + edge_node_conn, + input_core_dims=[["n_node"], ["n_edge", "two"]], + output_core_dims=[["n_edge"]], + dask="parallelized", + output_dtypes=[out_dtype], + dask_gufunc_kwargs={"output_sizes": {"n_edge": n_edge}}, + kwargs={ + "aggregation_func": aggregation_func, + "aggregation_func_kwargs": aggregation_func_kwargs, + }, + ) + + return uxarray.core.dataarray.UxDataArray(result, uxgrid=uxgrid, name=uxda.name) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index a2b55a796..b275272e6 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -245,14 +245,14 @@ def to_geodataframe( same name as the ``UxDataArray`` (or named ``var`` if no name exists) """ - if self.values.ndim > 1: + if self.ndim > 1: # data is multidimensional, must be a 1D slice raise ValueError( f"Data Variable must be 1-dimensional, with shape {self.uxgrid.n_face} " f"for face-centered data." ) - if self.values.size == self.uxgrid.n_face: + if self.size == self.uxgrid.n_face: gdf, non_nan_polygon_indices = self.uxgrid.to_geodataframe( periodic_elements=periodic_elements, projection=projection, @@ -289,21 +289,21 @@ def to_geodataframe( gdf[var_name] = _data - elif self.values.size == self.uxgrid.n_node: + elif self.size == self.uxgrid.n_node: raise ValueError( - f"Data Variable with size {self.values.size} does not match the number of faces " + f"Data Variable with size {self.size} does not match the number of faces " f"({self.uxgrid.n_face}. Current size matches the number of nodes. Consider running " f"``UxDataArray.topological_mean(destination='face') to aggregate the data onto the faces." ) - elif self.values.size == self.uxgrid.n_edge: + elif self.size == self.uxgrid.n_edge: raise ValueError( - f"Data Variable with size {self.values.size} does not match the number of faces " + f"Data Variable with size {self.size} does not match the number of faces " f"({self.uxgrid.n_face}. Current size matches the number of edges." ) else: # data is not mapped to raise ValueError( - f"Data Variable with size {self.values.size} does not match the number of faces " + f"Data Variable with size {self.size} does not match the number of faces " f"({self.uxgrid.n_face}." ) @@ -339,7 +339,7 @@ def to_polycollection( Flag to indicate whether to override a cached PolyCollection, if it exists """ # data is multidimensional, must be a 1D slice - if self.values.ndim > 1: + if self.ndim > 1: raise ValueError( f"Data Variable must be 1-dimensional, with shape {self.uxgrid.n_face} " f"for face-centered data." @@ -609,16 +609,21 @@ def integrate( >>> uxds = ux.open_dataset("grid.ug", "centroid_pressure_data_ug") >>> integral = uxds["psi"].integrate() """ - if self.values.shape[-1] == self.uxgrid.n_face: - face_areas = self.uxgrid.face_areas.values - - # perform dot product between face areas and last dimension of data - integral = np.einsum("i,...i", face_areas, self.values) + if self.shape[-1] == self.uxgrid.n_face: + # dot product between face areas and the face dimension of the data + if isinstance(self.data, np.ndarray): + # eager data: a direct einsum avoids xr.dot's per-call overhead + integral = np.einsum( + "i,...i", self.uxgrid.face_areas.values, self.values + ) + else: + # dask-backed data: xr.dot keeps the reduction lazy + integral = xr.dot(self, self.uxgrid.face_areas, dim="n_face") - elif self.values.shape[-1] == self.uxgrid.n_node: + elif self.shape[-1] == self.uxgrid.n_node: raise ValueError("Integrating data mapped to each node not yet supported.") - elif self.values.shape[-1] == self.uxgrid.n_edge: + elif self.shape[-1] == self.uxgrid.n_edge: raise ValueError("Integrating data mapped to each edge not yet supported.") else: @@ -626,7 +631,7 @@ def integrate( f"The final dimension of the data variable does not match the number of nodes, edges, " f"or faces. Expected one of " f"{self.uxgrid.n_node}, {self.uxgrid.n_edge}, or {self.uxgrid.n_face}, " - f"but received {self.values.shape[-1]}" + f"but received {self.shape[-1]}" ) # construct a uxda with integrated quantity @@ -1650,7 +1655,7 @@ def curl( ) # Compute curl = ∂v/∂x - ∂u/∂y - curl_values = grad_v_zonal.values - grad_u_meridional.values + curl_values = grad_v_zonal.data - grad_u_meridional.data u_units = self.attrs.get("units", "") has_sphere_radius = "sphere_radius" in self.uxgrid._ds.attrs @@ -2177,11 +2182,10 @@ def get_dual(self): # Get correct dimensions for the dual dims = [dim_map.get(dim, dim) for dim in self.dims] - # Get the values from the data array - data = np.array(self.values) - # Construct the new data array - uxda = uxarray.UxDataArray(uxgrid=dual, data=data, dims=dims, name=self.name) + uxda = uxarray.UxDataArray( + uxgrid=dual, data=self.data, dims=dims, name=self.name + ) return uxda diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index 4331dd30d..95bc6d8e2 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -6,7 +6,6 @@ from typing import IO, Any, Hashable, Mapping from warnings import warn -import numpy as np import xarray as xr from xarray.core import dtypes from xarray.core.options import OPTIONS @@ -715,11 +714,10 @@ def get_dual(self): # Get correct dimensions for the dual dims = [dim_map.get(dim, dim) for dim in self[var].dims] - # Get the values from the data array - data = np.array(self[var].values) - # Construct the new data array - uxda = uxarray.UxDataArray(uxgrid=dual, data=data, dims=dims, name=var) + uxda = uxarray.UxDataArray( + uxgrid=dual, data=self[var].data, dims=dims, name=var + ) # Add data array to dataset dataset[var] = uxda diff --git a/uxarray/cross_sections/dataarray_accessor.py b/uxarray/cross_sections/dataarray_accessor.py index b27e5c2b0..ac1f06fcf 100644 --- a/uxarray/cross_sections/dataarray_accessor.py +++ b/uxarray/cross_sections/dataarray_accessor.py @@ -8,7 +8,6 @@ from uxarray.constants import INT_DTYPE from .sample import ( - _fill_numba, sample_constant_latitude, sample_constant_longitude, sample_geodesic, @@ -114,22 +113,23 @@ def __call__( ) face_idx = np.array([row[0] if row else -1 for row in faces], dtype=INT_DTYPE) - orig_dims = list(self.uxda.dims) - face_axis = orig_dims.index("n_face") new_dim = "steps" - new_dims = [new_dim if d == "n_face" else d for d in orig_dims] - dim_axis = new_dims.index(new_dim) + new_dims = [new_dim if d == "n_face" else d for d in self.uxda.dims] - arr = np.moveaxis(self.uxda.compute().values, face_axis, -1) - M, Nf = arr.reshape(-1, arr.shape[-1]).shape - flat_orig = arr.reshape(M, Nf) - - # Fill along the arc with nearest‐neighbor - flat_filled = _fill_numba(flat_orig, face_idx, Nf, steps) - filled = flat_filled.reshape(*arr.shape[:-1], steps) + # Gather only the sampled columns (the nearest face for each step), + # lazily; steps with no containing face (-1) become NaN. + valid = face_idx >= 0 + valid_idx = xr.DataArray( + np.where(valid, face_idx, 0).astype(INT_DTYPE), dims=new_dim + ) - # Move steps axis back to its proper position - data = np.moveaxis(filled, -1, dim_axis) + data = ( + xr.DataArray(self.uxda.data, dims=self.uxda.dims) + .isel(n_face=valid_idx) + .where(xr.DataArray(valid, dims=new_dim)) + .transpose(*new_dims) + .data + ) # Build coords dict: keep everything except 'n_face' coords = { diff --git a/uxarray/cross_sections/sample.py b/uxarray/cross_sections/sample.py index 2b34b4e44..be5627448 100644 --- a/uxarray/cross_sections/sample.py +++ b/uxarray/cross_sections/sample.py @@ -1,19 +1,7 @@ import numpy as np -from numba import njit, prange from pyproj import Geod -@njit(parallel=True) -def _fill_numba(flat_orig, face_idx, n_face, n_steps): - M = flat_orig.shape[0] - out = np.full((M, n_steps), np.nan, np.float64) - for i in prange(n_steps): - f = face_idx[i] - if 0 <= f < n_face: - out[:, i] = flat_orig[:, f] - return out - - def sample_geodesic( start: tuple[float, float], end: tuple[float, float], steps: int ) -> tuple[np.ndarray, np.ndarray]: diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index 0385bed7b..bd69c991b 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -1241,7 +1241,7 @@ def face_node_connectivity(self) -> xr.DataArray: _populate_healpix_boundaries(self._ds) if self._ds["face_node_connectivity"].ndim == 1: - face_node_connectivity_1d = self._ds["face_node_connectivity"].values + face_node_connectivity_1d = self._ds["face_node_connectivity"].data face_node_connectivity_2d = np.expand_dims( face_node_connectivity_1d, axis=0 ) @@ -1969,7 +1969,7 @@ def calculate_total_face_area( and order == 4 and not latitude_adjusted_area ): - return np.sum(self.face_areas.values) + return float(self.face_areas.data.sum()) face_areas, _ = self._compute_face_areas( quadrature_rule, order, latitude_adjusted_area @@ -2490,7 +2490,7 @@ def get_dual(self, check_duplicate_nodes: bool = False): # Construct dual mesh dual = self.from_topology( - self.face_lon.values, self.face_lat.values, dual_node_face_conn + self.face_lon.data, self.face_lat.data, dual_node_face_conn ) return dual diff --git a/uxarray/grid/slice.py b/uxarray/grid/slice.py index 3bc97b9d1..56ad435c1 100644 --- a/uxarray/grid/slice.py +++ b/uxarray/grid/slice.py @@ -35,7 +35,9 @@ def _slice_node_indices( raise ValueError("Exclusive slicing is not yet supported.") # faces that saddle nodes given in 'indices' - face_indices = np.unique(grid.node_face_connectivity.values[indices].ravel()) + face_indices = np.unique( + grid.node_face_connectivity.isel(n_node=indices).values.ravel() + ) face_indices = face_indices[face_indices != INT_FILL_VALUE] return _slice_face_indices(grid, face_indices) @@ -65,7 +67,9 @@ def _slice_edge_indices( raise ValueError("Exclusive slicing is not yet supported.") # faces that saddle nodes given in 'indices' - face_indices = np.unique(grid.edge_face_connectivity.values[indices].ravel()) + face_indices = np.unique( + grid.edge_face_connectivity.isel(n_edge=indices).values.ravel() + ) face_indices = face_indices[face_indices != INT_FILL_VALUE] return _slice_face_indices(grid, face_indices) @@ -103,7 +107,9 @@ def _slice_face_indices( face_indices = np.atleast_1d(np.asarray(indices, dtype=INT_DTYPE)) # nodes of each face (inclusive) - node_indices = np.unique(grid.face_node_connectivity.values[face_indices].ravel()) + node_indices = np.unique( + grid.face_node_connectivity.isel(n_face=face_indices).values.ravel() + ) node_indices = node_indices[node_indices != INT_FILL_VALUE] # Index Node and Face variables @@ -113,7 +119,7 @@ def _slice_face_indices( # Only slice edge dimension if we have the face edge connectivity if "face_edge_connectivity" in ds: edge_indices = np.unique( - grid.face_edge_connectivity.values[face_indices].ravel() + grid.face_edge_connectivity.isel(n_face=face_indices).values.ravel() ) edge_indices = edge_indices[edge_indices != INT_FILL_VALUE] ds = ds.isel(n_edge=edge_indices) @@ -127,36 +133,30 @@ def _slice_face_indices( ds["subgrid_node_indices"] = xr.DataArray(node_indices, dims=["n_node"]) ds["subgrid_face_indices"] = xr.DataArray(face_indices, dims=["n_face"]) - # Construct updated Node Index Map - node_indices_dict = {orig: new for new, orig in enumerate(node_indices)} - node_indices_dict[INT_FILL_VALUE] = INT_FILL_VALUE + def _build_remap(orig_indices, size): + remap = np.full(size, INT_FILL_VALUE, dtype=INT_DTYPE) + remap[orig_indices] = np.arange(len(orig_indices), dtype=INT_DTYPE) + return remap - # Construct updated Edge Index Map - if edge_indices is not None: - edge_indices_dict = {orig: new for new, orig in enumerate(edge_indices)} - edge_indices_dict[INT_FILL_VALUE] = INT_FILL_VALUE - else: - edge_indices_dict = None - - def map_node_indices(i): - return node_indices_dict.get(i, INT_FILL_VALUE) + node_remap = _build_remap(node_indices, grid.n_node) + edge_remap = ( + _build_remap(edge_indices, grid.n_edge) if edge_indices is not None else None + ) - if edge_indices is not None: - - def map_edge_indices(i): - return edge_indices_dict.get(i, INT_FILL_VALUE) - else: - map_edge_indices = None + def _remap_indices(conn, remap): + # blockwise-safe: mask fill values before the lookup, restore them after + is_fill = conn == INT_FILL_VALUE + return np.where(is_fill, INT_FILL_VALUE, remap[np.where(is_fill, 0, conn)]) for conn_name in list(ds.data_vars): if conn_name.endswith("_node_connectivity"): - map_fn = map_node_indices + remap = node_remap elif conn_name.endswith("_edge_connectivity"): - if edge_indices_dict is None: + if edge_remap is None: ds = ds.drop_vars(conn_name) continue - map_fn = map_edge_indices + remap = edge_remap elif "_connectivity" in conn_name: # anything else we can't remap @@ -167,11 +167,14 @@ def map_edge_indices(i): # not a connectivity var, skip continue - # Apply Remapping - ds[conn_name] = xr.DataArray( - np.vectorize(map_fn, otypes=[INT_DTYPE])(ds[conn_name].values), - dims=ds[conn_name].dims, - attrs=ds[conn_name].attrs, + # Apply remapping (vectorized; stays lazy when the connectivity is dask) + ds[conn_name] = xr.apply_ufunc( + _remap_indices, + ds[conn_name], + kwargs={"remap": remap}, + dask="parallelized", + output_dtypes=[INT_DTYPE], + keep_attrs=True, ) if inverse_indices: diff --git a/uxarray/io/_geos.py b/uxarray/io/_geos.py index 227c85295..7b5c11605 100644 --- a/uxarray/io/_geos.py +++ b/uxarray/io/_geos.py @@ -12,8 +12,8 @@ def _read_geos_cs(in_ds: xr.Dataset): """ out_ds = xr.Dataset() - node_lon = in_ds["corner_lons"].values.ravel() - node_lat = in_ds["corner_lats"].values.ravel() + node_lon = in_ds["corner_lons"].data.ravel() + node_lat = in_ds["corner_lats"].data.ravel() out_ds["node_lon"] = xr.DataArray( data=node_lon, dims=ugrid.NODE_DIM, attrs=ugrid.NODE_LON_ATTRS @@ -24,8 +24,8 @@ def _read_geos_cs(in_ds: xr.Dataset): ) if "lons" in in_ds: - face_lon = in_ds["lons"].values.ravel() - face_lat = in_ds["lats"].values.ravel() + face_lon = in_ds["lons"].data.ravel() + face_lat = in_ds["lats"].data.ravel() out_ds["face_lon"] = xr.DataArray( data=face_lon, dims=ugrid.FACE_DIM, attrs=ugrid.FACE_LON_ATTRS diff --git a/uxarray/io/_scrip.py b/uxarray/io/_scrip.py index ca4c513fb..6e93d9644 100644 --- a/uxarray/io/_scrip.py +++ b/uxarray/io/_scrip.py @@ -10,7 +10,7 @@ def _values_in_degrees(data_array): - """Return data_array.values, converting to degrees if necessary, + """Return data_array's backing array, converting to degrees if necessary, i.e., if data_array.attrs["units"] in ["radians", "radian", "rad"]. Parameters @@ -22,10 +22,11 @@ def _values_in_degrees(data_array): Returns ------- - numpy.ndarray - data_array.values, converted to degrees if necessary. + numpy.ndarray or dask.array.Array + data_array's backing array (``.data``, so a lazily-opened array stays + lazy), converted to degrees if necessary. """ - values = data_array.values + values = data_array.data units = data_array.attrs.get("units", "").lower() # Check if units indicate radians @@ -49,8 +50,9 @@ def _to_ugrid(in_ds, out_ds): # Create node_lon & node_lat variables from grid_corner_lat/lon # Turn latitude and longitude scrip arrays into 1D # Convert to degrees if needed - corner_lat = _values_in_degrees(in_ds["grid_corner_lat"]).ravel() - corner_lon = _values_in_degrees(in_ds["grid_corner_lon"]).ravel() + # materialized here: fed to polars and numpy fancy-indexing below + corner_lat = np.asarray(_values_in_degrees(in_ds["grid_corner_lat"])).ravel() + corner_lon = np.asarray(_values_in_degrees(in_ds["grid_corner_lon"])).ravel() # Use Polars to find unique coordinate pairs df = pl.DataFrame({"lon": corner_lon, "lat": corner_lat}).with_row_count( diff --git a/uxarray/plot/matplotlib.py b/uxarray/plot/matplotlib.py index 5e43a30af..c20b7619f 100644 --- a/uxarray/plot/matplotlib.py +++ b/uxarray/plot/matplotlib.py @@ -202,7 +202,11 @@ def _nearest_neighbor_resample( # build an array of values for each valid point flat_vals = np.full(first_face.shape, np.nan, dtype=float) mask_has_face = first_face >= 0 - flat_vals[mask_has_face] = data.values[first_face[mask_has_face]] + # gather only the sampled faces (lazily for dask data): many pixels share a + # face, so deduplicate to materialize the minimal set rather than the whole + # field, then scatter back via the inverse map + unique_faces, inverse = np.unique(first_face[mask_has_face], return_inverse=True) + flat_vals[mask_has_face] = np.asarray(data.data[unique_faces])[inverse] # scatter back into a full raster via the valid mask res = np.full((ny, nx), np.nan, dtype=float) diff --git a/uxarray/remap/accessor.py b/uxarray/remap/accessor.py index 6415aa144..5306ab779 100644 --- a/uxarray/remap/accessor.py +++ b/uxarray/remap/accessor.py @@ -364,9 +364,9 @@ def apply_weights( Notes ----- - Dask-backed inputs are materialized in memory before the sparse - operator is applied. For lazy/chunked execution, prefer - ``nearest_neighbor`` or ``inverse_distance_weighted``. + Dask-backed inputs are remapped lazily: the sparse operator is applied + blockwise over the leading dimensions, keeping the source dimension in a + single chunk. Numpy-backed inputs are applied eagerly. """ return _apply_weights( diff --git a/uxarray/remap/apply_weights.py b/uxarray/remap/apply_weights.py index 9502ae228..dd8aa4685 100644 --- a/uxarray/remap/apply_weights.py +++ b/uxarray/remap/apply_weights.py @@ -58,9 +58,10 @@ def _apply_weights( ): """Apply a sparse remap operator to UXarray data. - Note: this path materializes dask-backed inputs eagerly when applying - the sparse operator. For lazy/chunked execution, use one of the other - remap methods (e.g., ``nearest_neighbor``, ``inverse_distance_weighted``). + Dask-backed inputs are remapped lazily: the sparse operator is applied + blockwise over the leading (non-source) dimensions via ``apply_ufunc``, with + the source dimension kept in a single chunk. Numpy-backed inputs are applied + eagerly. """ _assert_dimension(remap_to) @@ -87,7 +88,23 @@ def _apply_weights( remapped_any = True other_dims = [dim for dim in da.dims if dim != variable_source_dim] da_t = da.transpose(*other_dims, variable_source_dim) - remapped_values = weights_obj._apply(np.asarray(da_t.values)) + if isinstance(da_t.data, np.ndarray): + # eager: apply the sparse operator directly + remapped_values = weights_obj._apply(np.asarray(da_t.values)) + else: + # dask: apply blockwise over the leading dims; the sparse multiply + # spans the whole source dimension, so keep it in a single chunk + remapped_values = xr.apply_ufunc( + weights_obj._apply, + da_t.chunk({variable_source_dim: -1}), + input_core_dims=[[variable_source_dim]], + output_core_dims=[[destination_dim]], + dask="parallelized", + output_dtypes=[np.result_type(weights_obj.matrix.dtype, da_t.dtype)], + dask_gufunc_kwargs={ + "output_sizes": {destination_dim: weights_obj.destination_size} + }, + ).data other_dims_set = set(other_dims) coords = { diff --git a/uxarray/remap/structured.py b/uxarray/remap/structured.py index bc036a11d..01ec2ac3a 100644 --- a/uxarray/remap/structured.py +++ b/uxarray/remap/structured.py @@ -165,7 +165,7 @@ def _reshape_array_to_rectilinear( coords[spec.lon_name] = spec.lon return xr.DataArray( - np.asarray(da.values).reshape(shape), + da.data.reshape(shape), dims=dims, coords=coords, name=da.name,