Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add a function to fill NaNs in grids by interpolation #440

Open
wants to merge 8 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions verde/tests/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,27 @@
meshgrid_to_1d,
parse_engine,
partition_by_sum,
fill_nans
)


def test_fill_nans():
"""
Test filling NaNs on a small sample grid
"""
grid = xr.DataArray([[1, np.nan, 3],
[4, 5, np.nan],
[np.nan, 7, 8]])

filled_grid = fill_nans(grid)
expected_values = xr.DataArray([[1, 1, 3],
[4, 5, 3],
[4, 7, 8]])

Comment on lines +43 to +46
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The DataArrays should contain coordinates as well. They should be as close as possible to the format of a real dataset. That's how we make the tests more robust.

assert np.any(np.isnan(filled_grid))
assert np.allclose(filled_grid, expected_values)


def test_parse_engine():
"Check that it works for common input"
assert parse_engine("numba") == "numba"
Expand Down
33 changes: 33 additions & 0 deletions verde/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,10 @@
import numpy as np
import pandas as pd
import xarray as xr
import verde as vd
from scipy.spatial import cKDTree


try:
from pykdtree.kdtree import KDTree as pyKDTree
except ImportError:
Expand Down Expand Up @@ -681,6 +683,37 @@ def kdtree(coordinates, use_pykdtree=True, **kwargs):
return tree


def fill_nans(grid):
"""
Fill missing values in a grid by nearest neighbor interpolation

Parameters
----------
grid : :class:`xarray.DataArray`
A 2D grid with one or more data variables.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
A 2D grid with one or more data variables.
A 2D grid with missing values (NaNs).

Returns
-------
grid : :class:`xarray.DataArray`
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
grid : :class:`xarray.DataArray`
filled_grid : :class:`xarray.DataArray`

A 2D grid with the NaN values filled.
"""

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change

filled_grid = grid.copy()

not_nan_values = np.argwhere(~np.isnan(grid.values))
unknown_indices = np.argwhere(np.isnan(grid.values))
Comment on lines +702 to +703
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Use the verde.grid_to_table function and then drop the NaNs from it. It's easier and preserves the coordinates of the grid.


knn = vd.KNeighbors()
easting, northing = not_nan_values[:, 0], not_nan_values[:, 1]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Use the actual coordinates of the grid instead of generating indices. This makes the interpolation work even if the grid is not uniform.

knn.fit((easting, northing), grid.values[not_nan_values[:, 0],
not_nan_values[:, 1]])
predicted_values = knn_imputer.predict((easting, northing))

for i, idx in enumerate(unknown_indices):
filled_grid[tuple(idx)] = predicted_values[i]
Comment on lines +709 to +712
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instead of this, you could use knn.grid and pass in the coordinates of the original grid.


return filled_grid


def partition_by_sum(array, parts):
"""
Partition an array into parts of approximately equal sum.
Expand Down
Loading