From b690d715d8fc89058d91a6dbe248ce05975ef681 Mon Sep 17 00:00:00 2001 From: sidd3888 Date: Wed, 14 Aug 2024 22:26:58 +0530 Subject: [PATCH] added docstrings --- src/multinterp/backend/_cupy.py | 70 +++++++++++++++++++++ src/multinterp/backend/_jax.py | 57 +++++++++++++++++ src/multinterp/backend/_numba.py | 37 +++++++++++ src/multinterp/backend/_scipy.py | 70 +++++++++++++++++++++ src/multinterp/curvilinear/_scikit_image.py | 28 +++++++++ src/multinterp/grids.py | 1 + src/multinterp/rectilinear/_multi.py | 33 ++++++++++ src/multinterp/rectilinear/_utils.py | 11 ++++ 8 files changed, 307 insertions(+) diff --git a/src/multinterp/backend/_cupy.py b/src/multinterp/backend/_cupy.py index a989749..c783f3f 100644 --- a/src/multinterp/backend/_cupy.py +++ b/src/multinterp/backend/_cupy.py @@ -7,6 +7,25 @@ def cupy_multinterp(grids, values, args, options=None): + """ + Perform multivariate interpolation using CuPy. + + Parameters + ---------- + grids : array-like + Grid points in the domain. + values : array-like + Functional values at the grid points. + args : array-like + Points at which to interpolate data. + options : dict, optional + Additional options for interpolation. + + Returns + ------- + array-like + Interpolated values of the function. + """ mc_kwargs = update_mc_kwargs(options) args = cp.asarray(args) @@ -18,6 +37,27 @@ def cupy_multinterp(grids, values, args, options=None): def cupy_gradinterp(grids, values, args, axis=None, options=None): + """ + Computes the interpolated value of the gradient evaluated at specified points using CuPy. + + Parameters + ---------- + grids : list of array-like + Grid points in the domain. + values : array-like + Functional values at the grid points. + args : array-like + Points at which to interpolate data. + axis : int, optional + Axis along which to compute the gradient. + options : dict, optional + Additional options for interpolation. + + Returns + ------- + array-like + Interpolated values of the gradient. + """ mc_kwargs = update_mc_kwargs(options) eo = options.get("edge_order", 1) if options else 1 @@ -40,6 +80,21 @@ def cupy_gradinterp(grids, values, args, axis=None, options=None): def cupy_get_coordinates(grids, args): + """ + Takes input values and converts them to coordinates with respect to the specified grid. + + Parameters + ---------- + grids : cp.array + Grid points for each dimension. + args : cp.array + Points at which to interpolate data. + + Returns + ------- + cp.array + Coordinates with respect to the grid. + """ coords = cp.empty_like(args) for dim, grid in enumerate(grids): grid_size = cp.arange(grid.size) @@ -49,6 +104,21 @@ def cupy_get_coordinates(grids, args): def cupy_map_coordinates(values, coords, **kwargs): + """ + Run the map_coordinates function from the cupyx.scipy.ndimage module on the specified values. + + Parameters + ---------- + values : cp.array + Functional values from which to interpolate. + coords : cp.array + Coordinates at which to interpolate values. + + Returns + ------- + cp.array + Interpolated values. + """ original_shape = coords[0].shape coords = coords.reshape(len(values.shape), -1) output = map_coordinates(values, coords, **kwargs) diff --git a/src/multinterp/backend/_jax.py b/src/multinterp/backend/_jax.py index 4d59f45..e5f9b9e 100644 --- a/src/multinterp/backend/_jax.py +++ b/src/multinterp/backend/_jax.py @@ -41,6 +41,27 @@ def jax_multinterp(grids, values, args, options=None): def jax_gradinterp(grids, values, args, axis=None, options=None): + """ + Computes the interpolated value of the gradient evaluated at specified points using JAX. + + Parameters + ---------- + grids : list of array-like + Grid points in the domain. + values : array-like + Functional values at the grid points. + args : array-like + Points at which to interpolate data. + axis : int, optional + Axis along which to compute the gradient. + options : dict, optional + Additional options for interpolation. + + Returns + ------- + array-like + Interpolated values of the gradient. + """ mc_kwargs = update_mc_kwargs(options, jax=True) eo = options.get("edge_order", 1) if options else 1 @@ -64,6 +85,21 @@ def jax_gradinterp(grids, values, args, axis=None, options=None): @jit def jax_get_coordinates(grids, args): + """ + Takes input values and converts them to coordinates with respect to the specified grid. + + Parameters + ---------- + grids : jnp.array + Grid points for each dimension. + args : jnp.array + Points at which to interpolate data. + + Returns + ------- + jnp.array + Coordinates of the specified input points with respect to the grid. + """ grid_sizes = [jnp.arange(grid.size) for grid in grids] return jnp.array( [ @@ -75,6 +111,27 @@ def jax_get_coordinates(grids, args): @functools.partial(jit, static_argnums=(2, 3, 4)) def jax_map_coordinates(values, coords, order=None, mode=None, cval=None): + """ + Run the map_coordinates function from the jax.scipy.ndimage module on the specified values. + + Parameters + ---------- + values : jnp.array + The functional values from which to interpolate. + coords : jnp.array + The coordinates at which to interpolate the values. + order : int, optional + The order of interpolation, 0 for Nearest-Neighbour, 1 for Linear. + mode : str, optional + Method to handle extrapolation. See JAX documentation for options. + cval : float, optional + Value to use for extrapolation under 'constant' method. + + Returns + ------- + jnp.array + Interpolated values at specified coordinates. + """ original_shape = coords[0].shape coords = coords.reshape(len(values.shape), -1) output = map_coordinates(values, coords, order, mode, cval) diff --git a/src/multinterp/backend/_numba.py b/src/multinterp/backend/_numba.py index f246ad5..4d5ccbe 100644 --- a/src/multinterp/backend/_numba.py +++ b/src/multinterp/backend/_numba.py @@ -8,6 +8,25 @@ def numba_multinterp(grids, values, args, options=None): + """ + Perform multivariate interpolation using JIT-compiled functions with Numba. + + Parameters + ---------- + grids : array-like + Grid points in the domain. + values: array-like + Functional values at the grid points. + args: array-like + Points at which to interpolate data. + options: dict, optional + Additional options for interpolation. + + Returns + ------- + array-like + Interpolated values of the function. + """ mc_kwargs = update_mc_kwargs(options) args = np.asarray(args) @@ -20,6 +39,21 @@ def numba_multinterp(grids, values, args, options=None): @njit(parallel=True, cache=True, fastmath=True) def numba_get_coordinates(grids, args): + """ + Converts input arguments to coordinates with respect to the specified grid. JIT-compiled using Numba. + + Parameters + ---------- + grids : typed.List + Curvilinear grids for each dimension. + args : np.ndarray + Values in the domain at which the function is to be interpolated. + + Returns + ------- + np.ndarray + Coordinates of the input arguments. + """ coords = np.empty_like(args) for dim in prange(len(grids)): grid_size = np.arange(grids[dim].size) @@ -30,6 +64,9 @@ def numba_get_coordinates(grids, args): # same as scipy map_coordinates until replacement is found def numba_map_coordinates(values, coords, **kwargs): + """ + Identical to scipy_map_coordinates until a replacement is found. See documentation for scipy_map_coordinates. + """ original_shape = coords[0].shape coords = coords.reshape(len(values.shape), -1) output = map_coordinates(values, coords, **kwargs) diff --git a/src/multinterp/backend/_scipy.py b/src/multinterp/backend/_scipy.py index 9912f71..1769320 100644 --- a/src/multinterp/backend/_scipy.py +++ b/src/multinterp/backend/_scipy.py @@ -7,6 +7,25 @@ def scipy_multinterp(grids, values, args, options=None): + """ + Perform multivariate interpolation using SciPy. + + Parameters + ---------- + grids : list of array-like + Grid points in the domain. + values : array-like + Functional values at the grid points. + args : array-like + Points at which to interpolate data. + options : dict, optional + Additional options for interpolation. + + Returns + ------- + array-like + Interpolated values of the function. + """ mc_kwargs = update_mc_kwargs(options) args = np.asarray(args) @@ -18,6 +37,27 @@ def scipy_multinterp(grids, values, args, options=None): def scipy_gradinterp(grids, values, args, axis=None, options=None): + """ + Computes the interpolated value of the gradient evaluated at specified points using SciPy. + + Parameters + ---------- + grids : list of array-like + Grid points in the domain. + values : array-like + Functional values at the grid points. + args : array-like + Points at which to interpolate data. + axis : int, optional + Axis along which to compute the gradient. + options : dict, optional + Additional options for interpolation. + + Returns + ------- + array-like + Interpolated values of the gradient. + """ mc_kwargs = update_mc_kwargs(options) eo = options.get("edge_order", 1) if options else 1 @@ -40,6 +80,21 @@ def scipy_gradinterp(grids, values, args, axis=None, options=None): def scipy_get_coordinates(grids, args): + """ + Takes input values and converts them to coordinates with respect to the specified grid. + + Parameters + ---------- + grids : np.array + Grid points for each dimension. + args : np.array + Points at which to interpolate data. + + Returns + ------- + np.array + Coordinates with respect to the grid. + """ coords = np.empty_like(args) for dim, grid in enumerate(grids): grid_size = np.arange(grid.size) @@ -49,6 +104,21 @@ def scipy_get_coordinates(grids, args): def scipy_map_coordinates(values, coords, **kwargs): + """ + Run the map_coordinates function from the scipy.ndimage module on the specified values. + + Parameters + ---------- + values : np.array + Functional values from which to interpolate. + coords : np.array + Coordinates at which to interpolate values. + + Returns + ------- + np.array + Interpolated values of the function. + """ original_shape = coords[0].shape coords = coords.reshape(len(values.shape), -1) output = map_coordinates(values, coords, **kwargs) diff --git a/src/multinterp/curvilinear/_scikit_image.py b/src/multinterp/curvilinear/_scikit_image.py index 4b458a3..0c06acc 100644 --- a/src/multinterp/curvilinear/_scikit_image.py +++ b/src/multinterp/curvilinear/_scikit_image.py @@ -8,7 +8,22 @@ class PiecewiseAffineInterp(_CurvilinearGrid, MultivariateInterp): + """Curvilinear interpolator that uses the PiecewiseAffineTransform from scikit-image.""" + def __init__(self, values, grids, options=None): + """ + Initialize a PiecewiseAffineInterp object. + + Parameters + ---------- + values : np.ndarray + Functional values on a curvilinear grid. + grids : np.ndarray + Coordinates of the points in the curvilinear grid. + options : dict, optional + Additional keyword arguments to pass to the map_coordinates backend. + """ + super().__init__(values, grids, backend="scipy") self.mc_kwargs = update_mc_kwargs(options) @@ -25,6 +40,19 @@ def __init__(self, values, grids, options=None): self.interpolator = interpolator def _get_coordinates(self, args): + """Obtain the index coordinates for each of the arguments. + + Parameters + ---------- + args : np.ndarray + Arguments to be interpolated. + + Returns + ------- + np.ndarray + Index coordinates for each of the arguments. + """ + _input = args.reshape((self.ndim, -1)).T output = self.interpolator(_input).T.copy() return output.reshape(args.shape) diff --git a/src/multinterp/grids.py b/src/multinterp/grids.py index 967fc65..06383ac 100644 --- a/src/multinterp/grids.py +++ b/src/multinterp/grids.py @@ -33,6 +33,7 @@ def __init__(self, values, backend="scipy"): class _StructuredGrid(_AbstractGrid): + """Abstract class for interpolating on a structured grid. Serves as a base class for regular and abstract unstructured grid interpolators.""" def __init__(self, values, grids, backend="scipy"): super().__init__(values, backend=backend) diff --git a/src/multinterp/rectilinear/_multi.py b/src/multinterp/rectilinear/_multi.py index f8ce0e8..55ea239 100644 --- a/src/multinterp/rectilinear/_multi.py +++ b/src/multinterp/rectilinear/_multi.py @@ -91,6 +91,23 @@ def _map_coordinates(self, coords): return map_coords(self.values, coords, **self.mc_kwargs, backend=self.backend) def diff(self, axis=None, edge_order=1): + """ + Differentiates the interpolator along the specified axis. If axis is None, then returns a MultivaluedInterp object that approximates the partial derivative of the function across all axes. Otherwise, returns a MultivariateInterp object that approximates the partial derivative of the function along the specified axis. + + Parameters + ---------- + axis : int, optional + Axis along which to differentiate the function. + edge_order : int, optional + TODO: Add description + + Returns + ------- + MultivaluedInterp or MultivariateInterp + Interpolator object that approximates the partial derivative(s) of the function. + + """ + # if axis is not an integer less than or equal to the number # of dimensions of the input array, then a ValueError is raised. if axis is None: @@ -222,6 +239,22 @@ def _map_coordinates(self, coords): return asarray(fvals, backend=self.backend) def diff(self, axis=None, argnum=None, edge_order=1): + """ + Differentiates the interpolator along the specified axis. If both axis and argnum are specified, then returns the partial derivative of the specified function argument along the specified axis. If axis is None, then returns a MultivaluedInterp object that approximates the partial derivatives of the specified function argument along each axis. If argnum is None, then returns a MultivaluedInterp object that approximates the partial derivatives of all arguments of the function along the specified axes. + + Parameters + ---------- + axis : int, optional + Axis along which to differentiate the function. + edge_order : int, optional + TODO: Add description + + Returns + ------- + MultivaluedInterp or MultivariateInterp + Interpolator object that approximates the partial derivative(s) of the function. + + """ # if axis is not an integer less than or equal to the number # of dimensions of the input array, then a ValueError is raised. if axis is None: diff --git a/src/multinterp/rectilinear/_utils.py b/src/multinterp/rectilinear/_utils.py index 48e6376..2988eff 100644 --- a/src/multinterp/rectilinear/_utils.py +++ b/src/multinterp/rectilinear/_utils.py @@ -39,12 +39,23 @@ def get_coords(grids, args, backend="scipy"): + """ + Wrapper function for the get_coordinates function from the chosen backend. + """ return GET_COORDS[backend](grids, args) def map_coords(values, coords, backend="scipy", **kwargs): + """ + Wrapper function for the map_coordinates function from the chosen backend. + """ return MAP_COORDS[backend](values, coords, **kwargs) def get_grad(values, grids, axis=None, edge_order=None, backend="scipy"): + """ + Wrapper function for the gradient function from the chosen backend. + + TODO: use appropriate gradient functions from each backend. + """ return GET_GRAD[backend](values, grids, axis=axis, edge_order=edge_order)