diff --git a/regridding/_weights/_weights_conservative.py b/regridding/_weights/_weights_conservative.py index 67fc9c1..d5dcee1 100644 --- a/regridding/_weights/_weights_conservative.py +++ b/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 diff --git a/regridding/_weights/_weights_conservative_3d/__init__.py b/regridding/_weights/_weights_conservative_3d/__init__.py new file mode 100644 index 0000000..6c5336f --- /dev/null +++ b/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", +] diff --git a/regridding/_weights/_weights_conservative_3d/_arrays.py b/regridding/_weights/_weights_conservative_3d/_arrays.py new file mode 100644 index 0000000..07be0fc --- /dev/null +++ b/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 diff --git a/regridding/_weights/_weights_conservative_3d/_arrays_test.py b/regridding/_weights/_weights_conservative_3d/_arrays_test.py new file mode 100644 index 0000000..ad5f761 --- /dev/null +++ b/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) diff --git a/regridding/_weights/_weights_conservative_3d/_grids.py b/regridding/_weights/_weights_conservative_3d/_grids.py new file mode 100644 index 0000000..a0297cb --- /dev/null +++ b/regridding/_weights/_weights_conservative_3d/_grids.py @@ -0,0 +1,477 @@ +""" +Utilities for inspecting and searching logically-rectangular grids of +coordinates. +""" + +import sys +import numpy as np +import numba +import regridding as rg +from . import _arrays + +__all__ = [ + "shape_centers", + "cell_axes", + "cell_normals", + "cell_boundary", + "grid_boundary", + "index_of_point_brute", +] + + +@numba.njit(cache=True) +def shape_centers( + shape: tuple[int, int, int], +) -> tuple[int, int, int]: + """ + Given the shape of the grid of cell vertices, + compute the shape of the grid of cell centers. + + Parameters + ---------- + shape + The shape of the grid of cell vertices. + """ + + nx, ny, nz = shape + + return nx - 1, ny - 1, nz - 1 + + +@numba.njit(cache=True) +def grid_volume( + grid: tuple[np.ndarray, np.ndarray, np.ndarray], +) -> np.ndarray: + """ + Compute the volume of each cell in a logically-rectangular grid. + + Parameters + ---------- + grid + A 3D grid of cell vertices. + """ + + x, y, z = grid + + num_i, num_j, num_k = x.shape + + result = np.zeros((num_i - 1, num_j - 1, num_k - 1)) + + for axis in (0, 1, 2): + _grid_volume_sweep( + grid=grid, + out=result, + axis=axis, + ) + + return result + + +@numba.njit(cache=True) +def _grid_volume_sweep( + grid: tuple[np.ndarray, np.ndarray, np.ndarray], + out: np.ndarray, + axis: int, +) -> None: + """ + Compute the volume contribution of this axis. + + Parameters + ---------- + grid + A 3D grid of cell vertices. + out + An output to array to store the result. + axis + The axis along which to iterate. + """ + + x, y, z = grid + + x = _arrays.align_axis_right(x, axis) + y = _arrays.align_axis_right(y, axis) + z = _arrays.align_axis_right(z, axis) + + out = _arrays.align_axis_right(out, axis) + + num_i, num_j, num_k = x.shape + + for i in numba.prange(num_i): + + i = numba.types.int64(i) + + for j in range(num_j): + + for k in range(num_k - 1): + + i_left = i - 1 + i_right = i + + j_lower = j - 1 + j_upper = j + + k1 = k + k2 = k + 1 + + v1 = x[i, j, k1], y[i, j, k1], z[i, j, k1] + v2 = x[i, j, k2], y[i, j, k2], z[i, j, k2] + + if i_left >= 0: + + v0_left = x[i_left, j, k], y[i_left, j, k], z[i_left, j, k] + volume_left = rg.geometry.volume_tetrahedron(v0_left, v1, v2) + + if j_lower >= 0: + out[i_left, j_lower, k] -= volume_left + if j_upper < (num_j - 1): + out[i_left, j_upper, k] += volume_left + + if j_lower >= 0: + + v0_lower = x[i, j_lower, k], y[i, j_lower, k], z[i, j_lower, k] + volume_lower = -rg.geometry.volume_tetrahedron(v0_lower, v1, v2) + + if i_left >= 0: + out[i_left, j_lower, k] -= volume_lower + if i_right < (num_i - 1): + out[i_right, j_lower, k] += volume_lower + + +cell_axes = ( + 2, + 2, + 2, + 2, + 0, + 0, + 0, + 0, + 1, + 1, + 1, + 1, +) +""" +The index of the axis normal to each face in :func:`cell_boundary`. +""" + +cell_normals = ( + (0, 0, -1), + (0, 0, -1), + (0, 0, +1), + (0, 0, +1), + (-1, 0, 0), + (-1, 0, 0), + (+1, 0, 0), + (+1, 0, 0), + (0, -1, 0), + (0, -1, 0), + (0, +1, 0), + (0, +1, 0), +) +""" +Vectors normal to each face in :func:`cell_boundary`. +""" + + +@numba.njit(cache=True) +def cell_boundary( + index: tuple[int, int, int], + grid: tuple[np.ndarray, np.ndarray, np.ndarray], +) -> numba.typed.List[ + tuple[ + tuple[float, float, float], + tuple[float, float, float], + tuple[float, float, float], + ], +]: + """ + For a given cell in a grid of cell vertices, + compute the 12 triangles composing the boundary. + + Parameters + ---------- + index + The index of the cell to compute the boundary of. + grid + A grid of cell vertices. + + """ + i, j, k = index + + x, y, z = grid + + i_000 = i + 0, j + 0, k + 0 + i_001 = i + 0, j + 0, k + 1 + i_010 = i + 0, j + 1, k + 0 + i_011 = i + 0, j + 1, k + 1 + i_100 = i + 1, j + 0, k + 0 + i_101 = i + 1, j + 0, k + 1 + i_110 = i + 1, j + 1, k + 0 + i_111 = i + 1, j + 1, k + 1 + + indices_triangles = ( + (i_000, i_010, i_110), + (i_110, i_100, i_000), + (i_001, i_101, i_111), + (i_111, i_011, i_001), + (i_000, i_001, i_011), + (i_011, i_010, i_000), + (i_100, i_110, i_111), + (i_111, i_101, i_100), + (i_000, i_100, i_101), + (i_101, i_001, i_000), + (i_010, i_011, i_111), + (i_111, i_110, i_010), + ) + + polyhedron = numba.typed.List() + + for t in range(len(indices_triangles)): + + i0, i1, i2 = indices_triangles[t] + + v0 = x[i0], y[i0], z[i0] + v1 = x[i1], y[i1], z[i1] + v2 = x[i2], y[i2], z[i2] + + polyhedron.append((v0, v1, v2)) + + return polyhedron + + +@numba.njit(cache=True) +def grid_boundary( + grid: tuple[np.ndarray, np.ndarray, np.ndarray], +) -> tuple[ + numba.typed.List[tuple[int, int, int]], + numba.typed.List[ + tuple[ + tuple[float, float, float], + tuple[float, float, float], + tuple[float, float, float], + ], + ], +]: + """ + For a given grid, find the indices of the boundary cells + and the vertices of all the triangles on the boundary. + + find the set of indices that expresses the boundary + as a sequence of triangles. + + Parameters + ---------- + grid + A logically-rectangular grid of cell vertices. + """ + + x, y, z = grid + + shape_x, shape_y, shape_z = x.shape + + indices = numba.typed.List() + + triangles = numba.typed.List() + + for i0 in range(shape_x - 1): + for j0 in range(shape_y - 1): + + k0 = 0 + + i1 = i0 + 1 + j1 = j0 + 1 + + i_cell = i0, j0, k0 + + i_000 = i0, j0, k0 + i_010 = i0, j1, k0 + i_110 = i1, j1, k0 + i_100 = i1, j0, k0 + + v_000 = x[i_000], y[i_000], z[i_000] + v_010 = x[i_010], y[i_010], z[i_010] + v_110 = x[i_110], y[i_110], z[i_110] + v_100 = x[i_100], y[i_100], z[i_100] + + indices.append(i_cell) + indices.append(i_cell) + + triangles.append((v_000, v_010, v_110)) + triangles.append((v_110, v_100, v_000)) + + for i0 in range(shape_x - 1): + for j0 in range(shape_y - 1): + + k1 = shape_z - 1 + + i1 = i0 + 1 + j1 = j0 + 1 + + i_cell = i0, j0, k1 - 1 + + i_001 = i0, j0, k1 + i_101 = i1, j0, k1 + i_111 = i1, j1, k1 + i_011 = i0, j1, k1 + + v_001 = x[i_001], y[i_001], z[i_001] + v_101 = x[i_101], y[i_101], z[i_101] + v_111 = x[i_111], y[i_111], z[i_111] + v_011 = x[i_011], y[i_011], z[i_011] + + indices.append(i_cell) + indices.append(i_cell) + + triangles.append((v_001, v_101, v_111)) + triangles.append((v_111, v_011, v_001)) + + for j0 in range(shape_y - 1): + for k0 in range(shape_z - 1): + + i0 = 0 + + j1 = j0 + 1 + k1 = k0 + 1 + + i_cell = i0, j0, k0 + + i_000 = i0, j0, k0 + i_001 = i0, j0, k1 + i_011 = i0, j1, k1 + i_010 = i0, j1, k0 + + v_000 = x[i_000], y[i_000], z[i_000] + v_001 = x[i_001], y[i_001], z[i_001] + v_011 = x[i_011], y[i_011], z[i_011] + v_010 = x[i_010], y[i_010], z[i_010] + + indices.append(i_cell) + indices.append(i_cell) + + triangles.append((v_000, v_001, v_011)) + triangles.append((v_011, v_010, v_000)) + + for j0 in range(shape_y - 1): + for k0 in range(shape_z - 1): + + i1 = shape_x - 1 + + j1 = j0 + 1 + k1 = k0 + 1 + + i_cell = i1 - 1, j0, k0 + + i_100 = i1, j0, k0 + i_110 = i1, j1, k0 + i_111 = i1, j1, k1 + i_101 = i1, j0, k1 + + v_100 = x[i_100], y[i_100], z[i_100] + v_110 = x[i_110], y[i_110], z[i_110] + v_111 = x[i_111], y[i_111], z[i_111] + v_101 = x[i_101], y[i_101], z[i_101] + + indices.append(i_cell) + indices.append(i_cell) + + triangles.append((v_100, v_110, v_111)) + triangles.append((v_111, v_101, v_100)) + + for i0 in range(shape_x - 1): + for k0 in range(shape_z - 1): + + j0 = 0 + + i1 = i0 + 1 + k1 = k0 + 1 + + i_cell = i0, j0, k0 + + i_000 = i0, j0, k0 + i_100 = i1, j0, k0 + i_101 = i1, j0, k1 + i_001 = i0, j0, k1 + + v_000 = x[i_000], y[i_000], z[i_000] + v_100 = x[i_100], y[i_100], z[i_100] + v_101 = x[i_101], y[i_101], z[i_101] + v_001 = x[i_001], y[i_001], z[i_001] + + indices.append(i_cell) + indices.append(i_cell) + + triangles.append((v_000, v_100, v_101)) + triangles.append((v_101, v_001, v_000)) + + for i0 in range(shape_x - 1): + for k0 in range(shape_z - 1): + + j1 = shape_y - 1 + + i1 = i0 + 1 + k1 = k0 + 1 + + i_cell = i0, j1 - 1, k0 + + i_010 = i0, j1, k0 + i_011 = i0, j1, k1 + i_111 = i1, j1, k1 + i_110 = i1, j1, k0 + + v_010 = x[i_010], y[i_010], z[i_010] + v_011 = x[i_011], y[i_011], z[i_011] + v_111 = x[i_111], y[i_111], z[i_111] + v_110 = x[i_110], y[i_110], z[i_110] + + indices.append(i_cell) + indices.append(i_cell) + + triangles.append((v_010, v_011, v_111)) + triangles.append((v_111, v_110, v_010)) + + return indices, triangles + + +@numba.njit(cache=True) +def index_of_point_brute( + point: tuple[float, float, float], + grid: tuple[np.ndarray, np.ndarray, np.ndarray], +) -> tuple[int, int, int]: + """ + Find the index of the cell in the grid which contains the given point. + + This function uses brute force to search, + but this could be improved significantly by using the secant method + or possibly the bisection method. + + Parameters + ---------- + point + The query point. + grid + A logically-rectangular grid of cell vertices. + """ + + x, y, z = grid + + shape_x, shape_y, shape_z = x.shape + + for i in range(shape_x - 1): + for j in range(shape_y - 1): + for k in range(shape_z - 1): + + index = i, j, k + + polyhedron = cell_boundary( + index=index, + grid=grid, + ) + + if rg.geometry.point_is_inside_polyhedron( + point=point, + polyhedron=polyhedron, + ): + return index + + return sys.maxsize, sys.maxsize, sys.maxsize diff --git a/regridding/_weights/_weights_conservative_3d/_grids_test.py b/regridding/_weights/_weights_conservative_3d/_grids_test.py new file mode 100644 index 0000000..b09d81b --- /dev/null +++ b/regridding/_weights/_weights_conservative_3d/_grids_test.py @@ -0,0 +1,82 @@ +import pytest +import numpy as np +from . import _grids + + +@pytest.mark.parametrize( + argnames="grid,result_expected", + argvalues=[ + ( + np.meshgrid( + np.arange(2), + np.arange(2), + np.arange(2), + indexing="ij", + ), + np.ones((1, 1, 1)), + ), + ( + np.meshgrid( + np.arange(3), + np.arange(4), + np.arange(5), + indexing="ij", + ), + np.ones((2, 3, 4)), + ), + ( + np.meshgrid( + 2 * np.arange(3), + 2 * np.arange(4), + 2 * np.arange(5), + indexing="ij", + ), + 8 * np.ones((2, 3, 4)), + ), + ], +) +def test_volume_grid( + grid: tuple[np.ndarray, np.ndarray, np.ndarray], + result_expected: np.ndarray, +): + result = _grids.grid_volume(grid) + assert np.allclose(result, result_expected) + assert result.shape == result_expected.shape + + +@pytest.mark.parametrize( + argnames="point,grid,result_expected", + argvalues=[ + ( + (0.5, 0.5, 0.5), + np.meshgrid( + np.arange(3), + np.arange(4), + np.arange(5), + indexing="ij", + ), + (0, 0, 0), + ), + ( + (-0.5, -0.5, 3.5), + np.meshgrid( + -np.arange(3), + -np.arange(4), + np.arange(5), + indexing="ij", + ), + (0, 0, 3), + ), + ], +) +def test_index_of_point_brute( + point: tuple[float, float, float], + grid: tuple[np.ndarray, np.ndarray, np.ndarray], + result_expected: np.ndarray, +): + result = _grids.index_of_point_brute( + point=point, + grid=grid, + ) + + assert result == result_expected diff --git a/regridding/_weights/_weights_conservative_3d/_intercepts.py b/regridding/_weights/_weights_conservative_3d/_intercepts.py new file mode 100644 index 0000000..0e5e69f --- /dev/null +++ b/regridding/_weights/_weights_conservative_3d/_intercepts.py @@ -0,0 +1,374 @@ +""" +Create lists of intercepts generated during axis sweeps and determine their +contribution to the final list of weights. +""" + +import numpy as np +import numba +import regridding as rg +from . import _arrays +from . import _grids + + +@numba.njit(cache=True) +def empty( + shape_input: tuple[int, int, int], + shape_output: tuple[int, int, int], +) -> numba.typed.List: + """ + Create an empty list of intercepts between each plane in the + input grid and each plane in the output grid. + + There are three planes in the input grid and thre + + Parameters + ---------- + shape_input + The shape of the input grid. + shape_output + The shape of the output grid. + """ + intercepts = numba.typed.List() + for a in range(len(shape_input)): + intercepts_a = numba.typed.List() + for b in range(len(shape_output)): + intercepts_ab = numba.typed.List() + for i in range(shape_input[a]): + intercepts_abi = numba.typed.List() + for j in range(shape_output[b]): + intercepts_abij = numba.typed.List() + for x in range(0): + intercepts_abij.append((0, 0, (0.0, 0.0, 0.0))) + intercepts_abi.append(intercepts_abij) + intercepts_ab.append(intercepts_abi) + intercepts_a.append(intercepts_ab) + intercepts.append(intercepts_a) + return intercepts + + +@numba.njit(cache=True) +def insert_intercept( + intercepts: numba.typed.List[tuple[int, int, tuple[float, float, float]]], + intercept_new: tuple[int, int, tuple[float, float, float]], +) -> None: + """ + Insert a new intercept into the current list of intercepts at the correct + point so that the list maintains sorted order. + + Parameters + ---------- + intercepts + The current list of intercepts. + intercept_new + A new intercept to insert into the list. + """ + index = _bisect_intercepts(intercepts, intercept_new) + + intercepts.insert(index, intercept_new) + + +@numba.njit(cache=True) +def _bisect_intercepts( + intercepts: numba.typed.List[tuple[int, int, tuple[float, float, float]]], + intercept_new: tuple[int, int, tuple[float, float, float]], +) -> int: + """ + Given an ordered sequence of intercepts, + find the index for which a new intercept should be inserted to maintain + the ordering + + Parameters + ---------- + intercepts + The current list of intercepts. + intercept_new + A new intercept to insert into the list. + """ + + num_intercepts = len(intercepts) + + if num_intercepts < 2: + return num_intercepts + + _, _, intercept_new = intercept_new + + index_left = 0 + index_right = num_intercepts - 1 + + _, _, intercept_left = intercepts[index_left] + _, _, intercept_right = intercepts[index_right] + + t_start = _line_point_closest_approach_parameter( + line=(intercept_left, intercept_right), + point=intercept_new, + ) + + if t_start < 0: + return 0 + + if t_start >= 1: + return num_intercepts + + while (index_right - index_left) > 1: + + index_middle = (index_left + index_right) // 2 + + _, _, intercept_middle = intercepts[index_middle] + + t_left = _line_point_closest_approach_parameter( + line=(intercept_left, intercept_middle), + point=intercept_new, + ) + + if 0 <= t_left < 1: + index_right = index_middle + intercept_right = intercept_middle + + else: + index_left = index_middle + intercept_left = intercept_middle + + return index_right + + +@numba.njit(cache=True) +def _line_point_closest_approach_parameter( + line: tuple[ + tuple[float, float, float], + tuple[float, float, float], + ], + point: tuple[float, float, float], +) -> float: + """ + Compute the parameter describing the point of closest approach of line + segment to a point. + + Based on a `StackExchange answer by John Hughes `_. + + Parameters + ---------- + line + A line segment in 3D. + point + A point in 3D + """ + + a, b = line + p = point + + v = rg.math.difference_3d(b, a) + u = rg.math.difference_3d(a, p) + + numerator = -rg.math.dot_3d(v, u) + denominator = rg.math.dot_3d(v, v) + + return numerator / denominator + + +def _line_point_closest_approach( + line: tuple[ + tuple[float, float, float], + tuple[float, float, float], + ], + t: float, +) -> tuple[float, float, float]: + """ + Compute the point of closest approach of line segment to a point using + the parameter computed using :func:`_line_point_closest_approach_parameter`. + + Based on a `StackExchange answer by John Hughes `_. + + Parameters + ---------- + line + A line segment in 3D. + t + The parameter of closest approach. + """ + a, b = line + + s = 1 - t + + a2 = rg.math.multiply_3d(s, a) + b2 = rg.math.multiply_3d(t, b) + + return rg.math.sum_3d(a2, b2) + + +@numba.njit(cache=True) +def sweep( + intercepts: numba.typed.List, + weights: numba.typed.List, + grid_input: tuple[np.ndarray, np.ndarray, np.ndarray], + grid_output: tuple[np.ndarray, np.ndarray, np.ndarray], + volume_input: np.ndarray, +) -> None: + """ + Sweep through the list of intercepts and add the volume of each + segment to the list of weights. + + Parameters + ---------- + intercepts + A sequence of intercepts and the corresponding indices in the + input/output grids. + weights + The current list of weights. + New weights will be appended to this list. + grid_input + The vertices of the input grid. + grid_output + The vertices of the output grid. + volume_input + The volume of each cell in the input grid. + """ + + x_input, y_input, z_input = grid_input + x_output, y_output, z_output = grid_output + + shape_input = x_input.shape + shape_output = x_output.shape + + shape_centers_input = _grids.shape_centers(shape_input) + shape_centers_output = _grids.shape_centers(shape_output) + + for a in range(len(intercepts)): + + intercepts_a = intercepts[a] + + offset_a = [1 if axis == a else 0 for axis in range(len(intercepts))] + + for b in range(len(intercepts_a)): + + intercepts_ab = intercepts_a[b] + + offset_b = [1 if axis == b else 0 for axis in range(len(intercepts))] + + for i in range(len(intercepts_ab)): + + intercepts_abi = intercepts_ab[i] + + for j in range(len(intercepts_abi)): + + intercepts_abij = intercepts_abi[j] + + if not intercepts_abij: + continue + + i0_input, i0_output, p0 = intercepts_abij[0] + + i0_input = _arrays.index_3d(i0_input, shape_input) + i0_output = _arrays.index_3d(i0_output, shape_output) + + for t in range(1, len(intercepts_abij)): + + i1_input, i1_output, p1 = intercepts_abij[t] + + i1_input = _arrays.index_3d(i1_input, shape_input) + i1_output = _arrays.index_3d(i1_output, shape_output) + + i0_input_lower = rg.math.difference_3d(i0_input, offset_a) + i0_input_upper = i0_input + + i0_output_lower = rg.math.difference_3d(i0_output, offset_b) + i0_output_upper = i0_output + + apex_input = ( + x_input[i0_input], + y_input[i0_input], + z_input[i0_input], + ) + apex_output = ( + x_output[i0_output], + y_output[i0_output], + z_output[i0_output], + ) + + vol_input = rg.geometry.volume_tetrahedron( + apex_input, + p0, + p1, + ) + vol_output = -rg.geometry.volume_tetrahedron( + apex_output, + p0, + p1, + ) + + input_lower_in_bounds = _arrays.index_in_bounds( + index=i0_input_lower, + shape=shape_centers_input, + ) + input_upper_in_bounds = _arrays.index_in_bounds( + index=i0_input_upper, + shape=shape_centers_input, + ) + output_lower_in_bounds = _arrays.index_in_bounds( + index=i0_output_lower, + shape=shape_centers_output, + ) + output_upper_in_bounds = _arrays.index_in_bounds( + index=i0_output_upper, + shape=shape_centers_output, + ) + + n0_input_lower = _arrays.index_flat( + index=i0_input_lower, + shape=shape_centers_input, + ) + n0_input_upper = _arrays.index_flat( + index=i0_input_upper, + shape=shape_centers_input, + ) + n0_output_lower = _arrays.index_flat( + index=i0_output_lower, + shape=shape_centers_output, + ) + n0_output_upper = _arrays.index_flat( + index=i0_output_upper, + shape=shape_centers_output, + ) + + if input_lower_in_bounds: + + volume_input_lower = volume_input[i0_input_lower] + + if output_lower_in_bounds: + weight_lower_lower = ( + n0_input_lower, + n0_output_lower, + (vol_input + vol_output) / volume_input_lower, + ) + weights.append(weight_lower_lower) + + if output_upper_in_bounds: + weight_lower_upper = ( + n0_input_lower, + n0_output_upper, + (-vol_input - vol_output) / volume_input_lower, + ) + weights.append(weight_lower_upper) + + if input_upper_in_bounds: + + volume_input_upper = volume_input[i0_input_upper] + + if output_lower_in_bounds: + weight_upper_lower = ( + n0_input_upper, + n0_output_lower, + (-vol_input - vol_output) / volume_input_upper, + ) + weights.append(weight_upper_lower) + + if output_upper_in_bounds: + weight_upper_upper = ( + n0_input_upper, + n0_output_upper, + (vol_input + vol_output) / volume_input_upper, + ) + weights.append(weight_upper_upper) + + p0 = p1 + i0_input = i1_input + i0_output = i1_output diff --git a/regridding/_weights/_weights_conservative_3d/_intercepts_test.py b/regridding/_weights/_weights_conservative_3d/_intercepts_test.py new file mode 100644 index 0000000..1f293ac --- /dev/null +++ b/regridding/_weights/_weights_conservative_3d/_intercepts_test.py @@ -0,0 +1,209 @@ +import pytest +import numpy as np +import numba +from . import _intercepts + + +intercepts1 = numba.typed.List( + [ + (0, 0, (0, 0, 0)), + (0, 1, (1, 1, 1)), + (0, 2, (2, 2, 2)), + ] +) + +intercepts2 = numba.typed.List( + [ + (0, 0, (0, 0, 0)), + (0, 1, (-1, 1, -1)), + (0, 2, (-2, 2, -2)), + ] +) + + +@pytest.mark.parametrize( + argnames="intercepts,intercept_new", + argvalues=[ + ( + numba.typed.List( + [ + (0, 0, (0, 0, 0)), + (0, 1, (-1, 1, -1)), + (0, 2, (-2, 2, -2)), + ] + ), + (0, 0, (-1, -1, -1)), + ), + ( + _intercepts.empty((11, 11, 11), (11, 11, 11))[0][0][0][0], + (0, 0, (-1, -1, -1)), + ), + ], +) +def test_insert_intercept( + intercepts: numba.typed.List[tuple[int, int, tuple[float, float, float]]], + intercept_new: tuple[int, int, tuple[float, float, float]], +): + length_old = len(intercepts) + + _intercepts.insert_intercept(intercepts, intercept_new) + + assert len(intercepts) == length_old + 1 + + +@pytest.mark.parametrize( + argnames="intercepts,intercept_new,result_expected", + argvalues=[ + ( + numba.typed.List([(0, 0, (0.9, 0, 0)), (0, 1, (1.2, 0, 0))]), + (0, 2, (1.1, 0, 0)), + 1, + ), + ( + numba.typed.List([(0, 0, (0, 0.9, 0)), (0, 1, (0, 1.2, 0))]), + (0, 2, (0, 1.1, 0)), + 1, + ), + ( + numba.typed.List([(0, 0, (0, 0, 0.9)), (0, 1, (0, 0, 1.2))]), + (0, 2, (0, 0, 1.1)), + 1, + ), + ( + intercepts1, + (0, 0, (-1, -1, -1)), + 0, + ), + ( + intercepts1, + (0, 0, (0.5, 0.25, 0.5)), + 1, + ), + ( + intercepts1, + (0, 0, (1.5, 1.25, 1.75)), + 2, + ), + ( + intercepts1, + (0, 0, (2.5, 2.25, 2.75)), + 3, + ), + ( + intercepts2, + (0, 0, (1, -1, 1)), + 0, + ), + ( + intercepts2, + (0, 0, (-0.5, 0.25, -0.5)), + 1, + ), + ( + intercepts2, + (0, 0, (-1.5, 1.25, -1.75)), + 2, + ), + ( + intercepts2, + (0, 0, (-2.5, 2.25, -2.75)), + 3, + ), + ], +) +def test_bisect_intercepts( + intercepts: numba.typed.List[tuple[int, int, tuple[float, float, float]]], + intercept_new: tuple[int, int, tuple[float, float, float]], + result_expected: int, +): + result = _intercepts._bisect_intercepts( + intercepts=intercepts, + intercept_new=intercept_new, + ) + + assert result == result_expected + + +@pytest.mark.parametrize( + argnames="line, point, result_expected", + argvalues=[ + ( + ((-1, -1, 0), (1, 1, 0)), + (-1, 1, 0), + (0, 0, 0), + ) + ], +) +def test_line_point_closest_approach( + line: tuple[ + tuple[float, float, float], + tuple[float, float, float], + ], + point: tuple[float, float, float], + result_expected: tuple[float, float, float], +): + t = _intercepts._line_point_closest_approach_parameter(line, point) + + result = _intercepts._line_point_closest_approach(line, t) + + assert result == result_expected + + +@pytest.mark.parametrize( + argnames="grid_input,grid_output,volume_input", + argvalues=[ + ( + np.meshgrid( + np.arange(3), + np.arange(4), + np.arange(5), + indexing="ij", + ), + np.meshgrid( + np.arange(3) + 0.5, + np.arange(4) + 0.5, + np.arange(5) + 0.5, + indexing="ij", + ), + np.ones((2, 3, 4)), + ) + ], +) +def test_sweep( + grid_input: tuple[np.ndarray, np.ndarray, np.ndarray], + grid_output: tuple[np.ndarray, np.ndarray, np.ndarray], + volume_input: np.ndarray, +): + + weights = _test_sweep_compiled(grid_input, grid_output, volume_input) + + assert len(weights) > 0 + + +@numba.njit +def _test_sweep_compiled( + grid_input: tuple[np.ndarray, np.ndarray, np.ndarray], + grid_output: tuple[np.ndarray, np.ndarray, np.ndarray], + volume_input: np.ndarray, +): + intercepts = _intercepts.empty( + shape_input=grid_input[0].shape, + shape_output=grid_output[0].shape, + ) + + intercepts[0][0][0][0].append((0, 0, (0.05, 0.15, 0.35))) + intercepts[0][0][0][0].append((0, 1, (0.85, 0.25, 0.25))) + + weights = numba.typed.List() + for _ in range(0): + weights.append((0, 0, 0.0)) + + _intercepts.sweep( + intercepts=intercepts, + weights=weights, + grid_input=grid_input, + grid_output=grid_output, + volume_input=volume_input, + ) + + return weights diff --git a/regridding/_weights/_weights_conservative_3d/_weights_conservative_3d.py b/regridding/_weights/_weights_conservative_3d/_weights_conservative_3d.py new file mode 100644 index 0000000..236ba80 --- /dev/null +++ b/regridding/_weights/_weights_conservative_3d/_weights_conservative_3d.py @@ -0,0 +1,1036 @@ +import sys +import numpy as np +import numba +from numba import literal_unroll +import regridding as rg +from . import ( + _arrays, + _grids, + _intercepts, +) +from ._arrays import axis_x, axis_y, axis_z + +axes = ( + (axis_x,), + (axis_y,), + (axis_z,), + (axis_x, axis_y), + (axis_y, axis_z), + (axis_z, axis_x), +) + +__all__ = [ + "weights_conservative_3d", +] + + +@numba.njit(cache=True) +def weights_conservative_3d( + grid_input: tuple[np.ndarray, np.ndarray, np.ndarray], + grid_output: tuple[np.ndarray, np.ndarray, np.ndarray], +) -> numba.typed.List[tuple[int, int, float]]: + """ + For each cell of `grid_output`, + compute the fraction of volume shared with each cell of `grid_input` + and save as a list of weights. + + Parameters + ---------- + grid_input + The vertices of the old grid. + Every component must have the same 3D shape. + grid_output + The vertices of the new grid. + Every component must have the same 3D shape. + """ + + x_input, y_input, z_input = grid_input + x_output, y_output, z_output = grid_output + + shape_input = x_input.shape + shape_output = x_output.shape + + weights = numba.typed.List() + for x in range(0): + weights.append((0, 0, 0.0)) + + volume_input = _grids.grid_volume(grid_input) + + intercepts = _intercepts.empty(shape_input, shape_output) + + for sweep_input in (False, True): + _sweep_grid( + grid_input=grid_input, + grid_output=grid_output, + volume_input=volume_input, + weights=weights, + intercepts=intercepts, + sweep_input=sweep_input, + ) + + _intercepts.sweep( + intercepts=intercepts, + weights=weights, + grid_input=grid_input, + grid_output=grid_output, + volume_input=volume_input, + ) + + return weights + + +@numba.njit(cache=True) +def _sweep_grid( + grid_input: tuple[np.ndarray, np.ndarray, np.ndarray], + grid_output: tuple[np.ndarray, np.ndarray, np.ndarray], + volume_input: np.ndarray, + weights: numba.typed.List[tuple[int, int, float]], + intercepts: numba.typed.List, + sweep_input: bool, +) -> None: + """ + Sweep along each of the three axes and each of the three diagonals + of either `grid_input` or `grid_output` and compute the weights and + intercepts. + + Parameters + ---------- + grid_input + The vertices of the old grid. + grid_output + The vertices of the new grid. + volume_input + The area of the input grid which will used to compute the weight. + weights + The current list of weights. + New weights will be appended to this list. + intercepts + A sorted list of intercepts to be traversed later. + As new intercepts are found, they are inserted into the list. + sweep_input + A boolean flag indicating whether to iterate along `grid_input` + or `grid_output`. + If :obj:`True`, this function sweeps along `grid_input`. + """ + + shape_input = grid_input[0].shape + shape_output = grid_output[0].shape + + shape_cells_input = _grids.shape_centers(shape_input) + shape_cells_output = _grids.shape_centers(shape_output) + + grid_sweep, grid_static = _grid_sweep_static( + grid_input=grid_input, + grid_output=grid_output, + sweep_input=sweep_input, + ) + + x_sweep, y_sweep, z_sweep = grid_sweep + x_static, y_static, z_static = grid_static + + bbox_static = ( + (x_static.min(), y_static.min(), z_static.min()), + (x_static.max(), y_static.max(), z_static.max()), + ) + + _, boundary_static = _grids.grid_boundary(grid_static) + + for axis in literal_unroll(axes): + + axis_last = axis[~0] + + _sweep_along_axis( + grid_sweep=( + _arrays.align_axis_right(x_sweep, axis_last), + _arrays.align_axis_right(y_sweep, axis_last), + _arrays.align_axis_right(z_sweep, axis_last), + ), + grid_static=grid_static, + bbox_static=bbox_static, + boundary_static=boundary_static, + volume_input=volume_input, + shape_cells_input=shape_cells_input, + shape_cells_output=shape_cells_output, + weights=weights, + intercepts=intercepts, + sweep_input=sweep_input, + axis_sweep=axis, + ) + + +@numba.njit( + cache=True, + # parallel=True, +) +def _sweep_along_axis( + grid_sweep: tuple[np.ndarray, np.ndarray, np.ndarray], + grid_static: tuple[np.ndarray, np.ndarray, np.ndarray], + bbox_static: tuple[tuple[float, float, float], tuple[float, float, float]], + boundary_static: numba.typed.List[ + tuple[ + tuple[float, float, float], + tuple[float, float, float], + tuple[float, float, float], + ], + ], + volume_input: np.ndarray, + shape_cells_input: tuple[int, int, int], + shape_cells_output: tuple[int, int, int], + weights: numba.typed.List[tuple[int, int, float]], + intercepts: numba.typed.List, + sweep_input: bool, + axis_sweep: tuple[int] | tuple[int, int], +) -> None: + """ + Sweep along one logical axis of either `grid_input` or `grid_output`. + At each vertex, compute and save the weights and intercepts. + + Parameters + ---------- + grid_sweep + The vertices of the sweep grid. + The last axis of this grid must be the sweep axis. + grid_static + Coordinates of the static grid. + bbox_static + Two points defining the bounding box of the static grid. + boundary_static + A sequence of triangles defining the outer surface of the static grid. + volume_input + The volume of each cell in the input grid. + shape_cells_input + The number of cells along each axis of the input grid. + This could be calculated from the other arguments, + but we don't to save computation time. + shape_cells_output + The number of cells along each axis of the output grid. + Like `shape_cells_input`, + this could be calculated from the other arguments, + but we don't to save computation time. + weights + The current list of weights. + New weights will be appended to this list. + intercepts + A sorted list of intercepts to be traversed later. + As new intercepts are found, they are inserted into the list. + sweep_input + A boolean flag indicating whether to iterate along `grid_input` + or `grid_output`. + If :obj:`True`, this function sweeps along `grid_input`. + axis_sweep + The logical axis of the sweep grid to iterate along. + """ + + x_sweep, y_sweep, z_sweep = grid_sweep + x_static, y_static, z_static = grid_static + + shape_sweep = x_sweep.shape + shape_static = x_static.shape + + shape_cells_sweep = _grids.shape_centers(shape_sweep) + shape_cells_static = _grids.shape_centers(shape_static) + + shape_sweep_x, shape_sweep_y, shape_sweep_z = shape_sweep + shape_static_x, shape_static_y, shape_static_z = shape_static + + weight = numba.typed.List() + + for i in range(shape_sweep_x): + weight_i = numba.typed.List() + for _ in range(0): + weight_i.append((0, 0, 0.0)) + weight.append(weight_i) + + if len(axis_sweep) == 2: + d_start = -(shape_sweep_z - 2) + d_end = shape_sweep_y - 1 + step_index = (0, 1, 1) + else: + d_start = 0 + d_end = shape_sweep_y + step_index = (0, 0, 1) + + j_end = d_end + k_end = shape_cells_sweep[2] + + for i in range(shape_sweep_x): + + i = numba.types.int64(i) + + for d in range(d_start, d_end): + + if d < 0: + j = 0 + k = -d + else: + j = d + k = 0 + + index_sweep = i, j, k + + index_static = sys.maxsize, sys.maxsize, sys.maxsize + + sweep_is_outside_static = True + + point_1 = ( + x_sweep[index_sweep], + y_sweep[index_sweep], + z_sweep[index_sweep], + ) + + if rg.geometry.point_is_inside_box_3d( + point=point_1, + box=bbox_static, + ): + if rg.geometry.point_is_inside_polyhedron( + point=point_1, + polyhedron=boundary_static, + ): + index_static = _grids.index_of_point_brute( + point=point_1, + grid=grid_static, + ) + sweep_is_outside_static = False + + while (index_sweep[1] < j_end) and (index_sweep[2] < k_end): + + index_sweep_new = rg.math.sum_3d(index_sweep, step_index) + + point_2 = ( + x_sweep[index_sweep_new], + y_sweep[index_sweep_new], + z_sweep[index_sweep_new], + ) + + line = point_1, point_2 + + if sweep_is_outside_static: + + line, index_sweep, index_static = _step_outside_static( + line=line, + index_sweep=index_sweep, + index_static=index_static, + step_index=step_index, + grid_sweep=grid_sweep, + grid_static=grid_static, + shape_cells_input=shape_cells_input, + shape_cells_output=shape_cells_output, + intercepts=intercepts, + sweep_input=sweep_input, + axis_sweep=axis_sweep, + ) + + sweep_is_outside_static = index_static[0] == sys.maxsize + + else: + + line, index_sweep, index_static = _step_inside_static( + line=line, + index_sweep=index_sweep, + index_static=index_static, + step_index=step_index, + grid_sweep=grid_sweep, + grid_static=grid_static, + volume_input=volume_input, + shape_cells_input=shape_cells_input, + shape_cells_output=shape_cells_output, + weights=weights, + intercepts=intercepts, + sweep_input=sweep_input, + axis_sweep=axis_sweep, + ) + + if not _arrays.index_in_bounds( + index=index_static, + shape=shape_cells_static, + ): + break + + point_1 = line[1] + + for i in range(shape_sweep_x): + weight_i = weight[i] + for w in range(len(weight_i)): + weights.append(weight_i[w]) + + +@numba.njit(cache=True, error_model="numpy") +def _step_outside_static( + line: tuple[ + tuple[float, float, float], + tuple[float, float, float], + ], + index_sweep: tuple[int, int, int], + index_static: tuple[int, int, int], + step_index: tuple[int, int, int], + grid_sweep: tuple[np.ndarray, np.ndarray, np.ndarray], + grid_static: tuple[np.ndarray, np.ndarray, np.ndarray], + shape_cells_input: tuple[int, int, int], + shape_cells_output: tuple[int, int, int], + intercepts: numba.typed.List, + sweep_input: bool, + axis_sweep: tuple[int] | tuple[int, int], +) -> tuple[ + tuple[ + tuple[float, float, float], + tuple[float, float, float], + ], + tuple[int, int, int], + tuple[int, int, int], +]: + """ + Check if the current line segment crosses the boundary. + If it does, compute the intersection and return the intersection point + and the indices of the `grid_static` where the intersection occurs. + + Parameters + ---------- + line + The current line segment of the sweep grid. + index_sweep + The index of the current vertex in the sweep grid + This is assumed to be a valid, positive index. + index_static + The index of the current cell in the static grid. + This is assumed to be a valid, positive index. + step_index + A step size in index space on the sweep grid. + grid_sweep + The vertices of the sweep grid. + The last axis of this grid must be the sweep axis. + grid_static + The vertices of the static grid. + bbox_static + Two points defining the bounding box of the static grid. + shape_cells_input + The number of cells along each axis of the input grid. + This could be calculated from the other arguments, + but we don't to save computation time. + shape_cells_output + The number of cells along each axis of the output grid. + Like `shape_cells_input`, + this could be calculated from the other arguments, + but we don't to save computation time. + intercepts + A sorted list of intercepts to be traversed later. + As new intercepts are found, they are inserted into the list. + sweep_input + A boolean flag indicating whether to iterate along `grid_input` + or `grid_output`. + If :obj:`True`, this function sweeps along `grid_input`. + axis_sweep + The logical axis of the sweep grid to iterate along. + """ + + p1, p2 = line + + x, y, z = grid_static + + shape_cells_static = _grids.shape_centers(x.shape) + + for axis in _arrays.axes: + + x_axis = _arrays.align_axis_right(x, axis) + y_axis = _arrays.align_axis_right(y, axis) + z_axis = _arrays.align_axis_right(z, axis) + + for direction_face in (-1, 1): + + if direction_face > 0: + k_face = ~0 + else: + k_face = 0 + + x_face = x_axis[..., k_face] + y_face = y_axis[..., k_face] + z_face = z_axis[..., k_face] + + shape_face = x_face.shape + + num_i, num_j = shape_face + + for i0 in range(num_i - 1): + for j0 in range(num_j - 1): + + i1 = i0 + 1 + j1 = j0 + 1 + + for t in (-1, 1): + + if t > 0: + ind = ( + (i0, j0), + (i1, j0), + (i1, j1), + ) + else: + ind = ( + (i1, j1), + (i0, j1), + (i0, j0), + ) + + if direction_face < 0: + ind = ind[::-1] + + v0, v1, v2 = ind + + triangle = ( + (x_face[v0], y_face[v0], z_face[v0]), + (x_face[v1], y_face[v1], z_face[v1]), + (x_face[v2], y_face[v2], z_face[v2]), + ) + + tuv = rg.geometry.line_triangle_intersection_parameters( + line=line, + triangle=triangle, + ) + + if rg.geometry.line_intersects_triangle(tuv): + + if direction_face > 0: + index_static = i0, j0, shape_cells_static[axis] - 1 + else: + index_static = i0, j0, 0 + + i, j, k = index_static + + if axis == axis_x: + index_static = k, i, j + elif axis == axis_y: + index_static = j, k, i + + normal = [1 if ax == axis else 0 for ax in _arrays.axes] + normal = rg.math.multiply_3d(direction_face, normal) + + _, p2 = _calc_and_save_intercept( + line=line, + tuv=tuv, + index_sweep=index_sweep, + index_static=index_static, + grid_sweep=grid_sweep, + grid_static=grid_static, + shape_cells_input=shape_cells_input, + shape_cells_output=shape_cells_output, + intercepts=intercepts, + sweep_input=sweep_input, + axis_sweep=axis_sweep, + axis_static=axis, + offset_static=normal, + ) + + return (p1, p2), index_sweep, index_static + + index_sweep = rg.math.sum_3d(index_sweep, step_index) + + return (p1, p2), index_sweep, index_static + + +@numba.njit(cache=True, error_model="numpy") +def _step_inside_static( + line: tuple[ + tuple[float, float, float], + tuple[float, float, float], + ], + index_sweep: tuple[int, int, int], + index_static: tuple[int, int, int], + step_index: tuple[int, int, int], + grid_sweep: tuple[np.ndarray, np.ndarray, np.ndarray], + grid_static: tuple[np.ndarray, np.ndarray, np.ndarray], + volume_input: np.ndarray, + shape_cells_input: tuple[int, int, int], + shape_cells_output: tuple[int, int, int], + weights: numba.typed.List[tuple[int, int, float]], + intercepts: numba.typed.List, + sweep_input: bool, + axis_sweep: tuple[int] | tuple[int, int], +) -> tuple[ + tuple[ + tuple[float, float, float], + tuple[float, float, float], + ], + tuple[int, int, int], + tuple[int, int, int], +]: + """ + Check if the current line segment crosses any face of the current cell + in the static grid. + If it does, compute the intersection point and indices of the new + cell in the static grid. + Otherwise, increment the index of the sweep grid. + + Parameters + ---------- + line + The current line segment of the sweep grid. + index_sweep + The index of the current vertex in the sweep grid + This is assumed to be a valid, positive index. + index_static + The index of the current cell in the static grid. + This is assumed to be a valid, positive index. + step_index + A step size in index space on the sweep grid. + grid_sweep + The vertices of the sweep grid. + The last axis of this grid must be the sweep axis. + grid_static + The vertices of the static grid. + volume_input + The volume of each cell in the input grid. + shape_cells_input + The number of cells along each axis of the input grid. + This could be calculated from the other arguments, + but we don't to save computation time. + shape_cells_output + The number of cells along each axis of the output grid. + Like `shape_cells_input`, + this could be calculated from the other arguments, + but we don't to save computation time. + weights + The current list of weights. + New weights will be appended to this list. + intercepts + A sorted list of intercepts to be traversed later. + As new intercepts are found, they are inserted into the list. + sweep_input + A boolean flag indicating whether to iterate along `grid_input` + or `grid_output`. + If :obj:`True`, this function sweeps along `grid_input`. + axis_sweep + The logical axis of the sweep grid to iterate along. + """ + + p1, p2 = line + + polyhedron_cell = _grids.cell_boundary( + index=index_static, + grid=grid_static, + ) + + for t in range(len(polyhedron_cell)): + + triangle = polyhedron_cell[t] + + tuv = rg.geometry.line_triangle_intersection_parameters( + line=line, + triangle=triangle, + ) + + if rg.geometry.line_intersects_triangle(tuv): + + if tuv[0] > 1e-8: + + index_sweep_new = index_sweep + + index_static_new, p2 = _calc_and_save_intercept( + line=line, + tuv=tuv, + index_sweep=index_sweep, + index_static=index_static, + grid_sweep=grid_sweep, + grid_static=grid_static, + shape_cells_input=shape_cells_input, + shape_cells_output=shape_cells_output, + intercepts=intercepts, + sweep_input=sweep_input, + axis_sweep=axis_sweep, + axis_static=_grids.cell_axes[t], + offset_static=_grids.cell_normals[t], + ) + + break + + else: + index_sweep_new = rg.math.sum_3d(index_sweep, step_index) + index_static_new = index_static + + line = p1, p2 + + if len(axis_sweep) == 1: + _calc_and_save_weights( + line=line, + index_sweep=index_sweep, + index_static=index_static, + grid_sweep=grid_sweep, + volume_input=volume_input, + shape_cells_input=shape_cells_input, + shape_cells_output=shape_cells_output, + weights=weights, + sweep_input=sweep_input, + axis_sweep=axis_sweep, + ) + + return line, index_sweep_new, index_static_new + + +@numba.njit(cache=True) +def _calc_and_save_intercept( + line: tuple[ + tuple[float, float, float], + tuple[float, float, float], + ], + tuv: tuple[float, float, float], + index_sweep: tuple[int, int, int], + index_static: tuple[int, int, int], + grid_sweep: tuple[np.ndarray, np.ndarray, np.ndarray], + grid_static: tuple[np.ndarray, np.ndarray, np.ndarray], + shape_cells_input: tuple[int, int, int], + shape_cells_output: tuple[int, int, int], + intercepts: numba.typed.List, + sweep_input: bool, + axis_sweep: tuple[int] | tuple[int, int], + axis_static: int, + offset_static: tuple[int, int, int], +) -> tuple[ + tuple[int, int, int], + tuple[float, float, float], +]: + + intercept = rg.geometry.line_triangle_intersection( + line=line, + tuv=tuv, + ) + + point_sweep, point_sweep_new = line + + x_sweep, y_sweep, z_sweep = grid_sweep + x_static, y_static, z_static = grid_static + + vertex_static = ( + x_static[index_static], + y_static[index_static], + z_static[index_static], + ) + + shape_input = [n + 1 for n in shape_cells_input] + shape_output = [n + 1 for n in shape_cells_output] + + index_static_new = rg.math.sum_3d(index_static, offset_static) + + index_static_normal = index_static_new + if 0 <= index_static_normal[axis_static] < shape_output[axis_static]: + vertex_static_normal = ( + x_static[index_static_normal], + y_static[index_static_normal], + z_static[index_static_normal], + ) + normal_static = rg.math.difference_3d(vertex_static_normal, vertex_static) + else: + index_static_normal = rg.math.difference_3d(index_static, offset_static) + vertex_static_normal = ( + x_static[index_static_normal], + y_static[index_static_normal], + z_static[index_static_normal], + ) + normal_static = rg.math.difference_3d(vertex_static, vertex_static_normal) + + index_apex_static = index_static + if index_static_new[axis_static] > index_static[axis_static]: + index_apex_static = index_static_new + else: + normal_static = rg.math.negate_3d(normal_static) + + axes_orthogonal = [ax for ax in _arrays.axes if ax not in axis_sweep] + + for axis_orthogonal in axes_orthogonal: + + offset_axis = 2 - axis_sweep[~0] + + axis_orthogonal_apparent = (axis_orthogonal + offset_axis) % 3 + + offset_orthogonal = _arrays.vector_unit[axis_orthogonal_apparent] + + index_orthogonal = rg.math.difference_3d(index_sweep, offset_orthogonal) + + apex_sweep = ( + x_sweep[index_sweep], + y_sweep[index_sweep], + z_sweep[index_sweep], + ) + + if index_orthogonal[axis_orthogonal_apparent] >= 0: + point_orthogonal = ( + x_sweep[index_orthogonal], + y_sweep[index_orthogonal], + z_sweep[index_orthogonal], + ) + normal_sweep = rg.math.difference_3d(apex_sweep, point_orthogonal) + else: + index_orthogonal = rg.math.sum_3d(index_sweep, offset_orthogonal) + point_orthogonal = ( + x_sweep[index_orthogonal], + y_sweep[index_orthogonal], + z_sweep[index_orthogonal], + ) + normal_sweep = rg.math.difference_3d(point_orthogonal, apex_sweep) + + if sweep_input: + normal_input = normal_sweep + normal_output = normal_static + else: + normal_input = normal_static + normal_output = normal_sweep + + direction = rg.math.cross_3d(normal_input, normal_output) + + index_apex_sweep = index_sweep + if len(axis_sweep) == 1: + + axis_parallel_actual = [ + ax + for ax in _arrays.axes + if ax is not axis_orthogonal_apparent and ax is not _arrays.axis_z + ][0] + + offset_parallel = _arrays.vector_unit[axis_parallel_actual] + + index_parallel = rg.math.difference_3d(index_sweep, offset_parallel) + + if index_parallel[axis_parallel_actual] >= 0: + + point_parallel = ( + x_sweep[index_parallel], + y_sweep[index_parallel], + z_sweep[index_parallel], + ) + + parallel = rg.math.difference_3d(apex_sweep, point_parallel) + + sgn = rg.math.dot_3d(parallel, direction) + + if sgn < 0: + index_apex_sweep = index_parallel + + if sweep_input: + axis_input = axis_orthogonal + axis_output = axis_static + else: + axis_input = axis_static + axis_output = axis_orthogonal + + index_input, index_output = _index_input_output( + index_sweep=index_apex_sweep, + index_static=index_apex_static, + sweep_input=sweep_input, + axis_sweep=axis_sweep, + ) + + i_input = index_input[axis_input] + i_output = index_output[axis_output] + + intercept_new = ( + _arrays.index_flat(index_input, shape_input), + _arrays.index_flat(index_output, shape_output), + intercept, + ) + + intercepts_input_output = intercepts[axis_input][axis_output][i_input][i_output] + num_intercept = len(intercepts_input_output) + if num_intercept == 0: + index = 0 + + elif num_intercept == 1: + _, _, intercept_0 = intercepts_input_output[0] + direction_intercepts = rg.math.difference_3d(intercept, intercept_0) + sgn_intercepts = rg.math.dot_3d(direction_intercepts, direction) + if sgn_intercepts > 0: + index = 1 + else: + index = 0 + + else: + index = _intercepts._bisect_intercepts( + intercepts=intercepts_input_output, + intercept_new=intercept_new, + ) + + intercepts_input_output.insert(index, intercept_new) + + return index_static_new, intercept + + +@numba.njit(cache=True) +def _calc_and_save_weights( + line: tuple[ + tuple[float, float, float], + tuple[float, float, float], + ], + index_sweep: tuple[int, int, int], + index_static: tuple[int, int, int], + grid_sweep: tuple[np.ndarray, np.ndarray, np.ndarray], + volume_input: np.ndarray, + shape_cells_input: tuple[int, int, int], + shape_cells_output: tuple[int, int, int], + weights: numba.typed.List[tuple[int, int, float]], + sweep_input: bool, + axis_sweep: tuple[int] | tuple[int, int], +): + + x_sweep, y_sweep, z_sweep = grid_sweep + + shape_sweep = x_sweep.shape + + i, j, k = index_sweep + + i_left = i - 1 + i_right = i + + j_lower = j - 1 + j_upper = j + + index_sweep_left = i_left, j, k + index_sweep_lower = i, j_lower, k + + p0_left = ( + x_sweep[index_sweep_left], + y_sweep[index_sweep_left], + z_sweep[index_sweep_left], + ) + p0_lower = ( + x_sweep[index_sweep_lower], + y_sweep[index_sweep_lower], + z_sweep[index_sweep_lower], + ) + + p1, p2 = line + + volume_left = -rg.geometry.volume_tetrahedron(p0_left, p1, p2) + volume_lower = rg.geometry.volume_tetrahedron(p0_lower, p1, p2) + + if i_left >= 0: + if j_lower >= 0: + + index_sweep_00 = i_left, j_lower, k + + index_input_00, index_output_00 = _index_input_output( + index_sweep=index_sweep_00, + index_static=index_static, + sweep_input=sweep_input, + axis_sweep=axis_sweep, + ) + + volume_input_lower_left = volume_input[index_input_00] + + index_input_00 = _arrays.index_flat(index_input_00, shape_cells_input) + index_output_00 = _arrays.index_flat(index_output_00, shape_cells_output) + + volume_lower_left = (volume_lower + volume_left) / volume_input_lower_left + weights.append((index_input_00, index_output_00, volume_lower_left)) + + if j_upper < (shape_sweep[1] - 1): + + index_sweep_01 = i_left, j_upper, k + + index_input_01, index_output_01 = _index_input_output( + index_sweep=index_sweep_01, + index_static=index_static, + sweep_input=sweep_input, + axis_sweep=axis_sweep, + ) + + volume_input_upper_left = volume_input[index_input_01] + + index_input_01 = _arrays.index_flat(index_input_01, shape_cells_input) + index_output_01 = _arrays.index_flat(index_output_01, shape_cells_output) + + volume_upper_left = -volume_left / volume_input_upper_left + weights.append((index_input_01, index_output_01, volume_upper_left)) + + if i_right < (shape_sweep[0] - 1): + if j_lower >= 0: + + index_sweep_10 = i_right, j_lower, k + + index_input_10, index_output_10 = _index_input_output( + index_sweep=index_sweep_10, + index_static=index_static, + sweep_input=sweep_input, + axis_sweep=axis_sweep, + ) + + volume_input_lower_right = volume_input[index_input_10] + + index_input_10 = _arrays.index_flat(index_input_10, shape_cells_input) + index_output_10 = _arrays.index_flat(index_output_10, shape_cells_output) + + volume_lower_right = -volume_lower / volume_input_lower_right + weights.append((index_input_10, index_output_10, volume_lower_right)) + + +@numba.njit(cache=True) +def _grid_sweep_static( + grid_input: tuple[np.ndarray, np.ndarray, np.ndarray], + grid_output: tuple[np.ndarray, np.ndarray, np.ndarray], + sweep_input: bool, +) -> tuple[ + tuple[np.ndarray, np.ndarray, np.ndarray], + tuple[np.ndarray, np.ndarray, np.ndarray], +]: + """ + Given the input and output grids, + prepare the sweep and static grids. + + Parameters + ---------- + grid_input + The vertices of the old grid. + grid_output + The vertices of the new grid. + sweep_input + A boolean flag indicating whether to iterate along `grid_input` + or `grid_output`. + If :obj:`True`, this function sweeps along `grid_input`. + """ + if sweep_input: + grid_sweep = grid_input + grid_static = grid_output + else: + grid_sweep = grid_output + grid_static = grid_input + + return grid_sweep, grid_static + + +@numba.njit(cache=True) +def _index_input_output( + index_sweep: tuple[int, int, int], + index_static: tuple[int, int, int], + sweep_input: bool, + axis_sweep: tuple[int] | tuple[int, int], +) -> tuple[tuple[int, int, int], tuple[int, int, int]]: + """ + Convert indices on the sweep and static grids into indices on the + input/output grids. + + Parameters + ---------- + index_sweep + The index on the sweep grid to convert. + index_static + The index on the static grid to convert. + sweep_input + A boolean flag indicating whether to iterate along `grid_input` + or `grid_output`. + If :obj:`True`, this function sweeps along `grid_input`. + axis_sweep + The logical axis of the `grid_sweep` we're iterating along. + """ + + axis = axis_sweep[~0] + + i_sweep, j_sweep, k_sweep = index_sweep + + if axis == axis_x: + index_sweep = k_sweep, i_sweep, j_sweep + elif axis == axis_y: + index_sweep = j_sweep, k_sweep, i_sweep + + if sweep_input: + index_input = index_sweep + index_output = index_static + else: + index_input = index_static + index_output = index_sweep + + return index_input, index_output diff --git a/regridding/_weights/_weights_conservative_3d/_weights_conservative_3d_test.py b/regridding/_weights/_weights_conservative_3d/_weights_conservative_3d_test.py new file mode 100644 index 0000000..af73249 --- /dev/null +++ b/regridding/_weights/_weights_conservative_3d/_weights_conservative_3d_test.py @@ -0,0 +1,38 @@ +import pytest +import numpy as np +from . import _weights_conservative_3d + +_x = np.arange(3) +_y = np.arange(4) +_z = np.arange(5) + +_x, _y, _z = np.meshgrid(_x, _y, _z, indexing="ij") + + +@pytest.mark.parametrize( + argnames="grid_input,grid_output", + argvalues=[ + ( + np.meshgrid( + np.arange(2), + np.arange(2), + np.arange(2), + indexing="ij", + ), + np.meshgrid( + np.arange(2), + np.arange(2), + np.arange(2), + indexing="ij", + ), + ) + ], +) +def test_weights_conservative_3d( + grid_input: tuple[np.ndarray, np.ndarray, np.ndarray], + grid_output: tuple[np.ndarray, np.ndarray, np.ndarray], +): + result = _weights_conservative_3d.weights_conservative_3d( + grid_input=grid_input, + grid_output=grid_output, + )