diff --git a/docs/conf.py b/docs/conf.py index c4c2a97..48e5615 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -119,6 +119,7 @@ 'xarray': ('https://docs.xarray.dev/en/stable/', None), 'ndfilters': ('https://ndfilters.readthedocs.io/en/stable', None), 'colorsynth': ('https://colorsynth.readthedocs.io/en/stable', None), + 'regridding': ('https://regridding.readthedocs.io/en/stable', None), } # plt.Axes.__module__ = matplotlib.axes.__name__ diff --git a/named_arrays/__init__.py b/named_arrays/__init__.py index b564eab..4b07371 100644 --- a/named_arrays/__init__.py +++ b/named_arrays/__init__.py @@ -5,6 +5,7 @@ from . import transformations from . import ndfilters from . import colorsynth +from . import regridding from ._core import * from ._scalars.scalars import * from ._scalars.uncertainties.uncertainties import * diff --git a/named_arrays/_scalars/scalar_named_array_functions.py b/named_arrays/_scalars/scalar_named_array_functions.py index 3f3d7cf..f56b46a 100644 --- a/named_arrays/_scalars/scalar_named_array_functions.py +++ b/named_arrays/_scalars/scalar_named_array_functions.py @@ -9,6 +9,7 @@ import astroscrappy import ndfilters import colorsynth +import regridding import named_arrays as na from . import scalars @@ -1422,6 +1423,83 @@ def ndfilter( ) +@_implements(na.regridding.weights) +def regridding_weights( + coordinates_input: na.AbstractScalarArray | na.AbstractVectorArray, + coordinates_output: na.AbstractScalarArray | na.AbstractVectorArray, + axis_input: None | str | Sequence[str] = None, + axis_output: None | str | Sequence[str] = None, + method: Literal['multilinear', 'conservative'] = 'multilinear', +) -> tuple[na.AbstractScalar, dict[str, int], dict[str, int]]: + + if not isinstance(coordinates_output, na.AbstractVectorArray): + coordinates_output = na.CartesianNdVectorArray(dict(x=coordinates_output)) + + if not isinstance(coordinates_input, na.AbstractVectorArray): + coordinates_input = na.CartesianNdVectorArray(dict(x=coordinates_input)) + + return na.regridding.weights( + coordinates_input=coordinates_input, + coordinates_output=coordinates_output, + axis_input=axis_input, + axis_output=axis_output, + method=method, + ) + + +@_implements(na.regridding.regrid_from_weights) +def regridding_regrid_from_weights( + weights: na.AbstractScalarArray, + shape_input: dict[str, int], + shape_output: dict[str, int], + values_input: na.AbstractScalarArray, +) -> na.ScalarArray: + + try: + weights = scalars._normalize(weights) + values_input = scalars._normalize(values_input) + except scalars.ScalarTypeError: # pragma: nocover + return NotImplemented + + shape_weights = weights.shape + + axis_input = tuple(a for a in shape_input if a not in shape_weights) + axis_output = tuple(a for a in shape_output if a not in shape_weights) + + shape_values_input = values_input.shape + shape_orthogonal = { + a: shape_values_input[a] + for a in shape_values_input + if a not in axis_input + } + shape_orthogonal = na.broadcast_shapes(shape_orthogonal, shape_weights) + + shape_input = na.broadcast_shapes(shape_orthogonal, shape_input) + shape_output = na.broadcast_shapes(shape_orthogonal, shape_output) + + weights = weights.broadcast_to({ + a: shape_input[a] if a not in axis_input else 1 + for a in shape_input + }) + values_input = values_input.broadcast_to(shape_input) + + result = regridding.regrid_from_weights( + weights=weights.ndarray, + shape_input=tuple(shape_input.values()), + shape_output=tuple(shape_output.values()), + values_input=values_input.ndarray, + axis_input=tuple(tuple(shape_input).index(a) for a in axis_input), + axis_output=tuple(tuple(shape_output).index(a) for a in axis_output), + ) + + result = na.ScalarArray( + ndarray=result, + axes=tuple(shape_output), + ) + + return result + + @_implements(na.despike) def despike( array: na.AbstractScalarArray, diff --git a/named_arrays/_vectors/vector_named_array_functions.py b/named_arrays/_vectors/vector_named_array_functions.py index 4bbc7c5..3d511a3 100644 --- a/named_arrays/_vectors/vector_named_array_functions.py +++ b/named_arrays/_vectors/vector_named_array_functions.py @@ -1,8 +1,9 @@ -from typing import Callable, TypeVar +from typing import Callable, TypeVar, Sequence, Literal import numpy as np import numpy.typing as npt import matplotlib.axes import astropy.units as u +import regridding import named_arrays as na from named_arrays._scalars import scalars import named_arrays._scalars.scalar_named_array_functions @@ -603,3 +604,101 @@ def ndfilter( result = prototype.type_explicit.from_components(result) return result + + +@_implements(na.regridding.weights) +def regridding_weights( + coordinates_input: na.AbstractVectorArray, + coordinates_output: na.AbstractVectorArray, + axis_input: None | str | Sequence[str] = None, + axis_output: None | str | Sequence[str] = None, + method: Literal['multilinear', 'conservative'] = 'multilinear', +) -> tuple[na.AbstractScalar, dict[str, int], dict[str, int]]: + + try: + prototype = vectors._prototype(coordinates_input, coordinates_output) + coordinates_input = vectors._normalize(coordinates_input, prototype) + coordinates_output = vectors._normalize(coordinates_output, prototype) + except vectors.VectorTypeError: # pragma: nocover + return NotImplemented + + try: + coordinates_output = coordinates_output.components + coordinates_output = { + c: scalars._normalize(coordinates_output[c]) + for c in coordinates_output + if coordinates_output[c] is not None + } + coordinates_output = na.CartesianNdVectorArray(coordinates_output) + except scalars.ScalarTypeError: # pragma: nocover + return NotImplemented + + try: + coordinates_input = coordinates_input.components + coordinates_input = { + c: scalars._normalize(coordinates_input[c]) + for c in coordinates_output.components + } + coordinates_input = na.CartesianNdVectorArray(coordinates_input) + except scalars.ScalarTypeError: # pragma: nocover + return NotImplemented + + coordinates_output = coordinates_output.explicit + coordinates_input = coordinates_input.explicit + + shape_input = coordinates_input.shape + shape_output = coordinates_output.shape + + if axis_input is None: + axis_input = tuple(shape_input) + elif isinstance(axis_input, str): + axis_input = (axis_input,) + + if axis_output is None: + axis_output = tuple(shape_output) + elif isinstance(axis_output, str): + axis_output = (axis_output,) + + shape_orthogonal_input = { + a: shape_input[a] + for a in shape_input if a not in axis_input + } + shape_orthogonal_output = { + a: shape_output[a] + for a in shape_output if a not in axis_output + } + + shape_orthogonal = na.broadcast_shapes( + shape_orthogonal_input, + shape_orthogonal_output, + ) + + shape_input = na.broadcast_shapes(shape_orthogonal, shape_input) + shape_output = na.broadcast_shapes(shape_orthogonal, shape_output) + + coordinates_input = coordinates_input.broadcast_to(shape_input) + coordinates_output = coordinates_output.broadcast_to(shape_output) + + coordinates_input = coordinates_input.components + coordinates_output = coordinates_output.components + + coordinates_input = tuple(coordinates_input[c].ndarray for c in coordinates_input) + coordinates_output = tuple(coordinates_output[c].ndarray for c in coordinates_output) + + axis_input = tuple(tuple(shape_input).index(a) for a in axis_input) + axis_output = tuple(tuple(shape_output).index(a) for a in axis_output) + + result, _shape_input, _shape_output = regridding.weights( + coordinates_input=coordinates_input, + coordinates_output=coordinates_output, + axis_input=axis_input, + axis_output=axis_output, + method=method, + ) + + result = na.ScalarArray(result, tuple(shape_orthogonal)) + + shape_input = dict(zip(shape_input, _shape_input)) + shape_output = dict(zip(shape_output, _shape_output)) + + return result, shape_input, shape_output diff --git a/named_arrays/regridding.py b/named_arrays/regridding.py new file mode 100644 index 0000000..9671f6d --- /dev/null +++ b/named_arrays/regridding.py @@ -0,0 +1,218 @@ +""" +Array resampling and interpolation. + +A wrapper around the :mod:`regridding` module for named arrays. +""" + +from __future__ import annotations +from typing import Sequence, Literal +import astropy.units as u +import named_arrays as na + +__all__ = [ + "regrid", + "weights", + "regrid_from_weights", +] + + +def regrid( + coordinates_input: na.AbstractScalar | na.AbstractVectorArray, + coordinates_output: float | u.Quantity, + values_input: na.AbstractScalarArray, + axis_input: None | Sequence[str] = None, + axis_output: None | Sequence[str] = None, + method: Literal['multilinear', 'conservative'] = 'multilinear', +) -> na.AbstractScalarArray: + """ + Regrid an array of values defined on a logically-rectangular curvilinear + grid onto a new logically-rectangular curvilinear grid. + + Parameters + ---------- + coordinates_input + Coordinates of the input grid. + coordinates_output + Coordinates of the output grid. + Should have the same number of components as the input grid. + values_input + Input array of values to be resampled. + axis_input + Logical axes of the input grid to resample. + If :obj:`None`, resample all the axes of the input grid. + The number of axes should be equal to the number of + coordinates in the input grid. + axis_output + Logical axes of the output grid corresponding to the resampled axes + of the input grid. + If :obj:`None`, all the axes of the output grid correspond to resampled + axes in the input grid. + The number of axes should be equal to the number of + coordinates in the output grid. + method + The type of regridding to use. + + Examples + -------- + + Regrid a 2D array using conservative resampling. + + .. jupyter-execute:: + + import numpy as np + import matplotlib.pyplot as plt + import named_arrays as na + + # Define the number of edges in the input grid + num_x = 66 + num_y = 66 + + # Define a dummy linear grid + x = na.linspace(-5, 5, axis="x", num=num_x) + y = na.linspace(-5, 5, axis="y", num=num_y) + + # Define the curvilinear input grid using the dummy grid + angle = 0.4 + coordinates_input = na.Cartesian2dVectorArray( + x=x * np.cos(angle) - y * np.sin(angle) + 0.05 * x * x, + y=x * np.sin(angle) + y * np.cos(angle) + 0.05 * y * y, + ) + + # Define the test pattern + a_input = np.cos(np.square(x)) * np.cos(np.square(y)) + a_input = a_input.cell_centers() + + # Define a rectilinear output grid using the limits of the input grid + coordinates_output = na.Cartesian2dVectorLinearSpace( + start=coordinates_input.min(), + stop=coordinates_input.max(), + axis=na.Cartesian2dVectorArray("x2", "y2"), + num=66, + ) + + # Regrid the test pattern onto the new grid + a_output = na.regridding.regrid( + coordinates_input=coordinates_input, + coordinates_output=coordinates_output, + values_input=a_input, + method="conservative", + ) + + fig, ax = plt.subplots( + ncols=2, + sharex=True, + sharey=True, + figsize=(8, 4), + constrained_layout=True, + ); + na.plt.pcolormesh(coordinates_input, C=a_input, ax=ax[0]) + na.plt.pcolormesh(coordinates_output, C=a_output, ax=ax[1]) + ax[0].set_title("input array"); + ax[1].set_title("regridded array"); + """ + + _weights, shape_input, shape_output = weights( + coordinates_input=coordinates_input, + coordinates_output=coordinates_output, + axis_input=axis_input, + axis_output=axis_output, + method=method + ) + + result = regrid_from_weights( + weights=_weights, + shape_input=shape_input, + shape_output=shape_output, + values_input=values_input, + # values_output=values_output, + # axis_input=axis_input, + # axis_output=axis_output, + ) + + return result + + +def weights( + coordinates_input: na.AbstractScalar | na.AbstractVectorArray, + coordinates_output: na.AbstractScalar | na.AbstractVectorArray, + axis_input: None | str | Sequence[str] = None, + axis_output: None | str | Sequence[str] = None, + method: Literal['multilinear', 'conservative'] = 'multilinear', +) -> tuple[na.AbstractScalar, dict[str, int], dict[str, int]]: + """ + Save the results of a regridding operation as a sequence of weights, + which can be used in subsequent regridding operations on the same grid. + + The results of this function are designed to be used by + :func:`regrid_from_weights` + + This function returns a tuple containing a ragged array of weights, + the shape of the input coordinates, and the shape of the output coordinates. + + Parameters + ---------- + coordinates_input + Coordinates of the input grid. + coordinates_output + Coordinates of the output grid. + Should have the same number of coordinates as the input grid. + axis_input + Logical axes of the input grid to resample. + If :obj:`None`, resample all the axes of the input grid. + The number of axes should be equal to the number of + coordinates in the input grid. + axis_output + Logical axes of the output grid corresponding to the resampled axes + of the input grid. + If :obj:`None`, all the axes of the output grid correspond to resampled + axes in the input grid. + The number of axes should be equal to the number of + coordinates in the output grid. + method + The type of regridding to use. + + See Also + -------- + :func:`regridding.weights`: An equivalent function for instances of :class:`numpy.ndarray`. + :func:`regrid_from_weights`: A function designed to use the outputs of this function. + :func:`regrid`: Resample an array without saving the weights + + """ + return na._named_array_function( + func=weights, + coordinates_input=coordinates_input, + coordinates_output=coordinates_output, + axis_input=axis_input, + axis_output=axis_output, + method=method, + ) + + +def regrid_from_weights( + weights: na.AbstractScalar, + shape_input: dict[str, int], + shape_output: dict[str, int], + values_input: na.AbstractScalar, +) -> na.AbstractScalar: + """ + Regrid an array of values using weights computed by + :func:`weights`. + + Parameters + ---------- + weights + Ragged array of weights computed by :func:`weights`. + shape_input + Broadcasted shape of the input coordinates computed by :func:`weights`. + shape_output + Broadcasted shape of the output coordinates computed by :func:`weights`. + values_input + Input array of values to be resampled. + """ + return na._named_array_function( + func=regrid_from_weights, + weights=weights, + shape_input=shape_input, + shape_output=shape_output, + values_input=values_input, + ) diff --git a/named_arrays/tests/test_regridding.py b/named_arrays/tests/test_regridding.py new file mode 100644 index 0000000..311c322 --- /dev/null +++ b/named_arrays/tests/test_regridding.py @@ -0,0 +1,132 @@ +import pytest +import numpy as np +import named_arrays as na + +shape_vertices = dict(x=10, y=11) +shape_centers = {a: shape_vertices[a] - 1 for a in shape_vertices} + +x = na.linspace(-1, 1, axis="x", num=shape_vertices["x"]) +y = na.linspace(-1, 1, axis="y", num=shape_vertices["y"]) +z = na.linspace(-1, 1, axis="z", num=3) + +x_new = na.linspace(-1, 1, axis="x_new", num=5) +y_new = na.linspace(-1, 1, axis="y_new", num=6) + + +@pytest.mark.parametrize( + argnames="coordinates_input,coordinates_output,values_input,axis_input,axis_output,result_expected", + argvalues=[ + ( + na.linspace(-1, 1, axis="x_input", num=11), + na.linspace(-1, 1, axis="x_output", num=11), + np.square(na.linspace(-1, 1, axis="x_input", num=11)), + None, + None, + np.square(na.linspace(-1, 1, axis="x_output", num=11)), + ), + ( + y, + y_new, + x + y, + "y", + "y_new", + x + y_new, + ), + ( + x, + x_new, + x + y, + ("x",), + ("x_new",), + x_new + y, + ), + ( + x, + 0.1 * x_new + 0.001 * y_new, + x, + ("x",), + ("x_new",), + 0.1 * x_new + 0.001 * y_new, + ), + ], +) +def test_regrid_multilinear_1d( + coordinates_input: tuple[np.ndarray, ...], + coordinates_output: tuple[np.ndarray, ...], + values_input: np.ndarray, + axis_input: None | int | tuple[int, ...], + axis_output: None | int | tuple[int, ...], + result_expected: np.ndarray, +): + result = na.regridding.regrid( + coordinates_input=coordinates_input, + coordinates_output=coordinates_output, + values_input=values_input, + axis_input=axis_input, + axis_output=axis_output, + method="multilinear", + ) + assert isinstance(result, na.AbstractArray) + assert np.issubdtype(result.dtype, float) + assert np.allclose(result, result_expected) + + +@pytest.mark.parametrize( + argnames="coordinates_input, values_input, axis_input, coordinates_output, axis_output", + argvalues=[ + ( + na.Cartesian2dVectorArray(x, y), + na.random.normal(0, 1, shape_random=shape_centers), + None, + na.Cartesian2dVectorArray( + x=1.1 * x + 0.01, + y=1.2 * y + 0.01, + ), + None, + ), + ( + na.Cartesian2dVectorArray( + x=x + 0.01 * z, + y=y + 0.01 * z, + ), + na.random.normal(0, 1, shape_random=shape_centers | z.shape), + ("x", "y"), + na.Cartesian2dVectorArray( + x=1.1 * (x + 0.001 * z) + 0.01, + y=1.2 * (y + 0.01 * z) + 0.001, + ), + ("x", "y"), + ), + ], +) +def test_regrid_conservative_2d( + coordinates_input: tuple[np.ndarray, ...], + coordinates_output: tuple[np.ndarray, ...], + values_input: np.ndarray, + axis_input: None | int | tuple[int, ...], + axis_output: None | int | tuple[int, ...], +): + result = na.regridding.regrid( + coordinates_input=coordinates_input, + coordinates_output=coordinates_output, + values_input=values_input, + axis_input=axis_input, + axis_output=axis_output, + method="conservative", + ) + + if axis_output is None: + axis_output = tuple(coordinates_output.shape) + elif isinstance(axis_output, str): + axis_output = (axis_output, ) + + shape_result = coordinates_output.shape + shape_result = { + a: shape_result[a] - 1 if a in axis_output + else shape_result[a] + for a in shape_result + } + + assert np.issubdtype(result.dtype, float) + assert result.shape == shape_result + assert np.allclose(result.sum(), values_input.sum()) diff --git a/pyproject.toml b/pyproject.toml index 40306cc..f31cde4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ dependencies = [ "astroscrappy", "ndfilters==0.3.0", "colorsynth==0.1.5", + "regridding==0.1.1", ] dynamic = ["version"]