From 2e7db98f58e4b71db8db91c5ac195c815007b57d Mon Sep 17 00:00:00 2001
From: ahijevyc <ahijevyc@ucar.edu>
Date: Fri, 2 Aug 2024 09:15:12 -0600
Subject: [PATCH 1/7] remap.apply_func

Apply a function like np.mean or np.max to neighborhood around all
points. It is like inverse distance weighting with power=0 but there is
no weighting. All points in neighborhood are treated the same.

The neighborhood is at a fixed radius instead of a fixed number of
neighbors.

Eliminated depreciated destination_obj argument.
---
 uxarray/remap/apply_func.py         | 285 ++++++++++++++++++++++++++++
 uxarray/remap/dataarray_accessor.py |  52 +++--
 uxarray/remap/dataset_accessor.py   |  52 +++--
 3 files changed, 362 insertions(+), 27 deletions(-)
 create mode 100644 uxarray/remap/apply_func.py

diff --git a/uxarray/remap/apply_func.py b/uxarray/remap/apply_func.py
new file mode 100644
index 000000000..48d358fff
--- /dev/null
+++ b/uxarray/remap/apply_func.py
@@ -0,0 +1,285 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, Union
+
+if TYPE_CHECKING:
+    from uxarray.core.dataset import UxDataset
+    from uxarray.core.dataarray import UxDataArray
+
+import numpy as np
+from functools import partial
+
+import uxarray.core.dataarray
+import uxarray.core.dataset
+from uxarray.grid import Grid
+import warnings
+
+
+def _apply_func_remap(
+    source_grid: Grid,
+    destination_grid: Grid,
+    source_data: np.ndarray,
+    remap_to: str = "face centers",
+    coord_type: str = "spherical",
+    func: func = np.mean,
+    r: float = None,
+) -> np.array:
+    """Apply neighborhood function Remapping between two grids.
+
+    Parameters:
+    -----------
+    source_grid : Grid
+        Source grid that data is mapped from.
+    destination_grid : Grid
+        Destination grid to remap data to.
+    source_data : np.ndarray
+        Data variable to remap.
+    remap_to : str, default="nodes"
+        Location of where to map data, either "nodes", "edge centers", or "face centers".
+    coord_type: str, default="spherical"
+        Coordinate type to use for nearest neighbor query, either "spherical" or "Cartesian".
+    r : float, default=None
+        radius of neighborhoodFor spherical coordinates, the radius is in units of degrees,
+        and for cartesian coordinates, the radius is in meters.
+
+    Returns:
+    --------
+    destination_data : np.ndarray
+        Data mapped to the destination grid.
+    """
+
+
+    source_data = np.asarray(source_data)
+    n_elements = source_data.shape[-1]
+
+    if n_elements == source_grid.n_node:
+        source_data_mapping = "nodes"
+    elif n_elements == source_grid.n_face:
+        source_data_mapping = "face centers"
+    elif n_elements == source_grid.n_edge:
+        source_data_mapping = "edge centers"
+    else:
+        raise ValueError(
+            f"Invalid source_data shape. The final dimension should match the number of corner "
+            f"nodes ({source_grid.n_node}), edge nodes ({source_grid.n_edge}), or face centers ({source_grid.n_face}) "
+            f"in the source grid, but received: {source_data.shape}"
+        )
+
+    if coord_type == "spherical":
+        if remap_to == "nodes":
+            lon, lat = (
+                destination_grid.node_lon.values,
+                destination_grid.node_lat.values,
+            )
+        elif remap_to == "face centers":
+            lon, lat = (
+                destination_grid.face_lon.values,
+                destination_grid.face_lat.values,
+            )
+        elif remap_to == "edge centers":
+            lon, lat = (
+                destination_grid.edge_lon.values,
+                destination_grid.edge_lat.values,
+            )
+        else:
+            raise ValueError(
+                f"Invalid remap_to. Expected 'nodes', 'edge centers', or 'face centers', "
+                f"but received: {remap_to}"
+            )
+
+        _source_tree = source_grid.get_ball_tree(coordinates=source_data_mapping)
+
+        dest_coords = np.vstack([lon, lat]).T
+
+        neighbor_indices = _source_tree.query_radius(dest_coords, r=r)
+
+    elif coord_type == "cartesian":
+        if remap_to == "nodes":
+            x, y, z = (
+                destination_grid.node_x.values,
+                destination_grid.node_y.values,
+                destination_grid.node_z.values,
+            )
+        elif remap_to == "face centers":
+            x, y, z = (
+                destination_grid.face_x.values,
+                destination_grid.face_y.values,
+                destination_grid.face_z.values,
+            )
+        elif remap_to == "edge centers":
+            x, y, z = (
+                destination_grid.edge_x.values,
+                destination_grid.edge_y.values,
+                destination_grid.edge_z.values,
+            )
+        else:
+            raise ValueError(
+                f"Invalid remap_to. Expected 'nodes', 'edge centers', or 'face centers', "
+                f"but received: {remap_to}"
+            )
+
+        _source_tree = source_grid.get_ball_tree(
+            coordinates=source_data_mapping,
+            coordinate_system="cartesian",
+            distance_metric="minkowski",
+        )
+
+        dest_coords = np.vstack([x, y, z]).T
+
+        neighbor_indices = _source_tree.query_radius(dest_coords, r=r)
+
+    else:
+        raise ValueError(
+            f"Invalid coord_type. Expected either 'spherical' or 'cartesian', but received {coord_type}"
+        )
+
+    # make destination_shape a list instead of immutable tuple
+    destination_shape = list(source_data.shape)
+    # last dimension has same number of elements as neighbor_indices list
+    destination_shape[-1] = len(neighbor_indices)
+    destination_data = np.empty(destination_shape)
+    # Apply function to indices on last axis.
+    func = partial(func, axis=-1)
+    for i, idx in enumerate(neighbor_indices):
+        if len(idx):
+            destination_data[..., i] = func(source_data[..., idx])
+
+    return destination_data
+
+
+def _apply_func_remap_uxda(
+    source_uxda: UxDataArray,
+    destination_obj: Union[Grid, UxDataArray, UxDataset],
+    remap_to: str = "face centers",
+    coord_type: str = "spherical",
+    func: func = np.mean,
+    r=5,
+):
+    """Neighborhood function Remapping implementation for ``UxDataArray``.
+
+    Parameters
+    ---------
+    source_uxda : UxDataArray
+        Source UxDataArray for remapping
+    destination_obj : Grid, UxDataArray, UxDataset
+        Destination for remapping
+    remap_to : str, default="nodes"
+        Location of where to map data, either "nodes", "edge centers", or "face centers"
+    coord_type : str, default="spherical"
+        Indicates whether to remap using on Spherical or Cartesian coordinates for the computations when
+        remapping.
+    r : float, default=5
+        Radius of neighborhood.
+    """
+
+    # check dimensions remapped to and from
+    if (
+        (source_uxda._node_centered() and remap_to != "nodes")
+        or (source_uxda._face_centered() and remap_to != "face centers")
+        or (source_uxda._edge_centered() and remap_to != "edge centers")
+    ):
+        warnings.warn(
+            f"Your data is stored on {source_uxda.dims[-1]}, but you are remapping to {remap_to}"
+        )
+
+    # prepare dimensions
+    if remap_to == "nodes":
+        destination_dim = "n_node"
+    elif remap_to == "face centers":
+        destination_dim = "n_face"
+    else:
+        destination_dim = "n_edge"
+
+    destination_dims = list(source_uxda.dims)
+    destination_dims[-1] = destination_dim
+
+    if isinstance(destination_obj, Grid):
+        destination_grid = destination_obj
+    elif isinstance(
+        destination_obj,
+        (uxarray.core.dataarray.UxDataArray, uxarray.core.dataset.UxDataset),
+    ):
+        destination_grid = destination_obj.uxgrid
+    else:
+        raise ValueError("TODO: Invalid Input")
+
+    # perform remapping
+    destination_data = _apply_func_remap(
+        source_uxda.uxgrid,
+        destination_grid,
+        source_uxda.data,
+        remap_to,
+        coord_type,
+        func,
+        r,
+    )
+    # construct data array for remapping variable
+    uxda_remap = uxarray.core.dataarray.UxDataArray(
+        data=destination_data,
+        name=source_uxda.name,
+        coords=source_uxda.coords,
+        dims=destination_dims,
+        uxgrid=destination_grid,
+    )
+    # add remapped variable to existing UxDataset
+    if isinstance(destination_obj, uxarray.core.dataset.UxDataset):
+        uxds = destination_obj.copy()
+        uxds[source_uxda.name] = uxda_remap
+        return uxds
+
+    # construct a UxDataset from remapped variable and existing variable
+    elif isinstance(destination_obj, uxarray.core.dataset.UxDataArray):
+        uxds = destination_obj.copy().to_dataset()
+        uxds[source_uxda.name] = uxda_remap
+        return uxds
+
+    # return UxDataArray with remapped variable
+    else:
+        return uxda_remap
+
+
+def _apply_func_remap_uxds(
+    source_uxds: UxDataset,
+    destination_obj: Union[Grid, UxDataArray, UxDataset],
+    remap_to: str = "face centers",
+    coord_type: str = "spherical",
+    func: func = np.mean,
+    r: float = 5,
+):
+    """Neighboohood function implementation for ``UxDataset``.
+
+    Parameters
+    ---------
+    source_uxds : UxDataset
+        Source UxDataset for remapping
+    destination_obj : Grid, UxDataArray, UxDataset
+        Destination for remapping
+    remap_to : str, default="nodes"
+        Location of where to map data, either "nodes", "edge centers", or "face centers"
+    coord_type : str, default="spherical"
+        Indicates whether to remap using on Spherical or Cartesian coordinates
+    func : func = np.mean
+        function to apply to neighborhood
+    r : float, default=5
+        Radius of neighborhood
+    """
+
+    if isinstance(destination_obj, Grid):
+        destination_uxds = uxarray.core.dataset.UxDataset(uxgrid=destination_obj)
+    elif isinstance(destination_obj, uxarray.core.dataset.UxDataArray):
+        destination_uxds = destination_obj.to_dataset()
+    elif isinstance(destination_obj, uxarray.core.dataset.UxDataset):
+        destination_uxds = destination_obj
+    else:
+        raise ValueError
+
+    for var_name in source_uxds.data_vars:
+        destination_uxds = _apply_func_remap_uxda(
+            source_uxds[var_name],
+            destination_uxds,
+            remap_to,
+            coord_type,
+            func,
+            r,
+        )
+
+    return destination_uxds
diff --git a/uxarray/remap/dataarray_accessor.py b/uxarray/remap/dataarray_accessor.py
index 4a5f21dfe..f06196ce2 100644
--- a/uxarray/remap/dataarray_accessor.py
+++ b/uxarray/remap/dataarray_accessor.py
@@ -6,13 +6,14 @@
 from uxarray.remap.inverse_distance_weighted import (
     _inverse_distance_weighted_remap_uxda,
 )
