From b5cab8f9e045125d74719e8eda451a940169b3f3 Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Wed, 21 Oct 2026 18:14:06 -0600 Subject: [PATCH 1/5] Build multilinear weights on grids of any number of dimensions Work in progress: the locators and the cell math, with their tests. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01BYjDL98znSud1yFh9chnkP --- regridding/_regrid/_regrid_from_weights.py | 9 +- regridding/_util.py | 29 +- regridding/_weights/_weights.py | 4 +- regridding/_weights/_weights_multilinear.py | 206 ---- .../_weights/_weights_multilinear/__init__.py | 15 + .../_weights/_weights_multilinear/_cells.py | 490 ++++++++ .../_weights_multilinear/_cells_test.py | 197 ++++ .../_weights_multilinear/_curvilinear.py | 1003 +++++++++++++++++ .../_weights_multilinear/_curvilinear_test.py | 261 +++++ .../_weights_multilinear/_rectilinear.py | 144 +++ .../_weights_multilinear/_rectilinear_test.py | 110 ++ .../_weights_multilinear.py | 237 ++++ .../_weights_multilinear_test.py | 0 13 files changed, 2492 insertions(+), 213 deletions(-) delete mode 100644 regridding/_weights/_weights_multilinear.py create mode 100644 regridding/_weights/_weights_multilinear/__init__.py create mode 100644 regridding/_weights/_weights_multilinear/_cells.py create mode 100644 regridding/_weights/_weights_multilinear/_cells_test.py create mode 100644 regridding/_weights/_weights_multilinear/_curvilinear.py create mode 100644 regridding/_weights/_weights_multilinear/_curvilinear_test.py create mode 100644 regridding/_weights/_weights_multilinear/_rectilinear.py create mode 100644 regridding/_weights/_weights_multilinear/_rectilinear_test.py create mode 100644 regridding/_weights/_weights_multilinear/_weights_multilinear.py rename regridding/_weights/{ => _weights_multilinear}/_weights_multilinear_test.py (100%) diff --git a/regridding/_regrid/_regrid_from_weights.py b/regridding/_regrid/_regrid_from_weights.py index 1bf81bf..0c4b071 100644 --- a/regridding/_regrid/_regrid_from_weights.py +++ b/regridding/_regrid/_regrid_from_weights.py @@ -137,8 +137,11 @@ def regrid_from_weights( shape_output_tmp = values_output.shape - values_input = values_input.reshape(-1, *shape_input_numba) - values_output = values_output.reshape(-1, *shape_output_numba) + # The weights address each element by a flat index, so its resampled + # axes are flattened together. That also covers a single point, which + # has no resampled axes at all. + values_input = values_input.reshape(-1, np.prod(shape_input_numba, dtype=int)) + values_output = values_output.reshape(-1, np.prod(shape_output_numba, dtype=int)) flat = weights.reshape(-1) unit_weights = getattr(flat[0][2], "unit", None) if flat.size else None @@ -158,7 +161,7 @@ def regrid_from_weights( values_output=values_output, ) - values_output = values_output.reshape(*shape_output_tmp) + values_output = values_output.reshape(shape_output_tmp) values_output = np.moveaxis(values_output, axis_output_numba, axis_output) diff --git a/regridding/_util.py b/regridding/_util.py index db6e0b5..71396c4 100644 --- a/regridding/_util.py +++ b/regridding/_util.py @@ -110,7 +110,7 @@ def _normalize_axis( if axis is None: axis = tuple(range(ndim)) axis = np.lib.array_utils.normalize_axis_tuple(axis, ndim=ndim) - axis = tuple(~(~np.array(axis) % ndim)) + axis = tuple(~(~np.array(axis, dtype=int) % ndim)) return axis @@ -121,6 +121,7 @@ def _normalize_input_output_coordinates( axis_output: None | int | Sequence[int] = None, perturb: bool = False, seed: "None | int | np.random.Generator" = _seed_default, + match_axes: bool = True, ) -> tuple[ tuple[np.ndarray, ...], tuple[np.ndarray, ...], @@ -130,6 +131,30 @@ def _normalize_input_output_coordinates( tuple[int, ...], tuple[int, ...], ]: + """ + Broadcast the coordinates of an input and an output grid against each + other, apart from the axes which are resampled. + + Parameters + ---------- + coordinates_input + The coordinates of the input grid. + coordinates_output + The coordinates of the output grid. + axis_input + The resampled axes of the input grid. + axis_output + The resampled axes of the output grid. + perturb + Whether to perturb the output coordinates by a small random amount. + seed + The seed of the random perturbation. + match_axes + Whether `axis_output` must have as many axes as `axis_input`. + Multilinear interpolation does not need it, since its outputs are + points rather than cells, and they can be arranged in any number of + axes. + """ if isinstance(coordinates_input, np.ndarray): coordinates_input = (coordinates_input,) @@ -159,7 +184,7 @@ def _normalize_input_output_coordinates( axis_input = tuple(sorted(axis_input, reverse=True)) axis_output = tuple(sorted(axis_output, reverse=True)) - if len(axis_output) != len(axis_input): + if match_axes and 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}" diff --git a/regridding/_weights/_weights.py b/regridding/_weights/_weights.py index 3331ac6..81a20bd 100644 --- a/regridding/_weights/_weights.py +++ b/regridding/_weights/_weights.py @@ -1,7 +1,7 @@ from typing import Sequence, Literal import numpy as np from regridding import _util -from ._weights_multilinear import _weights_multilinear +from ._weights_multilinear import weights_multilinear from ._weights_conservative import _weights_conservative from ._weights_arrays import _weights_to_arrays @@ -339,7 +339,7 @@ def weights( clipped = False if method == "multilinear": - result = _weights_multilinear( + result = weights_multilinear( coordinates_input=coordinates_input, coordinates_output=coordinates_output, axis_input=axis_input, diff --git a/regridding/_weights/_weights_multilinear.py b/regridding/_weights/_weights_multilinear.py deleted file mode 100644 index 128e695..0000000 --- a/regridding/_weights/_weights_multilinear.py +++ /dev/null @@ -1,206 +0,0 @@ -from typing import Literal, Sequence -import numpy as np -import numba -from numba.typed.typedlist import List as TypedList -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, - weights_input: None | np.ndarray = None, - bounds: Literal["extrapolate", "nan", "raise"] = "extrapolate", - perturb: None | bool = False, - seed: "None | int | np.random.Generator" = _util._seed_default, -) -> 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, - method="searchsorted", - ) - 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, - weights_input=weights_input, - bounds=bounds, - perturb=perturb, - seed=seed, - ) - 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, - weights_input: None | np.ndarray = None, - bounds: Literal["extrapolate", "nan", "raise"] = "extrapolate", - perturb: None | bool = False, - seed: "None | int | np.random.Generator" = _util._seed_default, -) -> tuple[np.ndarray, tuple[int, ...], tuple[int, ...]]: - - if bounds not in ("extrapolate", "nan", "raise"): - raise ValueError( - f"Unrecognized {bounds=}, " - f"expected one of ('extrapolate', 'nan', 'raise')." - ) - - if perturb is None: - perturb = False - - ( - 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, - perturb=perturb, - seed=seed, - ) - - 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 weights_input is not None: - weights_input = np.broadcast_to(weights_input, shape_input) - weights_input = np.moveaxis(weights_input, axis_input, axis_input_numba) - weights_input = weights_input.reshape(-1, *shape_input_numba) - - if len(axis_input) == 1: - # `find_indices` marks an output point that falls outside the input - # grid with `fill_value`, which is not a usable index. Resolve those - # here so that the compiled kernel only ever sees a valid cell. - (index,) = indices_output - (x_input,) = coordinates_input - (x_output,) = coordinates_output - index_max = x_input.shape[~0] - 2 - outside = (index < 0) | (index > index_max) - if bounds == "raise" and outside.any(): - raise ValueError( - f"{outside.sum()} of the output points fall outside the input " - f"grid, and {bounds=}." - ) - # `fill_value` is a large positive number whichever side of the grid - # the point fell out on, so clamping it alone would send every such - # point to the last cell. Recover the side from the coordinate. - below = x_output < x_input[..., :1] - indices_output = (np.where(below, 0, np.clip(index, 0, index_max)),) - - weights_list = _weights_from_indices_multilinear_1d( - indices_output=indices_output, - coordinates_input=coordinates_input, - coordinates_output=coordinates_output, - weights_input=weights_input, - outside=outside if bounds == "nan" else np.zeros_like(outside), - ) - 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=TypedList) - 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=False) -def _weights_from_indices_multilinear_1d( - indices_output: tuple[np.ndarray, ...], - coordinates_input: tuple[np.ndarray, ...], - coordinates_output: tuple[np.ndarray, ...], - weights_input: None | np.ndarray, - outside: np.ndarray, -) -> TypedList: - (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 - - # `numba.typed.List()` is declared to return a plain `list` when Numba's - # JIT is disabled, a mode this library is not usable in. - weights: TypedList = TypedList() # type: ignore[assignment] - for _ in range(0): # pragma: nocover - weights.append( - ( - np.empty(0, dtype=np.int64), - np.empty(0, dtype=np.int64), - np.empty(0, dtype=np.float64), - ) - ) - - for d in range(num_d): - # each output vertex contributes exactly two weights, so the flat - # arrays for this element have a known size and need no compaction. - n = 2 * num_i_output - indices_input = np.empty(n, dtype=np.int64) - indices_output_d = np.empty(n, dtype=np.int64) - values = np.empty(n, dtype=np.float64) - - 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 - - if weights_input is not None: - w0 = w0 * weights_input[d, i0] - w1 = w1 * weights_input[d, i1] - - if outside[d, i]: - w0 = np.nan - w1 = np.nan - - indices_input[2 * i] = i0 - indices_output_d[2 * i] = i - values[2 * i] = w0 - - indices_input[2 * i + 1] = i1 - indices_output_d[2 * i + 1] = i - values[2 * i + 1] = w1 - - weights.append((indices_input, indices_output_d, values)) - - return weights diff --git a/regridding/_weights/_weights_multilinear/__init__.py b/regridding/_weights/_weights_multilinear/__init__.py new file mode 100644 index 0000000..c4202ed --- /dev/null +++ b/regridding/_weights/_weights_multilinear/__init__.py @@ -0,0 +1,15 @@ +""" +Multilinear weights for a logically-rectangular grid of any number of +dimensions. + +:mod:`._weights_multilinear` finds the cell of the input grid which contains +each output point, with :mod:`._rectilinear` if the grid is rectilinear and +:mod:`._curvilinear` otherwise, and turns the local coordinates of the point +in that cell into a weight for each corner of the cell with :mod:`._cells`. +""" + +from ._weights_multilinear import weights_multilinear + +__all__ = [ + "weights_multilinear", +] diff --git a/regridding/_weights/_weights_multilinear/_cells.py b/regridding/_weights/_weights_multilinear/_cells.py new file mode 100644 index 0000000..3a62c12 --- /dev/null +++ b/regridding/_weights/_weights_multilinear/_cells.py @@ -0,0 +1,490 @@ +r""" +The multilinear map from the unit cube onto a cell of a logically-rectangular +grid, in any number of dimensions. + +A grid of :math:`N` dimensions has :math:`N` coordinates, each sampled on the +same array of vertices. Here the vertices are stored flat, in row-major +order, as an array of shape ``(N, number of vertices)``, so that one compiled +function serves any number of dimensions. + +A cell is addressed by its first corner, the vertex with the smallest index +along every axis. Its :math:`2^N` corners are numbered so that bit :math:`m` +of a corner's number says whether that corner is one vertex further along +axis :math:`m`. A point of the cell has the local coordinates +:math:`\xi \in [0, 1]^N`, and corner :math:`b` contributes to it with the +weight + +.. math:: + + w_b(\xi) = \prod_{m=0}^{N-1} + \begin{cases} + \xi_m & \text{if bit } m \text{ of } b \text{ is set,} \\ + 1 - \xi_m & \text{otherwise,} + \end{cases} + +so that the position of the point is :math:`x(\xi) = \sum_b w_b(\xi) x_b`, +where :math:`x_b` is the position of corner :math:`b`. The same weights +interpolate any array sampled on the vertices, which is what +:func:`regridding.weights` saves. + +Positions are measured from the first corner of the cell, +:math:`x(\xi) - x_0 = \sum_b w_b(\xi) (x_b - x_0)`, since the weights sum to +one. Their rounding error is then a fraction of the size of the cell, rather +than of the distance of the cell from the origin, so that the local +coordinates of a point can be found to the same precision wherever the grid +is. + +Outside the unit cube, the same expressions continue the cell, which is how +a point outside the grid is extrapolated. +""" + +import numpy as np +import numba + +__all__ = [ + "num_work", + "strides", + "offsets", + "evaluate", + "position", + "target", + "distance_squared", + "solve", + "inverse", + "weights", +] + +num_work = 6 +""" +The number of rows of a workspace for :func:`inverse`, beyond the +:math:`2N` rows of the two matrices. + +A workspace has the shape ``(2 * N + num_work, N)``. +""" + +tolerance = 1e-13 +""" +The size of the last step, in units of the local coordinates, below which +Newton's method has converged. +""" + +num_iteration = 32 +"""The most steps Newton's method takes before it gives up.""" + +num_halving = 30 +"""The most times a step of Newton's method is halved before it gives up.""" + + +@numba.njit(cache=True) +def strides( + shape: np.ndarray, +) -> np.ndarray: + """ + The number of vertices between neighbors along each axis of a grid + stored in row-major order. + + Parameters + ---------- + shape + The number of vertices along each axis. + """ + num_axis = shape.shape[0] + result = np.empty(num_axis, dtype=np.int64) + stride = 1 + for m in range(num_axis - 1, -1, -1): + result[m] = stride + stride = stride * shape[m] + return result + + +@numba.njit(cache=True) +def offsets( + shape: np.ndarray, +) -> np.ndarray: + """ + The offset of each corner of a cell from its first corner, in the flat + index of the vertices. + + Parameters + ---------- + shape + The number of vertices along each axis. + """ + num_axis = shape.shape[0] + stride = strides(shape) + num_corner = 1 << num_axis + result = np.zeros(num_corner, dtype=np.int64) + for b in range(num_corner): + for m in range(num_axis): + if (b >> m) & 1: + result[b] += stride[m] + return result + + +@numba.njit(cache=True) +def evaluate( + vertices: np.ndarray, + base: int, + offset: np.ndarray, + xi: np.ndarray, + x: np.ndarray, + jacobian: np.ndarray, +) -> None: + r""" + The position of a point of a cell, measured from the first corner of the + cell, and its derivative with respect to the local coordinates. + + Parameters + ---------- + vertices + The coordinates of every vertex of the grid, of shape + ``(N, number of vertices)``. + base + The flat index of the first corner of the cell. + offset + The offset of each corner from the first, from :func:`offsets`. + xi + The local coordinates of the point. + x + Where to put the position of the point, measured from the first + corner of the cell. + jacobian + Where to put the derivative, an :math:`N \times N` array whose element + ``[k, m]`` is the derivative of coordinate :math:`k` with respect to + :math:`\xi_m`. + """ + num_axis = xi.shape[0] + x[:] = 0 + jacobian[:, :] = 0 + for b in range(offset.shape[0]): + index = base + offset[b] + weight = 1.0 + for m in range(num_axis): + if (b >> m) & 1: + weight *= xi[m] + else: + weight *= 1 - xi[m] + for k in range(num_axis): + x[k] += weight * (vertices[k, index] - vertices[k, base]) + for m in range(num_axis): + # the product of every factor but the one along axis m, which is + # computed directly rather than by dividing, since a factor can + # be zero + derivative = 1.0 if (b >> m) & 1 else -1.0 + for j in range(num_axis): + if j != m: + if (b >> j) & 1: + derivative *= xi[j] + else: + derivative *= 1 - xi[j] + for k in range(num_axis): + jacobian[k, m] += derivative * ( + vertices[k, index] - vertices[k, base] + ) + + +@numba.njit(cache=True) +def position( + vertices: np.ndarray, + base: int, + offset: np.ndarray, + xi: np.ndarray, + x: np.ndarray, +) -> None: + """ + The position of a point of a cell, measured from the first corner of the + cell. + + Parameters + ---------- + vertices + The coordinates of every vertex of the grid, of shape + ``(N, number of vertices)``. + base + The flat index of the first corner of the cell. + offset + The offset of each corner from the first, from :func:`offsets`. + xi + The local coordinates of the point. + x + Where to put the position of the point, measured from the first + corner of the cell. + """ + num_axis = xi.shape[0] + x[:] = 0 + for b in range(offset.shape[0]): + index = base + offset[b] + weight = 1.0 + for m in range(num_axis): + if (b >> m) & 1: + weight *= xi[m] + else: + weight *= 1 - xi[m] + for k in range(num_axis): + x[k] += weight * (vertices[k, index] - vertices[k, base]) + + +@numba.njit(cache=True) +def solve( + a: np.ndarray, + b: np.ndarray, + out: np.ndarray, +) -> bool: + """ + Solve the square system ``a @ out = b`` by Gaussian elimination with + partial pivoting. + + Both `a` and `b` are overwritten. Returns whether the system could be + solved, which it cannot if `a` is singular. + + Parameters + ---------- + a + The matrix of the system. + b + The right-hand side of the system. + out + Where to put the solution. + """ + num = b.shape[0] + for column in range(num): + pivot = column + largest = abs(a[column, column]) + for row in range(column + 1, num): + if abs(a[row, column]) > largest: + largest = abs(a[row, column]) + pivot = row + if not largest > 0: + return False + if pivot != column: + for j in range(num): + a[column, j], a[pivot, j] = a[pivot, j], a[column, j] + b[column], b[pivot] = b[pivot], b[column] + for row in range(column + 1, num): + factor = a[row, column] / a[column, column] + for j in range(column, num): + a[row, j] -= factor * a[column, j] + b[row] -= factor * b[column] + for row in range(num - 1, -1, -1): + total = b[row] + for j in range(row + 1, num): + total -= a[row, j] * out[j] + out[row] = total / a[row, row] + return True + + +@numba.njit(cache=True) +def target( + vertices: np.ndarray, + base: int, + point: np.ndarray, + out: np.ndarray, +) -> None: + """ + The position of a point measured from the first corner of a cell, as + :func:`evaluate` and :func:`position` measure it. + + Parameters + ---------- + vertices + The coordinates of every vertex of the grid. + base + The flat index of the first corner of the cell. + point + The point. + out + Where to put the result. + """ + for k in range(point.shape[0]): + out[k] = point[k] - vertices[k, base] + + +@numba.njit(cache=True) +def distance_squared( + x: np.ndarray, + y: np.ndarray, +) -> float: + """ + The square of the distance between two points. + + Parameters + ---------- + x + The first point. + y + The second point. + """ + result = 0.0 + for k in range(x.shape[0]): + result += (x[k] - y[k]) ** 2 + return result + + +@numba.njit(cache=True) +def inverse( + vertices: np.ndarray, + base: int, + offset: np.ndarray, + point: np.ndarray, + xi: np.ndarray, + work: np.ndarray, +) -> bool: + """ + Find the local coordinates of a point in a cell, which may be outside + the unit cube if the point is outside the cell. + + This uses Newton's method, starting from the value of `xi` it is given. + Each step is halved until it brings the cell's position closer to the + point, so that the method does not run away on a strongly curved cell. + It has converged once a step is shorter than :data:`tolerance`, which is + in units of the local coordinates and so does not depend on the units + or the size of the grid. + + Returns whether it converged. If it did not, `xi` is left wherever the + method stopped. + + Parameters + ---------- + vertices + The coordinates of every vertex of the grid, of shape + ``(N, number of vertices)``. + base + The flat index of the first corner of the cell. + offset + The offset of each corner from the first, from :func:`offsets`. + point + The point to find. + xi + The local coordinates to start from, which are replaced by the + result. + work + A workspace of shape ``(2 * N + num_work, N)``. + """ + num_axis = xi.shape[0] + + jacobian = work[:num_axis] + matrix = work[num_axis : 2 * num_axis] + x = work[2 * num_axis] + residual = work[2 * num_axis + 1] + step = work[2 * num_axis + 2] + trial = work[2 * num_axis + 3] + x_trial = work[2 * num_axis + 4] + goal = work[2 * num_axis + 5] + + target(vertices, base, point, goal) + evaluate(vertices, base, offset, xi, x, jacobian) + error = distance_squared(x, goal) + + for _ in range(num_iteration): + + if error == 0: + return True + + matrix[:, :] = jacobian + for k in range(num_axis): + residual[k] = x[k] - goal[k] + if not solve(matrix, residual, step): + return False + + size = 0.0 + scale = 1.0 + for m in range(num_axis): + size = max(size, abs(step[m])) + scale = max(scale, abs(xi[m])) + if not size < np.inf: + return False + + fraction = 1.0 + accepted = False + for _ in range(num_halving): + for m in range(num_axis): + trial[m] = xi[m] - fraction * step[m] + position(vertices, base, offset, trial, x_trial) + error_trial = distance_squared(x_trial, goal) + if error_trial < error: + accepted = True + break + fraction = fraction / 2 + + if not accepted: + # no point along the step is closer, which happens once the step + # is down to the rounding error of the position + return size <= tolerance * scale + + xi[:] = trial + + if fraction * size <= tolerance * scale: + return True + + evaluate(vertices, base, offset, xi, x, jacobian) + error = distance_squared(x, goal) + + return False + + +@numba.njit(cache=True, parallel=True) +def weights( + cell: np.ndarray, + xi: np.ndarray, + inside: np.ndarray, + shape: np.ndarray, + weights_input: np.ndarray, + nan_outside: bool, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + The weight of each corner of the cell containing each point. + + Returns the flat arrays ``(indices_input, indices_output, values)``, + with :math:`2^N` consecutive weights for each point, one for each corner + of its cell, in the order the corners are numbered. + + Parameters + ---------- + cell + The first corner of the cell of each point, of shape + ``(number of points, N)``. + xi + The local coordinates of each point in its cell, of the same shape. + inside + Whether each point is inside the grid. + shape + The number of vertices along each axis. + weights_input + Weights to multiply the weight of each vertex by, as a flat array, or + an empty array for none. + nan_outside + Whether to set the weights of a point outside the grid to NaN rather + than extrapolate. + """ + num_point, num_axis = cell.shape + stride = strides(shape) + offset = offsets(shape) + num_corner = offset.shape[0] + weighted = weights_input.shape[0] > 0 + + indices_input = np.empty(num_point * num_corner, dtype=np.int64) + indices_output = np.empty(num_point * num_corner, dtype=np.int64) + values = np.empty(num_point * num_corner, dtype=np.float64) + + for i in numba.prange(num_point): + base = 0 + for m in range(num_axis): + base += cell[i, m] * stride[m] + for b in range(num_corner): + index = base + offset[b] + weight = 1.0 + for m in range(num_axis): + if (b >> m) & 1: + weight *= xi[i, m] + else: + weight *= 1 - xi[i, m] + if weighted: + weight *= weights_input[index] + if nan_outside and not inside[i]: + weight = np.nan + w = i * num_corner + b + indices_input[w] = index + indices_output[w] = i + values[w] = weight + + return indices_input, indices_output, values diff --git a/regridding/_weights/_weights_multilinear/_cells_test.py b/regridding/_weights/_weights_multilinear/_cells_test.py new file mode 100644 index 0000000..cda016c --- /dev/null +++ b/regridding/_weights/_weights_multilinear/_cells_test.py @@ -0,0 +1,197 @@ +import pytest +import numpy as np +from . import _cells + + +def _cell( + num_axis: int, + seed: int = 0, + scale: float = 1, + shift: float = 0, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + A grid of a single cell, the unit cube with each corner moved at random, + as its vertices, its shape, and the offsets of its corners. + """ + rng = np.random.default_rng(seed) + shape = np.full(num_axis, 2, dtype=np.int64) + corners = np.stack(np.meshgrid(*[np.arange(2.0)] * num_axis, indexing="ij")) + corners = corners + 0.15 * rng.uniform(-1, 1, corners.shape) + vertices = shift + scale * corners.reshape(num_axis, -1) + return vertices, shape, _cells.offsets(shape) + + +def _work(num_axis: int) -> np.ndarray: + """A workspace for :func:`_cells.inverse`.""" + return np.empty((2 * num_axis + _cells.num_work, num_axis)) + + +def test_strides() -> None: + result = _cells.strides(np.array([4, 5, 6])) + np.testing.assert_array_equal(result, [30, 6, 1]) + + +def test_offsets() -> None: + result = _cells.offsets(np.array([4, 5, 6])) + np.testing.assert_array_equal(result, [0, 30, 6, 36, 1, 31, 7, 37]) + + +@pytest.mark.parametrize("num_axis", [1, 2, 3]) +class TestEvaluate: + + def test_corners(self, num_axis: int) -> None: + vertices, shape, offset = _cell(num_axis) + x = np.empty(num_axis) + jacobian = np.empty((num_axis, num_axis)) + result = np.empty((offset.size, num_axis)) + for b in range(offset.size): + xi = np.array([(b >> m) & 1 for m in range(num_axis)], dtype=float) + _cells.evaluate(vertices, 0, offset, xi, x, jacobian) + result[b] = x + expected = (vertices[:, offset] - vertices[:, :1]).T + np.testing.assert_allclose(result, expected, rtol=0, atol=1e-15) + + def test_position(self, num_axis: int) -> None: + vertices, shape, offset = _cell(num_axis) + xi = np.random.default_rng(1).uniform(-0.5, 1.5, num_axis) + x = np.empty(num_axis) + jacobian = np.empty((num_axis, num_axis)) + _cells.evaluate(vertices, 0, offset, xi, x, jacobian) + result = np.empty(num_axis) + _cells.position(vertices, 0, offset, xi, result) + np.testing.assert_array_equal(result, x) + + def test_jacobian(self, num_axis: int) -> None: + vertices, shape, offset = _cell(num_axis) + xi = np.random.default_rng(1).uniform(-0.5, 1.5, num_axis) + x = np.empty(num_axis) + jacobian = np.empty((num_axis, num_axis)) + _cells.evaluate(vertices, 0, offset, xi, x, jacobian) + h = 1e-6 + expected = np.empty((num_axis, num_axis)) + plus = np.empty(num_axis) + minus = np.empty(num_axis) + for m in range(num_axis): + step = np.zeros(num_axis) + step[m] = h + _cells.position(vertices, 0, offset, xi + step, plus) + _cells.position(vertices, 0, offset, xi - step, minus) + expected[:, m] = (plus - minus) / (2 * h) + np.testing.assert_allclose(jacobian, expected, rtol=0, atol=1e-9) + + +class TestSolve: + + @pytest.mark.parametrize("num", [1, 2, 3, 5]) + def test_random(self, num: int) -> None: + rng = np.random.default_rng(num) + a = rng.normal(size=(num, num)) + b = rng.normal(size=num) + expected = np.linalg.solve(a, b) + out = np.empty(num) + assert _cells.solve(a.copy(), b.copy(), out) + np.testing.assert_allclose(out, expected, rtol=1e-12) + + def test_pivot(self) -> None: + a = np.array([[0.0, 1.0], [1.0, 0.0]]) + b = np.array([2.0, 3.0]) + out = np.empty(2) + assert _cells.solve(a, b, out) + np.testing.assert_array_equal(out, [3.0, 2.0]) + + def test_singular(self) -> None: + a = np.array([[1.0, 2.0], [2.0, 4.0]]) + b = np.array([1.0, 1.0]) + out = np.empty(2) + assert not _cells.solve(a, b, out) + + +@pytest.mark.parametrize("num_axis", [1, 2, 3]) +class TestInverse: + + @pytest.mark.parametrize( + argnames="scale,shift", + argvalues=[ + (1, 0), + (1e-12, 0), + (1e9, 0), + (1e-3, 1e3), + ], + ) + def test_round_trip(self, num_axis: int, scale: float, shift: float) -> None: + vertices, shape, offset = _cell(num_axis, scale=scale, shift=shift) + rng = np.random.default_rng(2) + expected = rng.uniform(-0.5, 1.5, (20, num_axis)) + work = _work(num_axis) + point = np.empty(num_axis) + result = np.empty_like(expected) + for i in range(expected.shape[0]): + _cells.position(vertices, 0, offset, expected[i], point) + point = point + vertices[:, 0] + xi = np.full(num_axis, 0.5) + assert _cells.inverse(vertices, 0, offset, point, xi, work) + result[i] = xi + # the local coordinates are only as good as the point, whose rounding + # error grows with its distance from the origin + atol = 1e-12 + 1e-15 * shift / scale + np.testing.assert_allclose(result, expected, rtol=0, atol=atol) + + def test_unreachable(self, num_axis: int) -> None: + # a cell squashed flat along its first axis cannot reach a point off + # the plane it lies in + vertices, shape, offset = _cell(num_axis) + vertices[0] = 0 + point = np.full(num_axis, 0.5) + xi = np.full(num_axis, 0.5) + assert not _cells.inverse(vertices, 0, offset, point, xi, _work(num_axis)) + + +class TestWeights: + + def test_partition_of_unity(self) -> None: + rng = np.random.default_rng(3) + shape = np.array([4, 5, 3]) + cell = np.stack([rng.integers(0, n - 1, 10) for n in shape], axis=~0) + xi = rng.uniform(-0.5, 1.5, (10, 3)) + inside = np.ones(10, dtype=bool) + indices_input, indices_output, values = _cells.weights( + cell, xi, inside, shape, np.empty(0), False + ) + np.testing.assert_array_equal(indices_output, np.repeat(np.arange(10), 8)) + sums = np.bincount(indices_output, weights=values) + np.testing.assert_allclose(sums, 1, rtol=1e-14) + first = indices_input[::8] + np.testing.assert_array_equal(first, np.ravel_multi_index(cell.T, shape)) + + def test_corners_of_the_cell(self) -> None: + shape = np.array([3, 4]) + cell = np.array([[1, 2]]) + xi = np.array([[0.25, 0.5]]) + indices_input, indices_output, values = _cells.weights( + cell, xi, np.ones(1, dtype=bool), shape, np.empty(0), False + ) + np.testing.assert_array_equal(indices_input, [6, 10, 7, 11]) + np.testing.assert_allclose(values, [0.375, 0.125, 0.375, 0.125]) + + def test_weights_input(self) -> None: + shape = np.array([3, 4]) + cell = np.array([[1, 2]]) + xi = np.array([[0.25, 0.5]]) + weights_input = np.arange(12, dtype=float) + _, _, values = _cells.weights( + cell, xi, np.ones(1, dtype=bool), shape, weights_input, False + ) + expected = np.array([0.375, 0.125, 0.375, 0.125]) * [6, 10, 7, 11] + np.testing.assert_allclose(values, expected) + + @pytest.mark.parametrize("nan_outside", [False, True]) + def test_outside(self, nan_outside: bool) -> None: + shape = np.array([3, 4]) + cell = np.array([[1, 2], [0, 0]]) + xi = np.array([[1.5, 0.5], [0.5, 0.5]]) + inside = np.array([False, True]) + _, _, values = _cells.weights(cell, xi, inside, shape, np.empty(0), nan_outside) + first = [-0.25, 0.75, -0.25, 0.75] + if nan_outside: + first = [np.nan] * 4 + np.testing.assert_allclose(values, first + [0.25] * 4) diff --git a/regridding/_weights/_weights_multilinear/_curvilinear.py b/regridding/_weights/_weights_multilinear/_curvilinear.py new file mode 100644 index 0000000..a0aee83 --- /dev/null +++ b/regridding/_weights/_weights_multilinear/_curvilinear.py @@ -0,0 +1,1003 @@ +""" +Find the cell of a curvilinear grid which contains each of a set of points. + +The grid is searched through a hierarchy of bounding boxes. Level zero is +the cells themselves. Level one divides them into blocks of two cells along +every axis, or one at the far edge of an axis with an odd number of cells, +and each level above groups the blocks of the level below in the same way, +until one block holds the whole grid. Only the boxes of level one and above +are stored; the box of a cell is taken from its corners when it is needed. + +A point is looked for in the blocks whose boxes contain it, from the top +level down, and in each cell whose box contains it, its local coordinates +are found with Newton's method. The point is in the first cell where they +are inside the unit cube. Since the boxes of neighboring blocks overlap +little, a point reaches its cell in a number of steps which grows as the +logarithm of the size of the grid. + +Searching the boxes first means that a grid which curves back on itself, +such as an annulus, cannot lead the search astray. A search by Newton's +method across the whole grid can follow the continuation of a cell across +the hole of the annulus to a point on the other side, and conclude that the +point is outside the grid. It can also cycle between neighboring cells of +very different sizes. + +A point which is in no cell is outside the grid. To extrapolate to it, +Newton's method is run across the whole grid in index space, with the cells +along its edges continued beyond it, starting from the cell whose bounding +box is nearest to the point. On a grid of parallelograms, this finds the +same cell as clamping each index of the point to the grid. +""" + +import numpy as np +import numba +from . import _cells + +__all__ = [ + "tolerance", + "tree", + "locate", +] + +tolerance = 1e-9 +""" +How far outside the unit cube the local coordinates of a point may be, in +units of the cell, for the point to count as inside the cell. + +A point on the face shared by two cells, or on the boundary of the grid, +has local coordinates on the boundary of the unit cube, and finding them +numerically leaves them off it by a rounding error. Without some margin, +such a point could be outside every cell. +""" + +num_iteration_walk = 64 +""" +The most steps of Newton's method across the whole grid when extrapolating. +""" + + +@numba.njit(cache=True) +def _unravel( + index: int, + shape: np.ndarray, + out: np.ndarray, +) -> None: + """ + Convert a flat index into an index along each axis of an array stored in + row-major order. + + Parameters + ---------- + index + The flat index. + shape + The shape of the array. + out + Where to put the index along each axis. + """ + for m in range(shape.shape[0] - 1, -1, -1): + out[m] = index % shape[m] + index = index // shape[m] + + +@numba.njit(cache=True) +def _ravel( + index: np.ndarray, + shape: np.ndarray, +) -> int: + """ + Convert an index along each axis of an array stored in row-major order + into a flat index. + + Parameters + ---------- + index + The index along each axis. + shape + The shape of the array. + """ + result = 0 + for m in range(shape.shape[0]): + result = result * shape[m] + index[m] + return result + + +@numba.njit(cache=True, parallel=True) +def _boxes_of_cells( + vertices: np.ndarray, + shape: np.ndarray, + shape_block: np.ndarray, + lower: np.ndarray, + upper: np.ndarray, + num_chunk: int, +) -> None: + """ + The bounding boxes of the blocks of level one, each of two cells along + every axis, taken from their vertices. + + A vertex whose coordinate is NaN is left out of the box. + + Parameters + ---------- + vertices + The coordinates of every vertex of the grid, of shape + ``(N, number of vertices)``. + shape + The number of vertices along each axis. + shape_block + The number of blocks along each axis. + lower + Where to put the lower corner of the box of each block, of shape + ``(number of blocks, N)``. + upper + Where to put the upper corner of the box of each block. + num_chunk + The number of pieces to divide the work into. + """ + num_axis = shape.shape[0] + stride = _cells.strides(shape) + num_block = lower.shape[0] + + for chunk in numba.prange(num_chunk): + block = np.empty(num_axis, dtype=np.int64) + first = np.empty(num_axis, dtype=np.int64) + last = np.empty(num_axis, dtype=np.int64) + vertex = np.empty(num_axis, dtype=np.int64) + for k in range(chunk, num_block, num_chunk): + _unravel(k, shape_block, block) + for m in range(num_axis): + first[m] = 2 * block[m] + last[m] = min(2 * block[m] + 2, shape[m] - 1) + vertex[m] = first[m] + for c in range(num_axis): + lower[k, c] = np.inf + upper[k, c] = -np.inf + while True: + index = 0 + for m in range(num_axis): + index += vertex[m] * stride[m] + for c in range(num_axis): + value = vertices[c, index] + if value < lower[k, c]: + lower[k, c] = value + if value > upper[k, c]: + upper[k, c] = value + m = num_axis - 1 + while m >= 0: + vertex[m] += 1 + if vertex[m] <= last[m]: + break + vertex[m] = first[m] + m -= 1 + if m < 0: + break + + +@numba.njit(cache=True, parallel=True) +def _boxes_of_boxes( + lower_below: np.ndarray, + upper_below: np.ndarray, + shape_below: np.ndarray, + shape_block: np.ndarray, + lower: np.ndarray, + upper: np.ndarray, + num_chunk: int, +) -> None: + """ + The bounding boxes of the blocks of a level above level one, each of up + to two blocks of the level below along every axis. + + Parameters + ---------- + lower_below + The lower corner of the box of each block of the level below. + upper_below + The upper corner of the box of each block of the level below. + shape_below + The number of blocks along each axis of the level below. + shape_block + The number of blocks along each axis of this level. + lower + Where to put the lower corner of the box of each block of this level. + upper + Where to put the upper corner of the box of each block of this level. + num_chunk + The number of pieces to divide the work into. + """ + num_axis = shape_block.shape[0] + num_block = lower.shape[0] + + for chunk in numba.prange(num_chunk): + block = np.empty(num_axis, dtype=np.int64) + child = np.empty(num_axis, dtype=np.int64) + for k in range(chunk, num_block, num_chunk): + _unravel(k, shape_block, block) + for c in range(num_axis): + lower[k, c] = np.inf + upper[k, c] = -np.inf + for b in range(1 << num_axis): + valid = True + for m in range(num_axis): + child[m] = 2 * block[m] + ((b >> m) & 1) + if child[m] >= shape_below[m]: + valid = False + if not valid: + continue + j = _ravel(child, shape_below) + for c in range(num_axis): + if lower_below[j, c] < lower[k, c]: + lower[k, c] = lower_below[j, c] + if upper_below[j, c] > upper[k, c]: + upper[k, c] = upper_below[j, c] + + +def tree( + vertices: np.ndarray, + shape: np.ndarray, + num_chunk: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """ + Build the hierarchy of bounding boxes of a grid. + + Returns ``(lower, upper, start, shapes)``. ``shapes[level]`` is the + number of blocks along each axis of each level, starting with the cells + at level zero. The box of block ``k`` of a level above zero has the + lower corner ``lower[start[level] + k]`` and the upper corner + ``upper[start[level] + k]``. + + Parameters + ---------- + vertices + The coordinates of every vertex of the grid, of shape + ``(N, number of vertices)``, in row-major order. + shape + The number of vertices along each axis. + num_chunk + The number of pieces to divide the work into. + """ + num_axis = shape.size + + shapes = [shape - 1] + while np.any(shapes[~0] > 1): + shapes.append((shapes[~0] + 1) // 2) + shapes = np.array(shapes, dtype=np.int64) + num_level = shapes.shape[0] - 1 + + start = np.zeros(num_level + 2, dtype=np.int64) + start[2:] = np.cumsum(np.prod(shapes[1:], axis=~0)) + + lower = np.empty((start[~0], num_axis), dtype=np.float64) + upper = np.empty((start[~0], num_axis), dtype=np.float64) + + for level in range(1, num_level + 1): + a, b = start[level], start[level + 1] + if level == 1: + _boxes_of_cells( + vertices, + shape, + shapes[level], + lower[a:b], + upper[a:b], + num_chunk, + ) + else: + a_below, b_below = start[level - 1], start[level] + _boxes_of_boxes( + lower[a_below:b_below], + upper[a_below:b_below], + shapes[level - 1], + shapes[level], + lower[a:b], + upper[a:b], + num_chunk, + ) + + return lower, upper, start, shapes + + +@numba.njit(cache=True) +def _in_box( + lower: np.ndarray, + upper: np.ndarray, + index: int, + point: np.ndarray, +) -> bool: + """ + Whether a point is inside a bounding box, or on its boundary. + + Parameters + ---------- + lower + The lower corners of the boxes. + upper + The upper corners of the boxes. + index + Which box. + point + The point. + """ + for c in range(point.shape[0]): + if not (lower[index, c] <= point[c] <= upper[index, c]): + return False + return True + + +@numba.njit(cache=True) +def _in_cell( + vertices: np.ndarray, + base: int, + offset: np.ndarray, + point: np.ndarray, + xi: np.ndarray, + work: np.ndarray, +) -> bool: + """ + Whether a point is inside a cell, and if it is, its local coordinates. + + Parameters + ---------- + vertices + The coordinates of every vertex of the grid. + base + The flat index of the first corner of the cell. + offset + The offset of each corner from the first. + point + The point. + xi + Where to put the local coordinates of the point. + work + A workspace for :func:`~._cells.inverse`. + """ + num_axis = point.shape[0] + + for c in range(num_axis): + low = np.inf + high = -np.inf + for b in range(offset.shape[0]): + value = vertices[c, base + offset[b]] + if value < low: + low = value + if value > high: + high = value + if not (low <= point[c] <= high): + return False + + xi[:] = 0.5 + if not _cells.inverse(vertices, base, offset, point, xi, work): + return False + + for m in range(num_axis): + if not (-tolerance <= xi[m] <= 1 + tolerance): + return False + + return True + + +@numba.njit(cache=True) +def _search( + vertices: np.ndarray, + stride: np.ndarray, + offset: np.ndarray, + lower: np.ndarray, + upper: np.ndarray, + start: np.ndarray, + shapes: np.ndarray, + point: np.ndarray, + stack_level: np.ndarray, + stack_index: np.ndarray, + block: np.ndarray, + cell: np.ndarray, + xi: np.ndarray, + work: np.ndarray, +) -> bool: + """ + Find the cell which contains a point by descending the hierarchy of + bounding boxes. + + Returns whether the point is in a cell. If it is, the first corner of + the cell is left in `cell` and the local coordinates of the point in + `xi`. + + Parameters + ---------- + vertices + The coordinates of every vertex of the grid. + stride + The number of vertices between neighbors along each axis. + offset + The offset of each corner of a cell from the first. + lower + The lower corners of the boxes, from :func:`tree`. + upper + The upper corners of the boxes. + start + Where the boxes of each level start. + shapes + The number of blocks along each axis of each level. + point + The point to find. + stack_level + Room for the level of each block waiting to be searched. + stack_index + Room for the index of each block waiting to be searched. + block + Room for the index along each axis of a block. + cell + Where to put the first corner of the cell. + xi + Where to put the local coordinates of the point. + work + A workspace for :func:`~._cells.inverse`. + """ + num_axis = point.shape[0] + num_level = shapes.shape[0] - 1 + num_corner = offset.shape[0] + + if num_level == 0: + cell[:] = 0 + return _in_cell(vertices, 0, offset, point, xi, work) + + if not _in_box(lower, upper, int(start[num_level]), point): + return False + + stack_level[0] = num_level + stack_index[0] = 0 + size = 1 + + while size > 0: + size -= 1 + level = stack_level[size] + _unravel(stack_index[size], shapes[level], block) + below = shapes[level - 1] + + if level == 1: + # the children are cells, which are tried in order + for b in range(num_corner): + valid = True + base = 0 + for m in range(num_axis): + cell[m] = 2 * block[m] + ((b >> m) & 1) + if cell[m] >= below[m]: + valid = False + base += cell[m] * stride[m] + if valid and _in_cell(vertices, base, offset, point, xi, work): + return True + + else: + # the children are blocks, which are pushed in reverse so that + # they come off the stack in order + for b in range(num_corner - 1, -1, -1): + valid = True + for m in range(num_axis): + cell[m] = 2 * block[m] + ((b >> m) & 1) + if cell[m] >= below[m]: + valid = False + if not valid: + continue + index = _ravel(cell, below) + if _in_box(lower, upper, int(start[level - 1]) + index, point): + stack_level[size] = level - 1 + stack_index[size] = index + size += 1 + + return False + + +@numba.njit(cache=True) +def _distance( + low: float, + high: float, + value: float, +) -> float: + """ + The distance from a value to an interval, or zero if it is inside. + + Parameters + ---------- + low + The lower end of the interval. + high + The upper end of the interval. + value + The value. + """ + if value < low: + return low - value + if value > high: + return value - high + return 0.0 + + +@numba.njit(cache=True) +def _distance_box( + lower: np.ndarray, + upper: np.ndarray, + index: int, + point: np.ndarray, +) -> float: + """ + The square of the distance from a point to a bounding box. + + Parameters + ---------- + lower + The lower corners of the boxes. + upper + The upper corners of the boxes. + index + Which box. + point + The point. + """ + result = 0.0 + for c in range(point.shape[0]): + result += _distance(lower[index, c], upper[index, c], point[c]) ** 2 + return result + + +@numba.njit(cache=True) +def _distance_cell( + vertices: np.ndarray, + base: int, + offset: np.ndarray, + point: np.ndarray, +) -> float: + """ + The square of the distance from a point to the bounding box of a cell. + + Parameters + ---------- + vertices + The coordinates of every vertex of the grid. + base + The flat index of the first corner of the cell. + offset + The offset of each corner from the first. + point + The point. + """ + result = 0.0 + for c in range(point.shape[0]): + low = np.inf + high = -np.inf + for b in range(offset.shape[0]): + value = vertices[c, base + offset[b]] + if value < low: + low = value + if value > high: + high = value + result += _distance(low, high, point[c]) ** 2 + return result + + +@numba.njit(cache=True) +def _nearest( + vertices: np.ndarray, + stride: np.ndarray, + offset: np.ndarray, + lower: np.ndarray, + upper: np.ndarray, + start: np.ndarray, + shapes: np.ndarray, + point: np.ndarray, + stack_level: np.ndarray, + stack_index: np.ndarray, + block: np.ndarray, + child_index: np.ndarray, + child_distance: np.ndarray, + cell: np.ndarray, + nearest: np.ndarray, +) -> None: + """ + Find the cell whose bounding box is nearest to a point. + + This is a branch-and-bound search of the hierarchy of bounding boxes: + a block is skipped if its box is no nearer than the nearest cell found + so far, and the children of a block are searched nearest first, so that + a near cell is found early. The first corner of the cell is left in + `nearest`. + + Parameters + ---------- + vertices + The coordinates of every vertex of the grid. + stride + The number of vertices between neighbors along each axis. + offset + The offset of each corner of a cell from the first. + lower + The lower corners of the boxes, from :func:`tree`. + upper + The upper corners of the boxes. + start + Where the boxes of each level start. + shapes + The number of blocks along each axis of each level. + point + The point. + stack_level + Room for the level of each block waiting to be searched. + stack_index + Room for the index of each block waiting to be searched. + block + Room for the index along each axis of a block. + child_index + Room for the index of each child of a block. + child_distance + Room for the distance to each child of a block. + cell + Room for the index along each axis of a cell. + nearest + Where to put the first corner of the nearest cell. + """ + num_axis = point.shape[0] + num_level = shapes.shape[0] - 1 + num_corner = offset.shape[0] + + nearest[:] = 0 + if num_level == 0: + return + + best = np.inf + stack_level[0] = num_level + stack_index[0] = 0 + size = 1 + + while size > 0: + size -= 1 + level = stack_level[size] + index = stack_index[size] + if not _distance_box(lower, upper, int(start[level]) + index, point) < best: + continue + _unravel(index, shapes[level], block) + below = shapes[level - 1] + + if level == 1: + for b in range(num_corner): + valid = True + base = 0 + for m in range(num_axis): + cell[m] = 2 * block[m] + ((b >> m) & 1) + if cell[m] >= below[m]: + valid = False + base += cell[m] * stride[m] + if not valid: + continue + distance = _distance_cell(vertices, base, offset, point) + if distance < best: + best = distance + nearest[:] = cell + + else: + num_child = 0 + for b in range(num_corner): + valid = True + for m in range(num_axis): + cell[m] = 2 * block[m] + ((b >> m) & 1) + if cell[m] >= below[m]: + valid = False + if not valid: + continue + child = _ravel(cell, below) + distance = _distance_box( + lower, upper, int(start[level - 1]) + child, point + ) + if not distance < best: + continue + # insert in order of decreasing distance, so that the + # nearest child is pushed last and comes off the stack first + j = num_child + while j > 0 and child_distance[j - 1] < distance: + child_distance[j] = child_distance[j - 1] + child_index[j] = child_index[j - 1] + j -= 1 + child_distance[j] = distance + child_index[j] = child + num_child += 1 + for j in range(num_child): + stack_level[size] = level - 1 + stack_index[size] = child_index[j] + size += 1 + + +@numba.njit(cache=True) +def _clamp( + s: np.ndarray, + shape: np.ndarray, + stride: np.ndarray, + cell: np.ndarray, + xi: np.ndarray, +) -> int: + """ + Split a fractional index into the cell nearest to it and its local + coordinates in that cell, which are outside the unit cube if the index + is outside the grid. + + Returns the flat index of the first corner of the cell. + + Parameters + ---------- + s + The fractional index of a point, along each axis of the vertices. + shape + The number of vertices along each axis. + stride + The number of vertices between neighbors along each axis. + cell + Where to put the first corner of the cell. + xi + Where to put the local coordinates. + """ + base = 0 + for m in range(s.shape[0]): + value = s[m] + if not value >= 0: + value = 0.0 + elif value > shape[m] - 2: + value = shape[m] - 2.0 + cell[m] = int(np.floor(value)) + xi[m] = s[m] - cell[m] + base += cell[m] * stride[m] + return base + + +@numba.njit(cache=True) +def _walk( + vertices: np.ndarray, + shape: np.ndarray, + stride: np.ndarray, + offset: np.ndarray, + point: np.ndarray, + s: np.ndarray, + s_trial: np.ndarray, + cell: np.ndarray, + cell_trial: np.ndarray, + xi: np.ndarray, + xi_trial: np.ndarray, + work: np.ndarray, +) -> None: + """ + Extrapolate the grid to a point outside of it. + + This runs Newton's method across the whole grid in index space, from the + fractional index in `s`, continuing the cells along the edges of the grid + beyond it. Each step is halved until it brings the position closer to + the point. The first corner of the cell it stops in is left in `cell`, + and the local coordinates of the point in that cell in `xi`. + + Parameters + ---------- + vertices + The coordinates of every vertex of the grid. + shape + The number of vertices along each axis. + stride + The number of vertices between neighbors along each axis. + offset + The offset of each corner of a cell from the first. + point + The point. + s + The fractional index to start from, which is replaced by that of the + point. + s_trial + Room for a trial fractional index. + cell + Where to put the first corner of the cell. + cell_trial + Room for the first corner of a trial cell. + xi + Where to put the local coordinates of the point in the cell. + xi_trial + Room for the local coordinates in a trial cell. + work + A workspace of the same shape as one for :func:`~._cells.inverse`. + """ + num_axis = point.shape[0] + + jacobian = work[:num_axis] + matrix = work[num_axis : 2 * num_axis] + x = work[2 * num_axis] + residual = work[2 * num_axis + 1] + step = work[2 * num_axis + 2] + x_trial = work[2 * num_axis + 3] + goal = work[2 * num_axis + 4] + goal_trial = work[2 * num_axis + 5] + + base = _clamp(s, shape, stride, cell, xi) + _cells.target(vertices, base, point, goal) + _cells.evaluate(vertices, base, offset, xi, x, jacobian) + error = _cells.distance_squared(x, goal) + + for _ in range(num_iteration_walk): + + if not error > 0: + break + + matrix[:, :] = jacobian + for k in range(num_axis): + residual[k] = x[k] - goal[k] + if not _cells.solve(matrix, residual, step): + break + + size = 0.0 + scale = 1.0 + for m in range(num_axis): + size = max(size, abs(step[m])) + scale = max(scale, abs(s[m])) + if not size < np.inf: + break + + fraction = 1.0 + accepted = False + for _ in range(_cells.num_halving): + for m in range(num_axis): + s_trial[m] = s[m] - fraction * step[m] + base = _clamp(s_trial, shape, stride, cell_trial, xi_trial) + _cells.target(vertices, base, point, goal_trial) + _cells.position(vertices, base, offset, xi_trial, x_trial) + error_trial = _cells.distance_squared(x_trial, goal_trial) + if error_trial < error: + accepted = True + break + fraction = fraction / 2 + + if not accepted: + break + + s[:] = s_trial + + if fraction * size <= _cells.tolerance * scale: + break + + base = _clamp(s, shape, stride, cell, xi) + _cells.target(vertices, base, point, goal) + _cells.evaluate(vertices, base, offset, xi, x, jacobian) + error = _cells.distance_squared(x, goal) + + _clamp(s, shape, stride, cell, xi) + + +@numba.njit(cache=True, parallel=True) +def locate( + vertices: np.ndarray, + shape: np.ndarray, + lower: np.ndarray, + upper: np.ndarray, + start: np.ndarray, + shapes: np.ndarray, + points: np.ndarray, + extrapolate: bool, + num_chunk: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Find the cell of a curvilinear grid which contains each point. + + Returns ``(cell, xi, inside)``: the first corner of the cell along each + axis, the local coordinates of the point in the cell, and whether the + point is inside the grid. A point outside the grid, or with a + coordinate which is NaN, is given the first cell and local coordinates + of NaN, unless `extrapolate` is set and its coordinates are finite. + + Parameters + ---------- + vertices + The coordinates of every vertex of the grid, of shape + ``(N, number of vertices)``. + shape + The number of vertices along each axis. + lower + The lower corners of the boxes, from :func:`tree`. + upper + The upper corners of the boxes. + start + Where the boxes of each level start. + shapes + The number of blocks along each axis of each level. + points + The points to find, of shape ``(N, number of points)``. + extrapolate + Whether to find the cell to extrapolate from for a point outside the + grid. + num_chunk + The number of pieces to divide the work into. + """ + num_axis, num_point = points.shape + num_level = shapes.shape[0] - 1 + + stride = _cells.strides(shape) + offset = _cells.offsets(shape) + num_corner = offset.shape[0] + capacity = num_level * num_corner + 1 + + cell = np.zeros((num_point, num_axis), dtype=np.int64) + xi = np.empty((num_point, num_axis), dtype=np.float64) + inside = np.zeros(num_point, dtype=np.bool_) + + for chunk in numba.prange(num_chunk): + stack_level = np.empty(capacity, dtype=np.int64) + stack_index = np.empty(capacity, dtype=np.int64) + block = np.empty(num_axis, dtype=np.int64) + cell_i = np.empty(num_axis, dtype=np.int64) + cell_trial = np.empty(num_axis, dtype=np.int64) + point = np.empty(num_axis, dtype=np.float64) + xi_i = np.empty(num_axis, dtype=np.float64) + xi_trial = np.empty(num_axis, dtype=np.float64) + s = np.empty(num_axis, dtype=np.float64) + s_trial = np.empty(num_axis, dtype=np.float64) + child_index = np.empty(num_corner, dtype=np.int64) + child_distance = np.empty(num_corner, dtype=np.float64) + work = np.empty((2 * num_axis + _cells.num_work, num_axis)) + + for i in range(chunk, num_point, num_chunk): + finite = True + for c in range(num_axis): + point[c] = points[c, i] + if not np.isfinite(point[c]): + finite = False + + found = False + if finite: + found = _search( + vertices, + stride, + offset, + lower, + upper, + start, + shapes, + point, + stack_level, + stack_index, + block, + cell_i, + xi_i, + work, + ) + + if found: + inside[i] = True + elif finite and extrapolate: + _nearest( + vertices, + stride, + offset, + lower, + upper, + start, + shapes, + point, + stack_level, + stack_index, + block, + child_index, + child_distance, + cell_trial, + cell_i, + ) + for m in range(num_axis): + s[m] = cell_i[m] + 0.5 + _walk( + vertices, + shape, + stride, + offset, + point, + s, + s_trial, + cell_i, + cell_trial, + xi_i, + xi_trial, + work, + ) + else: + cell_i[:] = 0 + xi_i[:] = np.nan + + for m in range(num_axis): + cell[i, m] = cell_i[m] + xi[i, m] = xi_i[m] + + return cell, xi, inside diff --git a/regridding/_weights/_weights_multilinear/_curvilinear_test.py b/regridding/_weights/_weights_multilinear/_curvilinear_test.py new file mode 100644 index 0000000..1a59a8c --- /dev/null +++ b/regridding/_weights/_weights_multilinear/_curvilinear_test.py @@ -0,0 +1,261 @@ +import pytest +import numpy as np +import regridding +from . import _cells, _curvilinear + + +def _flat(grid: tuple[np.ndarray, ...]) -> tuple[np.ndarray, np.ndarray]: + """The vertices of a grid, flattened, and its shape.""" + vertices = np.stack([np.ravel(c) for c in grid]) + return vertices.astype(float), np.array(grid[0].shape) + + +def _annulus( + radius_inner: float = 1, + num_radius: int = 6, + num_angle: int = 37, +) -> tuple[np.ndarray, np.ndarray]: + """ + An annulus which wraps most of the way around, whose hole made the secant + locator of the conservative sweep report points inside it as outside. + """ + r, t = np.meshgrid( + np.linspace(radius_inner, 2, num_radius), + np.linspace(0, 0.9 * 2 * np.pi, num_angle), + indexing="ij", + ) + return r * np.cos(t), r * np.sin(t) + + +def _distorted(scale: float = 1) -> tuple[np.ndarray, np.ndarray]: + """A grid with the kind of smooth distortion an optical system has.""" + u, v = np.meshgrid(np.linspace(-1, 1, 9), np.linspace(-1, 1, 7), indexing="ij") + return ( + scale * (u + 0.05 * u * v + 0.03 * v**2), + scale * (v + 0.04 * u**2 - 0.02 * u * v), + ) + + +def _alternating() -> tuple[np.ndarray, np.ndarray]: + """ + A rotated grid whose cells alternate in width, on which Newton's method + across the whole grid cycles between cells. + """ + widths = np.tile([0.4, 1.0], 6) + a, b = np.meshgrid( + np.concatenate([[0], np.cumsum(widths)]), + np.linspace(0, 1, 4), + indexing="ij", + ) + angle = 0.3 + return ( + a * np.cos(angle) - b * np.sin(angle), + a * np.sin(angle) + b * np.cos(angle), + ) + + +def _distorted_3d() -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """A three-dimensional grid with a smooth distortion.""" + p, q, w = np.meshgrid( + np.linspace(-1, 1, 5), + np.linspace(-1, 1, 6), + np.linspace(-1, 1, 4), + indexing="ij", + ) + return ( + p + 0.05 * q * w, + q + 0.04 * p**2, + w + 0.03 * p * q + 0.02 * q**2, + ) + + +def _folded_1d() -> tuple[np.ndarray]: + """A one-dimensional grid which turns back on itself.""" + return (np.array([0.0, 1.0, 3.0, 2.0, 4.0]),) + + +grids = { + "annulus": _annulus(), + "disk-like annulus": _annulus(radius_inner=0.05), + "distorted": _distorted(), + "distorted, scale 1e-12": _distorted(scale=1e-12), + "distorted, scale 1e9": _distorted(scale=1e9), + "alternating": _alternating(), + "distorted 3D": _distorted_3d(), + "folded 1D": _folded_1d(), + "single cell": tuple(c[:2, :2] for c in _distorted()), +} + + +def _points( + grid: tuple[np.ndarray, ...], + num: int = 200, + seed: int = 0, +) -> np.ndarray: + """Random points over a region a little larger than the grid.""" + rng = np.random.default_rng(seed) + lower = np.array([c.min() for c in grid]) + upper = np.array([c.max() for c in grid]) + margin = 0.1 * (upper - lower) + return rng.uniform(lower - margin, upper + margin, (num, len(grid))).T + + +def _inside_brute( + grid: tuple[np.ndarray, ...], + points: np.ndarray, +) -> np.ndarray: + """ + Whether each point is inside the grid, by trying every cell. + + In two dimensions a cell is the quadrilateral between its corners, since + its edges are straight, so this tests each quadrilateral with + :func:`regridding.geometry.point_is_inside_polygon`, which shares nothing + with the code under test. In other dimensions it finds the local + coordinates of the point in every cell. + """ + vertices, shape = _flat(grid) + num_axis = len(grid) + offset = _cells.offsets(shape) + stride = _cells.strides(shape) + work = np.empty((2 * num_axis + _cells.num_work, num_axis)) + xi = np.empty(num_axis) + result = np.zeros(points.shape[1], dtype=bool) + for i in range(points.shape[1]): + for index in np.ndindex(*(shape - 1)): + base = int(np.dot(index, stride)) + if num_axis == 2: + corners = base + offset[[0, 1, 3, 2]] + inside = regridding.geometry.point_is_inside_polygon( + x=points[0, i], + y=points[1, i], + vertices_x=vertices[0, corners], + vertices_y=vertices[1, corners], + ) + else: + xi[:] = 0.5 + inside = _cells.inverse( + vertices, base, offset, points[:, i], xi, work + ) and np.all((-1e-9 <= xi) & (xi <= 1 + 1e-9)) + if inside: + result[i] = True + break + return result + + +def _locate( + grid: tuple[np.ndarray, ...], + points: np.ndarray, + extrapolate: bool, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Build the hierarchy of a grid and locate the points in it.""" + vertices, shape = _flat(grid) + tree = _curvilinear.tree(vertices, shape, num_chunk=3) + return _curvilinear.locate(vertices, shape, *tree, points, extrapolate, 3) + + +def _position( + grid: tuple[np.ndarray, ...], + cell: np.ndarray, + xi: np.ndarray, +) -> np.ndarray: + """The position of each point given by its cell and local coordinates.""" + vertices, shape = _flat(grid) + num_axis = len(grid) + offset = _cells.offsets(shape) + stride = _cells.strides(shape) + result = np.empty((num_axis, cell.shape[0])) + x = np.empty(num_axis) + for i in range(cell.shape[0]): + base = int(np.dot(cell[i], stride)) + _cells.position(vertices, base, offset, xi[i], x) + result[:, i] = x + vertices[:, base] + return result + + +class TestTree: + + @pytest.mark.parametrize("name", list(grids)) + def test_boxes(self, name: str) -> None: + grid = grids[name] + vertices, shape = _flat(grid) + lower, upper, start, shapes = _curvilinear.tree(vertices, shape, num_chunk=2) + np.testing.assert_array_equal(shapes[0], shape - 1) + np.testing.assert_array_equal(shapes[~0], np.ones(len(grid))) + num_level = shapes.shape[0] - 1 + for level in range(1, num_level + 1): + size = 2**level + for k, block in enumerate(np.ndindex(*shapes[level])): + index = tuple( + slice(size * b, min(size * (b + 1), n - 1) + 1) + for b, n in zip(block, shape) + ) + expected_lower = [c[index].min() for c in grid] + expected_upper = [c[index].max() for c in grid] + np.testing.assert_array_equal(lower[start[level] + k], expected_lower) + np.testing.assert_array_equal(upper[start[level] + k], expected_upper) + + +@pytest.mark.parametrize("name", list(grids)) +class TestLocate: + + def test_inside(self, name: str) -> None: + grid = grids[name] + points = _points(grid) + cell, xi, inside = _locate(grid, points, extrapolate=False) + np.testing.assert_array_equal(inside, _inside_brute(grid, points)) + + def test_local_coordinates(self, name: str) -> None: + grid = grids[name] + points = _points(grid) + cell, xi, inside = _locate(grid, points, extrapolate=False) + assert np.all((-1e-9 <= xi[inside]) & (xi[inside] <= 1 + 1e-9)) + scale = float(np.max([np.ptp(c) for c in grid])) + np.testing.assert_allclose( + _position(grid, cell[inside], xi[inside]), + points[:, inside], + rtol=0, + atol=1e-13 * scale, + ) + assert np.all(np.isnan(xi[~inside])) + assert np.all(cell[~inside] == 0) + + def test_extrapolate(self, name: str) -> None: + # every point the cells along the edge of the grid continue to is + # reached, so the position of each point is recovered exactly + grid = grids[name] + points = _points(grid) + cell, xi, inside = _locate(grid, points, extrapolate=True) + scale = float(np.max([np.ptp(c) for c in grid])) + np.testing.assert_allclose( + _position(grid, cell, xi), + points, + rtol=0, + atol=1e-12 * scale, + ) + assert not np.all(inside) + + +def test_extrapolate_clamps_parallelograms() -> None: + # on a grid of parallelograms the cell to extrapolate from is the one + # nearest in index space + u, v = np.meshgrid(np.arange(5.0), np.arange(4.0), indexing="ij") + grid = (u + 0.5 * v, v) + points = np.array([[-3.0, 10.0, 2.0, 1.0], [1.5, 1.5, -2.0, 10.0]]) + cell, xi, inside = _locate(grid, points, extrapolate=True) + np.testing.assert_array_equal(cell, [[0, 1], [3, 1], [3, 0], [0, 2]]) + np.testing.assert_allclose( + xi, + [[-3.75, 0.5], [6.25, 0.5], [0.0, -2.0], [-4.0, 8.0]], + rtol=0, + atol=1e-14, + ) + assert not np.any(inside) + + +def test_not_finite() -> None: + grid = _distorted() + points = np.array([[0.1, np.nan, np.inf], [np.nan, 0.1, 0.2]]) + cell, xi, inside = _locate(grid, points, extrapolate=True) + assert not np.any(inside) + assert np.all(np.isnan(xi)) + assert np.all(cell == 0) diff --git a/regridding/_weights/_weights_multilinear/_rectilinear.py b/regridding/_weights/_weights_multilinear/_rectilinear.py new file mode 100644 index 0000000..601d944 --- /dev/null +++ b/regridding/_weights/_weights_multilinear/_rectilinear.py @@ -0,0 +1,144 @@ +""" +Find the cell of a rectilinear grid which contains each of a set of points. + +A grid is rectilinear if each of its coordinates varies along one of its +axes only, a different axis for each coordinate, and is strictly monotonic +along it. The cell along each axis is then a binary search of that +coordinate alone. The search is exact, needs no preparation, and a point +outside the grid is extrapolated from the edge of the grid along each axis +on its own. +""" + +import numpy as np +import numba + +__all__ = [ + "profiles", + "locate", +] + + +def profiles( + vertices: np.ndarray, + shape: np.ndarray, +) -> None | tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """ + The coordinate along each axis of a grid, if the grid is rectilinear. + + Returns :obj:`None` if the grid is not rectilinear. Otherwise, returns + ``(profile, start, coordinate, descending)``, where ``profile`` holds the + values of the coordinate along each axis one after another, each sorted + in ascending order, ``profile[start[m]:start[m + 1]]`` is the one along + axis ``m``, ``coordinate[m]`` is which coordinate that is, and + ``descending[m]`` is whether it was reversed to sort it. + + Parameters + ---------- + vertices + The coordinates of every vertex of the grid, of shape + ``(N, number of vertices)``, in row-major order. + shape + The number of vertices along each axis. + """ + num_axis = shape.size + grid = vertices.reshape(num_axis, *shape) + + coordinate = np.full(num_axis, -1, dtype=np.int64) + for k in range(num_axis): + varies = [ + m + for m in range(num_axis) + if np.any(grid[k] != grid[k].take([0], axis=m)) + ] + if len(varies) != 1: + return None + (m,) = varies + if coordinate[m] >= 0: + return None + coordinate[m] = k + + profile = [] + descending = np.zeros(num_axis, dtype=np.bool_) + for m in range(num_axis): + index: list[int | slice] = [0] * num_axis + index[m] = slice(None) + values = grid[coordinate[m]][tuple(index)] + difference = np.diff(values) + if np.all(difference < 0): + values = values[::-1] + descending[m] = True + elif not np.all(difference > 0): + return None + profile.append(values) + + start = np.zeros(num_axis + 1, dtype=np.int64) + start[1:] = np.cumsum(shape) + + return ( + np.ascontiguousarray(np.concatenate(profile), dtype=np.float64), + start, + coordinate, + descending, + ) + + +@numba.njit(cache=True, parallel=True) +def locate( + profile: np.ndarray, + start: np.ndarray, + coordinate: np.ndarray, + descending: np.ndarray, + points: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Find the cell of a rectilinear grid which contains each point. + + Returns ``(cell, xi, inside)``: the first corner of the cell along each + axis, the local coordinates of the point in the cell, and whether the + point is inside the grid. A point outside the grid is given the cell at + the edge of the grid along each axis it is outside along, with local + coordinates outside of the unit interval, which extrapolate from that + cell. A point with a coordinate which is NaN is outside, and its local + coordinates are NaN. + + Parameters + ---------- + profile + The coordinate along each axis, from :func:`profiles`. + start + Where the coordinate along each axis starts in `profile`. + coordinate + Which coordinate varies along each axis. + descending + Whether the coordinate along each axis was reversed. + points + The points to find, of shape ``(N, number of points)``. + """ + num_axis, num_point = points.shape + + cell = np.empty((num_point, num_axis), dtype=np.int64) + xi = np.empty((num_point, num_axis), dtype=np.float64) + inside = np.empty(num_point, dtype=np.bool_) + + for i in numba.prange(num_point): + inside_i = True + for m in range(num_axis): + values = profile[start[m] : start[m + 1]] + num = values.shape[0] + value = points[coordinate[m], i] + if not (values[0] <= value <= values[num - 1]): + inside_i = False + j = np.searchsorted(values, value) - 1 + if j < 0: + j = 0 + elif j > num - 2: + j = num - 2 + fraction = (value - values[j]) / (values[j + 1] - values[j]) + if descending[m]: + j = num - 2 - j + fraction = 1 - fraction + cell[i, m] = j + xi[i, m] = fraction + inside[i] = inside_i + + return cell, xi, inside diff --git a/regridding/_weights/_weights_multilinear/_rectilinear_test.py b/regridding/_weights/_weights_multilinear/_rectilinear_test.py new file mode 100644 index 0000000..28ba0c2 --- /dev/null +++ b/regridding/_weights/_weights_multilinear/_rectilinear_test.py @@ -0,0 +1,110 @@ +import pytest +import numpy as np +from . import _rectilinear + + +def _flat(grid: tuple[np.ndarray, ...]) -> tuple[np.ndarray, np.ndarray]: + """The vertices of a grid, flattened, and its shape.""" + vertices = np.stack([np.ravel(c) for c in np.broadcast_arrays(*grid)]) + return vertices.astype(float), np.array(np.broadcast(*grid).shape) + + +x = np.array([0.0, 0.5, 2.0, 2.25]) +y = np.array([1.0, 3.0, 4.0]) +z = np.array([-2.0, -1.0, 0.0, 1.0, 5.0]) + + +class TestProfiles: + + def test_meshgrid(self) -> None: + vertices, shape = _flat(np.meshgrid(x, y, z, indexing="ij")) + result = _rectilinear.profiles(vertices, shape) + assert result is not None + profile, start, coordinate, descending = result + np.testing.assert_array_equal(profile, np.concatenate([x, y, z])) + np.testing.assert_array_equal(start, [0, 4, 7, 12]) + np.testing.assert_array_equal(coordinate, [0, 1, 2]) + np.testing.assert_array_equal(descending, [False, False, False]) + + def test_coordinates_out_of_order(self) -> None: + # the first coordinate varies along the second axis, and the second + # along the first + yy, xx = np.meshgrid(y, x, indexing="ij") + vertices, shape = _flat((xx, yy)) + result = _rectilinear.profiles(vertices, shape) + assert result is not None + profile, start, coordinate, descending = result + np.testing.assert_array_equal(profile, np.concatenate([y, x])) + np.testing.assert_array_equal(coordinate, [1, 0]) + + def test_descending(self) -> None: + vertices, shape = _flat(np.meshgrid(x, y[::-1], indexing="ij")) + result = _rectilinear.profiles(vertices, shape) + assert result is not None + profile, start, coordinate, descending = result + np.testing.assert_array_equal(profile, np.concatenate([x, y])) + np.testing.assert_array_equal(descending, [False, True]) + + @pytest.mark.parametrize( + argnames="grid", + argvalues=[ + # rotated + ( + np.add.outer(x, 0.1 * y), + np.add.outer(-0.1 * x, y), + ), + # not monotonic + np.meshgrid(np.array([0.0, 1.0, 0.5]), y, indexing="ij"), + # both coordinates vary along the first axis only + ( + np.broadcast_to(x[:, np.newaxis], (4, 3)), + np.broadcast_to(2 * x[:, np.newaxis], (4, 3)), + ), + # a coordinate which does not vary at all + ( + np.broadcast_to(x[:, np.newaxis], (4, 3)), + np.zeros((4, 3)), + ), + # a coordinate which is NaN somewhere + np.meshgrid(np.array([0.0, np.nan, 1.0]), y, indexing="ij"), + ], + ) + def test_not_rectilinear(self, grid: tuple[np.ndarray, ...]) -> None: + vertices, shape = _flat(grid) + assert _rectilinear.profiles(vertices, shape) is None + + +class TestLocate: + + @pytest.mark.parametrize("reverse", [False, True]) + def test_locate(self, reverse: bool) -> None: + # the second axis may run backwards, which has to give the same + # cells counted from the other end + y_ = y[::-1] if reverse else y + vertices, shape = _flat(np.meshgrid(x, y_, indexing="ij")) + result = _rectilinear.profiles(vertices, shape) + assert result is not None + points = np.array( + [ + # inside, on a vertex, on the edges, outside, NaN + [0.25, 2.0, 0.0, 2.25, -1.0, 3.0, 1.0], + [2.0, 3.0, 1.0, 4.0, 2.0, 5.0, np.nan], + ] + ) + cell, xi, inside = _rectilinear.locate(*result, points) + + expected_cell_x = [0, 1, 0, 2, 0, 2, 1] + expected_xi_x = [0.5, 1.0, 0.0, 1.0, -2.0, 4.0, 1 / 3] + expected_cell_y = np.array([0, 0, 0, 1, 0, 1, 1]) + expected_xi_y = np.array([0.5, 1.0, 0.0, 1.0, 0.5, 2.0, np.nan]) + if reverse: + expected_cell_y = 1 - expected_cell_y + expected_xi_y = 1 - expected_xi_y + + np.testing.assert_array_equal(cell[:, 0], expected_cell_x) + np.testing.assert_array_equal(cell[:, 1], expected_cell_y) + np.testing.assert_allclose(xi[:, 0], expected_xi_x, rtol=1e-15) + np.testing.assert_allclose(xi[:, 1], expected_xi_y, rtol=1e-15) + np.testing.assert_array_equal( + inside, [True, True, True, True, False, False, False] + ) diff --git a/regridding/_weights/_weights_multilinear/_weights_multilinear.py b/regridding/_weights/_weights_multilinear/_weights_multilinear.py new file mode 100644 index 0000000..b9da846 --- /dev/null +++ b/regridding/_weights/_weights_multilinear/_weights_multilinear.py @@ -0,0 +1,237 @@ +""" +Weights which interpolate an array sampled on a logically-rectangular grid +of any number of dimensions at any set of points. +""" + +from typing import Callable, Literal, Sequence +import numpy as np +import numba +from regridding import _util +from . import _cells, _rectilinear, _curvilinear + +__all__ = [ + "weights_multilinear", +] + + +def _index( + index: tuple[int, ...], + axis: tuple[int, ...], +) -> tuple[int | slice, ...]: + """ + The index of one orthogonal element of an array, which takes the element + `index` along the orthogonal axes and the whole of the resampled axes. + + Parameters + ---------- + index + The index of the element along the orthogonal axes. + axis + The resampled axes, as negative numbers. + """ + result: list[int | slice] = list(reversed(index)) + for ax in sorted(axis, reverse=True): + result.insert(~ax, slice(None)) + return tuple(reversed(result)) + + +def _locator( + vertices: np.ndarray, + shape: np.ndarray, + num_chunk: int, +) -> Callable[[np.ndarray, bool], tuple[np.ndarray, np.ndarray, np.ndarray]]: + """ + Prepare to find the cells of a grid which contain a set of points. + + A rectilinear grid is searched along each axis on its own, and any other + grid through a hierarchy of bounding boxes, which is built here. The + two give the same cells, but the first is exact and needs nothing built. + + Returns a function of the points, of shape ``(N, number of points)``, + and of whether to extrapolate, which returns the cell of each point, its + local coordinates in that cell, and whether it is inside the grid. + + Parameters + ---------- + vertices + The coordinates of every vertex of the grid, of shape + ``(N, number of vertices)``, in row-major order. + shape + The number of vertices along each axis. + num_chunk + The number of pieces to divide the work into. + """ + rectilinear = _rectilinear.profiles(vertices, shape) + + if rectilinear is not None: + profile, start, coordinate, descending = rectilinear + + def locate_rectilinear( + points: np.ndarray, + extrapolate: bool, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + # the search along each axis finds the cell to extrapolate from + # either way + del extrapolate + return _rectilinear.locate(profile, start, coordinate, descending, points) + + return locate_rectilinear + + lower, upper, start, shapes = _curvilinear.tree(vertices, shape, num_chunk) + + def locate_curvilinear( + points: np.ndarray, + extrapolate: bool, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + return _curvilinear.locate( + vertices, + shape, + lower, + upper, + start, + shapes, + points, + extrapolate, + num_chunk, + ) + + return locate_curvilinear + + +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, + weights_input: None | np.ndarray = None, + bounds: Literal["extrapolate", "nan", "raise"] = "extrapolate", + perturb: None | bool = False, + seed: "None | int | np.random.Generator" = _util._seed_default, +) -> tuple[np.ndarray, tuple[int, ...], tuple[int, ...]]: + """ + Weights which interpolate an array sampled on the vertices of a + logically-rectangular grid at a set of points. + + The grid may have any number of dimensions, and the points any number of + axes, which do not have to match. Each point is interpolated from the + corners of the cell which contains it, so it has :math:`2^N` weights. + + Parameters + ---------- + coordinates_input + The coordinates of the vertices of the input grid. + coordinates_output + The coordinates of the points. + axis_input + The logical axes of the input grid. + axis_output + The logical axes of the points. + weights_input + Weights to multiply each vertex of the input grid by. + bounds + How to treat a point outside the input grid. + perturb + Whether to perturb the points by a small random amount. + seed + The seed of the random perturbation. + """ + + if bounds not in ("extrapolate", "nan", "raise"): + raise ValueError( + f"Unrecognized {bounds=}, " + f"expected one of ('extrapolate', 'nan', 'raise')." + ) + + if perturb is None: + perturb = False + + ( + 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, + perturb=perturb, + seed=seed, + match_axes=False, + ) + + shape_input_broadcasted = coordinates_input[0].shape + + shape = np.array( + [shape_input_broadcasted[ax] for ax in sorted(axis_input_)], + dtype=np.int64, + ) + if np.any(shape < 2): + raise ValueError( + f"The input grid must have at least two vertices along each of " + f"its resampled axes, got {tuple(shape)}" + ) + + if weights_input is not None: + weights_input = np.broadcast_to(weights_input, shape_input_broadcasted) + + # The input grid is usually broadcast across the orthogonal axes, so + # every element shares it and the work of preparing to search it is done + # once. A zero stride guarantees that. + axis_input_orthogonal = tuple( + ax for ax in range(-len(shape_input_broadcasted), 0) if ax not in axis_input_ + ) + shared = all( + coordinate.strides[ax] == 0 + for coordinate in coordinates_input + for ax in axis_input_orthogonal + ) + + num_chunk = numba.get_num_threads() + + result = np.empty(shape_orthogonal, dtype=object) + locate = None + + for index in np.ndindex(*shape_orthogonal): + + index_input = _index(index, axis_input_) + index_output = _index(index, axis_output_) + + if locate is None or not shared: + vertices = np.stack( + [np.ravel(c[index_input]) for c in coordinates_input], + ).astype(np.float64) + locate = _locator(vertices, shape, num_chunk) + + points = np.stack( + [np.ravel(c[index_output]) for c in coordinates_output], + ).astype(np.float64) + + cell, xi, inside = locate(points, bounds == "extrapolate") + + if bounds == "raise" and not np.all(inside): + raise ValueError( + f"{np.count_nonzero(~inside)} of the output points fall outside " + f"the input grid, and {bounds=}." + ) + + if weights_input is None: + weights_input_index = np.empty(0, dtype=np.float64) + else: + weights_input_index = np.ravel(weights_input[index_input]).astype( + np.float64 + ) + + result[index] = _cells.weights( + cell, + xi, + inside, + shape, + weights_input_index, + bounds == "nan", + ) + + return result, shape_input, shape_output diff --git a/regridding/_weights/_weights_multilinear_test.py b/regridding/_weights/_weights_multilinear/_weights_multilinear_test.py similarity index 100% rename from regridding/_weights/_weights_multilinear_test.py rename to regridding/_weights/_weights_multilinear/_weights_multilinear_test.py From d307d7e7be00cca3c5175cf53eaac391319b8dcd Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Wed, 21 Oct 2026 18:25:09 -0600 Subject: [PATCH 2/5] Test multilinear weights end to end, and document them Recognize a rectilinear grid from the strides of its coordinates before comparing their values, return NaN where the continuation of a curvilinear grid cannot reach a point, and add a benchmark. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01BYjDL98znSud1yFh9chnkP --- README.md | 7 +- benchmarks/weights_multilinear.py | 66 +++++ docs/index.rst | 7 +- regridding/_regrid/_regrid.py | 53 +++- regridding/_weights/_weights.py | 18 +- .../_weights/_weights_multilinear/_cells.py | 6 +- .../_weights_multilinear/_cells_test.py | 8 + .../_weights_multilinear/_curvilinear.py | 26 +- .../_weights_multilinear/_curvilinear_test.py | 36 +++ .../_weights_multilinear/_rectilinear.py | 31 +- .../_weights_multilinear/_rectilinear_test.py | 38 ++- .../_weights_multilinear.py | 27 +- .../_weights_multilinear_test.py | 270 ++++++++++++++++++ 13 files changed, 538 insertions(+), 55 deletions(-) create mode 100644 benchmarks/weights_multilinear.py diff --git a/README.md b/README.md index 43b6ebd..2c9282a 100644 --- a/README.md +++ b/README.md @@ -39,7 +39,9 @@ pip install regridding * [`regrid()`](https://regridding.readthedocs.io/en/latest/_autosummary/regridding.regrid.html), which resamples an array onto a new grid using either of two methods: - * `"multilinear"`, linear interpolation along one axis; + * `"multilinear"`, linear interpolation on logically-rectangular grids of + any number of dimensions, rectilinear or curvilinear, at points which may + be arranged along any number of axes; * `"conservative"`, first-order conservative resampling of 1D grids and of 2D logically-rectangular curvilinear grids, using the algorithm described in [Ramshaw (1985)](https://doi.org/10.1016/0021-9991(85)90141-X). @@ -86,6 +88,9 @@ The remaining axes are orthogonal to the operation, and the resampling is repeated independently for every position along them. This is how a stack of images, or a spectrum for each pixel, is resampled in one call. +The `multilinear` method interpolates at points rather than resampling cells, +so `axis_output` may select any number of axes, such as a single axis of +scattered points to interpolate a 2D grid at. **Degenerate grids are perturbed.** Where a vertex of the output grid lands exactly on an edge of the input grid, diff --git a/benchmarks/weights_multilinear.py b/benchmarks/weights_multilinear.py new file mode 100644 index 0000000..0ed8817 --- /dev/null +++ b/benchmarks/weights_multilinear.py @@ -0,0 +1,66 @@ +import numpy as np +import regridding + +coordinates_input: tuple[np.ndarray, ...] = () +coordinates_output: tuple[np.ndarray, ...] = () + + +def setup_weights_multilinear_2d( + num_point: int, + grid: str, +) -> None: + """ + A grid of about a million vertices, and random points over it. + + The ``rectilinear`` grid is searched along each axis on its own, and the + ``curvilinear`` one, which is the same grid rotated and distorted, + through its hierarchy of bounding boxes. + """ + global coordinates_input + global coordinates_output + + x, y = np.meshgrid( + np.linspace(-1, 1, num=1001), + np.linspace(-1, 1, num=1000), + indexing="ij", + ) + + if grid == "curvilinear": + angle = 0.4 + x, y = ( + x * np.cos(angle) - y * np.sin(angle) + 0.05 * x * x, + x * np.sin(angle) + y * np.cos(angle) + 0.05 * y * y, + ) + + rng = np.random.default_rng(0) + coordinates_input = (x, y) + coordinates_output = ( + rng.uniform(x.min(), x.max(), size=num_point), + rng.uniform(y.min(), y.max(), size=num_point), + ) + + time_weights_multilinear_2d(num_point, grid) + + +def time_weights_multilinear_2d( + num_point: int, + grid: str, +) -> None: + regridding.weights( + coordinates_input=coordinates_input, + coordinates_output=coordinates_output, + method="multilinear", + coalesce=False, + ) + + +time_weights_multilinear_2d.setup = setup_weights_multilinear_2d +time_weights_multilinear_2d.params = ( + [1, 10_000, 1_000_000], + ["rectilinear", "curvilinear"], +) +time_weights_multilinear_2d.param_names = ( + "number of points", + "grid", +) +time_weights_multilinear_2d.timeout = 240 diff --git a/docs/index.rst b/docs/index.rst index 5532255..ebf08bb 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -36,7 +36,9 @@ Features * :func:`regridding.regrid`, which resamples an array onto a new grid using either of two methods: - * ``"multilinear"``, linear interpolation along one axis; + * ``"multilinear"``, linear interpolation on logically-rectangular grids + of any number of dimensions, rectilinear or curvilinear, at points + which may be arranged along any number of axes; * ``"conservative"``, first-order conservative resampling of 1D grids and of 2D logically-rectangular curvilinear grids, by one of two algorithms described below. @@ -98,6 +100,9 @@ The remaining axes are orthogonal to the operation, and the resampling is repeated independently for every position along them. This is how a stack of images, or a spectrum for each pixel, is resampled in one call. +The ``multilinear`` method interpolates at points rather than resampling cells, +so ``axis_output`` may select any number of axes, such as a single axis of +scattered points to interpolate a 2D grid at. **The conservative method has two algorithms, and picks between them.** Both compute the same weights: the fraction of each input cell shared with diff --git a/regridding/_regrid/_regrid.py b/regridding/_regrid/_regrid.py index 3c66b3c..cb51359 100644 --- a/regridding/_regrid/_regrid.py +++ b/regridding/_regrid/_regrid.py @@ -47,12 +47,17 @@ def regrid( 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. + For the ``conservative`` method, the number of axes should be equal + to the number of coordinates in the output grid. + The ``multilinear`` method interpolates at points rather than + resampling cells, so the points may be arranged along any number of + axes, or none at all for a single point. method The type of regridding to use. The ``multilinear`` method interprets `coordinates_input` as the - points where `values_input` is sampled. + points where `values_input` is sampled, and interpolates between the + corners of the cell which contains each output point, on a grid of + any number of dimensions, rectilinear or curvilinear. The ``conservative`` method interprets `coordinates_input` as the edges of the cells containing `values_input`, so `values_input` has one fewer element along each resampled axis, and the sum of the result matches the @@ -69,6 +74,12 @@ def regrid( ``extrapolate`` (the default) extrapolates linearly from the nearest cell of the input grid, ``nan`` sets those points to :obj:`numpy.nan`, and ``raise`` raises a :class:`ValueError`. + On a grid which is not rectilinear, the nearest cell is the one along + the edge of the grid whose continuation reaches the point, and a + point which no continuation reaches, as can happen beyond a strongly + curved edge, is set to :obj:`numpy.nan`. + A point with a coordinate which is not finite is set to + :obj:`numpy.nan`, or raises, whichever `bounds` is. perturb Whether to perturb `coordinates_output` by a small value to avoid degenerate grids. This is helpful for some methods, like ``conservative``, which @@ -139,6 +150,42 @@ def regrid( ax.stairs(values_output, x_output, label="output"); ax.legend(); + Interpolate an array sampled on a distorted 2D grid at scattered points. + + .. jupyter-execute:: + + # Define the vertices of the input grid + u, v = np.meshgrid( + np.linspace(-1, 1, num=21), + np.linspace(-1, 1, num=21), + indexing="ij", + ) + x_input = u + 0.2 * v * v + y_input = v - 0.2 * u * u + + # Sample a function on the vertices + values_input = np.cos(3 * x_input) * np.sin(2 * y_input) + + # Define a set of random points, some of them outside the grid + rng = np.random.default_rng(0) + x_output = rng.uniform(-1.3, 1.3, size=500) + y_output = rng.uniform(-1.3, 1.3, size=500) + + # Interpolate the function at the points + values_output = regridding.regrid( + coordinates_input=(x_input, y_input), + coordinates_output=(x_output, y_output), + values_input=values_input, + method="multilinear", + bounds="nan", + ) + + # Plot the input grid and the interpolated points + fig, ax = plt.subplots(constrained_layout=True); + ax.pcolormesh(x_input, y_input, values_input, shading="gouraud", vmin=-1, vmax=1); + ax.scatter(x_output, y_output, c=values_output, vmin=-1, vmax=1, edgecolors="black"); + ax.set_aspect("equal"); + References ---------- .. footbibliography:: diff --git a/regridding/_weights/_weights.py b/regridding/_weights/_weights.py index 81a20bd..462848a 100644 --- a/regridding/_weights/_weights.py +++ b/regridding/_weights/_weights.py @@ -141,12 +141,20 @@ def weights( 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. + For the ``conservative`` method, the number of axes should be equal + to the number of coordinates in the output grid. + The ``multilinear`` method interpolates at points rather than + resampling cells, so the points may be arranged along any number of + axes, or none at all for a single point. weights_input Weights applied to the values of the input grid before resampling. method The type of regridding to use. + The ``multilinear`` method interpolates between the corners of the + cell of the input grid which contains each output point, on a grid of + any number of dimensions, rectilinear or curvilinear. + The ``conservative`` method resamples the cells of the input grid + onto the cells of the output grid, in one or two dimensions. bounds How to treat output points that fall outside the input grid. Only applies when `method` is ``multilinear``; the ``conservative`` @@ -154,6 +162,12 @@ def weights( ``extrapolate`` (the default) extrapolates linearly from the nearest cell of the input grid, ``nan`` sets those points to :obj:`numpy.nan`, and ``raise`` raises a :class:`ValueError`. + On a grid which is not rectilinear, the nearest cell is the one along + the edge of the grid whose continuation reaches the point, and a + point which no continuation reaches, as can happen beyond a strongly + curved edge, is set to :obj:`numpy.nan`. + A point with a coordinate which is not finite is set to + :obj:`numpy.nan`, or raises, whichever `bounds` is. perturb Whether to perturb `coordinates_output` by a small value to avoid degenerate grids. This is helpful for some methods, like ``conservative``, which diff --git a/regridding/_weights/_weights_multilinear/_cells.py b/regridding/_weights/_weights_multilinear/_cells.py index 3a62c12..410f03c 100644 --- a/regridding/_weights/_weights_multilinear/_cells.py +++ b/regridding/_weights/_weights_multilinear/_cells.py @@ -178,9 +178,7 @@ def evaluate( else: derivative *= 1 - xi[j] for k in range(num_axis): - jacobian[k, m] += derivative * ( - vertices[k, index] - vertices[k, base] - ) + jacobian[k, m] += derivative * (vertices[k, index] - vertices[k, base]) @numba.njit(cache=True) @@ -419,7 +417,7 @@ def inverse( evaluate(vertices, base, offset, xi, x, jacobian) error = distance_squared(x, goal) - return False + return False # pragma: nocover @numba.njit(cache=True, parallel=True) diff --git a/regridding/_weights/_weights_multilinear/_cells_test.py b/regridding/_weights/_weights_multilinear/_cells_test.py index cda016c..832e27b 100644 --- a/regridding/_weights/_weights_multilinear/_cells_test.py +++ b/regridding/_weights/_weights_multilinear/_cells_test.py @@ -136,6 +136,14 @@ def test_round_trip(self, num_axis: int, scale: float, shift: float) -> None: atol = 1e-12 + 1e-15 * shift / scale np.testing.assert_allclose(result, expected, rtol=0, atol=atol) + @pytest.mark.filterwarnings("ignore:overflow encountered:RuntimeWarning") + def test_step_overflows(self, num_axis: int) -> None: + # a cell far smaller than its distance to the point + vertices, shape, offset = _cell(num_axis, scale=1e-300) + point = np.full(num_axis, 1e10) + xi = np.full(num_axis, 0.5) + assert not _cells.inverse(vertices, 0, offset, point, xi, _work(num_axis)) + def test_unreachable(self, num_axis: int) -> None: # a cell squashed flat along its first axis cannot reach a point off # the plane it lies in diff --git a/regridding/_weights/_weights_multilinear/_curvilinear.py b/regridding/_weights/_weights_multilinear/_curvilinear.py index a0aee83..aa6c3da 100644 --- a/regridding/_weights/_weights_multilinear/_curvilinear.py +++ b/regridding/_weights/_weights_multilinear/_curvilinear.py @@ -756,7 +756,7 @@ def _walk( xi: np.ndarray, xi_trial: np.ndarray, work: np.ndarray, -) -> None: +) -> bool: """ Extrapolate the grid to a point outside of it. @@ -766,6 +766,10 @@ def _walk( the point. The first corner of the cell it stops in is left in `cell`, and the local coordinates of the point in that cell in `xi`. + Returns whether it reached the point. It does not if the continuation + of the grid folds over before reaching it, as it can beyond a strongly + curved edge, in which case it stops where the continuation came closest. + Parameters ---------- vertices @@ -810,9 +814,12 @@ def _walk( _cells.evaluate(vertices, base, offset, xi, x, jacobian) error = _cells.distance_squared(x, goal) + reached = False + for _ in range(num_iteration_walk): if not error > 0: + reached = error == 0 break matrix[:, :] = jacobian @@ -844,11 +851,15 @@ def _walk( fraction = fraction / 2 if not accepted: + # no point along the step is closer, which happens once the step + # is down to the rounding error of the position + reached = size <= _cells.tolerance * scale break s[:] = s_trial if fraction * size <= _cells.tolerance * scale: + reached = True break base = _clamp(s, shape, stride, cell, xi) @@ -858,6 +869,8 @@ def _walk( _clamp(s, shape, stride, cell, xi) + return reached + @numba.njit(cache=True, parallel=True) def locate( @@ -877,8 +890,11 @@ def locate( Returns ``(cell, xi, inside)``: the first corner of the cell along each axis, the local coordinates of the point in the cell, and whether the point is inside the grid. A point outside the grid, or with a - coordinate which is NaN, is given the first cell and local coordinates - of NaN, unless `extrapolate` is set and its coordinates are finite. + coordinate which is not finite, is given the first cell and local + coordinates of NaN, unless `extrapolate` is set and its coordinates are + finite. An extrapolated point which the continuation of the grid does + not reach keeps the cell it was extrapolated from, with local + coordinates of NaN. Parameters ---------- @@ -978,7 +994,7 @@ def locate( ) for m in range(num_axis): s[m] = cell_i[m] + 0.5 - _walk( + reached = _walk( vertices, shape, stride, @@ -992,6 +1008,8 @@ def locate( xi_trial, work, ) + if not reached: + xi_i[:] = np.nan else: cell_i[:] = 0 xi_i[:] = np.nan diff --git a/regridding/_weights/_weights_multilinear/_curvilinear_test.py b/regridding/_weights/_weights_multilinear/_curvilinear_test.py index 1a59a8c..7e20ce1 100644 --- a/regridding/_weights/_weights_multilinear/_curvilinear_test.py +++ b/regridding/_weights/_weights_multilinear/_curvilinear_test.py @@ -252,6 +252,42 @@ def test_extrapolate_clamps_parallelograms() -> None: assert not np.any(inside) +def test_extrapolate_unreachable() -> None: + # the first cell has no length, so nothing to the left of the grid can + # be extrapolated to, while the right of the grid can + grid = (np.array([0.0, 0.0, 1.0, 2.0]),) + points = np.array([[-1.0, 3.0]]) + cell, xi, inside = _locate(grid, points, extrapolate=True) + assert not np.any(inside) + assert np.isnan(xi[0, 0]) + np.testing.assert_array_equal(cell[1], [2]) + np.testing.assert_allclose(xi[1], [2.0], rtol=1e-14) + + +@pytest.mark.filterwarnings("ignore:overflow encountered:RuntimeWarning") +def test_extrapolate_step_overflows() -> None: + # cells far smaller than the distance to the point + grid = (np.array([0.0, 1e-300, 3e-300]),) + cell, xi, inside = _locate(grid, np.array([[1e10]]), extrapolate=True) + assert not inside[0] + assert np.isnan(xi[0, 0]) + + +def test_degenerate_cell() -> None: + # the first row of cells has no height, so Newton's method cannot find + # a point on it in them, but the point is on the edge of the cell above + x, y = np.meshgrid( + np.array([0.0, 0.5, 1.0]), + np.array([1.0, 1.0, 2.0]), + indexing="ij", + ) + points = np.array([[0.1], [1.0]]) + cell, xi, inside = _locate((x, y), points, extrapolate=False) + assert inside[0] + np.testing.assert_array_equal(cell[0], [0, 1]) + np.testing.assert_allclose(xi[0], [0.2, 0.0], rtol=0, atol=1e-15) + + def test_not_finite() -> None: grid = _distorted() points = np.array([[0.1, np.nan, np.inf], [np.nan, 0.1, 0.2]]) diff --git a/regridding/_weights/_weights_multilinear/_rectilinear.py b/regridding/_weights/_weights_multilinear/_rectilinear.py index 601d944..3087ae0 100644 --- a/regridding/_weights/_weights_multilinear/_rectilinear.py +++ b/regridding/_weights/_weights_multilinear/_rectilinear.py @@ -11,6 +11,7 @@ import numpy as np import numba +from regridding import _util __all__ = [ "profiles", @@ -19,8 +20,7 @@ def profiles( - vertices: np.ndarray, - shape: np.ndarray, + coordinates: tuple[np.ndarray, ...], ) -> None | tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """ The coordinate along each axis of a grid, if the grid is rectilinear. @@ -34,21 +34,25 @@ def profiles( Parameters ---------- - vertices - The coordinates of every vertex of the grid, of shape - ``(N, number of vertices)``, in row-major order. - shape - The number of vertices along each axis. + coordinates + The coordinates of the vertices of the grid, each an array with an + axis for each axis of the grid. An array which has been broadcast + along an axis, and so has a stride of zero along it, is known not to + vary along it without its values being looked at, which makes the + grids that :func:`numpy.broadcast_to` or ``numpy.meshgrid(..., + sparse=True)`` give quick to recognize. """ - num_axis = shape.size - grid = vertices.reshape(num_axis, *shape) + num_axis = len(coordinates) + shape = np.array(coordinates[0].shape, dtype=np.int64) + grid = [] coordinate = np.full(num_axis, -1, dtype=np.int64) for k in range(num_axis): + (c,) = _util._unbroadcast((coordinates[k],)) varies = [ m for m in range(num_axis) - if np.any(grid[k] != grid[k].take([0], axis=m)) + if c.shape[m] > 1 and np.any(c != c.take([0], axis=m)) ] if len(varies) != 1: return None @@ -56,6 +60,7 @@ def profiles( if coordinate[m] >= 0: return None coordinate[m] = k + grid.append(c) profile = [] descending = np.zeros(num_axis, dtype=np.bool_) @@ -98,8 +103,8 @@ def locate( point is inside the grid. A point outside the grid is given the cell at the edge of the grid along each axis it is outside along, with local coordinates outside of the unit interval, which extrapolate from that - cell. A point with a coordinate which is NaN is outside, and its local - coordinates are NaN. + cell. A point with a coordinate which is not finite is outside, and its + local coordinate along that axis is NaN. Parameters ---------- @@ -134,6 +139,8 @@ def locate( elif j > num - 2: j = num - 2 fraction = (value - values[j]) / (values[j + 1] - values[j]) + if not np.isfinite(value): + fraction = np.nan if descending[m]: j = num - 2 - j fraction = 1 - fraction diff --git a/regridding/_weights/_weights_multilinear/_rectilinear_test.py b/regridding/_weights/_weights_multilinear/_rectilinear_test.py index 28ba0c2..7e45c2e 100644 --- a/regridding/_weights/_weights_multilinear/_rectilinear_test.py +++ b/regridding/_weights/_weights_multilinear/_rectilinear_test.py @@ -3,10 +3,9 @@ from . import _rectilinear -def _flat(grid: tuple[np.ndarray, ...]) -> tuple[np.ndarray, np.ndarray]: - """The vertices of a grid, flattened, and its shape.""" - vertices = np.stack([np.ravel(c) for c in np.broadcast_arrays(*grid)]) - return vertices.astype(float), np.array(np.broadcast(*grid).shape) +def _grid(grid: tuple[np.ndarray, ...]) -> tuple[np.ndarray, ...]: + """The coordinates of a grid, broadcast against each other.""" + return tuple(np.broadcast_arrays(*grid)) x = np.array([0.0, 0.5, 2.0, 2.25]) @@ -17,8 +16,8 @@ def _flat(grid: tuple[np.ndarray, ...]) -> tuple[np.ndarray, np.ndarray]: class TestProfiles: def test_meshgrid(self) -> None: - vertices, shape = _flat(np.meshgrid(x, y, z, indexing="ij")) - result = _rectilinear.profiles(vertices, shape) + coordinates = _grid(np.meshgrid(x, y, z, indexing="ij")) + result = _rectilinear.profiles(coordinates) assert result is not None profile, start, coordinate, descending = result np.testing.assert_array_equal(profile, np.concatenate([x, y, z])) @@ -26,20 +25,31 @@ def test_meshgrid(self) -> None: np.testing.assert_array_equal(coordinate, [0, 1, 2]) np.testing.assert_array_equal(descending, [False, False, False]) + def test_broadcast(self) -> None: + # a grid broadcast from one dimension along each axis is recognized + # from its strides, without its values being compared + coordinates = _grid(np.meshgrid(x, y, indexing="ij", sparse=True)) + assert all(0 in c.strides for c in coordinates) + result = _rectilinear.profiles(coordinates) + assert result is not None + profile, start, coordinate, descending = result + np.testing.assert_array_equal(profile, np.concatenate([x, y])) + np.testing.assert_array_equal(coordinate, [0, 1]) + def test_coordinates_out_of_order(self) -> None: # the first coordinate varies along the second axis, and the second # along the first yy, xx = np.meshgrid(y, x, indexing="ij") - vertices, shape = _flat((xx, yy)) - result = _rectilinear.profiles(vertices, shape) + coordinates = _grid((xx, yy)) + result = _rectilinear.profiles(coordinates) assert result is not None profile, start, coordinate, descending = result np.testing.assert_array_equal(profile, np.concatenate([y, x])) np.testing.assert_array_equal(coordinate, [1, 0]) def test_descending(self) -> None: - vertices, shape = _flat(np.meshgrid(x, y[::-1], indexing="ij")) - result = _rectilinear.profiles(vertices, shape) + coordinates = _grid(np.meshgrid(x, y[::-1], indexing="ij")) + result = _rectilinear.profiles(coordinates) assert result is not None profile, start, coordinate, descending = result np.testing.assert_array_equal(profile, np.concatenate([x, y])) @@ -70,8 +80,8 @@ def test_descending(self) -> None: ], ) def test_not_rectilinear(self, grid: tuple[np.ndarray, ...]) -> None: - vertices, shape = _flat(grid) - assert _rectilinear.profiles(vertices, shape) is None + coordinates = _grid(grid) + assert _rectilinear.profiles(coordinates) is None class TestLocate: @@ -81,8 +91,8 @@ def test_locate(self, reverse: bool) -> None: # the second axis may run backwards, which has to give the same # cells counted from the other end y_ = y[::-1] if reverse else y - vertices, shape = _flat(np.meshgrid(x, y_, indexing="ij")) - result = _rectilinear.profiles(vertices, shape) + coordinates = _grid(np.meshgrid(x, y_, indexing="ij")) + result = _rectilinear.profiles(coordinates) assert result is not None points = np.array( [ diff --git a/regridding/_weights/_weights_multilinear/_weights_multilinear.py b/regridding/_weights/_weights_multilinear/_weights_multilinear.py index b9da846..48b54df 100644 --- a/regridding/_weights/_weights_multilinear/_weights_multilinear.py +++ b/regridding/_weights/_weights_multilinear/_weights_multilinear.py @@ -36,8 +36,7 @@ def _index( def _locator( - vertices: np.ndarray, - shape: np.ndarray, + coordinates: tuple[np.ndarray, ...], num_chunk: int, ) -> Callable[[np.ndarray, bool], tuple[np.ndarray, np.ndarray, np.ndarray]]: """ @@ -53,15 +52,13 @@ def _locator( Parameters ---------- - vertices - The coordinates of every vertex of the grid, of shape - ``(N, number of vertices)``, in row-major order. - shape - The number of vertices along each axis. + coordinates + The coordinates of the vertices of the grid, each an array with an + axis for each axis of the grid. num_chunk The number of pieces to divide the work into. """ - rectilinear = _rectilinear.profiles(vertices, shape) + rectilinear = _rectilinear.profiles(coordinates) if rectilinear is not None: profile, start, coordinate, descending = rectilinear @@ -77,6 +74,8 @@ def locate_rectilinear( return locate_rectilinear + shape = np.array(coordinates[0].shape, dtype=np.int64) + vertices = np.stack([np.ravel(c) for c in coordinates]).astype(np.float64) lower, upper, start, shapes = _curvilinear.tree(vertices, shape, num_chunk) def locate_curvilinear( @@ -180,12 +179,12 @@ def weights_multilinear( # The input grid is usually broadcast across the orthogonal axes, so # every element shares it and the work of preparing to search it is done - # once. A zero stride guarantees that. + # once. A zero stride, or a single element, guarantees that. axis_input_orthogonal = tuple( ax for ax in range(-len(shape_input_broadcasted), 0) if ax not in axis_input_ ) shared = all( - coordinate.strides[ax] == 0 + coordinate.strides[ax] == 0 or coordinate.shape[ax] == 1 for coordinate in coordinates_input for ax in axis_input_orthogonal ) @@ -201,10 +200,10 @@ def weights_multilinear( index_output = _index(index, axis_output_) if locate is None or not shared: - vertices = np.stack( - [np.ravel(c[index_input]) for c in coordinates_input], - ).astype(np.float64) - locate = _locator(vertices, shape, num_chunk) + locate = _locator( + tuple(c[index_input] for c in coordinates_input), + num_chunk, + ) points = np.stack( [np.ravel(c[index_output]) for c in coordinates_output], diff --git a/regridding/_weights/_weights_multilinear/_weights_multilinear_test.py b/regridding/_weights/_weights_multilinear/_weights_multilinear_test.py index 8be2543..3809ccf 100644 --- a/regridding/_weights/_weights_multilinear/_weights_multilinear_test.py +++ b/regridding/_weights/_weights_multilinear/_weights_multilinear_test.py @@ -1,5 +1,7 @@ +from typing import Any import pytest import numpy as np +import scipy.interpolate import astropy.units as u import regridding @@ -147,3 +149,271 @@ def test_weights_multilinear_1d( assert isinstance(result, np.ndarray) assert np.issubdtype(result.dtype, float) assert np.allclose(result, result_expected) + + +def _regrid( + grid: tuple[np.ndarray, ...], + values: np.ndarray, + points: tuple[np.ndarray, ...], + **kwargs: Any, +) -> np.ndarray: + """Interpolate an array sampled on a grid at a set of points.""" + return regridding.regrid( + coordinates_input=grid, + coordinates_output=points, + values_input=values, + method="multilinear", + **kwargs, + ) + + +def _affine(coordinates: tuple[np.ndarray, ...]) -> np.ndarray: + """ + An affine function of the coordinates, which multilinear interpolation + reproduces exactly on any grid, since the weights of the corners of a + cell sum to one and average the corners to the point. + """ + result = np.full(np.broadcast(*coordinates).shape, 0.3) + for k, c in enumerate(coordinates): + result = result + (0.7 - 0.4 * k) * c + return result + + +def _annulus() -> tuple[np.ndarray, np.ndarray]: + """An annulus which wraps most of the way around.""" + r, t = np.meshgrid( + np.linspace(1, 2, 5), + np.linspace(0, 0.9 * 2 * np.pi, 31), + indexing="ij", + ) + return r * np.cos(t), r * np.sin(t) + + +def _distorted() -> tuple[np.ndarray, np.ndarray]: + """A grid with the kind of smooth distortion an optical system has.""" + u, v = np.meshgrid(np.linspace(-1, 1, 7), np.linspace(-1, 1, 6), indexing="ij") + return u + 0.05 * u * v + 0.03 * v**2, v + 0.04 * u**2 - 0.02 * u * v + + +def _distorted_3d() -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """A three-dimensional grid with a smooth distortion.""" + p, q, w = np.meshgrid( + np.linspace(-1, 1, 4), + np.linspace(-1, 1, 5), + np.linspace(-1, 1, 3), + indexing="ij", + ) + return p + 0.05 * q * w, q + 0.04 * p**2, w + 0.03 * p * q + + +grids_curvilinear = { + "annulus": _annulus(), + "distorted": _distorted(), + "distorted 3D": _distorted_3d(), +} + + +def _points( + grid: tuple[np.ndarray, ...], + num: int = 100, + seed: int = 0, +) -> tuple[np.ndarray, ...]: + """Random points over a region a little larger than the grid.""" + rng = np.random.default_rng(seed) + lower = np.array([c.min() for c in grid]) + upper = np.array([c.max() for c in grid]) + margin = 0.1 * (upper - lower) + return tuple(rng.uniform(lower - margin, upper + margin, (num, len(grid))).T) + + +class TestRectilinear: + + @pytest.mark.parametrize("num_axis", [1, 2, 3]) + @pytest.mark.parametrize("bounds", ["extrapolate", "nan"]) + def test_matches_scipy(self, num_axis: int, bounds: str) -> None: + rng = np.random.default_rng(num_axis) + axes = [np.sort(rng.uniform(-1, 1, 3 + m)) for m in range(num_axis)] + # one of the axes runs backwards + axes[0] = axes[0][::-1] + grid = tuple(np.meshgrid(*axes, indexing="ij")) + values = rng.normal(size=grid[0].shape) + points = rng.uniform(-1.3, 1.3, (num_axis, 30)) + result = _regrid(grid, values, tuple(points), bounds=bounds) + expected = scipy.interpolate.RegularGridInterpolator( + points=axes, + values=values, + bounds_error=False, + fill_value=None if bounds == "extrapolate" else np.nan, + )(points.T) + np.testing.assert_allclose(result, expected, rtol=0, atol=1e-13) + assert np.any(np.isnan(result)) == (bounds == "nan") + + def test_descending_1d(self) -> None: + # a grid which runs backwards used to be searched as if it ran forwards + x = np.array([3.0, 2.0, 0.5, 0.0]) + result = _regrid((x,), x**2, (np.array([0.25, 1.0, 2.5, 4.0]),)) + np.testing.assert_allclose(result, [0.125, 1.5, 6.5, 14.0]) + + +@pytest.mark.parametrize("name", list(grids_curvilinear)) +class TestCurvilinear: + + def test_affine(self, name: str) -> None: + grid = grids_curvilinear[name] + points = _points(grid) + result = _regrid(grid, _affine(grid), points) + np.testing.assert_allclose(result, _affine(points), rtol=0, atol=1e-12) + + def test_nan(self, name: str) -> None: + grid = grids_curvilinear[name] + points = _points(grid) + result = _regrid(grid, _affine(grid), points, bounds="nan") + inside = np.isfinite(result) + assert 0 < inside.sum() < inside.size + np.testing.assert_allclose( + result[inside], + _affine(points)[inside], + rtol=0, + atol=1e-12, + ) + + def test_raise(self, name: str) -> None: + grid = grids_curvilinear[name] + with pytest.raises(ValueError, match="fall outside the input grid"): + _regrid(grid, _affine(grid), _points(grid), bounds="raise") + + def test_raise_inside_grid(self, name: str) -> None: + grid = grids_curvilinear[name] + points = tuple(c[1:-1, 1:-1] for c in grid) + result = _regrid(grid, _affine(grid), points, bounds="raise") + np.testing.assert_allclose(result, _affine(points), rtol=0, atol=1e-12) + + +class TestPoints: + """ + Multilinear interpolation weights points, rather than cells, so the + points may be arranged in any number of axes. + """ + + def test_grid_2d_points_1d(self) -> None: + grid = _distorted() + points = _points(grid) + result = _regrid(grid, _affine(grid), points, axis_output=(0,)) + assert result.shape == points[0].shape + np.testing.assert_allclose(result, _affine(points), rtol=0, atol=1e-12) + + def test_grid_3d_points_2d(self) -> None: + grid = _distorted_3d() + points = tuple(c.reshape(10, 10) for c in _points(grid)) + result = _regrid(grid, _affine(grid), points) + assert result.shape == (10, 10) + np.testing.assert_allclose(result, _affine(points), rtol=0, atol=1e-12) + + def test_single_point(self) -> None: + grid = _distorted() + point = (np.array(0.1), np.array(-0.2)) + result = _regrid(grid, _affine(grid), point) + assert result.shape == () + np.testing.assert_allclose(result, _affine(point), rtol=0, atol=1e-12) + + def test_conservative_axes_must_match(self) -> None: + grid = _distorted() + with pytest.raises(ValueError, match="number of axes in `axis_output`"): + regridding.weights( + coordinates_input=grid, + coordinates_output=_points(grid), + axis_output=(0,), + method="conservative", + ) + + +class TestOrthogonal: + + def test_grid_per_element(self) -> None: + # two grids, rotated by different angles, each with its own points + angle = np.array([0.1, 0.7])[:, np.newaxis, np.newaxis] + u, v = _distorted() + grid = ( + u * np.cos(angle) - v * np.sin(angle), + u * np.sin(angle) + v * np.cos(angle), + ) + values = np.stack([_affine(grid)[0], 2 * _affine(grid)[1]]) + points = tuple(np.stack([c, c[::-1]]) for c in _points(_distorted())) + result = _regrid(grid, values, points, axis_input=(1, 2), axis_output=(1,)) + expected = np.stack( + [ + _regrid( + tuple(c[i] for c in grid), + values[i], + tuple(c[i] for c in points), + ) + for i in range(2) + ] + ) + np.testing.assert_allclose(result, expected, rtol=0, atol=1e-14) + np.testing.assert_allclose(result[0], _affine(points)[0], atol=1e-12) + + def test_grid_shared(self) -> None: + # one grid, broadcast over a batch of points, which gives the same + # result as taking every point at once + grid = tuple(c[np.newaxis] for c in _annulus()) + points = tuple(c.reshape(4, 25) for c in _points(_annulus())) + values = _affine(_annulus()) + result = _regrid(grid, values, points, axis_input=(1, 2), axis_output=(1,)) + expected = _regrid(_annulus(), values, points) + assert result.shape == expected.shape == (4, 25) + np.testing.assert_array_equal(result, expected) + + +class TestWeightsInput: + + @pytest.mark.parametrize("name", list(grids_curvilinear)) + def test_weights_input(self, name: str) -> None: + grid = grids_curvilinear[name] + points = _points(grid) + values = _affine(grid) + weights_input = np.random.default_rng(1).uniform(0.5, 1.5, values.shape) + weights = regridding.weights( + coordinates_input=grid, + coordinates_output=points, + weights_input=weights_input, + method="multilinear", + ) + result = regridding.regrid_from_weights(*weights, values_input=values) + expected = _regrid(grid, values * weights_input, points) + np.testing.assert_allclose(result, expected, rtol=1e-14, atol=1e-14) + + +def test_units() -> None: + grid = tuple(c * u.mm for c in _distorted()) + points = tuple(c * u.m for c in _points(_distorted())) + result = _regrid(grid, _affine(_distorted()), points, bounds="nan") + expected = _regrid( + _distorted(), + _affine(_distorted()), + tuple(1000 * c.value for c in points), + bounds="nan", + ) + np.testing.assert_allclose(result, expected, rtol=0, atol=1e-12) + + +@pytest.mark.parametrize( + argnames="grid", + argvalues=[ + _distorted(), + tuple(np.meshgrid(np.linspace(-1, 1, 4), np.linspace(-1, 1, 3), indexing="ij")), + ], +) +def test_points_not_finite(grid: tuple[np.ndarray, ...]) -> None: + points = (np.array([0.1, np.nan, np.inf]), np.array([np.nan, 0.2, -0.3])) + result = _regrid(grid, _affine(grid), points) + assert np.all(np.isnan(result)) + with pytest.raises(ValueError, match="fall outside the input grid"): + _regrid(grid, _affine(grid), points, bounds="raise") + + +def test_too_few_vertices() -> None: + x = np.linspace(0, 1, 4) + grid = np.meshgrid(x, x[:1], indexing="ij") + with pytest.raises(ValueError, match="at least two vertices"): + _regrid(tuple(grid), np.zeros((4, 1)), (np.array([0.5]), np.array([0.0]))) From c66b0d3649c49faae86338924bd860d6cde703b7 Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Wed, 21 Oct 2026 20:38:32 -0600 Subject: [PATCH 3/5] Address the review of the multilinear weights - Newton's method has converged once the distance from the position to the point is down to its rounding error, as well as once a step is short. In a thin cell at an angle to the axes, rounding along its length kept the step across its width above the tolerance, and points well inside the cell were reported outside it. One routine now serves both the search within a cell and the extrapolating walk across the grid, which had copied it. - Without `axis_output`, the points must have as many axes as the grid again, so that a forgotten `axis_output` raises rather than taking the orthogonal axes for points. Points arranged in other numbers of axes pass `axis_output` explicitly. - Every orthogonal element is searched in one call: rectilinear grids together, even if each element has its own, and the points of every element at once on a curvilinear grid they share. A spectrum for each of many pixels had become several times slower than on `main`. - A rectilinear grid is compared in floating point, since the difference of two unsigned integers wraps around instead of going negative, which took a descending grid for an ascending one. - `regrid_from_weights()` gives the number of elements when it flattens them, which it cannot infer when there are no points, and its docstring says what `axis_output` may be. - The tests have thin cells, an oracle for three dimensions which shares nothing with the code under test, and the boundary of the grid under `bounds="raise"`. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01BYjDL98znSud1yFh9chnkP --- README.md | 4 +- benchmarks/weights_multilinear.py | 1 + docs/index.rst | 4 +- regridding/_regrid/_regrid.py | 8 +- regridding/_regrid/_regrid_from_weights.py | 17 +- regridding/_weights/_weights.py | 7 +- .../_weights/_weights_multilinear/_cells.py | 253 +++++++++++++---- .../_weights_multilinear/_cells_test.py | 171 ++++++++++-- .../_weights_multilinear/_curvilinear.py | 238 +++------------- .../_weights_multilinear/_curvilinear_test.py | 146 ++++++++-- .../_weights_multilinear/_rectilinear.py | 102 ++++--- .../_weights_multilinear/_rectilinear_test.py | 82 +++++- .../_weights_multilinear.py | 257 ++++++++++-------- .../_weights_multilinear_test.py | 192 +++++++++++-- 14 files changed, 967 insertions(+), 515 deletions(-) diff --git a/README.md b/README.md index 2c9282a..f7d8615 100644 --- a/README.md +++ b/README.md @@ -89,8 +89,8 @@ repeated independently for every position along them. This is how a stack of images, or a spectrum for each pixel, is resampled in one call. The `multilinear` method interpolates at points rather than resampling cells, -so `axis_output` may select any number of axes, such as a single axis of -scattered points to interpolate a 2D grid at. +so an explicit `axis_output` may select any number of axes, such as a single +axis of scattered points to interpolate a 2D grid at. **Degenerate grids are perturbed.** Where a vertex of the output grid lands exactly on an edge of the input grid, diff --git a/benchmarks/weights_multilinear.py b/benchmarks/weights_multilinear.py index 0ed8817..1a6f7fa 100644 --- a/benchmarks/weights_multilinear.py +++ b/benchmarks/weights_multilinear.py @@ -49,6 +49,7 @@ def time_weights_multilinear_2d( regridding.weights( coordinates_input=coordinates_input, coordinates_output=coordinates_output, + axis_output=0, method="multilinear", coalesce=False, ) diff --git a/docs/index.rst b/docs/index.rst index ebf08bb..c8f9c28 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -101,8 +101,8 @@ repeated independently for every position along them. This is how a stack of images, or a spectrum for each pixel, is resampled in one call. The ``multilinear`` method interpolates at points rather than resampling cells, -so ``axis_output`` may select any number of axes, such as a single axis of -scattered points to interpolate a 2D grid at. +so an explicit ``axis_output`` may select any number of axes, such as a single +axis of scattered points to interpolate a 2D grid at. **The conservative method has two algorithms, and picks between them.** Both compute the same weights: the fraction of each input cell shared with diff --git a/regridding/_regrid/_regrid.py b/regridding/_regrid/_regrid.py index cb51359..9dabf4b 100644 --- a/regridding/_regrid/_regrid.py +++ b/regridding/_regrid/_regrid.py @@ -48,10 +48,11 @@ def regrid( If :obj:`None`, all the axes of the output grid correspond to resampled axes in the input grid. For the ``conservative`` method, the number of axes should be equal - to the number of coordinates in the output grid. + to the number of coordinates in the output grid, as it should be for + the ``multilinear`` method if `axis_output` is :obj:`None`. The ``multilinear`` method interpolates at points rather than - resampling cells, so the points may be arranged along any number of - axes, or none at all for a single point. + resampling cells, so if `axis_output` is given, the points may be + arranged along any number of axes, or none at all for a single point. method The type of regridding to use. The ``multilinear`` method interprets `coordinates_input` as the @@ -176,6 +177,7 @@ def regrid( coordinates_input=(x_input, y_input), coordinates_output=(x_output, y_output), values_input=values_input, + axis_output=0, method="multilinear", bounds="nan", ) diff --git a/regridding/_regrid/_regrid_from_weights.py b/regridding/_regrid/_regrid_from_weights.py index 0c4b071..ffeed12 100644 --- a/regridding/_regrid/_regrid_from_weights.py +++ b/regridding/_regrid/_regrid_from_weights.py @@ -51,8 +51,9 @@ def regrid_from_weights( 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`. + These should be the axes passed as `axis_output` to + :func:`regridding.weights`, which for the ``multilinear`` method may + be any number of axes, or none for a single point. Notes ----- @@ -139,9 +140,15 @@ def regrid_from_weights( # The weights address each element by a flat index, so its resampled # axes are flattened together. That also covers a single point, which - # has no resampled axes at all. - values_input = values_input.reshape(-1, np.prod(shape_input_numba, dtype=int)) - values_output = values_output.reshape(-1, np.prod(shape_output_numba, dtype=int)) + # has no resampled axes at all. The number of elements is given rather + # than inferred, since it cannot be when there are no points. + num_orthogonal = int(np.prod(shape_orthogonal, dtype=int)) + values_input = values_input.reshape( + num_orthogonal, np.prod(shape_input_numba, dtype=int) + ) + values_output = values_output.reshape( + num_orthogonal, np.prod(shape_output_numba, dtype=int) + ) flat = weights.reshape(-1) unit_weights = getattr(flat[0][2], "unit", None) if flat.size else None diff --git a/regridding/_weights/_weights.py b/regridding/_weights/_weights.py index 462848a..5756894 100644 --- a/regridding/_weights/_weights.py +++ b/regridding/_weights/_weights.py @@ -142,10 +142,11 @@ def weights( If :obj:`None`, all the axes of the output grid correspond to resampled axes in the input grid. For the ``conservative`` method, the number of axes should be equal - to the number of coordinates in the output grid. + to the number of coordinates in the output grid, as it should be for + the ``multilinear`` method if `axis_output` is :obj:`None`. The ``multilinear`` method interpolates at points rather than - resampling cells, so the points may be arranged along any number of - axes, or none at all for a single point. + resampling cells, so if `axis_output` is given, the points may be + arranged along any number of axes, or none at all for a single point. weights_input Weights applied to the values of the input grid before resampling. method diff --git a/regridding/_weights/_weights_multilinear/_cells.py b/regridding/_weights/_weights_multilinear/_cells.py index 410f03c..bd8a101 100644 --- a/regridding/_weights/_weights_multilinear/_cells.py +++ b/regridding/_weights/_weights_multilinear/_cells.py @@ -50,13 +50,13 @@ "target", "distance_squared", "solve", - "inverse", + "newton", "weights", ] -num_work = 6 +num_work = 7 """ -The number of rows of a workspace for :func:`inverse`, beyond the +The number of rows of a workspace for :func:`newton`, beyond the :math:`2N` rows of the two matrices. A workspace has the shape ``(2 * N + num_work, N)``. @@ -64,11 +64,19 @@ tolerance = 1e-13 """ -The size of the last step, in units of the local coordinates, below which +The size of a step, in units of the local coordinates, below which Newton's method has converged. """ -num_iteration = 32 +rounding = 64 * np.finfo(np.float64).eps +""" +The rounding error of a position, for each corner of a cell, as a fraction +of the largest of the terms it is the sum of. + +See :func:`newton`. +""" + +num_iteration = 64 """The most steps Newton's method takes before it gives up.""" num_halving = 30 @@ -319,42 +327,165 @@ def distance_squared( @numba.njit(cache=True) -def inverse( +def _move( + shape: np.ndarray, + stride: np.ndarray, + clamp: bool, + cell: np.ndarray, + xi: np.ndarray, +) -> int: + """ + Move a point given by a cell and its local coordinates in that cell to + the cell its local coordinates fall in, if `clamp` is set. + + The cell is clamped to the grid, so a point beyond the grid stays in the + cell along its edge, with local coordinates outside the unit cube. The + local coordinates shift by whole cells, which is exact. Returns the flat + index of the first corner of the cell, moved or not. + + Parameters + ---------- + shape + The number of vertices along each axis. + stride + The number of vertices between neighbors along each axis. + clamp + Whether to move the point to the cell its local coordinates fall in. + cell + The first corner of the cell, which is replaced by that of the cell + the point is moved to. + xi + The local coordinates of the point, which are replaced by those in + the cell it is moved to. + """ + base = 0 + for m in range(cell.shape[0]): + if clamp: + shift = np.floor(xi[m]) + low = -float(cell[m]) + high = float(shape[m] - 2 - cell[m]) + if shift != shift: + shift = 0.0 + elif shift < low: + shift = low + elif shift > high: + shift = high + cell[m] += int(shift) + xi[m] -= shift + base += cell[m] * stride[m] + return base + + +@numba.njit(cache=True) +def _floor( vertices: np.ndarray, base: int, offset: np.ndarray, + xi: np.ndarray, + goal: np.ndarray, +) -> float: + """ + The square of the distance between the position of a point of a cell + and the point it should be at, below which the difference is no more + than rounding error. + + Each coordinate of the position, measured from the first corner of the + cell, is a sum over the corners of their weights times their distances + from the first corner, and the point is measured from the first corner + too. Each rounds to a multiple of :data:`rounding` of the largest of + its terms: the sum of the magnitudes of the weights times the extent of + the cell along that coordinate, and the distance of the point from the + first corner. + + Parameters + ---------- + vertices + The coordinates of every vertex of the grid. + base + The flat index of the first corner of the cell. + offset + The offset of each corner from the first. + xi + The local coordinates of the point. + goal + The point, measured from the first corner of the cell. + """ + num_axis = xi.shape[0] + total = 1.0 + for m in range(num_axis): + total *= abs(xi[m]) + abs(1 - xi[m]) + factor = (1 << num_axis) * rounding + result = 0.0 + for k in range(num_axis): + extent = 0.0 + for b in range(offset.shape[0]): + extent = max(extent, abs(vertices[k, base + offset[b]] - vertices[k, base])) + error = factor * (total * extent + abs(goal[k])) + result += error * error + return result + + +@numba.njit(cache=True) +def newton( + vertices: np.ndarray, + shape: np.ndarray, + stride: np.ndarray, + offset: np.ndarray, point: np.ndarray, + clamp: bool, + cell: np.ndarray, xi: np.ndarray, + cell_trial: np.ndarray, work: np.ndarray, ) -> bool: """ - Find the local coordinates of a point in a cell, which may be outside - the unit cube if the point is outside the cell. - - This uses Newton's method, starting from the value of `xi` it is given. - Each step is halved until it brings the cell's position closer to the - point, so that the method does not run away on a strongly curved cell. - It has converged once a step is shorter than :data:`tolerance`, which is - in units of the local coordinates and so does not depend on the units - or the size of the grid. - - Returns whether it converged. If it did not, `xi` is left wherever the - method stopped. + Find the local coordinates of a point by Newton's method. + + The iterate is a cell and local coordinates in it, starting from those + given. Without `clamp`, it stays in its cell, whose continuation is + searched beyond it, so the local coordinates come out outside the unit + cube if the point is outside the cell. With `clamp`, it moves after + every step to the cell its local coordinates fall in, clamped to the + grid, so it can cross the grid and continue past its edges, which is how + a point outside the grid is extrapolated. + + Each step is halved until it brings the position closer to the point, so + that the method does not run away on a strongly curved cell. It has + converged once a step is shorter than :data:`tolerance`, in units of the + local coordinates, or once the distance from the position to the point + is down to the rounding error of the position (see :func:`_floor`). The + second is needed for a cell much longer than it is wide, if it is not + aligned with the axes: rounding along its length moves the local + coordinate across its width by much more than the first allows, and no + step then brings the position any closer. Neither depends on the units, + the size, or the position of the grid. + + Returns whether it converged. The cell and local coordinates are left + wherever it stopped. Parameters ---------- vertices The coordinates of every vertex of the grid, of shape ``(N, number of vertices)``. - base - The flat index of the first corner of the cell. + shape + The number of vertices along each axis. + stride + The number of vertices between neighbors along each axis. offset The offset of each corner from the first, from :func:`offsets`. point The point to find. + clamp + Whether the iterate moves between cells. + cell + The first corner of the cell to start in, which is replaced by that + of the cell it stops in. xi - The local coordinates to start from, which are replaced by the - result. + The local coordinates to start from, which are replaced by those + where it stops. + cell_trial + Room for the first corner of the cell of a trial step. work A workspace of shape ``(2 * N + num_work, N)``. """ @@ -365,17 +496,19 @@ def inverse( x = work[2 * num_axis] residual = work[2 * num_axis + 1] step = work[2 * num_axis + 2] - trial = work[2 * num_axis + 3] - x_trial = work[2 * num_axis + 4] - goal = work[2 * num_axis + 5] + goal = work[2 * num_axis + 3] + xi_trial = work[2 * num_axis + 4] + x_trial = work[2 * num_axis + 5] + goal_trial = work[2 * num_axis + 6] + base = _move(shape, stride, clamp, cell, xi) target(vertices, base, point, goal) evaluate(vertices, base, offset, xi, x, jacobian) error = distance_squared(x, goal) for _ in range(num_iteration): - if error == 0: + if error <= _floor(vertices, base, offset, xi, goal): return True matrix[:, :] = jacobian @@ -394,11 +527,15 @@ def inverse( fraction = 1.0 accepted = False + base_trial = base for _ in range(num_halving): + cell_trial[:] = cell for m in range(num_axis): - trial[m] = xi[m] - fraction * step[m] - position(vertices, base, offset, trial, x_trial) - error_trial = distance_squared(x_trial, goal) + xi_trial[m] = xi[m] - fraction * step[m] + base_trial = _move(shape, stride, clamp, cell_trial, xi_trial) + target(vertices, base_trial, point, goal_trial) + position(vertices, base_trial, offset, xi_trial, x_trial) + error_trial = distance_squared(x_trial, goal_trial) if error_trial < error: accepted = True break @@ -409,11 +546,14 @@ def inverse( # is down to the rounding error of the position return size <= tolerance * scale - xi[:] = trial + cell[:] = cell_trial + xi[:] = xi_trial + base = base_trial if fraction * size <= tolerance * scale: return True + target(vertices, base, point, goal) evaluate(vertices, base, offset, xi, x, jacobian) error = distance_squared(x, goal) @@ -430,59 +570,66 @@ def weights( nan_outside: bool, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """ - The weight of each corner of the cell containing each point. + The weight of each corner of the cell containing each point, for each + of a number of elements. - Returns the flat arrays ``(indices_input, indices_output, values)``, - with :math:`2^N` consecutive weights for each point, one for each corner - of its cell, in the order the corners are numbered. + Returns the arrays ``(indices_input, indices_output, values)``, each with + a row for each element, which holds :math:`2^N` consecutive weights for + each point, one for each corner of its cell, in the order the corners + are numbered. Parameters ---------- cell The first corner of the cell of each point, of shape - ``(number of points, N)``. + ``(number of elements, number of points, N)``. xi The local coordinates of each point in its cell, of the same shape. inside - Whether each point is inside the grid. + Whether each point is inside the grid, of shape + ``(number of elements, number of points)``. shape The number of vertices along each axis. weights_input - Weights to multiply the weight of each vertex by, as a flat array, or - an empty array for none. + Weights to multiply the weight of each vertex by, flattened, of shape + ``(number of elements, number of vertices)``, or with a single row + which every element shares, or with none for no weights. nan_outside Whether to set the weights of a point outside the grid to NaN rather than extrapolate. """ - num_point, num_axis = cell.shape + num_element, num_point, num_axis = cell.shape stride = strides(shape) offset = offsets(shape) num_corner = offset.shape[0] - weighted = weights_input.shape[0] > 0 + num_weighted = weights_input.shape[0] - indices_input = np.empty(num_point * num_corner, dtype=np.int64) - indices_output = np.empty(num_point * num_corner, dtype=np.int64) - values = np.empty(num_point * num_corner, dtype=np.float64) + size = num_point * num_corner + indices_input = np.empty((num_element, size), dtype=np.int64) + indices_output = np.empty((num_element, size), dtype=np.int64) + values = np.empty((num_element, size), dtype=np.float64) - for i in numba.prange(num_point): + for t in numba.prange(num_element * num_point): + e = t // num_point + i = t - e * num_point base = 0 for m in range(num_axis): - base += cell[i, m] * stride[m] + base += cell[e, i, m] * stride[m] for b in range(num_corner): index = base + offset[b] weight = 1.0 for m in range(num_axis): if (b >> m) & 1: - weight *= xi[i, m] + weight *= xi[e, i, m] else: - weight *= 1 - xi[i, m] - if weighted: - weight *= weights_input[index] - if nan_outside and not inside[i]: + weight *= 1 - xi[e, i, m] + if num_weighted > 0: + weight *= weights_input[min(e, num_weighted - 1), index] + if nan_outside and not inside[e, i]: weight = np.nan w = i * num_corner + b - indices_input[w] = index - indices_output[w] = i - values[w] = weight + indices_input[e, w] = index + indices_output[e, w] = i + values[e, w] = weight return indices_input, indices_output, values diff --git a/regridding/_weights/_weights_multilinear/_cells_test.py b/regridding/_weights/_weights_multilinear/_cells_test.py index 832e27b..0137ff2 100644 --- a/regridding/_weights/_weights_multilinear/_cells_test.py +++ b/regridding/_weights/_weights_multilinear/_cells_test.py @@ -106,8 +106,37 @@ def test_singular(self) -> None: assert not _cells.solve(a, b, out) +def _newton( + vertices: np.ndarray, + shape: np.ndarray, + point: np.ndarray, + cell: tuple[int, ...], + clamp: bool = False, +) -> tuple[bool, np.ndarray, np.ndarray]: + """ + Run Newton's method from the middle of a cell, and return whether it + converged, the cell it stopped in, and its local coordinates there. + """ + num_axis = shape.size + cell_ = np.array(cell, dtype=np.int64) + xi = np.full(num_axis, 0.5) + converged = _cells.newton( + vertices, + shape, + _cells.strides(shape), + _cells.offsets(shape), + point, + clamp, + cell_, + xi, + np.empty(num_axis, dtype=np.int64), + _work(num_axis), + ) + return converged, cell_, xi + + @pytest.mark.parametrize("num_axis", [1, 2, 3]) -class TestInverse: +class TestNewton: @pytest.mark.parametrize( argnames="scale,shift", @@ -122,15 +151,15 @@ def test_round_trip(self, num_axis: int, scale: float, shift: float) -> None: vertices, shape, offset = _cell(num_axis, scale=scale, shift=shift) rng = np.random.default_rng(2) expected = rng.uniform(-0.5, 1.5, (20, num_axis)) - work = _work(num_axis) point = np.empty(num_axis) result = np.empty_like(expected) for i in range(expected.shape[0]): _cells.position(vertices, 0, offset, expected[i], point) - point = point + vertices[:, 0] - xi = np.full(num_axis, 0.5) - assert _cells.inverse(vertices, 0, offset, point, xi, work) - result[i] = xi + converged, cell, result[i] = _newton( + vertices, shape, point + vertices[:, 0], (0,) * num_axis + ) + assert converged + assert np.all(cell == 0) # the local coordinates are only as good as the point, whose rounding # error grows with its distance from the origin atol = 1e-12 + 1e-15 * shift / scale @@ -141,8 +170,8 @@ def test_step_overflows(self, num_axis: int) -> None: # a cell far smaller than its distance to the point vertices, shape, offset = _cell(num_axis, scale=1e-300) point = np.full(num_axis, 1e10) - xi = np.full(num_axis, 0.5) - assert not _cells.inverse(vertices, 0, offset, point, xi, _work(num_axis)) + converged, _, _ = _newton(vertices, shape, point, (0,) * num_axis) + assert not converged def test_unreachable(self, num_axis: int) -> None: # a cell squashed flat along its first axis cannot reach a point off @@ -150,8 +179,59 @@ def test_unreachable(self, num_axis: int) -> None: vertices, shape, offset = _cell(num_axis) vertices[0] = 0 point = np.full(num_axis, 0.5) - xi = np.full(num_axis, 0.5) - assert not _cells.inverse(vertices, 0, offset, point, xi, _work(num_axis)) + converged, _, _ = _newton(vertices, shape, point, (0,) * num_axis) + assert not converged + + def test_clamp(self, num_axis: int) -> None: + # starting in the first cell of a sheared grid, the iterate crosses + # the grid and stops in the cell at its far corner, continued past it + axes = [np.arange(4.0 + m) for m in range(num_axis)] + grid = list(np.meshgrid(*axes, indexing="ij")) + grid[0] = grid[0] + 0.3 * grid[~0] + shape = np.array(grid[0].shape, dtype=np.int64) + vertices = np.stack([np.ravel(c) for c in grid]) + index = shape + 0.5 + point = index.copy() + point[0] = point[0] + 0.3 * index[~0] + converged, cell, xi = _newton( + vertices, shape, point, (0,) * num_axis, clamp=True + ) + assert converged + np.testing.assert_array_equal(cell, shape - 2) + # 1.5 vertices past the last one, so 2.5 from the start of the last cell + np.testing.assert_allclose(xi, 2.5, rtol=0, atol=1e-12) + + +@pytest.mark.parametrize("num_axis", [2, 3]) +@pytest.mark.parametrize("aspect", [1e4, 1e6]) +def test_newton_thin_cell(num_axis: int, aspect: float) -> None: + # a cell much longer than it is wide, at an angle to the axes. Rounding + # along its length moves the local coordinates across its width by far + # more than the tolerance on a step, which once left Newton's method + # unconverged for points well inside the cell. + corners = np.stack(np.meshgrid(*[np.arange(2.0)] * num_axis, indexing="ij")) + corners[1:] = corners[1:] / aspect + angle = 0.6 + rotation = np.eye(num_axis) + rotation[:2, :2] = [ + [np.cos(angle), -np.sin(angle)], + [np.sin(angle), np.cos(angle)], + ] + corners = np.einsum("ij,j...->i...", rotation, corners) + vertices = corners.reshape(num_axis, -1) + shape = np.full(num_axis, 2, dtype=np.int64) + offset = _cells.offsets(shape) + rng = np.random.default_rng(0) + expected = rng.uniform(0.05, 0.95, (50, num_axis)) + point = np.empty(num_axis) + result = np.empty_like(expected) + for i in range(expected.shape[0]): + _cells.position(vertices, 0, offset, expected[i], point) + converged, _, result[i] = _newton( + vertices, shape, point + vertices[:, 0], (0,) * num_axis + ) + assert converged + np.testing.assert_allclose(result, expected, rtol=0, atol=1e-13 * aspect) class TestWeights: @@ -163,43 +243,74 @@ def test_partition_of_unity(self) -> None: xi = rng.uniform(-0.5, 1.5, (10, 3)) inside = np.ones(10, dtype=bool) indices_input, indices_output, values = _cells.weights( - cell, xi, inside, shape, np.empty(0), False + cell[np.newaxis], + xi[np.newaxis], + inside[np.newaxis], + shape, + np.empty((0, 0)), + False, ) - np.testing.assert_array_equal(indices_output, np.repeat(np.arange(10), 8)) - sums = np.bincount(indices_output, weights=values) + np.testing.assert_array_equal(indices_output, [np.repeat(np.arange(10), 8)]) + sums = np.bincount(indices_output[0], weights=values[0]) np.testing.assert_allclose(sums, 1, rtol=1e-14) - first = indices_input[::8] + first = indices_input[0, ::8] np.testing.assert_array_equal(first, np.ravel_multi_index(cell.T, shape)) def test_corners_of_the_cell(self) -> None: shape = np.array([3, 4]) - cell = np.array([[1, 2]]) - xi = np.array([[0.25, 0.5]]) + cell = np.array([[[1, 2]]]) + xi = np.array([[[0.25, 0.5]]]) indices_input, indices_output, values = _cells.weights( - cell, xi, np.ones(1, dtype=bool), shape, np.empty(0), False + cell, xi, np.ones((1, 1), dtype=bool), shape, np.empty((0, 0)), False ) - np.testing.assert_array_equal(indices_input, [6, 10, 7, 11]) - np.testing.assert_allclose(values, [0.375, 0.125, 0.375, 0.125]) + np.testing.assert_array_equal(indices_input, [[6, 10, 7, 11]]) + np.testing.assert_allclose(values, [[0.375, 0.125, 0.375, 0.125]]) - def test_weights_input(self) -> None: + @pytest.mark.parametrize("per_element", [False, True]) + def test_weights_input(self, per_element: bool) -> None: + # two elements, which share the weights of the vertices or have + # their own shape = np.array([3, 4]) - cell = np.array([[1, 2]]) - xi = np.array([[0.25, 0.5]]) - weights_input = np.arange(12, dtype=float) + cell = np.array([[[1, 2]], [[0, 0]]]) + xi = np.array([[[0.25, 0.5]], [[0.5, 0.5]]]) + weights_input = np.arange(12, dtype=float)[np.newaxis] + if per_element: + weights_input = np.concatenate([weights_input, 2 * weights_input]) _, _, values = _cells.weights( - cell, xi, np.ones(1, dtype=bool), shape, weights_input, False + cell, xi, np.ones((2, 1), dtype=bool), shape, weights_input, False ) - expected = np.array([0.375, 0.125, 0.375, 0.125]) * [6, 10, 7, 11] + expected = np.array( + [ + np.array([0.375, 0.125, 0.375, 0.125]) * [6, 10, 7, 11], + np.array([0.25, 0.25, 0.25, 0.25]) * [0, 4, 1, 5], + ] + ) + if per_element: + expected[1] *= 2 np.testing.assert_allclose(values, expected) @pytest.mark.parametrize("nan_outside", [False, True]) def test_outside(self, nan_outside: bool) -> None: shape = np.array([3, 4]) - cell = np.array([[1, 2], [0, 0]]) - xi = np.array([[1.5, 0.5], [0.5, 0.5]]) - inside = np.array([False, True]) - _, _, values = _cells.weights(cell, xi, inside, shape, np.empty(0), nan_outside) + cell = np.array([[[1, 2], [0, 0]]]) + xi = np.array([[[1.5, 0.5], [0.5, 0.5]]]) + inside = np.array([[False, True]]) + _, _, values = _cells.weights( + cell, xi, inside, shape, np.empty((0, 0)), nan_outside + ) first = [-0.25, 0.75, -0.25, 0.75] if nan_outside: first = [np.nan] * 4 - np.testing.assert_allclose(values, first + [0.25] * 4) + np.testing.assert_allclose(values, [first + [0.25] * 4]) + + def test_no_points(self) -> None: + shape = np.array([3, 4]) + result = _cells.weights( + np.empty((2, 0, 2), dtype=np.int64), + np.empty((2, 0, 2)), + np.empty((2, 0), dtype=bool), + shape, + np.empty((0, 0)), + False, + ) + assert [a.shape for a in result] == [(2, 0)] * 3 diff --git a/regridding/_weights/_weights_multilinear/_curvilinear.py b/regridding/_weights/_weights_multilinear/_curvilinear.py index aa6c3da..6013056 100644 --- a/regridding/_weights/_weights_multilinear/_curvilinear.py +++ b/regridding/_weights/_weights_multilinear/_curvilinear.py @@ -50,11 +50,6 @@ such a point could be outside every cell. """ -num_iteration_walk = 64 -""" -The most steps of Newton's method across the whole grid when extrapolating. -""" - @numba.njit(cache=True) def _unravel( @@ -325,10 +320,13 @@ def _in_box( @numba.njit(cache=True) def _in_cell( vertices: np.ndarray, - base: int, + shape: np.ndarray, + stride: np.ndarray, offset: np.ndarray, point: np.ndarray, + cell: np.ndarray, xi: np.ndarray, + cell_trial: np.ndarray, work: np.ndarray, ) -> bool: """ @@ -338,19 +336,29 @@ def _in_cell( ---------- vertices The coordinates of every vertex of the grid. - base - The flat index of the first corner of the cell. + shape + The number of vertices along each axis. + stride + The number of vertices between neighbors along each axis. offset The offset of each corner from the first. point The point. + cell + The first corner of the cell. xi Where to put the local coordinates of the point. + cell_trial + Room for :func:`~._cells.newton`. work - A workspace for :func:`~._cells.inverse`. + A workspace for :func:`~._cells.newton`. """ num_axis = point.shape[0] + base = 0 + for m in range(num_axis): + base += cell[m] * stride[m] + for c in range(num_axis): low = np.inf high = -np.inf @@ -364,7 +372,9 @@ def _in_cell( return False xi[:] = 0.5 - if not _cells.inverse(vertices, base, offset, point, xi, work): + if not _cells.newton( + vertices, shape, stride, offset, point, False, cell, xi, cell_trial, work + ): return False for m in range(num_axis): @@ -377,6 +387,7 @@ def _in_cell( @numba.njit(cache=True) def _search( vertices: np.ndarray, + shape: np.ndarray, stride: np.ndarray, offset: np.ndarray, lower: np.ndarray, @@ -389,6 +400,7 @@ def _search( block: np.ndarray, cell: np.ndarray, xi: np.ndarray, + cell_trial: np.ndarray, work: np.ndarray, ) -> bool: """ @@ -403,6 +415,8 @@ def _search( ---------- vertices The coordinates of every vertex of the grid. + shape + The number of vertices along each axis. stride The number of vertices between neighbors along each axis. offset @@ -427,8 +441,10 @@ def _search( Where to put the first corner of the cell. xi Where to put the local coordinates of the point. + cell_trial + Room for :func:`~._cells.newton`. work - A workspace for :func:`~._cells.inverse`. + A workspace for :func:`~._cells.newton`. """ num_axis = point.shape[0] num_level = shapes.shape[0] - 1 @@ -436,7 +452,9 @@ def _search( if num_level == 0: cell[:] = 0 - return _in_cell(vertices, 0, offset, point, xi, work) + return _in_cell( + vertices, shape, stride, offset, point, cell, xi, cell_trial, work + ) if not _in_box(lower, upper, int(start[num_level]), point): return False @@ -455,13 +473,13 @@ def _search( # the children are cells, which are tried in order for b in range(num_corner): valid = True - base = 0 for m in range(num_axis): cell[m] = 2 * block[m] + ((b >> m) & 1) if cell[m] >= below[m]: valid = False - base += cell[m] * stride[m] - if valid and _in_cell(vertices, base, offset, point, xi, work): + if valid and _in_cell( + vertices, shape, stride, offset, point, cell, xi, cell_trial, work + ): return True else: @@ -701,177 +719,6 @@ def _nearest( size += 1 -@numba.njit(cache=True) -def _clamp( - s: np.ndarray, - shape: np.ndarray, - stride: np.ndarray, - cell: np.ndarray, - xi: np.ndarray, -) -> int: - """ - Split a fractional index into the cell nearest to it and its local - coordinates in that cell, which are outside the unit cube if the index - is outside the grid. - - Returns the flat index of the first corner of the cell. - - Parameters - ---------- - s - The fractional index of a point, along each axis of the vertices. - shape - The number of vertices along each axis. - stride - The number of vertices between neighbors along each axis. - cell - Where to put the first corner of the cell. - xi - Where to put the local coordinates. - """ - base = 0 - for m in range(s.shape[0]): - value = s[m] - if not value >= 0: - value = 0.0 - elif value > shape[m] - 2: - value = shape[m] - 2.0 - cell[m] = int(np.floor(value)) - xi[m] = s[m] - cell[m] - base += cell[m] * stride[m] - return base - - -@numba.njit(cache=True) -def _walk( - vertices: np.ndarray, - shape: np.ndarray, - stride: np.ndarray, - offset: np.ndarray, - point: np.ndarray, - s: np.ndarray, - s_trial: np.ndarray, - cell: np.ndarray, - cell_trial: np.ndarray, - xi: np.ndarray, - xi_trial: np.ndarray, - work: np.ndarray, -) -> bool: - """ - Extrapolate the grid to a point outside of it. - - This runs Newton's method across the whole grid in index space, from the - fractional index in `s`, continuing the cells along the edges of the grid - beyond it. Each step is halved until it brings the position closer to - the point. The first corner of the cell it stops in is left in `cell`, - and the local coordinates of the point in that cell in `xi`. - - Returns whether it reached the point. It does not if the continuation - of the grid folds over before reaching it, as it can beyond a strongly - curved edge, in which case it stops where the continuation came closest. - - Parameters - ---------- - vertices - The coordinates of every vertex of the grid. - shape - The number of vertices along each axis. - stride - The number of vertices between neighbors along each axis. - offset - The offset of each corner of a cell from the first. - point - The point. - s - The fractional index to start from, which is replaced by that of the - point. - s_trial - Room for a trial fractional index. - cell - Where to put the first corner of the cell. - cell_trial - Room for the first corner of a trial cell. - xi - Where to put the local coordinates of the point in the cell. - xi_trial - Room for the local coordinates in a trial cell. - work - A workspace of the same shape as one for :func:`~._cells.inverse`. - """ - num_axis = point.shape[0] - - jacobian = work[:num_axis] - matrix = work[num_axis : 2 * num_axis] - x = work[2 * num_axis] - residual = work[2 * num_axis + 1] - step = work[2 * num_axis + 2] - x_trial = work[2 * num_axis + 3] - goal = work[2 * num_axis + 4] - goal_trial = work[2 * num_axis + 5] - - base = _clamp(s, shape, stride, cell, xi) - _cells.target(vertices, base, point, goal) - _cells.evaluate(vertices, base, offset, xi, x, jacobian) - error = _cells.distance_squared(x, goal) - - reached = False - - for _ in range(num_iteration_walk): - - if not error > 0: - reached = error == 0 - break - - matrix[:, :] = jacobian - for k in range(num_axis): - residual[k] = x[k] - goal[k] - if not _cells.solve(matrix, residual, step): - break - - size = 0.0 - scale = 1.0 - for m in range(num_axis): - size = max(size, abs(step[m])) - scale = max(scale, abs(s[m])) - if not size < np.inf: - break - - fraction = 1.0 - accepted = False - for _ in range(_cells.num_halving): - for m in range(num_axis): - s_trial[m] = s[m] - fraction * step[m] - base = _clamp(s_trial, shape, stride, cell_trial, xi_trial) - _cells.target(vertices, base, point, goal_trial) - _cells.position(vertices, base, offset, xi_trial, x_trial) - error_trial = _cells.distance_squared(x_trial, goal_trial) - if error_trial < error: - accepted = True - break - fraction = fraction / 2 - - if not accepted: - # no point along the step is closer, which happens once the step - # is down to the rounding error of the position - reached = size <= _cells.tolerance * scale - break - - s[:] = s_trial - - if fraction * size <= _cells.tolerance * scale: - reached = True - break - - base = _clamp(s, shape, stride, cell, xi) - _cells.target(vertices, base, point, goal) - _cells.evaluate(vertices, base, offset, xi, x, jacobian) - error = _cells.distance_squared(x, goal) - - _clamp(s, shape, stride, cell, xi) - - return reached - - @numba.njit(cache=True, parallel=True) def locate( vertices: np.ndarray, @@ -939,9 +786,6 @@ def locate( cell_trial = np.empty(num_axis, dtype=np.int64) point = np.empty(num_axis, dtype=np.float64) xi_i = np.empty(num_axis, dtype=np.float64) - xi_trial = np.empty(num_axis, dtype=np.float64) - s = np.empty(num_axis, dtype=np.float64) - s_trial = np.empty(num_axis, dtype=np.float64) child_index = np.empty(num_corner, dtype=np.int64) child_distance = np.empty(num_corner, dtype=np.float64) work = np.empty((2 * num_axis + _cells.num_work, num_axis)) @@ -957,6 +801,7 @@ def locate( if finite: found = _search( vertices, + shape, stride, offset, lower, @@ -969,6 +814,7 @@ def locate( block, cell_i, xi_i, + cell_trial, work, ) @@ -992,20 +838,20 @@ def locate( cell_trial, cell_i, ) - for m in range(num_axis): - s[m] = cell_i[m] + 0.5 - reached = _walk( + # Newton's method across the whole grid, with the cells + # along its edges continued beyond it, from the middle of + # the nearest cell + xi_i[:] = 0.5 + reached = _cells.newton( vertices, shape, stride, offset, point, - s, - s_trial, + True, cell_i, - cell_trial, xi_i, - xi_trial, + cell_trial, work, ) if not reached: diff --git a/regridding/_weights/_weights_multilinear/_curvilinear_test.py b/regridding/_weights/_weights_multilinear/_curvilinear_test.py index 7e20ce1..c18a6e0 100644 --- a/regridding/_weights/_weights_multilinear/_curvilinear_test.py +++ b/regridding/_weights/_weights_multilinear/_curvilinear_test.py @@ -1,5 +1,6 @@ import pytest import numpy as np +import scipy.optimize import regridding from . import _cells, _curvilinear @@ -69,6 +70,40 @@ def _distorted_3d() -> tuple[np.ndarray, np.ndarray, np.ndarray]: ) +def _thin(aspect: float) -> tuple[np.ndarray, np.ndarray]: + """ + A grid of cells much longer than they are wide, rotated and slightly + curved, in which rounding along the length of a cell once kept Newton's + method from converging on points well inside it. + """ + a, b = np.meshgrid( + np.linspace(0, 1, 6), + np.linspace(0, 1 / aspect, 6), + indexing="ij", + ) + angle = 0.6 + return ( + a * np.cos(angle) - b * np.sin(angle) + 0.001 * a * a, + a * np.sin(angle) + b * np.cos(angle), + ) + + +def _thin_3d(aspect: float) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """A three-dimensional grid of cells much longer than they are wide.""" + a, b, c = np.meshgrid( + np.linspace(0, 1, 4), + np.linspace(0, 1 / aspect, 4), + np.linspace(0, 1 / aspect, 3), + indexing="ij", + ) + angle = 0.6 + return ( + a * np.cos(angle) - b * np.sin(angle) + 0.001 * a * a, + a * np.sin(angle) + b * np.cos(angle), + c + 0.3 * a, + ) + + def _folded_1d() -> tuple[np.ndarray]: """A one-dimensional grid which turns back on itself.""" return (np.array([0.0, 1.0, 3.0, 2.0, 4.0]),) @@ -82,6 +117,9 @@ def _folded_1d() -> tuple[np.ndarray]: "distorted, scale 1e9": _distorted(scale=1e9), "alternating": _alternating(), "distorted 3D": _distorted_3d(), + "thin, aspect 1e4": _thin(1e4), + "thin, aspect 1e6": _thin(1e6), + "thin 3D, aspect 1e4": _thin_3d(1e4), "folded 1D": _folded_1d(), "single cell": tuple(c[:2, :2] for c in _distorted()), } @@ -100,42 +138,93 @@ def _points( return rng.uniform(lower - margin, upper + margin, (num, len(grid))).T +def _multilinear( + corners: np.ndarray, + xi: np.ndarray, +) -> np.ndarray: + """ + The position of a point of a cell, by interpolating between its corners + along one axis after another. + + Parameters + ---------- + corners + The position of each corner, of shape ``(2,) * N + (N,)``. + xi + The local coordinates of the point. + """ + result = corners + for m in range(xi.size): + result = (1 - xi[m]) * result[0] + xi[m] * result[1] + return result + + +def _interior( + grid: tuple[np.ndarray, ...], + num: int = 100, + seed: int = 1, +) -> np.ndarray: + """Random points well inside random cells of a grid.""" + rng = np.random.default_rng(seed) + stacked = np.stack(grid, axis=~0) + num_axis = len(grid) + result = np.empty((num_axis, num)) + for i in range(num): + cell = [rng.integers(0, n - 1) for n in grid[0].shape] + corners = stacked[tuple(slice(j, j + 2) for j in cell)] + result[:, i] = _multilinear(corners, rng.uniform(0.05, 0.95, num_axis)) + return result + + def _inside_brute( grid: tuple[np.ndarray, ...], points: np.ndarray, ) -> np.ndarray: """ - Whether each point is inside the grid, by trying every cell. + Whether each point is inside the grid, by trying every cell, in ways + which share nothing with the code under test. In two dimensions a cell is the quadrilateral between its corners, since - its edges are straight, so this tests each quadrilateral with - :func:`regridding.geometry.point_is_inside_polygon`, which shares nothing - with the code under test. In other dimensions it finds the local - coordinates of the point in every cell. + its edges are straight, so each quadrilateral is tested with + :func:`regridding.geometry.point_is_inside_polygon`. In other + dimensions :func:`scipy.optimize.root` finds the local coordinates of the + point in each cell whose bounding box holds it. """ - vertices, shape = _flat(grid) num_axis = len(grid) - offset = _cells.offsets(shape) - stride = _cells.strides(shape) - work = np.empty((2 * num_axis + _cells.num_work, num_axis)) - xi = np.empty(num_axis) + shape = grid[0].shape + stacked = np.stack(grid, axis=~0) result = np.zeros(points.shape[1], dtype=bool) for i in range(points.shape[1]): - for index in np.ndindex(*(shape - 1)): - base = int(np.dot(index, stride)) + point = points[:, i] + for cell in np.ndindex(*(n - 1 for n in shape)): + corners = stacked[tuple(slice(j, j + 2) for j in cell)] if num_axis == 2: - corners = base + offset[[0, 1, 3, 2]] + x = corners[..., 0][[0, 1, 1, 0], [0, 0, 1, 1]] + y = corners[..., 1][[0, 1, 1, 0], [0, 0, 1, 1]] inside = regridding.geometry.point_is_inside_polygon( - x=points[0, i], - y=points[1, i], - vertices_x=vertices[0, corners], - vertices_y=vertices[1, corners], + x=point[0], + y=point[1], + vertices_x=x, + vertices_y=y, ) else: - xi[:] = 0.5 - inside = _cells.inverse( - vertices, base, offset, points[:, i], xi, work - ) and np.all((-1e-9 <= xi) & (xi <= 1 + 1e-9)) + flat = corners.reshape(-1, num_axis) + if np.any(point < flat.min(0)) or np.any(point > flat.max(0)): + continue + solution = scipy.optimize.root( + lambda xi: _multilinear(corners, xi) - point, + x0=np.full(num_axis, 0.5), + method="hybr", + options=dict(xtol=1e-14), + ) + # judged by the residual, since asked for this tolerance the + # solver reports failure once it cannot improve on a + # solution, however good + xi = solution.x + residual = np.abs(_multilinear(corners, xi) - point).max() + inside = bool(residual <= 1e-12) and bool( + np.all((-1e-9 <= xi) & (xi <= 1 + 1e-9)) + ) if inside: result[i] = True break @@ -200,9 +289,10 @@ class TestLocate: def test_inside(self, name: str) -> None: grid = grids[name] - points = _points(grid) + points = np.concatenate([_points(grid), _interior(grid)], axis=~0) cell, xi, inside = _locate(grid, points, extrapolate=False) np.testing.assert_array_equal(inside, _inside_brute(grid, points)) + assert np.all(inside[-100:]) def test_local_coordinates(self, name: str) -> None: grid = grids[name] @@ -221,17 +311,15 @@ def test_local_coordinates(self, name: str) -> None: def test_extrapolate(self, name: str) -> None: # every point the cells along the edge of the grid continue to is - # reached, so the position of each point is recovered exactly + # reached, so the position of each point is recovered, to the + # rounding error of a position that many cells along the continuation grid = grids[name] points = _points(grid) cell, xi, inside = _locate(grid, points, extrapolate=True) scale = float(np.max([np.ptp(c) for c in grid])) - np.testing.assert_allclose( - _position(grid, cell, xi), - points, - rtol=0, - atol=1e-12 * scale, - ) + reach = np.prod(np.abs(xi) + np.abs(1 - xi), axis=~0) + error = np.abs(_position(grid, cell, xi) - points).max(axis=0) + np.testing.assert_array_less(error, 1e-12 * scale * reach) assert not np.all(inside) diff --git a/regridding/_weights/_weights_multilinear/_rectilinear.py b/regridding/_weights/_weights_multilinear/_rectilinear.py index 3087ae0..4146770 100644 --- a/regridding/_weights/_weights_multilinear/_rectilinear.py +++ b/regridding/_weights/_weights_multilinear/_rectilinear.py @@ -7,6 +7,9 @@ coordinate alone. The search is exact, needs no preparation, and a point outside the grid is extrapolated from the edge of the grid along each axis on its own. + +The grids of many elements, such as a spectrum for each pixel of an image, +are searched at once, each with its own coordinates. """ import numpy as np @@ -23,27 +26,33 @@ def profiles( coordinates: tuple[np.ndarray, ...], ) -> None | tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """ - The coordinate along each axis of a grid, if the grid is rectilinear. - - Returns :obj:`None` if the grid is not rectilinear. Otherwise, returns - ``(profile, start, coordinate, descending)``, where ``profile`` holds the - values of the coordinate along each axis one after another, each sorted - in ascending order, ``profile[start[m]:start[m + 1]]`` is the one along - axis ``m``, ``coordinate[m]`` is which coordinate that is, and - ``descending[m]`` is whether it was reversed to sort it. + The coordinate along each axis of the grid of each element, if every + grid is rectilinear in the same way. + + Returns :obj:`None` unless the same coordinate varies along the same + axis of every grid, and is strictly monotonic along it. Otherwise, + returns ``(profile, start, coordinate, descending)``, where ``profile`` + holds, for each element, the values of the coordinate along each axis + one after another, each sorted in ascending order, + ``profile[:, start[m]:start[m + 1]]`` is the one along axis ``m``, + ``coordinate[m]`` is which coordinate that is, and ``descending[:, m]`` + is whether it was reversed to sort it. If every element shares one + grid, ``profile`` and ``descending`` have a single row. Parameters ---------- coordinates - The coordinates of the vertices of the grid, each an array with an - axis for each axis of the grid. An array which has been broadcast - along an axis, and so has a stride of zero along it, is known not to - vary along it without its values being looked at, which makes the - grids that :func:`numpy.broadcast_to` or ``numpy.meshgrid(..., - sparse=True)`` give quick to recognize. + The coordinates of the vertices of the grid of each element, each an + array with an axis for the elements followed by an axis for each axis + of the grid. An array which has been broadcast along an axis, and + so has a stride of zero along it, is known not to vary along it + without its values being looked at, which makes the grids that + :func:`numpy.broadcast_to` or ``numpy.meshgrid(..., sparse=True)`` + give quick to recognize. """ num_axis = len(coordinates) - shape = np.array(coordinates[0].shape, dtype=np.int64) + num_element = coordinates[0].shape[0] + shape = np.array(coordinates[0].shape[1:], dtype=np.int64) grid = [] coordinate = np.full(num_axis, -1, dtype=np.int64) @@ -52,7 +61,7 @@ def profiles( varies = [ m for m in range(num_axis) - if c.shape[m] > 1 and np.any(c != c.take([0], axis=m)) + if c.shape[1 + m] > 1 and np.any(c != c.take([0], axis=1 + m)) ] if len(varies) != 1: return None @@ -62,25 +71,30 @@ def profiles( coordinate[m] = k grid.append(c) + num_row = 1 if all(c.shape[0] == 1 for c in grid) else num_element + profile = [] - descending = np.zeros(num_axis, dtype=np.bool_) + descending = np.zeros((num_row, num_axis), dtype=np.bool_) for m in range(num_axis): - index: list[int | slice] = [0] * num_axis - index[m] = slice(None) - values = grid[coordinate[m]][tuple(index)] - difference = np.diff(values) - if np.all(difference < 0): - values = values[::-1] - descending[m] = True - elif not np.all(difference > 0): + index: list[int | slice] = [slice(None)] + [0] * num_axis + index[1 + m] = slice(None) + # compared as floating point numbers, since the difference of two + # unsigned integers wraps around instead of going negative + values = np.asarray(grid[coordinate[m]][tuple(index)], dtype=np.float64) + values = np.broadcast_to(values, (num_row, shape[m])) + difference = np.diff(values, axis=~0) + reverse = np.all(difference < 0, axis=~0) + if not np.all(reverse | np.all(difference > 0, axis=~0)): return None + values = np.where(reverse[:, np.newaxis], values[:, ::-1], values) + descending[:, m] = reverse profile.append(values) start = np.zeros(num_axis + 1, dtype=np.int64) start[1:] = np.cumsum(shape) return ( - np.ascontiguousarray(np.concatenate(profile), dtype=np.float64), + np.ascontiguousarray(np.concatenate(profile, axis=~0)), start, coordinate, descending, @@ -96,7 +110,8 @@ def locate( points: np.ndarray, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """ - Find the cell of a rectilinear grid which contains each point. + Find the cell of the rectilinear grid of each element which contains + each of its points. Returns ``(cell, xi, inside)``: the first corner of the cell along each axis, the local coordinates of the point in the cell, and whether the @@ -109,28 +124,33 @@ def locate( Parameters ---------- profile - The coordinate along each axis, from :func:`profiles`. + The coordinate along each axis of each grid, from :func:`profiles`. start Where the coordinate along each axis starts in `profile`. coordinate Which coordinate varies along each axis. descending - Whether the coordinate along each axis was reversed. + Whether the coordinate along each axis of each grid was reversed. points - The points to find, of shape ``(N, number of points)``. + The points to find, of shape + ``(number of elements, N, number of points)``. """ - num_axis, num_point = points.shape + num_element, num_axis, num_point = points.shape + num_row = profile.shape[0] - cell = np.empty((num_point, num_axis), dtype=np.int64) - xi = np.empty((num_point, num_axis), dtype=np.float64) - inside = np.empty(num_point, dtype=np.bool_) + cell = np.empty((num_element, num_point, num_axis), dtype=np.int64) + xi = np.empty((num_element, num_point, num_axis), dtype=np.float64) + inside = np.empty((num_element, num_point), dtype=np.bool_) - for i in numba.prange(num_point): + for t in numba.prange(num_element * num_point): + e = t // num_point + i = t - e * num_point + row = min(e, num_row - 1) inside_i = True for m in range(num_axis): - values = profile[start[m] : start[m + 1]] + values = profile[row, start[m] : start[m + 1]] num = values.shape[0] - value = points[coordinate[m], i] + value = points[e, coordinate[m], i] if not (values[0] <= value <= values[num - 1]): inside_i = False j = np.searchsorted(values, value) - 1 @@ -141,11 +161,11 @@ def locate( fraction = (value - values[j]) / (values[j + 1] - values[j]) if not np.isfinite(value): fraction = np.nan - if descending[m]: + if descending[row, m]: j = num - 2 - j fraction = 1 - fraction - cell[i, m] = j - xi[i, m] = fraction - inside[i] = inside_i + cell[e, i, m] = j + xi[e, i, m] = fraction + inside[e, i] = inside_i return cell, xi, inside diff --git a/regridding/_weights/_weights_multilinear/_rectilinear_test.py b/regridding/_weights/_weights_multilinear/_rectilinear_test.py index 7e45c2e..e91cdea 100644 --- a/regridding/_weights/_weights_multilinear/_rectilinear_test.py +++ b/regridding/_weights/_weights_multilinear/_rectilinear_test.py @@ -4,8 +4,11 @@ def _grid(grid: tuple[np.ndarray, ...]) -> tuple[np.ndarray, ...]: - """The coordinates of a grid, broadcast against each other.""" - return tuple(np.broadcast_arrays(*grid)) + """ + The coordinates of one grid, broadcast against each other, with an axis + for a single element in front. + """ + return tuple(c[np.newaxis] for c in np.broadcast_arrays(*grid)) x = np.array([0.0, 0.5, 2.0, 2.25]) @@ -20,10 +23,10 @@ def test_meshgrid(self) -> None: result = _rectilinear.profiles(coordinates) assert result is not None profile, start, coordinate, descending = result - np.testing.assert_array_equal(profile, np.concatenate([x, y, z])) + np.testing.assert_array_equal(profile, [np.concatenate([x, y, z])]) np.testing.assert_array_equal(start, [0, 4, 7, 12]) np.testing.assert_array_equal(coordinate, [0, 1, 2]) - np.testing.assert_array_equal(descending, [False, False, False]) + np.testing.assert_array_equal(descending, [[False, False, False]]) def test_broadcast(self) -> None: # a grid broadcast from one dimension along each axis is recognized @@ -33,7 +36,7 @@ def test_broadcast(self) -> None: result = _rectilinear.profiles(coordinates) assert result is not None profile, start, coordinate, descending = result - np.testing.assert_array_equal(profile, np.concatenate([x, y])) + np.testing.assert_array_equal(profile, [np.concatenate([x, y])]) np.testing.assert_array_equal(coordinate, [0, 1]) def test_coordinates_out_of_order(self) -> None: @@ -44,7 +47,7 @@ def test_coordinates_out_of_order(self) -> None: result = _rectilinear.profiles(coordinates) assert result is not None profile, start, coordinate, descending = result - np.testing.assert_array_equal(profile, np.concatenate([y, x])) + np.testing.assert_array_equal(profile, [np.concatenate([y, x])]) np.testing.assert_array_equal(coordinate, [1, 0]) def test_descending(self) -> None: @@ -52,8 +55,38 @@ def test_descending(self) -> None: result = _rectilinear.profiles(coordinates) assert result is not None profile, start, coordinate, descending = result - np.testing.assert_array_equal(profile, np.concatenate([x, y])) - np.testing.assert_array_equal(descending, [False, True]) + np.testing.assert_array_equal(profile, [np.concatenate([x, y])]) + np.testing.assert_array_equal(descending, [[False, True]]) + + def test_descending_unsigned(self) -> None: + # the difference of two unsigned integers wraps around instead of + # going negative, which once passed this for an ascending grid + coordinates = (np.array([[5, 4, 3, 2, 1]], dtype=np.uint8),) + result = _rectilinear.profiles(coordinates) + assert result is not None + profile, start, coordinate, descending = result + np.testing.assert_array_equal(profile, [[1.0, 2.0, 3.0, 4.0, 5.0]]) + np.testing.assert_array_equal(descending, [[True]]) + + def test_elements(self) -> None: + # a grid for each of three elements, the second running backwards + # along its second axis, and every one sharing its first coordinate + xx = np.broadcast_to(x[np.newaxis, :, np.newaxis], (3, 4, 3)) + yy = np.stack([y, 10 - y, 2 * y])[:, np.newaxis, :] + yy = np.broadcast_to(yy, (3, 4, 3)) + result = _rectilinear.profiles((xx, yy)) + assert result is not None + profile, start, coordinate, descending = result + expected = [ + np.concatenate([x, y]), + np.concatenate([x, 10 - y[::-1]]), + np.concatenate([x, 2 * y]), + ] + np.testing.assert_array_equal(profile, expected) + np.testing.assert_array_equal( + descending, + [[False, False], [False, True], [False, False]], + ) @pytest.mark.parametrize( argnames="grid", @@ -83,6 +116,11 @@ def test_not_rectilinear(self, grid: tuple[np.ndarray, ...]) -> None: coordinates = _grid(grid) assert _rectilinear.profiles(coordinates) is None + def test_elements_disagree(self) -> None: + # rectilinear for one element but not monotonic for the other + xx = np.stack([x, np.array([0.0, 1.0, 0.5, 2.0])]) + assert _rectilinear.profiles((xx,)) is None + class TestLocate: @@ -101,7 +139,7 @@ def test_locate(self, reverse: bool) -> None: [2.0, 3.0, 1.0, 4.0, 2.0, 5.0, np.nan], ] ) - cell, xi, inside = _rectilinear.locate(*result, points) + cell, xi, inside = _rectilinear.locate(*result, points[np.newaxis]) expected_cell_x = [0, 1, 0, 2, 0, 2, 1] expected_xi_x = [0.5, 1.0, 0.0, 1.0, -2.0, 4.0, 1 / 3] @@ -111,10 +149,26 @@ def test_locate(self, reverse: bool) -> None: expected_cell_y = 1 - expected_cell_y expected_xi_y = 1 - expected_xi_y - np.testing.assert_array_equal(cell[:, 0], expected_cell_x) - np.testing.assert_array_equal(cell[:, 1], expected_cell_y) - np.testing.assert_allclose(xi[:, 0], expected_xi_x, rtol=1e-15) - np.testing.assert_allclose(xi[:, 1], expected_xi_y, rtol=1e-15) + np.testing.assert_array_equal(cell[0, :, 0], expected_cell_x) + np.testing.assert_array_equal(cell[0, :, 1], expected_cell_y) + np.testing.assert_allclose(xi[0, :, 0], expected_xi_x, rtol=1e-15) + np.testing.assert_allclose(xi[0, :, 1], expected_xi_y, rtol=1e-15) np.testing.assert_array_equal( - inside, [True, True, True, True, False, False, False] + inside, [[True, True, True, True, False, False, False]] + ) + + def test_elements(self) -> None: + # each element is searched in its own grid, the second of which is + # [4.5, 4, 1, 0], running backwards + xx = np.stack([x, 2 * x[::-1]]) + result = _rectilinear.profiles((xx,)) + assert result is not None + points = np.array([[[0.25, 3.0]], [[0.5, 4.25]]]) + cell, xi, inside = _rectilinear.locate(*result, points) + np.testing.assert_array_equal(cell[..., 0], [[0, 2], [2, 0]]) + np.testing.assert_allclose( + xi[..., 0], + [[0.5, 4.0], [0.5, 0.5]], + rtol=1e-15, ) + np.testing.assert_array_equal(inside, [[True, False], [True, True]]) diff --git a/regridding/_weights/_weights_multilinear/_weights_multilinear.py b/regridding/_weights/_weights_multilinear/_weights_multilinear.py index 48b54df..7200b45 100644 --- a/regridding/_weights/_weights_multilinear/_weights_multilinear.py +++ b/regridding/_weights/_weights_multilinear/_weights_multilinear.py @@ -3,7 +3,7 @@ of any number of dimensions at any set of points. """ -from typing import Callable, Literal, Sequence +from typing import Literal, Sequence import numpy as np import numba from regridding import _util @@ -14,87 +14,69 @@ ] -def _index( - index: tuple[int, ...], - axis: tuple[int, ...], -) -> tuple[int | slice, ...]: - """ - The index of one orthogonal element of an array, which takes the element - `index` along the orthogonal axes and the whole of the resampled axes. - - Parameters - ---------- - index - The index of the element along the orthogonal axes. - axis - The resampled axes, as negative numbers. - """ - result: list[int | slice] = list(reversed(index)) - for ax in sorted(axis, reverse=True): - result.insert(~ax, slice(None)) - return tuple(reversed(result)) - - -def _locator( - coordinates: tuple[np.ndarray, ...], +def _locate_curvilinear( + grids: tuple[np.ndarray, ...], + points: np.ndarray, + extrapolate: bool, num_chunk: int, -) -> Callable[[np.ndarray, bool], tuple[np.ndarray, np.ndarray, np.ndarray]]: +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """ - Prepare to find the cells of a grid which contain a set of points. + Find the cell of a curvilinear grid which contains each point, for each + element. - A rectilinear grid is searched along each axis on its own, and any other - grid through a hierarchy of bounding boxes, which is built here. The - two give the same cells, but the first is exact and needs nothing built. - - Returns a function of the points, of shape ``(N, number of points)``, - and of whether to extrapolate, which returns the cell of each point, its - local coordinates in that cell, and whether it is inside the grid. + Each grid is searched through its own hierarchy of bounding boxes. If + every element shares one grid, the points of every element are searched + for together. Parameters ---------- - coordinates - The coordinates of the vertices of the grid, each an array with an - axis for each axis of the grid. + grids + The coordinates of the vertices of the grid of each element, each an + array with an axis for the elements, of length one if every element + shares the grid, followed by an axis for each axis of the grid. + points + The points to find, of shape + ``(number of elements, N, number of points)``. + extrapolate + Whether to find the cell to extrapolate from for a point outside the + grid. num_chunk The number of pieces to divide the work into. """ - rectilinear = _rectilinear.profiles(coordinates) - - if rectilinear is not None: - profile, start, coordinate, descending = rectilinear - - def locate_rectilinear( - points: np.ndarray, - extrapolate: bool, - ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - # the search along each axis finds the cell to extrapolate from - # either way - del extrapolate - return _rectilinear.locate(profile, start, coordinate, descending, points) - - return locate_rectilinear - - shape = np.array(coordinates[0].shape, dtype=np.int64) - vertices = np.stack([np.ravel(c) for c in coordinates]).astype(np.float64) - lower, upper, start, shapes = _curvilinear.tree(vertices, shape, num_chunk) - - def locate_curvilinear( - points: np.ndarray, - extrapolate: bool, - ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - return _curvilinear.locate( + num_element, num_axis, num_point = points.shape + num_grid = grids[0].shape[0] + shape = np.array(grids[0].shape[1:], dtype=np.int64) + + cell = np.empty((num_element, num_point, num_axis), dtype=np.int64) + xi = np.empty((num_element, num_point, num_axis), dtype=np.float64) + inside = np.empty((num_element, num_point), dtype=np.bool_) + + for g in range(num_grid): + vertices = np.stack([np.ravel(c[g]) for c in grids], dtype=np.float64) + tree = _curvilinear.tree(vertices, shape, num_chunk) + if num_grid == 1: + index = slice(None) + num_row = num_element + points_g = points.transpose(1, 0, 2).reshape( + num_axis, num_element * num_point + ) + else: + index = slice(g, g + 1) + num_row = 1 + points_g = points[g] + cell_g, xi_g, inside_g = _curvilinear.locate( vertices, shape, - lower, - upper, - start, - shapes, - points, + *tree, + np.ascontiguousarray(points_g), extrapolate, num_chunk, ) + cell[index] = cell_g.reshape(num_row, num_point, num_axis) + xi[index] = xi_g.reshape(num_row, num_point, num_axis) + inside[index] = inside_g.reshape(num_row, num_point) - return locate_curvilinear + return cell, xi, inside def weights_multilinear( @@ -115,6 +97,13 @@ def weights_multilinear( axes, which do not have to match. Each point is interpolated from the corners of the cell which contains it, so it has :math:`2^N` weights. + A rectilinear grid is searched along each axis on its own, and any other + grid through a hierarchy of bounding boxes. The two find the same + cells, but the first is exact and needs nothing built. The elements + along the orthogonal axes are all searched in the same call, so that + many small grids, such as a spectrum for each pixel of an image, cost + little more than one large one. + Parameters ---------- coordinates_input @@ -124,7 +113,10 @@ def weights_multilinear( axis_input The logical axes of the input grid. axis_output - The logical axes of the points. + The logical axes of the points. If :obj:`None`, every axis of the + points, and there must be as many as in `axis_input`, which keeps a + forgotten `axis_output` from passing for points arranged in a + different number of axes. weights_input Weights to multiply each vertex of the input grid by. bounds @@ -159,78 +151,105 @@ def weights_multilinear( axis_output=axis_output, perturb=perturb, seed=seed, - match_axes=False, + match_axes=axis_output is None, ) shape_input_broadcasted = coordinates_input[0].shape - - shape = np.array( - [shape_input_broadcasted[ax] for ax in sorted(axis_input_)], - dtype=np.int64, + num_orthogonal = len(shape_orthogonal) + num_element = int(np.prod(shape_orthogonal, dtype=int)) + + # Move the resampled axes behind the orthogonal ones, in ascending order, + # which is the order of the flat indices of the weights. + source_input = sorted(axis_input_) + source_output = sorted(axis_output_) + destination_input = range(-len(source_input), 0) + destination_output = range(-len(source_output), 0) + coordinates_input = tuple( + np.moveaxis(c, source_input, destination_input) for c in coordinates_input + ) + coordinates_output = tuple( + np.moveaxis(c, source_output, destination_output) for c in coordinates_output ) + + shape = np.array(coordinates_input[0].shape[num_orthogonal:], dtype=np.int64) if np.any(shape < 2): raise ValueError( f"The input grid must have at least two vertices along each of " f"its resampled axes, got {tuple(shape)}" ) - if weights_input is not None: - weights_input = np.broadcast_to(weights_input, shape_input_broadcasted) + num_point = int(np.prod(coordinates_output[0].shape[num_orthogonal:], dtype=int)) + + result = np.empty(shape_orthogonal, dtype=object) + if num_element == 0: + return result, shape_input, shape_output # The input grid is usually broadcast across the orthogonal axes, so - # every element shares it and the work of preparing to search it is done - # once. A zero stride, or a single element, guarantees that. - axis_input_orthogonal = tuple( - ax for ax in range(-len(shape_input_broadcasted), 0) if ax not in axis_input_ - ) + # that every element shares it and it is searched once. A zero stride, + # or a single element, guarantees that. shared = all( - coordinate.strides[ax] == 0 or coordinate.shape[ax] == 1 - for coordinate in coordinates_input - for ax in axis_input_orthogonal + c.strides[ax] == 0 or c.shape[ax] == 1 + for c in coordinates_input + for ax in range(num_orthogonal) + ) + if shared: + grids = tuple(c[(0,) * num_orthogonal][np.newaxis] for c in coordinates_input) + else: + grids = tuple(np.reshape(c, (num_element, *shape)) for c in coordinates_input) + + points = np.stack( + [np.reshape(c, (num_element, num_point)) for c in coordinates_output], + axis=1, + dtype=np.float64, ) num_chunk = numba.get_num_threads() - result = np.empty(shape_orthogonal, dtype=object) - locate = None - - for index in np.ndindex(*shape_orthogonal): - - index_input = _index(index, axis_input_) - index_output = _index(index, axis_output_) - - if locate is None or not shared: - locate = _locator( - tuple(c[index_input] for c in coordinates_input), - num_chunk, - ) - - points = np.stack( - [np.ravel(c[index_output]) for c in coordinates_output], - ).astype(np.float64) + rectilinear = _rectilinear.profiles(grids) + if rectilinear is not None: + cell, xi, inside = _rectilinear.locate(*rectilinear, points) + else: + cell, xi, inside = _locate_curvilinear( + grids=grids, + points=points, + extrapolate=bounds == "extrapolate", + num_chunk=num_chunk, + ) - cell, xi, inside = locate(points, bounds == "extrapolate") + if bounds == "raise" and not np.all(inside): + raise ValueError( + f"{np.count_nonzero(~inside)} of the output points fall outside " + f"the input grid, and {bounds=}." + ) - if bounds == "raise" and not np.all(inside): - raise ValueError( - f"{np.count_nonzero(~inside)} of the output points fall outside " - f"the input grid, and {bounds=}." + if weights_input is None: + weights_input = np.empty((0, 0), dtype=np.float64) + else: + num_vertex = int(np.prod(shape)) + weights_input = np.broadcast_to(weights_input, shape_input_broadcasted) + weights_input = np.moveaxis(weights_input, source_input, destination_input) + if all( + weights_input.strides[ax] == 0 or weights_input.shape[ax] == 1 + for ax in range(num_orthogonal) + ): + weights_input = np.reshape( + weights_input[(0,) * num_orthogonal], (1, num_vertex) ) - - if weights_input is None: - weights_input_index = np.empty(0, dtype=np.float64) else: - weights_input_index = np.ravel(weights_input[index_input]).astype( - np.float64 - ) + weights_input = np.reshape(weights_input, (num_element, num_vertex)) + weights_input = np.ascontiguousarray(weights_input, dtype=np.float64) + + indices_input, indices_output, values = _cells.weights( + cell, + xi, + inside, + shape, + weights_input, + bounds == "nan", + ) - result[index] = _cells.weights( - cell, - xi, - inside, - shape, - weights_input_index, - bounds == "nan", - ) + flat = result.reshape(-1) + for e in range(num_element): + flat[e] = (indices_input[e], indices_output[e], values[e]) return result, shape_input, shape_output diff --git a/regridding/_weights/_weights_multilinear/_weights_multilinear_test.py b/regridding/_weights/_weights_multilinear/_weights_multilinear_test.py index 3809ccf..1447ffd 100644 --- a/regridding/_weights/_weights_multilinear/_weights_multilinear_test.py +++ b/regridding/_weights/_weights_multilinear/_weights_multilinear_test.py @@ -238,7 +238,7 @@ def test_matches_scipy(self, num_axis: int, bounds: str) -> None: grid = tuple(np.meshgrid(*axes, indexing="ij")) values = rng.normal(size=grid[0].shape) points = rng.uniform(-1.3, 1.3, (num_axis, 30)) - result = _regrid(grid, values, tuple(points), bounds=bounds) + result = _regrid(grid, values, tuple(points), axis_output=0, bounds=bounds) expected = scipy.interpolate.RegularGridInterpolator( points=axes, values=values, @@ -254,6 +254,50 @@ def test_descending_1d(self) -> None: result = _regrid((x,), x**2, (np.array([0.25, 1.0, 2.5, 4.0]),)) np.testing.assert_allclose(result, [0.125, 1.5, 6.5, 14.0]) + @pytest.mark.parametrize("bounds", ["extrapolate", "nan"]) + def test_descending_unsigned(self, bounds: str) -> None: + # the difference of two unsigned integers wraps around instead of + # going negative, which once took this grid for an ascending one + x = np.array([5, 4, 3, 2, 1], dtype=np.uint8) + result = _regrid( + (x,), + x.astype(float) ** 2, + (np.array([4.5, 2.5, 1.5]),), + bounds=bounds, + ) + np.testing.assert_allclose(result, [20.5, 6.5, 2.5]) + + @pytest.mark.parametrize("bounds", ["extrapolate", "nan"]) + def test_grid_per_element(self, bounds: str) -> None: + # a spectrum for each of many pixels, each on its own grid, some + # running backwards, which are all searched in one call + rng = np.random.default_rng(5) + num = 50 + x = np.sort(rng.uniform(0, 1, (num, 8)), axis=~0) + x[::3] = x[::3, ::-1] + values = rng.normal(size=(num, 8)) + points = rng.uniform(-0.1, 1.1, 20) + result = _regrid( + (x,), + values, + (points,), + axis_input=~0, + axis_output=~0, + bounds=bounds, + ) + expected = np.stack( + [ + scipy.interpolate.RegularGridInterpolator( + points=(x[i],), + values=values[i], + bounds_error=False, + fill_value=None if bounds == "extrapolate" else np.nan, + )(points[:, np.newaxis]) + for i in range(num) + ] + ) + np.testing.assert_allclose(result, expected, rtol=0, atol=1e-13) + @pytest.mark.parametrize("name", list(grids_curvilinear)) class TestCurvilinear: @@ -261,13 +305,13 @@ class TestCurvilinear: def test_affine(self, name: str) -> None: grid = grids_curvilinear[name] points = _points(grid) - result = _regrid(grid, _affine(grid), points) + result = _regrid(grid, _affine(grid), points, axis_output=0) np.testing.assert_allclose(result, _affine(points), rtol=0, atol=1e-12) def test_nan(self, name: str) -> None: grid = grids_curvilinear[name] points = _points(grid) - result = _regrid(grid, _affine(grid), points, bounds="nan") + result = _regrid(grid, _affine(grid), points, axis_output=0, bounds="nan") inside = np.isfinite(result) assert 0 < inside.sum() < inside.size np.testing.assert_allclose( @@ -280,19 +324,47 @@ def test_nan(self, name: str) -> None: def test_raise(self, name: str) -> None: grid = grids_curvilinear[name] with pytest.raises(ValueError, match="fall outside the input grid"): - _regrid(grid, _affine(grid), _points(grid), bounds="raise") + _regrid(grid, _affine(grid), _points(grid), axis_output=0, bounds="raise") - def test_raise_inside_grid(self, name: str) -> None: + def test_raise_vertices(self, name: str) -> None: + # every vertex of the grid, including those on its boundary, is + # inside it grid = grids_curvilinear[name] - points = tuple(c[1:-1, 1:-1] for c in grid) - result = _regrid(grid, _affine(grid), points, bounds="raise") - np.testing.assert_allclose(result, _affine(points), rtol=0, atol=1e-12) + result = _regrid(grid, _affine(grid), grid, bounds="raise") + np.testing.assert_allclose(result, _affine(grid), rtol=0, atol=1e-12) + + +def test_thin_cells() -> None: + # rotated cells much longer than they are wide, in which rounding along + # the length of a cell once left points well inside it outside the grid. + # The grid is not curved, so that the points made by the same map as its + # vertices are inside its cells. + aspect = 1e4 + a, b = np.meshgrid( + np.linspace(0, 1, 6), + np.linspace(0, 1 / aspect, 6), + indexing="ij", + ) + angle = 0.6 + grid = ( + a * np.cos(angle) - b * np.sin(angle), + a * np.sin(angle) + b * np.cos(angle), + ) + rng = np.random.default_rng(0) + pa = rng.uniform(0.05, 0.95, 200) + pb = rng.uniform(0.05, 0.95, 200) / aspect + points = ( + pa * np.cos(angle) - pb * np.sin(angle), + pa * np.sin(angle) + pb * np.cos(angle), + ) + result = _regrid(grid, _affine(grid), points, axis_output=0, bounds="raise") + np.testing.assert_allclose(result, _affine(points), rtol=0, atol=1e-12) class TestPoints: """ - Multilinear interpolation weights points, rather than cells, so the - points may be arranged in any number of axes. + Multilinear interpolation weights points, rather than cells, so given + explicitly, `axis_output` may have any number of axes. """ def test_grid_2d_points_1d(self) -> None: @@ -305,17 +377,55 @@ def test_grid_2d_points_1d(self) -> None: def test_grid_3d_points_2d(self) -> None: grid = _distorted_3d() points = tuple(c.reshape(10, 10) for c in _points(grid)) - result = _regrid(grid, _affine(grid), points) + result = _regrid(grid, _affine(grid), points, axis_output=(0, 1)) assert result.shape == (10, 10) np.testing.assert_allclose(result, _affine(points), rtol=0, atol=1e-12) def test_single_point(self) -> None: grid = _distorted() point = (np.array(0.1), np.array(-0.2)) - result = _regrid(grid, _affine(grid), point) + result = _regrid(grid, _affine(grid), point, axis_output=()) assert result.shape == () np.testing.assert_allclose(result, _affine(point), rtol=0, atol=1e-12) + @pytest.mark.parametrize("name", ["distorted", "rectilinear"]) + def test_no_points(self, name: str) -> None: + grid = _distorted() + if name == "rectilinear": + x = np.linspace(-1, 1, 5) + grid = tuple(np.meshgrid(x, x, indexing="ij")) + empty = (np.zeros(0), np.zeros(0)) + result = _regrid(grid, _affine(grid), empty, axis_output=0) + assert result.shape == (0,) + + def test_axis_output_must_match_if_not_given(self) -> None: + # without `axis_output`, every axis of the points is taken to be + # resampled, and there must be as many as the grid has, so that a + # forgotten `axis_output` is not taken for points arranged in fewer + # axes + grid = _distorted() + with pytest.raises(ValueError, match="number of axes in `axis_output`"): + _regrid(grid, _affine(grid), _points(grid)) + + def test_forgotten_axis_output(self) -> None: + # a batch of five grids along the second axis, with points for each, + # but no `axis_output`, which would otherwise interpolate every + # element at the points of every other one + x = np.linspace(-1, 1, 10) + y = np.linspace(-1, 1, 5) + values = x[:, np.newaxis] + y + points = np.linspace(-1, 1, 6)[:, np.newaxis] + 0 * y + with pytest.raises(ValueError, match="number of axes in `axis_output`"): + _regrid((x[:, np.newaxis],), values, (points,), axis_input=0) + result = _regrid( + (x[:, np.newaxis],), + values, + (points,), + axis_input=0, + axis_output=0, + ) + np.testing.assert_allclose(result, points + y, rtol=0, atol=1e-14) + def test_conservative_axes_must_match(self) -> None: grid = _distorted() with pytest.raises(ValueError, match="number of axes in `axis_output`"): @@ -346,6 +456,7 @@ def test_grid_per_element(self) -> None: tuple(c[i] for c in grid), values[i], tuple(c[i] for c in points), + axis_output=0, ) for i in range(2) ] @@ -376,22 +487,62 @@ def test_weights_input(self, name: str) -> None: weights = regridding.weights( coordinates_input=grid, coordinates_output=points, + axis_output=0, weights_input=weights_input, method="multilinear", ) - result = regridding.regrid_from_weights(*weights, values_input=values) - expected = _regrid(grid, values * weights_input, points) + result = regridding.regrid_from_weights( + *weights, + values_input=values, + axis_output=0, + ) + expected = _regrid(grid, values * weights_input, points, axis_output=0) + np.testing.assert_allclose(result, expected, rtol=1e-14, atol=1e-14) + + def test_weights_input_per_element(self) -> None: + # weights which differ from one element to the next, on a grid which + # every element shares + grid = tuple(np.broadcast_to(c, (3, *c.shape)) for c in _distorted()) + points = tuple(c[np.newaxis] for c in _points(_distorted())) + values = np.stack([_affine(_distorted())] * 3) + weights_input = np.random.default_rng(2).uniform(0.5, 1.5, values.shape) + weights = regridding.weights( + coordinates_input=grid, + coordinates_output=points, + axis_input=(1, 2), + axis_output=(1,), + weights_input=weights_input, + method="multilinear", + ) + result = regridding.regrid_from_weights( + *weights, + values_input=values, + axis_input=(1, 2), + axis_output=(1,), + ) + expected = np.stack( + [ + _regrid( + _distorted(), + values[i] * weights_input[i], + tuple(c[0] for c in points), + axis_output=0, + ) + for i in range(3) + ] + ) np.testing.assert_allclose(result, expected, rtol=1e-14, atol=1e-14) def test_units() -> None: grid = tuple(c * u.mm for c in _distorted()) points = tuple(c * u.m for c in _points(_distorted())) - result = _regrid(grid, _affine(_distorted()), points, bounds="nan") + result = _regrid(grid, _affine(_distorted()), points, axis_output=0, bounds="nan") expected = _regrid( _distorted(), _affine(_distorted()), tuple(1000 * c.value for c in points), + axis_output=0, bounds="nan", ) np.testing.assert_allclose(result, expected, rtol=0, atol=1e-12) @@ -406,14 +557,19 @@ def test_units() -> None: ) def test_points_not_finite(grid: tuple[np.ndarray, ...]) -> None: points = (np.array([0.1, np.nan, np.inf]), np.array([np.nan, 0.2, -0.3])) - result = _regrid(grid, _affine(grid), points) + result = _regrid(grid, _affine(grid), points, axis_output=0) assert np.all(np.isnan(result)) with pytest.raises(ValueError, match="fall outside the input grid"): - _regrid(grid, _affine(grid), points, bounds="raise") + _regrid(grid, _affine(grid), points, axis_output=0, bounds="raise") def test_too_few_vertices() -> None: x = np.linspace(0, 1, 4) grid = np.meshgrid(x, x[:1], indexing="ij") with pytest.raises(ValueError, match="at least two vertices"): - _regrid(tuple(grid), np.zeros((4, 1)), (np.array([0.5]), np.array([0.0]))) + _regrid( + tuple(grid), + np.zeros((4, 1)), + (np.array([0.5]), np.array([0.0])), + axis_output=0, + ) From 3b150469dc30fece098f33279a6b957afa523268 Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Wed, 21 Oct 2026 20:43:26 -0600 Subject: [PATCH 4/5] Name the axes in the axis-count error as plain integers Without `axis_output`, a scattered set of points now trips this error, so it reads better without NumPy's integer types in it. Also give scipy its options as a literal, which is what its type annotations ask for. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01BYjDL98znSud1yFh9chnkP --- regridding/_util.py | 6 ++++-- .../_weights/_weights_multilinear/_curvilinear_test.py | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/regridding/_util.py b/regridding/_util.py index 71396c4..3dec29d 100644 --- a/regridding/_util.py +++ b/regridding/_util.py @@ -186,8 +186,10 @@ def _normalize_input_output_coordinates( if match_axes and 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}" + f"The number of axes in `axis_output`, " + f"{tuple(int(ax) for ax in axis_output)}, " + f"must match the number of axes in `axis_input`, " + f"{tuple(int(ax) for ax in axis_input)}" ) if len(coordinates_input) != len(axis_input): diff --git a/regridding/_weights/_weights_multilinear/_curvilinear_test.py b/regridding/_weights/_weights_multilinear/_curvilinear_test.py index c18a6e0..78080d3 100644 --- a/regridding/_weights/_weights_multilinear/_curvilinear_test.py +++ b/regridding/_weights/_weights_multilinear/_curvilinear_test.py @@ -215,7 +215,7 @@ def _inside_brute( lambda xi: _multilinear(corners, xi) - point, x0=np.full(num_axis, 0.5), method="hybr", - options=dict(xtol=1e-14), + options={"xtol": 1e-14}, ) # judged by the residual, since asked for this tolerance the # solver reports failure once it cannot improve on a From 74934fe8d4bdc723e5336fafed9b5fcbadc8e1c5 Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Wed, 21 Oct 2026 20:46:43 -0600 Subject: [PATCH 5/5] Apply weights for no elements, and test Newton's halved steps `regrid_from_weights()` handed the kernel an empty typed list, which has no type to compile it for, when an orthogonal axis had a length of zero. With no elements the output is already all zeros, so the kernel is skipped. Newton's method halves a step which does not bring the position closer, as crossing from a narrow cell into a wide one does, and a test now does that. The exits which the test of the rounding error of the position now reaches first are marked as such. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01BYjDL98znSud1yFh9chnkP --- regridding/_regrid/_regrid_from_weights.py | 13 ++++++++----- .../_weights/_weights_multilinear/_cells.py | 10 ++++++---- .../_weights_multilinear/_cells_test.py | 18 ++++++++++++++++++ .../_weights_multilinear_test.py | 17 +++++++++++++++++ 4 files changed, 49 insertions(+), 9 deletions(-) diff --git a/regridding/_regrid/_regrid_from_weights.py b/regridding/_regrid/_regrid_from_weights.py index ffeed12..8a01b81 100644 --- a/regridding/_regrid/_regrid_from_weights.py +++ b/regridding/_regrid/_regrid_from_weights.py @@ -162,11 +162,14 @@ def regrid_from_weights( values_input = np.ascontiguousarray(values_input) values_output = np.ascontiguousarray(values_output) - _regrid_from_weights( - weights=weights_numba, - values_input=values_input, - values_output=values_output, - ) + # an empty typed list has no type for the kernel to be compiled for, and + # with no elements there is nothing to add to the zeros of the output + if flat.size: + _regrid_from_weights( + weights=weights_numba, + values_input=values_input, + values_output=values_output, + ) values_output = values_output.reshape(shape_output_tmp) diff --git a/regridding/_weights/_weights_multilinear/_cells.py b/regridding/_weights/_weights_multilinear/_cells.py index bd8a101..7e77ad1 100644 --- a/regridding/_weights/_weights_multilinear/_cells.py +++ b/regridding/_weights/_weights_multilinear/_cells.py @@ -364,7 +364,7 @@ def _move( shift = np.floor(xi[m]) low = -float(cell[m]) high = float(shape[m] - 2 - cell[m]) - if shift != shift: + if shift != shift: # pragma: nocover shift = 0.0 elif shift < low: shift = low @@ -541,16 +541,18 @@ def newton( break fraction = fraction / 2 - if not accepted: + if not accepted: # pragma: nocover # no point along the step is closer, which happens once the step - # is down to the rounding error of the position + # is down to the rounding error of the position, if the test of + # the distance above has not already stopped the iteration, or + # where the continuation of a cell stops short of the point return size <= tolerance * scale cell[:] = cell_trial xi[:] = xi_trial base = base_trial - if fraction * size <= tolerance * scale: + if fraction * size <= tolerance * scale: # pragma: nocover return True target(vertices, base, point, goal) diff --git a/regridding/_weights/_weights_multilinear/_cells_test.py b/regridding/_weights/_weights_multilinear/_cells_test.py index 0137ff2..fc53013 100644 --- a/regridding/_weights/_weights_multilinear/_cells_test.py +++ b/regridding/_weights/_weights_multilinear/_cells_test.py @@ -202,6 +202,24 @@ def test_clamp(self, num_axis: int) -> None: np.testing.assert_allclose(xi, 2.5, rtol=0, atol=1e-12) +def test_newton_halves_steps() -> None: + # crossing from a cell one wide, past one a thousandth wide, into one a + # hundred wide: the full step from the first cell lands far beyond the + # point in the continuation of the last, and has to be halved + x, y = np.meshgrid( + np.array([0.0, 1.0, 1.001, 100.0]), + np.array([0.0, 1.0]), + indexing="ij", + ) + vertices = np.stack([x.ravel(), y.ravel()]) + shape = np.array(x.shape, dtype=np.int64) + point = np.array([50.0, 0.5]) + converged, cell, xi = _newton(vertices, shape, point, (0, 0), clamp=True) + assert converged + np.testing.assert_array_equal(cell, [2, 0]) + np.testing.assert_allclose(xi, [(50 - 1.001) / 98.999, 0.5], rtol=1e-14) + + @pytest.mark.parametrize("num_axis", [2, 3]) @pytest.mark.parametrize("aspect", [1e4, 1e6]) def test_newton_thin_cell(num_axis: int, aspect: float) -> None: diff --git a/regridding/_weights/_weights_multilinear/_weights_multilinear_test.py b/regridding/_weights/_weights_multilinear/_weights_multilinear_test.py index 1447ffd..fe95013 100644 --- a/regridding/_weights/_weights_multilinear/_weights_multilinear_test.py +++ b/regridding/_weights/_weights_multilinear/_weights_multilinear_test.py @@ -476,6 +476,23 @@ def test_grid_shared(self) -> None: np.testing.assert_array_equal(result, expected) +def test_no_elements() -> None: + # an orthogonal axis of length zero gives no elements, and nothing to do + x = np.linspace(-1, 1, 5) + grid = tuple( + np.broadcast_to(c, (0, 5, 5)) for c in np.meshgrid(x, x, indexing="ij") + ) + points = (np.zeros((0, 7)), np.zeros((0, 7))) + result = _regrid( + grid, + np.zeros((0, 5, 5)), + points, + axis_input=(1, 2), + axis_output=(1,), + ) + assert result.shape == (0, 7) + + class TestWeightsInput: @pytest.mark.parametrize("name", list(grids_curvilinear))