Added support for 3D conservative resampling to regridding.weights().

regridding/_weights/_weights_conservative.py

@@ -6,6 +6,7 @@
 from regridding import _util
 from ._weights_conservative_1d import weights_conservative_1d
 from regridding._conservative_ramshaw import _conservative_ramshaw
+from ._weights_conservative_3d import weights_conservative_3d
 
 
 def _weights_conservative(
@@ -127,9 +128,25 @@ def _weights_conservative(
                     weights_input=weights_input_index,
                 )
 
+            elif len(axis_input) == 3:
+                x_input, y_input, z_input = coordinates_input
+                x_output, y_output, z_output = coordinates_output
+                weights[index] = weights_conservative_3d(
+                    grid_input=(
+                        x_input[index_vertices_input],
+                        y_input[index_vertices_input],
+                        z_input[index_vertices_input],
+                    ),
+                    grid_output=(
+                        x_output[index_vertices_output],
+                        y_output[index_vertices_output],
+                        z_output[index_vertices_output],
+                    ),
+                )
+
             else:  # pragma: nocover
                 raise NotImplementedError(
-                    "Regridding operations greater than 2D are not supported"
+                    "Regridding operations greater than 3D are not supported"
                 )
 
     return weights, shape_values_input, shape_values_output

regridding/_weights/_weights_conservative_3d/__init__.py

@@ -0,0 +1,9 @@
+"""
+A subpackage for performing 3D conservative resampling
+"""
+
+from ._weights_conservative_3d import weights_conservative_3d
+
+__all__ = [
+    "weights_conservative_3d",
+]

regridding/_weights/_weights_conservative_3d/_arrays.py

@@ -0,0 +1,147 @@
+"""
+3-dimensional array manipulation routines.
+"""
+
+import numpy as np
+import numba
+
+__all__ = [
+    "axis_x",
+    "axis_y",
+    "axis_z",
+    "axes",
+    "index_in_bounds",
+    "index_flat",
+    "index_3d",
+    "align_axis_right",
+]
+
+axis_x = 0
+axis_y = 1
+axis_z = 2
+
+axes = (axis_x, axis_y, axis_z)
+
+vector_unit = (
+    (1, 0, 0),
+    (0, 1, 0),
+    (0, 0, 1),
+)
+
+
+@numba.njit(cache=True)
+def index_in_bounds(
+    index: tuple[int, int, int],
+    shape: tuple[int, int, int],
+):
+    """
+    Check if a 3D index is within array bounds specified by `shape`.
+
+    Parameters
+    ----------
+    index
+        The 3D index to check.
+    shape
+        The shape of the array to use as the bounds.
+    """
+    i, j, k = index
+
+    if i < 0:
+        return False
+    if j < 0:
+        return False
+    if k < 0:
+        return False
+
+    nx, ny, nz = shape
+
+    if i >= nx:
+        return False
+    if j >= ny:
+        return False
+    if k >= nz:
+        return False
+
+    return True
+
+
+@numba.njit(cache=True)
+def index_flat(
+    index: tuple[int, int, int],
+    shape: tuple[int, int, int],
+) -> int:
+    """
+    Convert a 3D index to a flat index.
+
+    Parameters
+    ----------
+    index
+        A 3D index to convert.
+    shape
+        The sizes of each axis of the array.
+    """
+
+    i, j, k = index
+
+    num_x, num_y, num_z = shape
+
+    result = i * num_y * num_z + j * num_z + k
+
+    return result
+
+
+@numba.njit(cache=True)
+def index_3d(
+    index: int,
+    shape: tuple[int, int, int],
+) -> tuple[int, int, int]:
+    """
+    Convert a flat index to a 3D index.
+
+    Parameters
+    ----------
+    index
+        A flat index to convert.
+    shape
+        The sizes of each axis of the array.
+    """
+
+    num_x, num_y, num_z = shape
+
+    k = index % num_z
+
+    j = (index // num_z) % num_y
+
+    i = index // (num_z * num_y)
+
+    return i, j, k
+
+
+@numba.njit(cache=True)
+def align_axis_right(
+    a: np.ndarray,
+    axis: int,
+) -> np.ndarray:
+    """
+    Roll all the axes of a 3D array to the right until `axis` is the last axis.
+
+    This function is needed to permute the axes in such a way to retain
+    their right-handedness.
+
+    Parameters
+    ----------
+    a
+        The array to modify the axes of.
+    axis
+        The axis to set as the last axis.
+    """
+    axis = axis % a.ndim
+
+    if axis == axis_x:
+        result = a.transpose((axis_y, axis_z, axis_x))
+    elif axis == axis_y:
+        result = a.transpose((axis_z, axis_x, axis_y))
+    else:
+        result = a
+
+    return result

regridding/_weights/_weights_conservative_3d/_arrays_test.py

@@ -0,0 +1,85 @@
+import pytest
+import numpy as np
+from . import _arrays
+
+
+@pytest.mark.parametrize(
+    argnames="index,shape,result_expected",
+    argvalues=[
+        (
+            (0, 0, 0),
+            (11, 12, 13),
+            True,
+        ),
+        (
+            (5, 5, 5),
+            (11, 12, 13),
+            True,
+        ),
+        (
+            (-1, 0, 0),
+            (11, 12, 13),
+            False,
+        ),
+        (
+            (0, 12, 0),
+            (11, 12, 13),
+            False,
+        ),
+    ],
+)
+def test_index_in_bounds(
+    index: tuple[int, int, int],
+    shape: tuple[int, int, int],
+    result_expected: bool,
+):
+    result = _arrays.index_in_bounds(index=index, shape=shape)
+
+    assert result == result_expected
+
+
+@pytest.mark.parametrize(
+    argnames="index",
+    argvalues=[
+        (1, 1, 1),
+        (5, 6, 7),
+    ],
+)
+@pytest.mark.parametrize(
+    argnames="shape",
+    argvalues=[
+        (11, 12, 13),
+    ],
+)
+def test_index_flat(
+    index: tuple[int, int, int],
+    shape: tuple[int, int, int],
+):
+    result = _arrays.index_flat(index=index, shape=shape)
+    result_expected = np.ravel_multi_index(index, shape)
+
+    assert result == result_expected
+
+
+@pytest.mark.parametrize(
+    argnames="index",
+    argvalues=[0, 11, 25],
+)
+@pytest.mark.parametrize(
+    argnames="shape",
+    argvalues=[
+        (11, 12, 13),
+    ],
+)
+def test_index_3d(
+    index: int,
+    shape: tuple[int, int, int],
+):
+    result = _arrays.index_3d(index=index, shape=shape)
+    result_expected = np.unravel_index(index, shape)
+
+    print(f"{result=}")
+    print(f"{result_expected=}")
+
+    assert result == result_expected
+    assert index == np.ravel_multi_index(result, shape)