From 89dea250d837fc29590bd237e08de0ee339d3457 Mon Sep 17 00:00:00 2001 From: jacobdparker Date: Wed, 17 Jan 2024 15:21:04 -0500 Subject: [PATCH 1/8] changing from shape_output to shape_output_numba --- regridding/_find_indices/_find_indices.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/regridding/_find_indices/_find_indices.py b/regridding/_find_indices/_find_indices.py index 5668696..518051d 100644 --- a/regridding/_find_indices/_find_indices.py +++ b/regridding/_find_indices/_find_indices.py @@ -86,7 +86,7 @@ def find_indices( indices_output = tuple( np.moveaxis( - a=i.reshape(*shape_orthogonal, *shape_output), + a=i.reshape(*shape_orthogonal, *shape_output_numba), source=axis_output_numba, destination=axis_output, ) From f9f31f039917c7242d8627d3b07a2b4907646f7b Mon Sep 17 00:00:00 2001 From: jacobdparker Date: Wed, 17 Jan 2024 15:23:41 -0500 Subject: [PATCH 2/8] modifying normalized_input_output_coordinates to handle negative axes --- regridding/_util.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/regridding/_util.py b/regridding/_util.py index 1002787..5aba591 100644 --- a/regridding/_util.py +++ b/regridding/_util.py @@ -75,12 +75,14 @@ def _normalize_input_output_coordinates( ) shape_input = list(shape_orthogonal) - for ax in axis_input: + for ax in reversed(axis_input): + ax = ax % ndim_input shape_input.insert(ax, shape_coordinates_input[ax]) shape_input = tuple(shape_input) shape_output = list(shape_orthogonal) - for ax in axis_output: + for ax in reversed(axis_output): + ax = ax % ndim_input shape_output.insert(ax, shape_coordinates_output[ax]) shape_output = tuple(shape_output) From 156dd3b464219fe874676e2419f2c0b374faef48 Mon Sep 17 00:00:00 2001 From: jacobdparker Date: Wed, 17 Jan 2024 15:28:06 -0500 Subject: [PATCH 3/8] addressing possible race condition in _weights_from_indices_multilinear_1d --- regridding/_weights/_weights_multilinear.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/regridding/_weights/_weights_multilinear.py b/regridding/_weights/_weights_multilinear.py index 2960bce..4b09ddb 100644 --- a/regridding/_weights/_weights_multilinear.py +++ b/regridding/_weights/_weights_multilinear.py @@ -88,7 +88,7 @@ def _weights_from_indices_multilinear( return weights, shape_input, shape_output -@numba.njit(parallel=True) +@numba.njit(parallel=False) def _weights_from_indices_multilinear_1d( indices_output: tuple[np.ndarray], coordinates_input: tuple[np.ndarray], @@ -103,7 +103,7 @@ def _weights_from_indices_multilinear_1d( weights = numba.typed.List() - for d in numba.prange(num_d): + for d in range(num_d): weights_d = numba.typed.List() for _ in range(0): weights_d.append((0.0, 0.0, 0.0)) From 0eb6116b1d7f918c61c60a03968e60a7abec2a03 Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Wed, 6 Mar 2024 08:10:19 -0700 Subject: [PATCH 4/8] Fixing Black errors. --- regridding/_interp_ndarray.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/regridding/_interp_ndarray.py b/regridding/_interp_ndarray.py index 948c83d..8af03d3 100644 --- a/regridding/_interp_ndarray.py +++ b/regridding/_interp_ndarray.py @@ -108,9 +108,11 @@ def ndarray_linear_interpolation( for ax in range(-ndim_broadcasted_a, 0) ) shape_broadcasted_indices = tuple( - shape_indices[ax] - if ax in axis_indices - else shape_orthogonal[axis_orthogonal_indices.index(ax)] + ( + shape_indices[ax] + if ax in axis_indices + else shape_orthogonal[axis_orthogonal_indices.index(ax)] + ) for ax in range(-ndim_broadcasted_indices, 0) ) @@ -125,9 +127,11 @@ def index_a_indices(index: tuple[int, ...]) -> tuple[tuple, tuple]: for ax in range(-ndim_broadcasted_a, 0) ) index_indices = tuple( - slice(None) - if ax in axis_indices - else index[axis_orthogonal_indices.index(ax)] + ( + slice(None) + if ax in axis_indices + else index[axis_orthogonal_indices.index(ax)] + ) for ax in range(-ndim_broadcasted_indices, 0) ) return index_a, index_indices From 2be6e5d2c6160c1b5e695f819094687a02718cea Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Wed, 6 Mar 2024 08:38:33 -0700 Subject: [PATCH 5/8] Pin asv to 0.6.1 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 29053e1..0402e67 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,7 @@ doc = [ "sphinx-favicon", ] benchmark = [ - "asv", + "asv<=0.6.1", ] [project.urls] From a83e9eab92f4006e0ba83643e46fae53cb42ccf0 Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Wed, 6 Mar 2024 08:41:06 -0700 Subject: [PATCH 6/8] Pin asv to 0.6.1 --- .github/workflows/benchmarks.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml index 81ab7c1..2bfd0b8 100644 --- a/.github/workflows/benchmarks.yml +++ b/.github/workflows/benchmarks.yml @@ -39,7 +39,7 @@ jobs: run: | python -m pip install --upgrade pip pip install setuptools wheel - pip install asv + pip install asv<=0.6.1 - name: Setup Pages id: pages uses: actions/configure-pages@v3 From 9135aee4b10440e83895e1f0d700ed8bdaf0b1f8 Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Wed, 6 Mar 2024 08:44:21 -0700 Subject: [PATCH 7/8] Pin asv to 0.6.1 --- .github/workflows/benchmarks.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml index 2bfd0b8..a2d207d 100644 --- a/.github/workflows/benchmarks.yml +++ b/.github/workflows/benchmarks.yml @@ -39,7 +39,7 @@ jobs: run: | python -m pip install --upgrade pip pip install setuptools wheel - pip install asv<=0.6.1 + pip install asv==0.6.1 - name: Setup Pages id: pages uses: actions/configure-pages@v3 From 1a589156e08d27e4cf7f07a5730a97e51addc136 Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Wed, 6 Mar 2024 12:54:31 -0700 Subject: [PATCH 8/8] Added `regridding.fill()` function to fill in missing data using interpolation (#3) Added `regridding.fill()` function to fill in missing data using interpolation. --- regridding/__init__.py | 1 + regridding/_fill/__init__.py | 1 + regridding/_fill/_fill.py | 99 ++++++++++++++++++++++++++ regridding/_fill/_fill_test.py | 48 +++++++++++++ regridding/_fill/_gauss_seidel.py | 111 ++++++++++++++++++++++++++++++ 5 files changed, 260 insertions(+) create mode 100644 regridding/_fill/__init__.py create mode 100644 regridding/_fill/_fill.py create mode 100644 regridding/_fill/_fill_test.py create mode 100644 regridding/_fill/_gauss_seidel.py diff --git a/regridding/__init__.py b/regridding/__init__.py index 1bf4a8b..2e20e47 100644 --- a/regridding/__init__.py +++ b/regridding/__init__.py @@ -4,3 +4,4 @@ from ._weights import * from ._interp_ndarray import * from ._regrid import * +from ._fill import * diff --git a/regridding/_fill/__init__.py b/regridding/_fill/__init__.py new file mode 100644 index 0000000..50da3c4 --- /dev/null +++ b/regridding/_fill/__init__.py @@ -0,0 +1 @@ +from ._fill import * diff --git a/regridding/_fill/_fill.py b/regridding/_fill/_fill.py new file mode 100644 index 0000000..578f4b7 --- /dev/null +++ b/regridding/_fill/_fill.py @@ -0,0 +1,99 @@ +from typing import Sequence, Literal +import numpy as np +from ._gauss_seidel import fill_gauss_seidel + +__all__ = [ + "fill", +] + + +def fill( + a: np.ndarray, + where: None | np.ndarray = None, + axis: None | int | Sequence[int] = None, + method: Literal["gauss_seidel"] = "gauss_seidel", + **kwargs, +) -> np.ndarray: + """ + Fill an array with missing values by interpolating from the valid points. + + Parameters + ---------- + a + The array with missing values to be filled + where + Boolean array of missing values. + If :obj:`None` (the default), all NaN values will be filled. + axis + The axes to use for interpolation. + If :obj:`None` (the default), interpolate along all the axes of `a`. + method + The interpolation method to use. + The only option is "gauss_seidel", which uses the Gauss-Seidel relaxation + technique to interpolate the valid data points. + kwargs + Additional method-specific keyword arguments. + For the Gauss-Seidel method, the valid keyword arguments are: + - ``num_iterations=100``, the number of red-black Gauss-Seidel iterations to perform. + + Examples + -------- + + Set random elements of an array to NaN, and then fill in the missing elements + using the Gauss-Seidel relaxation method. + + .. jupyter-execute:: + + import numpy as np + import matplotlib.pyplot as plt + import regridding + + # Define the independent variables + x = 3 * np.pi * np.linspace(-1, 1, num=51) + y = 3 * np.pi * np.linspace(-1, 1, num=51) + x, y = np.meshgrid(x, y, indexing="ij") + + # Define the array to remove elements from + a = np.cos(x) * np.cos(y) + + # Define the elements of the array to remove + where = np.random.uniform(0, 1, size=a.shape) > 0.9 + + # Set random elements of the array to NaN + a_missing = a.copy() + a_missing[where] = np.nan + + # Fill the missing elements using Gauss-Seidel relaxation + b = regridding.fill(a_missing, method="gauss_seidel", num_iterations=11) + + # Plot the results + fig, axs = plt.subplots( + ncols=3, + figsize=(6, 3), + sharey=True, + constrained_layout=True, + ) + kwargs_imshow = dict( + vmin=a.min(), + vmax=a.max(), + ) + axs[0].imshow(a_missing, **kwargs_imshow); + axs[1].imshow(b, **kwargs_imshow); + axs[2].imshow(a - b, **kwargs_imshow); + axs[0].set_title("original array"); + axs[1].set_title("filled array"); + axs[2].set_title("difference"); + """ + + if where is None: + where = np.isnan(a) + + if method == "gauss_seidel": + return fill_gauss_seidel( + a=a, + where=where, + axis=axis, + **kwargs, + ) + else: # pragma: nocover + raise ValueError("Unrecognized method '{method}'") diff --git a/regridding/_fill/_fill_test.py b/regridding/_fill/_fill_test.py new file mode 100644 index 0000000..901b76c --- /dev/null +++ b/regridding/_fill/_fill_test.py @@ -0,0 +1,48 @@ +import pytest +import numpy as np +import regridding + +_num_x = 11 +_num_y = 12 +_num_t = 13 + + +@pytest.mark.parametrize( + argnames="a,where,axis", + argvalues=[ + ( + np.random.uniform(0, 1, size=(_num_x, _num_y)), + np.random.uniform(0, 1, size=(_num_x, _num_y)) > 0.9, + None, + ), + ( + np.random.uniform(0, 1, size=(_num_t, _num_x, _num_y)), + np.random.uniform(0, 1, size=(_num_t, _num_x, _num_y)) > 0.9, + (~1, ~0), + ), + ( + np.sqrt(np.random.uniform(-0.1, 1, size=(_num_x, _num_t, _num_y))), + None, + (0, ~0), + ), + ], +) +@pytest.mark.parametrize("num_iterations", [11]) +def test_fill_gauss_sidel_2d( + a: np.ndarray, + where: np.ndarray, + axis: None | tuple[int, ...], + num_iterations: int, +): + result = regridding.fill( + a=a, + where=where, + axis=axis, + method="gauss_seidel", + num_iterations=num_iterations, + ) + if where is None: + where = np.isnan(a) + + assert np.allclose(result[~where], a[~where]) + assert np.all(result[where] != 0) diff --git a/regridding/_fill/_gauss_seidel.py b/regridding/_fill/_gauss_seidel.py new file mode 100644 index 0000000..9d24658 --- /dev/null +++ b/regridding/_fill/_gauss_seidel.py @@ -0,0 +1,111 @@ +from typing import Sequence +import numpy as np +import numba +import regridding._util + +__all__ = [ + "fill_gauss_seidel", +] + + +def fill_gauss_seidel( + a: np.ndarray, + where: np.ndarray, + axis: None | int | Sequence[int], + num_iterations: int = 100, +) -> np.ndarray: + + a = a.copy() + + a, where = np.broadcast_arrays(a, where, subok=True) + + a[where] = 0 + + axis = regridding._util._normalize_axis(axis=axis, ndim=a.ndim) + axis_numba = ~np.arange(len(axis))[::-1] + + shape = a.shape + shape_numba = tuple(shape[ax] for ax in axis) + + a = np.moveaxis(a, axis, axis_numba) + where = np.moveaxis(where, axis, axis_numba) + + shape_moved = a.shape + + a = a.reshape(-1, *shape_numba) + where = where.reshape(-1, *shape_numba) + + if len(axis) == 2: + result = _fill_gauss_seidel_2d( + a=a, + where=where, + num_iterations=num_iterations, + ) + else: # pragma: nocover + raise ValueError( + f"The number of interpolation axes, {len(axis)}," f"is not supported" + ) + + result = result.reshape(shape_moved) + result = np.moveaxis(result, axis_numba, axis) + + return result + + +@numba.njit(parallel=True) +def _fill_gauss_seidel_2d( + a: np.ndarray, + where: np.ndarray, + num_iterations: int, +) -> np.ndarray: + + num_t, num_y, num_x = a.shape + + for t in numba.prange(num_t): + for k in range(num_iterations): + for is_odd in [False, True]: + _iteration_gauss_seidel_2d( + a=a, + where=where, + t=t, + num_x=num_x, + num_y=num_y, + is_odd=is_odd, + ) + + return a + + +@numba.njit(fastmath=True) +def _iteration_gauss_seidel_2d( + a: np.ndarray, + where: np.ndarray, + t: int, + num_x: int, + num_y: int, + is_odd: bool, +) -> None: + + xmin, xmax = -1, 1 + ymin, ymax = -1, 1 + + dx = (xmax - xmin) / (num_x - 1) + dy = (ymax - ymin) / (num_y - 1) + + dxxinv = 1 / (dx * dx) + dyyinv = 1 / (dy * dy) + + dcent = 1 / (2 * (dxxinv + dyyinv)) + + for j in range(num_y): + for i in range(num_x): + if (i + j) & 1 == is_odd: + if where[t, j, i]: + i9 = (i - 1) % num_x + i1 = (i + 1) % num_x + j9 = (j - 1) % num_y + j1 = (j + 1) % num_y + + xterm = dxxinv * (a[t, j, i9] + a[t, j, i1]) + yterm = dyyinv * (a[t, j9, i] + a[t, j1, i]) + a[t, j, i] = (xterm + yterm) * dcent