diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml index 4a39fe5..81ab7c1 100644 --- a/.github/workflows/benchmarks.yml +++ b/.github/workflows/benchmarks.yml @@ -5,7 +5,6 @@ on: push: branches: - main - pull_request: permissions: pages: write diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index b56ee39..869c458 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -1,7 +1,11 @@ name: tests -on: [push, workflow_dispatch, pull_request] +on: + push: + branches: + - main + pull_request: jobs: build: diff --git a/docs/refs.bib b/docs/refs.bib index 1020aed..35c88d3 100644 --- a/docs/refs.bib +++ b/docs/refs.bib @@ -33,3 +33,16 @@ @article{Kumar2018 keywords = {Axis-crossing, Computational Geometry, Polygons, Point-in-Polygon test, Winding Number}, abstract = {This work is an extension of an axis-crossing algorithm to compute winding number for solving point in polygon for an arbitary polygon. Polygons are popular drawings in computer graphics to represent different types of structures with approximations. Solutions for point-in-polygons are many, like even-odd rule, positive-negative number, and winding number. This paper mainly deals with improvements of ‘A winding number and point in polygon algorithm’. Point in polygon is a fundamental problem and has various applications in ray tracing, computer graphics, image processing, gaming applications, robotics, acoustics, geo-science etc. The main focus of this paper explains about winding number for a closed polygon ‘S’, to test whether point ‘P’ lies either inside or outside with respect to positive and negative axis-crossing algorithm method.} } +@article{Ramshaw1985, + title = {Conservative rezoning algorithm for generalized two-dimensional meshes}, + journal = {Journal of Computational Physics}, + volume = {59}, + number = {2}, + pages = {193-199}, + year = {1985}, + issn = {0021-9991}, + doi = {https://doi.org/10.1016/0021-9991(85)90141-X}, + url = {https://www.sciencedirect.com/science/article/pii/002199918590141X}, + author = {John D Ramshaw}, + abstract = {A method is presented for transferring a conserved quantity Q from one generalized mesh to another when the volumetric density of Q is uniform within each cell of the original mesh.} +} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 0637925..29053e1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ test = [ ] doc = [ "pytest", + "scipy", "matplotlib", "graphviz", "sphinx-autodoc-typehints", diff --git a/regridding/__init__.py b/regridding/__init__.py index 186360a..1bf4a8b 100644 --- a/regridding/__init__.py +++ b/regridding/__init__.py @@ -1,5 +1,6 @@ from . import math from . import geometry from ._find_indices import * +from ._weights import * from ._interp_ndarray import * from ._regrid import * diff --git a/regridding/_regrid.py b/regridding/_regrid.py deleted file mode 100644 index 26e6d45..0000000 --- a/regridding/_regrid.py +++ /dev/null @@ -1,408 +0,0 @@ -from __future__ import annotations - -import dataclasses -from typing import Sequence -import numpy as np -import numba -from ._conservative_ramshaw import _conservative_ramshaw -from ._util import _normalize_axis - -__all__ = [ - "regrid", - "calc_weights", - "regrid_from_weights", -] - - -def regrid( - vertices_input: tuple[np.ndarray, ...], - vertices_output: tuple[np.ndarray, ...], - values_input: np.ndarray, - values_output: None | np.ndarray = None, - axis_input: None | int | tuple[int, ...] = None, - axis_output: None | int | tuple[int, ...] = None, - method: str = "conservative", - order: int = 1, -) -> np.ndarray: - """ - Transfer a histogram defined on any curvilinear grid onto a new curvilinear grid. - - Parameters - ---------- - vertices_input - The vertices of each bin in the input histogram. - The number of elements in ``vertices``, ``len(vertices)``, - should match the number of regridding axes, ``len(axis)``. - All elements of ``vertices`` should be broadcastable with the shape :math:`(...,M,...,N,...)`, - where :math:`M` and :math:`N` are the number of elements along each regridding axis. - vertices_output - The vertices of each new bin in the output histogram. - The number of elements in ``vertices``, ``len(vertices)``, - should match the number of regridding axes, ``len(axis)``. - values_input - The value of each bin in the input histogram. - Should be broadcastable with :math:`(...,M-1,...,N-1,...)`. - values_output - An alternative output array to place the result. - It must have the same shape as the expected output. - axis_input - The axes of the input histogram to regrid. - If :class:`None`, regrid over all the axes of the input histogram. - axis_output - The axes of ``vertices_output`` corresponding to the axes in ``axis``. - method - The type of regridding to use. Currently, the only valid option is ``conservative``. - order - The order of the regridding operation. Currently, only first-order regridding (``order=1``) is supported - - Returns - ------- - The regridded histogram - - Examples - -------- - Define an input curvilinear grid - - .. jupyter-execute:: - - import numpy as np - import matplotlib.pyplot as plt - import regridding - - num_x = 66 - num_y = 66 - - x = np.linspace(-5, 5, num=num_x) - y = np.linspace(-5, 5, num=num_y) - - x, y = np.meshgrid(x, y, indexing="ij") - - angle = 0.4 - x_input = x * np.cos(angle) - y * np.sin(angle) + 0.05 * x * x - y_input = x * np.sin(angle) + y * np.cos(angle) + 0.05 * y * y - - Define a test patten - - .. jupyter-execute:: - - pitch = 16 - a_input = 0 * x[:~0,:~0] - a_input[::pitch, :] = 1 - a_input[:, ::pitch] = 1 - a_input[pitch//2::pitch, pitch//2::pitch] = 1 - - plt.figure(); - plt.pcolormesh(x_input, y_input, a_input); - - Define a new grid - - .. jupyter-execute:: - - x_output = np.linspace(x_input.min(), x_input.max(), num_x // 2) - y_output = np.linspace(y_input.min(), y_input.max(), num_y // 2) - - x_output, y_output = np.meshgrid(x_output, y_output, indexing="ij") - - Regrid the test pattern onto the new grid - - .. jupyter-execute:: - - a_output = regridding.regrid( - vertices_input=(x_input, y_input), - vertices_output=(x_output, y_output), - values_input=a_input, - ) - - plt.figure(); - plt.pcolormesh(x_output, y_output, a_output); - """ - weights = calc_weights( - vertices_input=vertices_input, - vertices_output=vertices_output, - axis_input=axis_input, - axis_output=axis_output, - method=method, - order=order, - ) - result = regrid_from_weights( - weights=weights, - values_input=values_input, - values_output=values_output, - axis_input=axis_input, - axis_output=axis_output, - ) - return result - - -def calc_weights( - vertices_input: tuple[np.ndarray, ...], - vertices_output: tuple[np.ndarray, ...], - # values_input: np.ndarray, - # values_output: None | np.ndarray = None, - axis_input: None | int | tuple[int, ...] = None, - axis_output: None | int | tuple[int, ...] = None, - method: str = "conservative", - order: int = 1, -) -> tuple[np.ndarray, tuple[int, ...]]: - # shape_values_input = values_input.shape - shape_vertices_input = np.broadcast(*vertices_input).shape - shape_vertices_output = np.broadcast(*vertices_output).shape - - # ndim_values_input = len(shape_values_input) - # ndim_vertices_input = len(shape_vertices_input) - ndim_input = len(shape_vertices_input) - ndim_output = len(shape_vertices_output) - - axis_input = _normalize_axis(axis_input, ndim=ndim_input) - axis_output = _normalize_axis(axis_output, ndim=ndim_output) - - axis_input = sorted(axis_input, reverse=True) - axis_output = sorted(axis_output, reverse=True) - - if len(axis_output) != len(axis_input): - raise ValueError( - f"The number of axes in `axis_output`, {axis_output}, " - f"must match the number of axes in `axis_input`, {axis_input}" - ) - - if len(vertices_input) != len(axis_input): - raise ValueError( - f"The number of elements in `vertices_input`, {len(vertices_input)}, " - f"should match the number of axes in `axis_input`, {axis_input}" - ) - - if len(vertices_output) != len(vertices_input): - raise ValueError( - f"The number of elements in `vertices_output`, {len(vertices_output)}, " - f"should match the number of elements in `vertices_input`, {len(vertices_input)}" - ) - - # shape_input = np.broadcast_shapes( - # tuple(1 if ax in axis_input else shape_values_input[ax] for ax in _normalize_axis(None, ndim_values_input)), - # shape_vertices_input, - # ) - - axis_input_orthogonal = tuple( - ax for ax in _normalize_axis(None, ndim_input) if ax not in axis_input - ) - axis_output_orthogonal = tuple( - ax for ax in _normalize_axis(None, ndim_output) if ax not in axis_output - ) - - shape_input_orthogonal = tuple( - shape_vertices_input[ax] for ax in axis_input_orthogonal - ) - shape_output_orthogonal = tuple( - shape_vertices_output[ax] for ax in axis_output_orthogonal - ) - - shape_orthogonal = np.broadcast_shapes( - shape_input_orthogonal, shape_output_orthogonal - ) - - shape_input = list(shape_orthogonal) - for ax in axis_input: - shape_input.insert(ax, shape_vertices_input[ax]) - shape_input = tuple(shape_input) - - shape_output = list(shape_orthogonal) - for ax in axis_output: - shape_output.insert(ax, shape_vertices_output[ax]) - shape_output = tuple(shape_output) - - shape_centers_output = list(shape_output) - for ax in axis_output: - shape_centers_output[ax] -= 1 - shape_centers_output = tuple(shape_centers_output) - - # bshape_values_input = list(shape_orthogonal) - # for ax in axis_input: - # bshape_values_input.insert(ax, shape_values_input[ax]) - # bshape_values_input = tuple(bshape_values_input) - - # if values_output is None: - # bshape_values_output = list(shape_orthogonal) - # for ax in axis_output: - # bshape_values_output.insert(ax, shape_output[ax] - 1) - # bshape_values_output = tuple(bshape_values_output) - # - # values_output = np.zeros(bshape_values_output, dtype=values_input.dtype) - - vertices_input = tuple( - np.broadcast_to(component, shape_input) for component in vertices_input - ) - vertices_output = tuple( - np.broadcast_to(component, shape_output) for component in vertices_output - ) - # values_input = np.broadcast_to(values_input, bshape_values_input) - - weights = np.empty(shape_orthogonal, dtype=numba.typed.List) - - for index in np.ndindex(*shape_orthogonal): - index_vertices_input = list(index) - for ax in axis_input: - index_vertices_input.insert(ax, slice(None)) - index_vertices_input = tuple(index_vertices_input) - - index_vertices_output = list(index) - for ax in axis_output: - index_vertices_output.insert(ax, slice(None)) - index_vertices_output = tuple(index_vertices_output) - - # index_values_input = list(index) - # for ax in axis_input: - # index_values_input.insert(ax, slice(None)) - # index_values_input = tuple(index_values_input) - - # index_values_output = list(index) - # for ax in axis_output: - # index_values_output.insert(ax, slice(None)) - # index_values_output = tuple(index_values_output) - - if len(axis_input) == 1: - raise NotImplementedError("1D regridding not supported") - - elif len(axis_input) == 2: - vertices_input_x, vertices_input_y = vertices_input - vertices_output_x, vertices_output_y = vertices_output - - if method == "conservative": - if order == 1: - weights[index] = _conservative_ramshaw( - # values_input=values_input[index_values_input], - # values_output=values_output[index_values_output], - grid_input=( - vertices_input_x[index_vertices_input], - vertices_input_y[index_vertices_input], - ), - grid_output=( - vertices_output_x[index_vertices_output], - vertices_output_y[index_vertices_output], - ), - ) - else: - raise NotImplementedError(f"order {order} not supported") - else: - raise NotImplementedError(f"method {method} not supported") - - else: - raise NotImplementedError( - "Regridding operations greater than 2D are not supported" - ) - - return weights, shape_centers_output - - -def regrid_from_weights( - weights: tuple[np.ndarray, tuple[int, ...]], - values_input: np.ndarray, - values_output: None | np.ndarray = None, - axis_input: None | int | tuple[int, ...] = None, - axis_output: None | int | tuple[int, ...] = None, -): - weights, shape_centers_output = weights - - shape_weights = weights.shape - - shape_values_input = values_input.shape - if values_output is not None: - shape_values_output = values_output.shape - else: - shape_values_output = shape_centers_output - - ndim_input = len(shape_values_input) - ndim_output = len(shape_values_output) - - axis_input = _normalize_axis(axis_input, ndim=ndim_input) - axis_output = _normalize_axis(axis_output, ndim=ndim_output) - - axis_input = sorted(axis_input, reverse=True) - axis_output = sorted(axis_output, reverse=True) - - if len(axis_output) != len(axis_input): - raise ValueError( - f"The number of axes in `axis_output`, {axis_output}, " - f"must match the number of axes in `axis_input`, {axis_input}" - ) - - axis_input_orthogonal = tuple( - ax for ax in _normalize_axis(None, ndim_input) if ax not in axis_input - ) - axis_output_orthogonal = tuple( - ax for ax in _normalize_axis(None, ndim_output) if ax not in axis_output - ) - - shape_input_orthogonal = tuple( - shape_values_input[ax] for ax in axis_input_orthogonal - ) - shape_output_orthogonal = tuple( - shape_values_output[ax] for ax in axis_output_orthogonal - ) - - shape_orthogonal = np.broadcast_shapes( - shape_input_orthogonal, shape_output_orthogonal, shape_weights - ) - - weights = np.broadcast_to(weights, shape_orthogonal) - - shape_input = list(shape_orthogonal) - for ax in axis_input: - shape_input.insert(ax, shape_values_input[ax]) - shape_input = tuple(shape_input) - - values_input = np.broadcast_to(values_input, shape_input) - - shape_output = list(shape_orthogonal) - for ax in axis_output: - shape_output.insert(ax, shape_values_output[ax]) - shape_output = tuple(shape_output) - - if values_output is None: - values_output = np.zeros(shape_output) - else: - if values_output.shape != shape_output: - raise ValueError( - f"provided output array has the wrong shape. Expected {shape_output}, got {values_output.shape}" - ) - - for index in np.ndindex(*shape_orthogonal): - index_values_input = list(index) - for ax in axis_input: - index_values_input.insert(ax, slice(None)) - index_values_input = tuple(index_values_input) - - index_values_output = list(index) - for ax in axis_output: - index_values_output.insert(ax, slice(None)) - index_values_output = tuple(index_values_output) - - _regrid_from_weights_2d( - weights=weights[index], - values_input=values_input[index_values_input], - values_output=values_output[index_values_output], - ) - - return values_output - - -@numba.njit(error_model="numpy") -def _regrid_from_weights_2d( - weights: np.ndarray, - values_input: np.ndarray, - values_output: np.ndarray, -) -> None: - values_input = values_input.reshape(-1) - values_output = values_output.reshape(-1) - - for i in range(len(weights)): - index_input, index_output, weight = weights[i] - values_output[int(index_output)] += weight * values_input[int(index_input)] - - -@dataclasses.dataclass -class Regridder: - index_input: np.ndarray - index_output: np.ndarray - shape_input: tuple[int, ...] - shape_output: tuple[int, ...] - weights: np.ndarray diff --git a/regridding/_regrid/__init__.py b/regridding/_regrid/__init__.py new file mode 100644 index 0000000..7c9c066 --- /dev/null +++ b/regridding/_regrid/__init__.py @@ -0,0 +1,2 @@ +from ._regrid_from_weights import * +from ._regrid import * diff --git a/regridding/_regrid/_regrid.py b/regridding/_regrid/_regrid.py new file mode 100644 index 0000000..2676b5e --- /dev/null +++ b/regridding/_regrid/_regrid.py @@ -0,0 +1,165 @@ +from typing import Sequence, Literal +import numpy as np +import regridding +from . import regrid_from_weights + +__all__ = [ + "regrid", +] + + +def regrid( + coordinates_input: tuple[np.ndarray, ...], + coordinates_output: tuple[np.ndarray, ...], + values_input: np.ndarray, + values_output: None | np.ndarray = None, + axis_input: None | int | Sequence[int] = None, + axis_output: None | int | Sequence[int] = None, + method: Literal["multilinear", "conservative"] = "multilinear", +) -> np.ndarray: + """ + 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 coordinates as the input grid. + values_input + Input array of values to be resampled. + values_output + Optional array in which to place the output. + 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. + The ``conservative`` method uses the algorithm described in + :footcite:t:`Ramshaw1985` + + See Also + -------- + :func:`regridding.regrid` + :func:`regridding.regrid_from_weights` + + Examples + -------- + + Regrid a 1D array using multilinear interpolation. + + .. jupyter-execute:: + + import numpy as np + import matplotlib.pyplot as plt + import regridding + + # Define the input grid + x_input = np.linspace(-1, 1, num=11) + + # Define the input array + values_input = np.square(x_input) + + # Define the output grid + x_output = np.linspace(-1, 1, num=51) + + # Regrid the input array onto the output grid + values_output = regridding.regrid( + coordinates_input=(x_input,), + coordinates_output=(x_output,), + values_input=values_input, + method="multilinear", + ) + + # Plot the results + plt.figure(figsize=(6, 3)); + plt.scatter(x_input, values_input, s=100, label="input", zorder=1); + plt.scatter(x_output, values_output, label="interpolated", zorder=0); + plt.legend(); + + | + + Regrid a 2D array using conservative resampling. + + .. jupyter-execute:: + + # Define the number of edges in the input grid + num_x = 66 + num_y = 66 + + # Define a dummy linear grid + x = np.linspace(-5, 5, num=num_x) + y = np.linspace(-5, 5, num=num_y) + x, y = np.meshgrid(x, y, indexing="ij") + + # Define the curvilinear input grid using the dummy grid + angle = 0.4 + x_input = x * np.cos(angle) - y * np.sin(angle) + 0.05 * x * x + y_input = x * np.sin(angle) + y * np.cos(angle) + 0.05 * y * y + + # Define the test pattern + pitch = 16 + a_input = 0 * x[:~0,:~0] + a_input[::pitch, :] = 1 + a_input[:, ::pitch] = 1 + a_input[pitch//2::pitch, pitch//2::pitch] = 1 + + # Define a rectilinear output grid using the limits of the input grid + x_output = np.linspace(x_input.min(), x_input.max(), num_x // 2) + y_output = np.linspace(y_input.min(), y_input.max(), num_y // 2) + x_output, y_output = np.meshgrid(x_output, y_output, indexing="ij") + + # Regrid the test pattern onto the new grid + a_output = regridding.regrid( + coordinates_input=(x_input, y_input), + coordinates_output=(x_output, y_output), + values_input=a_input, + method="conservative", + ) + + fig, axs = plt.subplots( + ncols=2, + sharex=True, + sharey=True, + figsize=(8, 4), + constrained_layout=True, + ); + axs[0].pcolormesh(x_input, y_input, a_input); + axs[0].set_title("input array"); + axs[1].pcolormesh(x_output, y_output, a_output); + axs[1].set_title("regridded array"); + + | + + References + ---------- + .. footbibliography:: + """ + weights, 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 = 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 diff --git a/regridding/_regrid/_regrid_from_weights.py b/regridding/_regrid/_regrid_from_weights.py new file mode 100644 index 0000000..a952185 --- /dev/null +++ b/regridding/_regrid/_regrid_from_weights.py @@ -0,0 +1,112 @@ +from typing import Sequence +import numpy as np +import numba +from regridding import _util + +__all__ = [ + "regrid_from_weights", +] + + +def regrid_from_weights( + weights: np.ndarray, + shape_input: tuple[int, ...], + shape_output: tuple[int, ...], + values_input: np.ndarray, + values_output: None | np.ndarray = None, + axis_input: None | int | Sequence[int] = None, + axis_output: None | int | Sequence[int] = None, +) -> np.ndarray: + """ + Regrid an array of values using weights computed by + :func:`regridding.weights`. + + Parameters + ---------- + weights + Ragged array of weights computed by :func:`regridding.weights`. + shape_input + Broadcasted shape of the input coordinates computed by :func:`regridding.weights`. + shape_output + Broadcasted shape of the output coordinates computed by :func:`regridding.weights`. + values_input + Input array of values to be resampled. + values_output + Optional array in which to place the output. + axis_input + Logical axes of the input array to resample. + If :obj:`None`, resample all the axes of the input array. + The number of axes should be equal to the number of + coordinates in the original input grid passed to :func:`regridding.weights`. + axis_output + Logical axes of the output array corresponding to the resampled axes + of the input array. + If :obj:`None`, all the axes of the output array correspond to resampled + axes in the input grid. + The number of axes should be equal to the original number of + coordinates in the output grid passed to :func:`regridding.weights`. + + See Also + -------- + :func:`regridding.regrid` + :func:`regridding.regrid_from_weights` + """ + + shape_input = np.broadcast_shapes(shape_input, values_input.shape) + values_input = np.broadcast_to(values_input, shape=shape_input, subok=True) + ndim_input = len(shape_input) + axis_input = _util._normalize_axis(axis_input, ndim=ndim_input) + + if values_output is None: + shape_output = np.broadcast_shapes( + shape_output, + tuple( + shape_input[ax] if ax not in axis_input else 1 + for ax in _util._normalize_axis(None, ndim_input) + ), + ) + values_output = np.zeros_like(values_input, shape=shape_output) + else: + if values_output.shape != shape_output: + raise ValueError(f"") + values_output.fill(0) + values_output_copy = values_output + ndim_output = len(shape_output) + axis_output = _util._normalize_axis(axis_output, ndim=ndim_output) + + axis_input_numba = ~np.arange(len(axis_input))[::-1] + axis_output_numba = ~np.arange(len(axis_output))[::-1] + + shape_input_numba = tuple(shape_input[ax] for ax in axis_input) + shape_output_numba = tuple(shape_output[ax] for ax in axis_output) + + values_input = np.moveaxis(values_input, axis_input, axis_input_numba) + values_output = np.moveaxis(values_output, axis_output, axis_output_numba) + + weights = numba.typed.List(weights.reshape(-1)) + values_input = values_input.reshape(-1, *shape_input_numba) + values_output = values_output.reshape(-1, *shape_output_numba) + + _regrid_from_weights( + weights=weights, + values_input=values_input, + values_output=values_output, + ) + + return values_output_copy + + +@numba.njit() +def _regrid_from_weights( + weights: numba.typed.List, + values_input: np.ndarray, + values_output: np.ndarray, +) -> None: + for d in numba.prange(len(weights)): + weights_d = weights[d] + values_input_d = values_input[d].reshape(-1) + values_output_d = values_output[d].reshape(-1) + + for w in range(len(weights_d)): + i_input, i_output, weight = weights_d[w] + values_output_d[int(i_output)] += weight * values_input_d[int(i_input)] diff --git a/regridding/_regrid/_tests/__init__.py b/regridding/_regrid/_tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/regridding/_regrid/_tests/test_regrid.py b/regridding/_regrid/_tests/test_regrid.py new file mode 100644 index 0000000..8d3870d --- /dev/null +++ b/regridding/_regrid/_tests/test_regrid.py @@ -0,0 +1,100 @@ +import pytest +import numpy as np +import regridding + + +@pytest.mark.parametrize( + argnames="coordinates_input,coordinates_output,values_input,values_output,axis_input,axis_output,result_expected", + argvalues=[ + ( + (np.linspace(-1, 1, num=11),), + (np.linspace(-1, 1, num=11),), + np.square(np.linspace(-1, 1, num=11)), + None, + None, + None, + np.square(np.linspace(-1, 1, num=11)), + ), + ( + (np.linspace(-1, 1, num=11),), + (np.linspace(-1, 1, num=11),), + np.square(np.linspace(-1, 1, num=11)), + np.empty(shape=11), + None, + None, + np.square(np.linspace(-1, 1, num=11)), + ), + ], +) +def test_regrid_multilinear_1d( + coordinates_input: tuple[np.ndarray, ...], + coordinates_output: tuple[np.ndarray, ...], + values_input: np.ndarray, + values_output: None | np.ndarray, + axis_input: None | int | tuple[int, ...], + axis_output: None | int | tuple[int, ...], + result_expected: np.ndarray, +): + result = regridding.regrid( + coordinates_input=coordinates_input, + coordinates_output=coordinates_output, + values_input=values_input, + values_output=values_output, + axis_input=axis_input, + axis_output=axis_output, + method="multilinear", + ) + assert isinstance(result, np.ndarray) + assert np.issubdtype(result.dtype, float) + assert np.all(result == result_expected) + + +@pytest.mark.parametrize( + argnames="coordinates_input, values_input, axis_input", + argvalues=[ + ( + np.meshgrid( + np.linspace(-1, 1, num=10), + np.linspace(-1, 1, num=11), + indexing="ij", + ), + np.random.normal(size=(10 - 1, 11 - 1)), + None, + ), + ], +) +@pytest.mark.parametrize( + argnames="coordinates_output, values_output, axis_output", + argvalues=[ + ( + np.meshgrid( + 1.1 * np.linspace(-1, 1, num=10) + 0.001, + 1.2 * np.linspace(-1, 1, num=11) + 0.001, + indexing="ij", + ), + None, + None, + ) + ], +) +def test_regrid_conservative_2d( + coordinates_input: tuple[np.ndarray, ...], + coordinates_output: tuple[np.ndarray, ...], + values_input: np.ndarray, + values_output: None | np.ndarray, + axis_input: None | int | tuple[int, ...], + axis_output: None | int | tuple[int, ...], +): + result = regridding.regrid( + coordinates_input=coordinates_input, + coordinates_output=coordinates_output, + values_input=values_input, + values_output=values_output, + axis_input=axis_input, + axis_output=axis_output, + method="conservative", + ) + + assert np.issubdtype(result.dtype, float) + assert result.shape == tuple(np.array(np.broadcast(*coordinates_output).shape) - 1) + assert np.isclose(result.sum(), values_input.sum()) diff --git a/regridding/_weights/__init__.py b/regridding/_weights/__init__.py new file mode 100644 index 0000000..71ecb36 --- /dev/null +++ b/regridding/_weights/__init__.py @@ -0,0 +1 @@ +from ._weights import * diff --git a/regridding/_weights/_weights.py b/regridding/_weights/_weights.py new file mode 100644 index 0000000..4d4877d --- /dev/null +++ b/regridding/_weights/_weights.py @@ -0,0 +1,136 @@ +from typing import Sequence, Literal +import numpy as np +from ._weights_multilinear import _weights_multilinear +from ._weights_conservative import _weights_conservative + +__all__ = [ + "weights", +] + + +def weights( + coordinates_input: tuple[np.ndarray, ...], + coordinates_output: tuple[np.ndarray, ...], + axis_input: None | int | Sequence[int] = None, + axis_output: None | int | Sequence[int] = None, + method: Literal["multilinear", "conservative"] = "multilinear", +) -> tuple[np.ndarray, tuple[int, ...], tuple[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:`regridding.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.regrid` + :func:`regridding.regrid_from_weights` + + Examples + -------- + + Regrid two arrays of values defined on the same grid using saved weights. + + .. jupyter-execute:: + + import numpy as np + import scipy.signal + import matplotlib.pyplot as plt + import regridding + + # Define input grid + x_input = np.linspace(-4, 4, num=101) + y_input = np.linspace(-4, 4, num=101) + x_input, y_input = np.meshgrid(x_input, y_input, indexing="ij") + + # Define rotated output grid + angle = 0.2 + x_output = x_input * np.cos(angle) - y_input * np.sin(angle) + y_output = x_input * np.sin(angle) + y_input * np.cos(angle) + + # Define two arrays of values defined on the same grid + values_input_1 = np.cos(np.square(x_input)) * np.cos(np.square(y_input)) + values_input_2 = np.sin(np.square(x_input) + np.square(y_input)) + + # Convolve with a 2x2 uniform kernel to simulate values defined on cell centers + values_input_1 = scipy.signal.convolve(values_input_1, np.ones((2, 2)), mode="valid") + values_input_2 = scipy.signal.convolve(values_input_2, np.ones((2, 2)), mode="valid") + + # Save regridding weights relating the input and output grids + weights = regridding.weights( + coordinates_input=(x_input, y_input), + coordinates_output=(x_output, y_output), + method="conservative", + ) + + # Regrid the first array of values using the saved weights + values_output_1 = regridding.regrid_from_weights( + *weights, + values_input=values_input_1, + ) + + # Regrid the second array of values using the saved weights + values_output_2 = regridding.regrid_from_weights( + *weights, + values_input=values_input_2, + ) + + # Plot the original and regridded arrays of values + fig, axs = plt.subplots( + nrows=2, + ncols=2, + sharex=True, + sharey=True, + constrained_layout=True, + ) + axs[0, 0].pcolormesh(x_input, y_input, values_input_1); + axs[0, 0].set_title(r"values_input_1"); + axs[0, 1].pcolormesh(x_input, y_input, values_input_2); + axs[0, 1].set_title(r"values_input_2"); + axs[1, 0].pcolormesh(x_output, y_output, values_output_1); + axs[1, 0].set_title(r"values_output_1"); + axs[1, 1].pcolormesh(x_output, y_output, values_output_2); + axs[1, 1].set_title(r"values_output_2"); + """ + if method == "multilinear": + return _weights_multilinear( + coordinates_input=coordinates_input, + coordinates_output=coordinates_output, + axis_input=axis_input, + axis_output=axis_output, + ) + elif method == "conservative": + return _weights_conservative( + coordinates_input=coordinates_input, + coordinates_output=coordinates_output, + axis_input=axis_input, + axis_output=axis_output, + ) + else: + raise ValueError(f"unrecognized method '{method}'") diff --git a/regridding/_weights/_weights_conservative.py b/regridding/_weights/_weights_conservative.py new file mode 100644 index 0000000..a70c009 --- /dev/null +++ b/regridding/_weights/_weights_conservative.py @@ -0,0 +1,75 @@ +from typing import Sequence +import numpy as np +import numba +from regridding import _util +from regridding._conservative_ramshaw import _conservative_ramshaw + + +def _weights_conservative( + coordinates_input: tuple[np.ndarray, ...], + coordinates_output: tuple[np.ndarray, ...], + axis_input: None | int | Sequence[int] = None, + axis_output: None | int | Sequence[int] = None, +) -> tuple[np.ndarray, tuple[int, ...], tuple[int, ...]]: + ( + coordinates_input, + coordinates_output, + axis_input, + axis_output, + shape_input, + shape_output, + shape_orthogonal, + ) = _util._normalize_input_output_coordinates( + coordinates_input=coordinates_input, + coordinates_output=coordinates_output, + axis_input=axis_input, + axis_output=axis_output, + ) + + shape_values_input = list(shape_input) + for ax in axis_input: + shape_values_input[ax] -= 1 + shape_values_input = tuple(shape_values_input) + + shape_values_output = list(shape_output) + for ax in axis_output: + shape_values_output[ax] -= 1 + shape_values_output = tuple(shape_values_output) + + weights = np.empty(shape_orthogonal, dtype=numba.typed.List) + + for index in np.ndindex(*shape_orthogonal): + index_vertices_input = list(index) + for ax in axis_input: + index_vertices_input.insert(ax, slice(None)) + index_vertices_input = tuple(index_vertices_input) + + index_vertices_output = list(index) + for ax in axis_output: + index_vertices_output.insert(ax, slice(None)) + index_vertices_output = tuple(index_vertices_output) + + if len(axis_input) == 1: + raise NotImplementedError("1D regridding not supported") + + elif len(axis_input) == 2: + coordinates_input_x, coordinates_input_y = coordinates_input + coordinates_output_x, coordinates_output_y = coordinates_output + + weights[index] = _conservative_ramshaw( + grid_input=( + coordinates_input_x[index_vertices_input], + coordinates_input_y[index_vertices_input], + ), + grid_output=( + coordinates_output_x[index_vertices_output], + coordinates_output_y[index_vertices_output], + ), + ) + + else: + raise NotImplementedError( + "Regridding operations greater than 2D are not supported" + ) + + return weights, shape_values_input, shape_values_output diff --git a/regridding/_weights/_weights_multilinear.py b/regridding/_weights/_weights_multilinear.py new file mode 100644 index 0000000..2960bce --- /dev/null +++ b/regridding/_weights/_weights_multilinear.py @@ -0,0 +1,127 @@ +from typing import Sequence +import numpy as np +import numba +import regridding +from regridding import _util + + +def _weights_multilinear( + coordinates_input: tuple[np.ndarray, ...], + coordinates_output: tuple[np.ndarray, ...], + axis_input: None | int | Sequence[int] = None, + axis_output: None | int | Sequence[int] = None, +) -> tuple[np.ndarray, tuple[int, ...], tuple[int, ...]]: + indices_output = regridding.find_indices( + coordinates_input=coordinates_input, + coordinates_output=coordinates_output, + axis_input=axis_input, + axis_output=axis_output, + ) + result = _weights_from_indices_multilinear( + indices_output=indices_output, + coordinates_input=coordinates_input, + coordinates_output=coordinates_output, + axis_input=axis_input, + axis_output=axis_output, + ) + return result + + +def _weights_from_indices_multilinear( + indices_output: tuple[np.ndarray, ...], + coordinates_input: tuple[np.ndarray, ...], + coordinates_output: tuple[np.ndarray, ...], + axis_input: None | int | Sequence[int] = None, + axis_output: None | int | Sequence[int] = None, +) -> tuple[np.ndarray, tuple[int, ...], tuple[int, ...]]: + ( + coordinates_input, + coordinates_output, + axis_input, + axis_output, + shape_input, + shape_output, + shape_orthogonal, + ) = _util._normalize_input_output_coordinates( + coordinates_input=coordinates_input, + coordinates_output=coordinates_output, + axis_input=axis_input, + axis_output=axis_output, + ) + + axis_input_numba = ~np.arange(len(axis_input))[::-1] + axis_output_numba = ~np.arange(len(axis_output))[::-1] + + shape_input_numba = tuple(shape_input[ax] for ax in axis_input) + shape_output_numba = tuple(shape_output[ax] for ax in axis_output) + + indices_output = tuple( + np.moveaxis(v, axis_output, axis_output_numba).reshape(-1, *shape_output_numba) + for v in indices_output + ) + coordinates_input = tuple( + np.moveaxis(v, axis_input, axis_input_numba).reshape(-1, *shape_input_numba) + for v in coordinates_input + ) + coordinates_output = tuple( + np.moveaxis(v, axis_output, axis_output_numba).reshape(-1, *shape_output_numba) + for v in coordinates_output + ) + + if len(axis_input) == 1: + weights_list = _weights_from_indices_multilinear_1d( + indices_output=indices_output, + coordinates_input=coordinates_input, + coordinates_output=coordinates_output, + ) + else: + raise ValueError( + f"{len(axis_input)}-dimensional multilinear interpolation is not supported" + ) + + num_d = len(weights_list) + weights = np.empty(shape=num_d, dtype=numba.typed.List) + for d in range(num_d): + weights[d] = weights_list[d] + weights = weights.reshape(shape_orthogonal) + + return weights, shape_input, shape_output + + +@numba.njit(parallel=True) +def _weights_from_indices_multilinear_1d( + indices_output: tuple[np.ndarray], + coordinates_input: tuple[np.ndarray], + coordinates_output: tuple[np.ndarray], +) -> np.ndarray: + (i_output,) = indices_output + (x_input,) = coordinates_input + (x_output,) = coordinates_output + + num_d, num_i_input = x_input.shape + num_d, num_i_output = x_output.shape + + weights = numba.typed.List() + + for d in numba.prange(num_d): + weights_d = numba.typed.List() + for _ in range(0): + weights_d.append((0.0, 0.0, 0.0)) + + for i in numba.prange(num_i_output): + i0 = i_output[d, i] + i1 = i0 + 1 + + x0 = x_input[d, i0] + x1 = x_input[d, i1] + x = x_output[d, i] + + w1 = (x - x0) / (x1 - x0) + w0 = 1 - w1 + + weights_d.append((i0, i, w0)) + weights_d.append((i1, i, w1)) + + weights.append(weights_d) + + return weights diff --git a/regridding/tests/test_regrid.py b/regridding/tests/test_regrid.py deleted file mode 100644 index 5dbe118..0000000 --- a/regridding/tests/test_regrid.py +++ /dev/null @@ -1,57 +0,0 @@ -import pytest -import numpy as np -import regridding - - -@pytest.mark.parametrize( - argnames="vertices_input, values_input, axis_input", - argvalues=[ - ( - np.meshgrid( - np.linspace(-1, 1, num=10), - np.linspace(-1, 1, num=11), - indexing="ij", - ), - np.random.normal(size=(10 - 1, 11 - 1)), - None, - ), - ], -) -@pytest.mark.parametrize( - argnames="vertices_output, values_output, axis_output", - argvalues=[ - ( - np.meshgrid( - 1.1 * np.linspace(-1, 1, num=10) + 0.001, - 1.2 * np.linspace(-1, 1, num=11) + 0.001, - indexing="ij", - ), - None, - None, - ) - ], -) -@pytest.mark.parametrize("order", [1]) -def test_regrid_conservative_2d( - vertices_input: tuple[np.ndarray, ...], - vertices_output: tuple[np.ndarray, ...], - values_input: np.ndarray, - values_output: None | np.ndarray, - axis_input: None | int | tuple[int, ...], - axis_output: None | int | tuple[int, ...], - order: int, -): - result = regridding.regrid( - vertices_input=vertices_input, - vertices_output=vertices_output, - values_input=values_input, - values_output=values_output, - axis_input=axis_input, - axis_output=axis_output, - method="conservative", - order=order, - ) - - assert np.issubdtype(result.dtype, float) - assert result.shape == tuple(np.array(np.broadcast(*vertices_output).shape) - 1) - assert np.isclose(result.sum(), values_input.sum())