diff --git a/src/parcels/interpolators.py b/src/parcels/interpolators.py index dd22d7cdf..45dda5bea 100644 --- a/src/parcels/interpolators.py +++ b/src/parcels/interpolators.py @@ -22,6 +22,7 @@ "UXPiecewiseLinearNode", "XFreeslip", "XLinear", + "XLinearInvdistLandTracer", "XNearest", "XPartialslip", "ZeroInterpolator", @@ -75,11 +76,11 @@ def _get_corner_data_Agrid( # Y coordinates: [yi, yi, yi+1, yi+1] for each spatial point, repeated for time/z yi_1 = np.clip(yi + 1, 0, data.shape[2] - 1) - yi = np.tile(np.repeat(np.column_stack([yi, yi_1]), 2), (lenT) * (lenZ)) + yi = np.tile(np.array([yi, yi, yi_1, yi_1]).flatten(), lenT * lenZ) # X coordinates: [xi, xi+1, xi, xi+1] for each spatial point, repeated for time/z xi_1 = np.clip(xi + 1, 0, data.shape[3] - 1) - xi = np.tile(np.column_stack([xi, xi_1, xi, xi_1]).flatten(), (lenT) * (lenZ)) + xi = np.tile(np.array([xi, xi_1]).flatten(), lenT * lenZ * 2) # Create DataArrays for indexing selection_dict = { @@ -91,7 +92,7 @@ def _get_corner_data_Agrid( if "time" in data.dims: selection_dict["time"] = xr.DataArray(ti, dims=("points")) - return data.isel(selection_dict).data.reshape(lenT, lenZ, npart, 4) + return data.isel(selection_dict).data.reshape(lenT, lenZ, 2, 2, npart) def XLinear( @@ -114,22 +115,22 @@ def XLinear( corner_data = _get_corner_data_Agrid(data, ti, zi, yi, xi, lenT, lenZ, len(xsi), axis_dim) if lenT == 2: - tau = tau[np.newaxis, :, np.newaxis] - corner_data = corner_data[0, :, :, :] * (1 - tau) + corner_data[1, :, :, :] * tau + tau = tau[np.newaxis, :] + corner_data = corner_data[0, :] * (1 - tau) + corner_data[1, :] * tau else: - corner_data = corner_data[0, :, :, :] + corner_data = corner_data[0, :] if lenZ == 2: - zeta = zeta[:, np.newaxis] - corner_data = corner_data[0, :, :] * (1 - zeta) + corner_data[1, :, :] * zeta + zeta = zeta[np.newaxis, :] + corner_data = corner_data[0, :] * (1 - zeta) + corner_data[1, :] * zeta else: - corner_data = corner_data[0, :, :] + corner_data = corner_data[0, :] value = ( - (1 - xsi) * (1 - eta) * corner_data[:, 0] - + xsi * (1 - eta) * corner_data[:, 1] - + (1 - xsi) * eta * corner_data[:, 2] - + xsi * eta * corner_data[:, 3] + (1 - xsi) * (1 - eta) * corner_data[0, 0, :] + + xsi * (1 - eta) * corner_data[0, 1, :] + + (1 - xsi) * eta * corner_data[1, 0, :] + + xsi * eta * corner_data[1, 1, :] ) return value.compute() if is_dask_collection(value) else value @@ -409,8 +410,8 @@ def _Spatialslip( corner_dataV = _get_corner_data_Agrid(vectorfield.V.data, ti, zi, yi, xi, lenT, lenZ, npart, axis_dim) def is_land(ti: int, zi: int, yi: int, xi: int): - uval = corner_dataU[ti, zi, :, xi + 2 * yi] - vval = corner_dataV[ti, zi, :, xi + 2 * yi] + uval = corner_dataU[ti, zi, yi, xi, :] + vval = corner_dataV[ti, zi, yi, xi, :] return np.where(np.isclose(uval, 0.0) & np.isclose(vval, 0.0), True, False) f_u = np.ones_like(xsi) @@ -571,6 +572,52 @@ def XNearest( return value.compute() if is_dask_collection(value) else value +def XLinearInvdistLandTracer( + particle_positions: dict[str, float | np.ndarray], + grid_positions: dict[_XGRID_AXES, dict[str, int | float | np.ndarray]], + field: Field, +): + """Linear spatial interpolation on a regular grid, where points on land are not used.""" + values = XLinear(particle_positions, grid_positions, field) + + xi, xsi = grid_positions["X"]["index"], grid_positions["X"]["bcoord"] + yi, eta = grid_positions["Y"]["index"], grid_positions["Y"]["bcoord"] + zi, zeta = grid_positions["Z"]["index"], grid_positions["Z"]["bcoord"] + ti, tau = grid_positions["T"]["index"], grid_positions["T"]["bcoord"] + + axis_dim = field.grid.get_axis_dim_mapping(field.data.dims) + lenT = 2 if np.any(tau > 0) else 1 + lenZ = 2 if np.any(zeta > 0) else 1 + + corner_data = _get_corner_data_Agrid(field.data, ti, zi, yi, xi, lenT, lenZ, len(xsi), axis_dim) + + land_mask = np.isnan(corner_data) + nb_land = np.sum(land_mask, axis=(0, 1, 2, 3)) + + if np.any(nb_land): + all_land_mask = nb_land == 4 * lenZ * lenT + values[all_land_mask] = 0.0 + + not_all_land = ~all_land_mask + if np.any(not_all_land): + i_grid = np.arange(2)[None, None, None, :, None] + j_grid = np.arange(2)[None, None, :, None, None] + eta_b = eta[None, None, None, None, :] + xsi_b = xsi[None, None, None, None, :] + + inv_dist = 1.0 / ((eta_b - j_grid) ** 2 + (xsi_b - i_grid) ** 2) + + valid_mask = ~land_mask + weighted = np.where(valid_mask, corner_data * inv_dist, 0.0) + + val = np.sum(weighted, axis=(0, 1, 2, 3)) + w_sum = np.sum(np.where(valid_mask, inv_dist, 0.0), axis=(0, 1, 2, 3)) + + values[not_all_land] = val[not_all_land] / w_sum[not_all_land] + + return values.compute() if is_dask_collection(values) else values + + def UXPiecewiseConstantFace( particle_positions: dict[str, float | np.ndarray], grid_positions: dict[_UXGRID_AXES, dict[str, int | float | np.ndarray]], diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py index db421f713..f270dd24b 100644 --- a/tests/test_interpolation.py +++ b/tests/test_interpolation.py @@ -21,6 +21,7 @@ UXPiecewiseLinearNode, XFreeslip, XLinear, + XLinearInvdistLandTracer, XNearest, XPartialslip, ZeroInterpolator, @@ -68,9 +69,19 @@ def field(): [0.49, 0.49], [0.51, 0.51], [1.49, 6.49], - id="Linear", + id="Linear-1", ), pytest.param(XLinear, np.timedelta64(1, "s"), 2.5, 0.49, 0.51, 13.99, id="Linear-2"), + pytest.param( + XLinear, + [np.timedelta64(0, "s"), np.timedelta64(1, "s"), np.timedelta64(1, "s")], + [0, 0, 2.5], + [0.49, 0.49, 0.49], + [0.51, 0.51, 0.51], + [1.49, 6.49, 13.99], + id="Linear-3", + ), + pytest.param(XLinearInvdistLandTracer, np.timedelta64(1, "s"), 2.5, 0.49, 0.51, 13.99, id="LinearInvDistLand"), pytest.param( XNearest, [np.timedelta64(0, "s"), np.timedelta64(3, "s")], @@ -122,6 +133,38 @@ def test_spatial_slip_interpolation(field, func, t, z, y, x, expected): np.testing.assert_array_almost_equal(velocities, expected) +@pytest.mark.parametrize( + "func, t, z, y, x, expected", + [ + (XLinearInvdistLandTracer, np.timedelta64(1, "s"), 0, 0.5, 0.5, 1.0), + (XLinearInvdistLandTracer, np.timedelta64(1, "s"), 0, 1.5, 1.5, 0.0), + ( + XLinearInvdistLandTracer, + [np.timedelta64(0, "s"), np.timedelta64(1, "s")], + [0, 2], + [0.5, 0.5], + [0.5, 0.5], + 1.0, + ), + ( + XLinearInvdistLandTracer, + [np.timedelta64(0, "s"), np.timedelta64(1, "s")], + [0, 2], + [0.5, 1.5], + [0.5, 1.5], + [1.0, 0.0], + ), + ], +) +def test_invdistland_interpolation(field, func, t, z, y, x, expected): + field.data[:] = 1.0 + field.data[:, :, 1:3, 1:3] = np.nan # Set NaN land value to test inv_dist + field.interp_method = func + + value = field[t, z, y, x] + np.testing.assert_array_almost_equal(value, expected) + + @pytest.mark.parametrize("mesh", ["spherical", "flat"]) def test_interpolation_mesh_type(mesh, npart=10): ds = simple_UV_dataset(mesh=mesh)