+from uxarray.remap.apply_func import _apply_func_remap_uxda
 
 if TYPE_CHECKING:
     from uxarray.core.dataset import UxDataset
     from uxarray.core.dataarray import UxDataArray
 
 from uxarray.grid import Grid
-
+import numpy as np
 
 class UxDataArrayRemapAccessor:
     def __init__(self, uxda: UxDataArray):
@@ -26,13 +27,13 @@ def __repr__(self):
             "  * nearest_neighbor(destination_obj, remap_to, coord_type)\n"
         )
         methods_heading += "  * inverse_distance_weighted(destination_obj, remap_to, coord_type, power, k)\n"
+        methods_heading += "  * apply_func(destination_grid, remap_to, coord_type, func, r)\n"
 
         return prefix + methods_heading
 
     def nearest_neighbor(
         self,
-        destination_grid: Optional[Grid] = None,
-        destination_obj: Optional[Grid, UxDataArray, UxDataset] = None,
+        destination_grid: Grid = None,
         remap_to: str = "face centers",
         coord_type: str = "spherical",
     ):
@@ -43,19 +44,12 @@ def nearest_neighbor(
         ---------
         destination_grid : Grid
             Destination Grid for remapping
-        destination_obj : Grid, UxDataArray, UxDataset
-            Optional destination for remapping, deprecating
         remap_to : str, default="nodes"
-            Location of where to map data, either "nodes" or "face centers"
+            Location of where to map data, either "nodes", "edge centers", or "face centers"
         coord_type : str, default="spherical"
             Indicates whether to remap using on spherical or cartesian coordinates
         """
-        if destination_grid is not None and destination_obj is not None:
-            raise ValueError(
-                "Only one destination allowed, "
-                "please remove either `destination_grid` or `destination_obj`."
-            )
-        elif destination_grid is None and destination_obj is None:
+        if destination_grid is None:
             raise ValueError("Destination needed for remap.")
 
         if destination_grid is not None:
@@ -71,6 +65,37 @@ def nearest_neighbor(
                 self.uxda, destination_obj, remap_to, coord_type
             )
 
+    def apply_func(
+        self,
+        destination_grid: Grid = None,
+        remap_to: str = "face centers",
+        coord_type: str = "spherical",
+        func: func = np.mean,
+        r=1,
+    ):
+        """Neighborhood function Remapping between a source
+        (``UxDataArray``) and destination.`.
+
+        Parameters
+        ---------
+        destination_grid : Grid
+            Destination Grid for remapping
+        remap_to : str, default="nodes"
+            Location of where to map data, either "nodes", "edge centers", or "face centers"
+        coord_type : str, default="spherical"
+            Indicates whether to remap using on spherical or cartesian coordinates
+        func : func, default = np.mean
+            Function to apply to neighborhood
+        r : float, default=1
+            Radius of neighborhood in deg
+        """
+        if destination_grid is None:
+            raise ValueError("Destination needed for remap.")
+
+        return _apply_func_remap_uxda(
+            self.uxda, destination_grid, remap_to, coord_type, func, r
+        )
+
     def inverse_distance_weighted(
         self,
         destination_grid: Optional[Grid] = None,
@@ -90,7 +115,7 @@ def inverse_distance_weighted(
         destination_obj : Grid, UxDataArray, UxDataset
             Optional destination for remapping, deprecating
         remap_to : str, default="nodes"
-            Location of where to map data, either "nodes" or "face centers"
+            Location of where to map data, either "nodes", "edge centers", or "face centers"
         coord_type : str, default="spherical"
             Indicates whether to remap using on spherical or cartesian coordinates
         power : int, default=2
@@ -99,6 +124,7 @@ def inverse_distance_weighted(
         k : int, default=8
             Number of nearest neighbors to consider in the weighted calculation.
         """
+
         if destination_grid is not None and destination_obj is not None:
             raise ValueError(
                 "Only one destination allowed, "
diff --git a/uxarray/remap/dataset_accessor.py b/uxarray/remap/dataset_accessor.py
index d5a9edf3f..20ccd9fb1 100644
--- a/uxarray/remap/dataset_accessor.py
+++ b/uxarray/remap/dataset_accessor.py
@@ -6,13 +6,14 @@
 from uxarray.remap.inverse_distance_weighted import (
     _inverse_distance_weighted_remap_uxds,
 )
+from uxarray.remap.apply_func import _apply_func_remap_uxds
 
 if TYPE_CHECKING:
     from uxarray.core.dataset import UxDataset
     from uxarray.core.dataarray import UxDataArray
 
 from uxarray.grid import Grid
-
+import numpy as np
 
 class UxDatasetRemapAccessor:
     def __init__(self, uxds: UxDataset):
@@ -26,13 +27,13 @@ def __repr__(self):
             "  * nearest_neighbor(destination_obj, remap_to, coord_type)\n"
         )
         methods_heading += "  * inverse_distance_weighted(destination_obj, remap_to, coord_type, power, k)\n"
+        methods_heading += "  * apply_func(destination_grid, remap_to, coord_type, func, r)\n"
 
         return prefix + methods_heading
 
     def nearest_neighbor(
         self,
-        destination_grid: Optional[Grid] = None,
-        destination_obj: Optional[Grid, UxDataArray, UxDataset] = None,
+        destination_grid: Grid = None,
         remap_to: str = "face centers",
         coord_type: str = "spherical",
     ):
@@ -43,20 +44,12 @@ def nearest_neighbor(
         ---------
         destination_grid : Grid
             Destination Grid for remapping
-        destination_obj : Grid, UxDataArray, UxDataset
-            Optional destination for remapping, deprecating
         remap_to : str, default="nodes"
             Location of where to map data, either "nodes", "edge centers", or "face centers"
         coord_type : str, default="spherical"
             Indicates whether to remap using on spherical or cartesian coordinates
         """
-
-        if destination_grid is not None and destination_obj is not None:
-            raise ValueError(
-                "Only one destination allowed, "
-                "please remove either `destination_grid` or `destination_obj`."
-            )
-        elif destination_grid is None and destination_obj is None:
+        if destination_grid is None:
             raise ValueError("Destination needed for remap.")
 
         if destination_grid is not None:
@@ -72,6 +65,37 @@ def nearest_neighbor(
                 self.uxds, destination_obj, remap_to, coord_type
             )
 
+    def apply_func(
+        self,
+        destination_grid: Grid = None,
+        remap_to: str = "face centers",
+        coord_type: str = "spherical",
+        func: func = np.mean,
+        r=1,
+    ):
+        """Neighborhood function Remapping between a source
+        (``UxDataset``) and destination.`.
+
+        Parameters
+        ---------
+        destination_grid : Grid
+            Destination Grid for remapping
+        remap_to : str, default="nodes"
+            Location of where to map data, either "nodes", "edge centers", or "face centers"
+        coord_type : str, default="spherical"
+            Indicates whether to remap using on spherical or cartesian coordinates
+        func : func, default = np.mean
+            Function to apply to neighborhood
+        r : float, default=1
+            Radius of neighborhood in deg
+        """
+        if destination_grid is None:
+            raise ValueError("Destination needed for remap.")
+
+        return _apply_func_remap_uxds(
+            self.uxds, destination_grid, remap_to, coord_type, func, r
+        )
+
     def inverse_distance_weighted(
         self,
         destination_grid: Optional[Grid] = None,
@@ -81,8 +105,8 @@ def inverse_distance_weighted(
         power=2,
         k=8,
     ):
-        """Inverse Distance Weighted Remapping between a source (``UxDataset``)
-        and destination.`.
+        """Inverse Distance Weighted Remapping between a source
+        (``UxDataset``) and destination.`.
 
         Parameters
         ---------

From af2d13b35830476e514f19529510d50c619530ad Mon Sep 17 00:00:00 2001
From: ahijevyc <ahijevyc@ucar.edu>
Date: Fri, 2 Aug 2024 12:14:39 -0600
Subject: [PATCH 2/7] Cleaned apply_func.py

default r=1 (degrees)

No more destination_obj, only destination_grid

Don't think we need to add remapped variable to existing UxDataset or
construct a UxDataset from remapped variable and existing variable.

Use [var_name] when constructing destination_uxds.
---
 uxarray/remap/apply_func.py | 64 ++++++++++---------------------------
 1 file changed, 16 insertions(+), 48 deletions(-)

diff --git a/uxarray/remap/apply_func.py b/uxarray/remap/apply_func.py
index 48d358fff..170673c97 100644
--- a/uxarray/remap/apply_func.py
+++ b/uxarray/remap/apply_func.py
@@ -21,7 +21,7 @@ def _apply_func_remap(
     remap_to: str = "face centers",
     coord_type: str = "spherical",
     func: func = np.mean,
-    r: float = None,
+    r: float = 1.,
 ) -> np.array:
     """Apply neighborhood function Remapping between two grids.
 
@@ -37,7 +37,7 @@ def _apply_func_remap(
         Location of where to map data, either "nodes", "edge centers", or "face centers".
     coord_type: str, default="spherical"
         Coordinate type to use for nearest neighbor query, either "spherical" or "Cartesian".
-    r : float, default=None
+    r : float, default=1.
         radius of neighborhoodFor spherical coordinates, the radius is in units of degrees,
         and for cartesian coordinates, the radius is in meters.
 
@@ -148,11 +148,11 @@ def _apply_func_remap(
 
 def _apply_func_remap_uxda(
     source_uxda: UxDataArray,
-    destination_obj: Union[Grid, UxDataArray, UxDataset],
+    destination_grid: Grid,
     remap_to: str = "face centers",
     coord_type: str = "spherical",
     func: func = np.mean,
-    r=5,
+    r=1.,
 ):
     """Neighborhood function Remapping implementation for ``UxDataArray``.
 
@@ -160,14 +160,14 @@ def _apply_func_remap_uxda(
     ---------
     source_uxda : UxDataArray
         Source UxDataArray for remapping
-    destination_obj : Grid, UxDataArray, UxDataset
-        Destination for remapping
+    destination_grid : Grid
+        Destination grid for remapping
     remap_to : str, default="nodes"
         Location of where to map data, either "nodes", "edge centers", or "face centers"
     coord_type : str, default="spherical"
         Indicates whether to remap using on Spherical or Cartesian coordinates for the computations when
         remapping.
-    r : float, default=5
+    r : float, default=1.
         Radius of neighborhood.
     """
 
@@ -192,16 +192,6 @@ def _apply_func_remap_uxda(
     destination_dims = list(source_uxda.dims)
     destination_dims[-1] = destination_dim
 
-    if isinstance(destination_obj, Grid):
-        destination_grid = destination_obj
-    elif isinstance(
-        destination_obj,
-        (uxarray.core.dataarray.UxDataArray, uxarray.core.dataset.UxDataset),
-    ):
-        destination_grid = destination_obj.uxgrid
-    else:
-        raise ValueError("TODO: Invalid Input")
-
     # perform remapping
     destination_data = _apply_func_remap(
         source_uxda.uxgrid,
@@ -220,30 +210,16 @@ def _apply_func_remap_uxda(
         dims=destination_dims,
         uxgrid=destination_grid,
     )
-    # add remapped variable to existing UxDataset
-    if isinstance(destination_obj, uxarray.core.dataset.UxDataset):
-        uxds = destination_obj.copy()
-        uxds[source_uxda.name] = uxda_remap
-        return uxds
-
-    # construct a UxDataset from remapped variable and existing variable
-    elif isinstance(destination_obj, uxarray.core.dataset.UxDataArray):
-        uxds = destination_obj.copy().to_dataset()
-        uxds[source_uxda.name] = uxda_remap
-        return uxds
-
-    # return UxDataArray with remapped variable
-    else:
-        return uxda_remap
+    return uxda_remap
 
 
 def _apply_func_remap_uxds(
     source_uxds: UxDataset,
-    destination_obj: Union[Grid, UxDataArray, UxDataset],
+    destination_grid: Grid,
     remap_to: str = "face centers",
     coord_type: str = "spherical",
     func: func = np.mean,
-    r: float = 5,
+    r: float = 1.,
 ):
     """Neighboohood function implementation for ``UxDataset``.
 
@@ -251,29 +227,21 @@ def _apply_func_remap_uxds(
     ---------
     source_uxds : UxDataset
         Source UxDataset for remapping
-    destination_obj : Grid, UxDataArray, UxDataset
-        Destination for remapping
+    destination_grid : Grid
+        Destination grid for remapping
     remap_to : str, default="nodes"
         Location of where to map data, either "nodes", "edge centers", or "face centers"
     coord_type : str, default="spherical"
         Indicates whether to remap using on Spherical or Cartesian coordinates
     func : func = np.mean
         function to apply to neighborhood
-    r : float, default=5
-        Radius of neighborhood
+    r : float, default=1.
+        Radius of neighborhood in deg
     """
 
-    if isinstance(destination_obj, Grid):
-        destination_uxds = uxarray.core.dataset.UxDataset(uxgrid=destination_obj)
-    elif isinstance(destination_obj, uxarray.core.dataset.UxDataArray):
-        destination_uxds = destination_obj.to_dataset()
-    elif isinstance(destination_obj, uxarray.core.dataset.UxDataset):
-        destination_uxds = destination_obj
-    else:
-        raise ValueError
-
+    destination_uxds = uxarray.core.dataset.UxDataset(uxgrid=destination_grid)
     for var_name in source_uxds.data_vars:
-        destination_uxds = _apply_func_remap_uxda(
+        destination_uxds[var_name] = _apply_func_remap_uxda(
             source_uxds[var_name],
             destination_uxds,
             remap_to,

From d0cb074bf9430577edc5d3dc5b6efb4d5d6d9191 Mon Sep 17 00:00:00 2001
From: ahijevyc <ahijevyc@ucar.edu>
Date: Thu, 29 Aug 2024 11:50:23 -0600
Subject: [PATCH 3/7] don't think you need axis=-1 arg

---
 uxarray/remap/apply_func.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/uxarray/remap/apply_func.py b/uxarray/remap/apply_func.py
index 170673c97..763f52398 100644
--- a/uxarray/remap/apply_func.py
+++ b/uxarray/remap/apply_func.py
@@ -6,7 +6,6 @@
     from uxarray.core.dataarray import UxDataArray
 
 import numpy as np
-from functools import partial
 
 import uxarray.core.dataarray
 import uxarray.core.dataset
@@ -138,7 +137,6 @@ def _apply_func_remap(
     destination_shape[-1] = len(neighbor_indices)
     destination_data = np.empty(destination_shape)
     # Apply function to indices on last axis.
-    func = partial(func, axis=-1)
     for i, idx in enumerate(neighbor_indices):
         if len(idx):
             destination_data[..., i] = func(source_data[..., idx])

From 236dc793e4c1e0fb04eaa6ffc8e7420707b303da Mon Sep 17 00:00:00 2001
From: ahijevyc <ahijevyc@ucar.edu>
Date: Fri, 30 Aug 2024 12:01:02 -0600
Subject: [PATCH 4/7] not depreciating destination_obj yet

---
 uxarray/remap/dataarray_accessor.py | 13 ++++++++++---
 uxarray/remap/dataset_accessor.py   | 13 +++++++++++--
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/uxarray/remap/dataarray_accessor.py b/uxarray/remap/dataarray_accessor.py
index f06196ce2..be730a718 100644
--- a/uxarray/remap/dataarray_accessor.py
+++ b/uxarray/remap/dataarray_accessor.py
@@ -33,7 +33,8 @@ def __repr__(self):
 
     def nearest_neighbor(
         self,
-        destination_grid: Grid = None,
+        destination_grid: Optional[Grid] = None,
+        destination_obj: Optional[Grid, UxDataArray, UxDataset] = None,
         remap_to: str = "face centers",
         coord_type: str = "spherical",
     ):
@@ -44,12 +45,19 @@ def nearest_neighbor(
         ---------
         destination_grid : Grid
             Destination Grid for remapping
+        destination_obj : Grid, UxDataArray, UxDataset
+            Optional destination for remapping, deprecating
         remap_to : str, default="nodes"
             Location of where to map data, either "nodes", "edge centers", or "face centers"
         coord_type : str, default="spherical"
             Indicates whether to remap using on spherical or cartesian coordinates
         """
-        if destination_grid is None:
+        if destination_grid is not None and destination_obj is not None:
+            raise ValueError(
+                "Only one destination allowed, "
+                "please remove either `destination_grid` or `destination_obj`."
+            )
+        elif destination_grid is None and destination_obj is None:
             raise ValueError("Destination needed for remap.")
 
         if destination_grid is not None:
@@ -124,7 +132,6 @@ def inverse_distance_weighted(
         k : int, default=8
             Number of nearest neighbors to consider in the weighted calculation.
         """
-
         if destination_grid is not None and destination_obj is not None:
             raise ValueError(
                 "Only one destination allowed, "
diff --git a/uxarray/remap/dataset_accessor.py b/uxarray/remap/dataset_accessor.py
index 20ccd9fb1..502cd108b 100644
--- a/uxarray/remap/dataset_accessor.py
+++ b/uxarray/remap/dataset_accessor.py
@@ -33,7 +33,8 @@ def __repr__(self):
 
     def nearest_neighbor(
         self,
-        destination_grid: Grid = None,
+        destination_grid: Optional[Grid] = None,
+        destination_obj: Optional[Grid, UxDataArray, UxDataset] = None,
         remap_to: str = "face centers",
         coord_type: str = "spherical",
     ):
@@ -44,12 +45,20 @@ def nearest_neighbor(
         ---------
         destination_grid : Grid
             Destination Grid for remapping
+        destination_obj : Grid, UxDataArray, UxDataset
+            Optional destination for remapping, deprecating
         remap_to : str, default="nodes"
             Location of where to map data, either "nodes", "edge centers", or "face centers"
         coord_type : str, default="spherical"
             Indicates whether to remap using on spherical or cartesian coordinates
         """
-        if destination_grid is None:
+
+        if destination_grid is not None and destination_obj is not None:
+            raise ValueError(
+                "Only one destination allowed, "
+                "please remove either `destination_grid` or `destination_obj`."
+            )
+        elif destination_grid is None and destination_obj is None:
             raise ValueError("Destination needed for remap.")
 
         if destination_grid is not None:

From b1575470284f2dde9628e3141ffcc3c6671929c0 Mon Sep 17 00:00:00 2001
From: ahijevyc <ahijevyc@ucar.edu>
Date: Fri, 30 Aug 2024 12:26:13 -0600
Subject: [PATCH 5/7] pre-commit format fixes

---
 uxarray/remap/apply_func.py         |  8 ++++----
 uxarray/remap/dataarray_accessor.py |  9 ++++++---
 uxarray/remap/dataset_accessor.py   | 13 ++++++++-----
 3 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/uxarray/remap/apply_func.py b/uxarray/remap/apply_func.py
index 763f52398..cf54f86a1 100644
--- a/uxarray/remap/apply_func.py
+++ b/uxarray/remap/apply_func.py
@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from uxarray.core.dataset import UxDataset
@@ -20,7 +20,7 @@ def _apply_func_remap(
     remap_to: str = "face centers",
     coord_type: str = "spherical",
     func: func = np.mean,
-    r: float = 1.,
+    r: float = 1.0,
 ) -> np.array:
     """Apply neighborhood function Remapping between two grids.
 
@@ -150,7 +150,7 @@ def _apply_func_remap_uxda(
     remap_to: str = "face centers",
     coord_type: str = "spherical",
     func: func = np.mean,
-    r=1.,
+    r=1.0,
 ):
     """Neighborhood function Remapping implementation for ``UxDataArray``.
 
@@ -217,7 +217,7 @@ def _apply_func_remap_uxds(
     remap_to: str = "face centers",
     coord_type: str = "spherical",
     func: func = np.mean,
-    r: float = 1.,
+    r: float = 1.0,
 ):
     """Neighboohood function implementation for ``UxDataset``.
 
diff --git a/uxarray/remap/dataarray_accessor.py b/uxarray/remap/dataarray_accessor.py
index be730a718..8f20ddbc5 100644
--- a/uxarray/remap/dataarray_accessor.py
+++ b/uxarray/remap/dataarray_accessor.py
@@ -15,6 +15,7 @@
 from uxarray.grid import Grid
 import numpy as np
 
+
 class UxDataArrayRemapAccessor:
     def __init__(self, uxda: UxDataArray):
         self.uxda = uxda
@@ -27,7 +28,9 @@ def __repr__(self):
             "  * nearest_neighbor(destination_obj, remap_to, coord_type)\n"
         )
         methods_heading += "  * inverse_distance_weighted(destination_obj, remap_to, coord_type, power, k)\n"
-        methods_heading += "  * apply_func(destination_grid, remap_to, coord_type, func, r)\n"
+        methods_heading += (
+                "  * apply_func(destination_grid, remap_to, coord_type, func, r)\n"
+        )
 
         return prefix + methods_heading
 
@@ -81,8 +84,8 @@ def apply_func(
         func: func = np.mean,
         r=1,
     ):
-        """Neighborhood function Remapping between a source
-        (``UxDataArray``) and destination.`.
+        """Neighborhood function Remapping between a source (``UxDataArray``)
+        and destination.`.
 
         Parameters
         ---------
diff --git a/uxarray/remap/dataset_accessor.py b/uxarray/remap/dataset_accessor.py
index 502cd108b..2d954388d 100644
--- a/uxarray/remap/dataset_accessor.py
+++ b/uxarray/remap/dataset_accessor.py
@@ -15,6 +15,7 @@
 from uxarray.grid import Grid
 import numpy as np
 
+
 class UxDatasetRemapAccessor:
     def __init__(self, uxds: UxDataset):
         self.uxds = uxds
@@ -27,7 +28,9 @@ def __repr__(self):
             "  * nearest_neighbor(destination_obj, remap_to, coord_type)\n"
         )
         methods_heading += "  * inverse_distance_weighted(destination_obj, remap_to, coord_type, power, k)\n"
-        methods_heading += "  * apply_func(destination_grid, remap_to, coord_type, func, r)\n"
+        methods_heading += (
+                "  * apply_func(destination_grid, remap_to, coord_type, func, r)\n"
+        )
 
         return prefix + methods_heading
 
@@ -82,8 +85,8 @@ def apply_func(
         func: func = np.mean,
         r=1,
     ):
-        """Neighborhood function Remapping between a source
-        (``UxDataset``) and destination.`.
+        """Neighborhood function Remapping between a source (``UxDataset``) and
+        destination.`.
 
         Parameters
         ---------
@@ -114,8 +117,8 @@ def inverse_distance_weighted(
         power=2,
         k=8,
     ):
-        """Inverse Distance Weighted Remapping between a source
-        (``UxDataset``) and destination.`.
+        """Inverse Distance Weighted Remapping between a source (``UxDataset``)
+        and destination.`.
 
         Parameters
         ---------

From 90003968f0edd88c655645d15e6e785512b64545 Mon Sep 17 00:00:00 2001
From: ahijevyc <ahijevyc@ucar.edu>
Date: Fri, 30 Aug 2024 12:32:14 -0600
Subject: [PATCH 6/7] pre-commit ruff format fixes

---
 uxarray/remap/apply_func.py         | 1 -
 uxarray/remap/dataarray_accessor.py | 2 +-
 uxarray/remap/dataset_accessor.py   | 2 +-
 3 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/uxarray/remap/apply_func.py b/uxarray/remap/apply_func.py
index cf54f86a1..1d5fe097b 100644
--- a/uxarray/remap/apply_func.py
+++ b/uxarray/remap/apply_func.py
@@ -46,7 +46,6 @@ def _apply_func_remap(
         Data mapped to the destination grid.
     """
 
-
     source_data = np.asarray(source_data)
     n_elements = source_data.shape[-1]
 
diff --git a/uxarray/remap/dataarray_accessor.py b/uxarray/remap/dataarray_accessor.py
index 8f20ddbc5..a3ee83a9f 100644
--- a/uxarray/remap/dataarray_accessor.py
+++ b/uxarray/remap/dataarray_accessor.py
@@ -29,7 +29,7 @@ def __repr__(self):
         )
         methods_heading += "  * inverse_distance_weighted(destination_obj, remap_to, coord_type, power, k)\n"
         methods_heading += (
-                "  * apply_func(destination_grid, remap_to, coord_type, func, r)\n"
+            "  * apply_func(destination_grid, remap_to, coord_type, func, r)\n"
         )
 
         return prefix + methods_heading
diff --git a/uxarray/remap/dataset_accessor.py b/uxarray/remap/dataset_accessor.py
index 2d954388d..6346b4e15 100644
--- a/uxarray/remap/dataset_accessor.py
+++ b/uxarray/remap/dataset_accessor.py
@@ -29,7 +29,7 @@ def __repr__(self):
         )
         methods_heading += "  * inverse_distance_weighted(destination_obj, remap_to, coord_type, power, k)\n"
         methods_heading += (
-                "  * apply_func(destination_grid, remap_to, coord_type, func, r)\n"
+            "  * apply_func(destination_grid, remap_to, coord_type, func, r)\n"
         )
 
         return prefix + methods_heading

From 22f98e572b0b5e95801065d65f6b96c47ba97cde Mon Sep 17 00:00:00 2001
From: ahijevyc <ahijevyc@ucar.edu>
Date: Fri, 30 Aug 2024 12:44:34 -0600
Subject: [PATCH 7/7] pre-commit ruff Callable TYPING instead of func

---
 uxarray/remap/apply_func.py         | 9 +++++----
 uxarray/remap/dataarray_accessor.py | 5 +++--
 uxarray/remap/dataset_accessor.py   | 5 +++--
 3 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/uxarray/remap/apply_func.py b/uxarray/remap/apply_func.py
index 1d5fe097b..e5e914856 100644
--- a/uxarray/remap/apply_func.py
+++ b/uxarray/remap/apply_func.py
@@ -1,4 +1,5 @@
 from __future__ import annotations
+from collections.abc import Callable
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
@@ -19,7 +20,7 @@ def _apply_func_remap(
     source_data: np.ndarray,
     remap_to: str = "face centers",
     coord_type: str = "spherical",
-    func: func = np.mean,
+    func: Callable = np.mean,
     r: float = 1.0,
 ) -> np.array:
     """Apply neighborhood function Remapping between two grids.
@@ -148,7 +149,7 @@ def _apply_func_remap_uxda(
     destination_grid: Grid,
     remap_to: str = "face centers",
     coord_type: str = "spherical",
-    func: func = np.mean,
+    func: Callable = np.mean,
     r=1.0,
 ):
     """Neighborhood function Remapping implementation for ``UxDataArray``.
@@ -215,7 +216,7 @@ def _apply_func_remap_uxds(
     destination_grid: Grid,
     remap_to: str = "face centers",
     coord_type: str = "spherical",
-    func: func = np.mean,
+    func: Callable = np.mean,
     r: float = 1.0,
 ):
     """Neighboohood function implementation for ``UxDataset``.
@@ -230,7 +231,7 @@ def _apply_func_remap_uxds(
         Location of where to map data, either "nodes", "edge centers", or "face centers"
     coord_type : str, default="spherical"
         Indicates whether to remap using on Spherical or Cartesian coordinates
-    func : func = np.mean
+    func : Callable = np.mean
         function to apply to neighborhood
     r : float, default=1.
         Radius of neighborhood in deg
diff --git a/uxarray/remap/dataarray_accessor.py b/uxarray/remap/dataarray_accessor.py
index a3ee83a9f..5bfb6a06d 100644
--- a/uxarray/remap/dataarray_accessor.py
+++ b/uxarray/remap/dataarray_accessor.py
@@ -1,4 +1,5 @@
 from __future__ import annotations
+from collections.abc import Callable
 from typing import TYPE_CHECKING, Optional
 from warnings import warn
 
@@ -81,7 +82,7 @@ def apply_func(
         destination_grid: Grid = None,
         remap_to: str = "face centers",
         coord_type: str = "spherical",
-        func: func = np.mean,
+        func: Callable = np.mean,
         r=1,
     ):
         """Neighborhood function Remapping between a source (``UxDataArray``)
@@ -95,7 +96,7 @@ def apply_func(
             Location of where to map data, either "nodes", "edge centers", or "face centers"
         coord_type : str, default="spherical"
             Indicates whether to remap using on spherical or cartesian coordinates
-        func : func, default = np.mean
+        func : Callable, default = np.mean
             Function to apply to neighborhood
         r : float, default=1
             Radius of neighborhood in deg
diff --git a/uxarray/remap/dataset_accessor.py b/uxarray/remap/dataset_accessor.py
index 6346b4e15..4ab30ebf8 100644
--- a/uxarray/remap/dataset_accessor.py
+++ b/uxarray/remap/dataset_accessor.py
@@ -1,4 +1,5 @@
 from __future__ import annotations
+from collections.abc import Callable
 from typing import TYPE_CHECKING, Optional
 from warnings import warn
 
@@ -82,7 +83,7 @@ def apply_func(
         destination_grid: Grid = None,
         remap_to: str = "face centers",
         coord_type: str = "spherical",
-        func: func = np.mean,
+        func: Callable = np.mean,
         r=1,
     ):
         """Neighborhood function Remapping between a source (``UxDataset``) and
@@ -96,7 +97,7 @@ def apply_func(
             Location of where to map data, either "nodes", "edge centers", or "face centers"
         coord_type : str, default="spherical"
             Indicates whether to remap using on spherical or cartesian coordinates
-        func : func, default = np.mean
+        func : Callable, default = np.mean
             Function to apply to neighborhood
         r : float, default=1
             Radius of neighborhood in deg