diff --git a/docs/user_guide/examples/explanation_interpolation.md b/docs/user_guide/examples/explanation_interpolation.md new file mode 100644 index 0000000000..a1f615ae38 --- /dev/null +++ b/docs/user_guide/examples/explanation_interpolation.md @@ -0,0 +1,76 @@ +# Interpolator explanation + +Interpolation is an important functionality of Parcels. On this page we will discuss the way it is +implemented in **Parcels** and how to write a custom interpolator function. + +```{note} +TODO: expand explanation (similar to Kernel loop explanation) +``` + +When we want to know the state of particles in an environmental field, such as temperature or velocity, +we _evaluate_ the `parcels.Field` at the particles real position in time and space (`t`, `z`, `lat`, `lon`). +In Parcels we can do this using square brackets: + +``` +particles.temperature = fieldset.temperature[particles] +``` + +````{note} +The statement above is shorthand for +```python +particles.temperature = fieldset.temperature[particles.time, particles.z, particles.lat, particles.lon, particles] +``` +where the `particles` argument at the end provides the grid search algorithm with a first guess for the element indices to interpolate on. + +If you want to sample at a different location, or time, that is not necessarily close to the particles location, you can use +```python +particles.temperature = fieldset.temperature[time, depth, lat, lon] +``` +but this could be slower for curvilinear and unstructured because the entire grid needs to be searched. +```` + +The values of the `temperature` field at the particles' positions are determined using an interpolation +method. This interpolation method defines how the discretized values of the `parcels.Field` should +relate to the value at any point within a grid cell. + +Each `parcels.Field` is defined on a (structured) `parcels.XGrid` or (unstructured) `parcels.UXGrid`. +The interpolation function takes information about the particles position relative to this grid (`grid_positions`), +as well as the values of the grid points of the `parcels.Field` in time and space, to calculate +the requested value at the particles location. Note that all grid values are available so that higher-order interpolation is possible. + +## Interpolator API + +The interpolators included in Parcels are designed for common interpolation schemes in Parcels simulations. +If we want to add a custom interpolation method, we need to look at the interpolator API: + +We can write an interpolator function that takes a `parcels.Field` (or `parcels.VectorField`), a dictionary with the `particle_positions` +in real space and time, and a dictionary with the `grid_positions`. + +The `particle_positions` dictionary contains: + +``` +particle_positions = {"time", time, "z", z, "lat", lat, "lon", lon} +``` + +For structured (`X`) grids, the `grid_positions` dictionary contains: + +``` +grid_positions = { + "T": {"index": ti, "bcoord": tau}, + "Z": {"index": zi, "bcoord": zeta}, + "Y": {"index": yi, "bcoord": eta}, + "X": {"index": xi, "bcoord", xsi}, +} +``` + +where `index` is the grid index in the corresponding dimension, and `bcoord` is the barycentric coordinate in the grid cell. + +For unstructured (`UX`) grids, the same dictionary is defined as: + +``` +grid_positions = { + "T": {"index": ti, "bcoord": tau}, + "Z": {"index": zi, "bcoord": zeta}, + "FACE": {"index": fi, "bcoord": bcoord} +} +``` diff --git a/docs/user_guide/examples/tutorial_interpolation.ipynb b/docs/user_guide/examples/tutorial_interpolation.ipynb new file mode 100644 index 0000000000..6180c70cef --- /dev/null +++ b/docs/user_guide/examples/tutorial_interpolation.ipynb @@ -0,0 +1,294 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Using `parcels.interpolators`\n", + "Parcels comes with a number of different interpolation methods for tracer fields, such as temperature. Here, we will look at a few common `parcels.interpolators` for structured (`X`) grids, and how to configure them in an idealised example. For more guidance on the sampling of such fields, check out the [sampling tutorial](./tutorial_sampling).\n", + "\n", + "We first import the relevant modules" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from matplotlib import cm\n", + "\n", + "import parcels" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We create a small 2D structured grid, where `P` is a tracer that we want to interpolate. In each grid cell, `P` has a random value between 0.1 and 1.1. We then set `P[1,1]` to `0`, which for Parcels specifies that this is a land cell\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from parcels._datasets.structured.generated import simple_UV_dataset\n", + "\n", + "ds = simple_UV_dataset(dims=(1, 1, 5, 4), mesh=\"flat\").isel(time=0, depth=0)\n", + "ds[\"lat\"][:] = np.linspace(0.0, 1.0, len(ds.YG))\n", + "ds[\"lon\"][:] = np.linspace(0.0, 1.0, len(ds.XG))\n", + "dx, dy = 1.0 / len(ds.XG), 1.0 / len(ds.YG)\n", + "ds[\"P\"] = ds[\"U\"] + np.random.rand(5, 4) + 0.1\n", + "ds[\"P\"][1, 1] = 0\n", + "ds" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "From this dataset we create a `parcels.FieldSet`. Parcels requires an interpolation method to be set for each `parcels.Field`, which we will later adapt to see the effects of the different interpolators. A common interpolator for fields on structured grids is (tri)linear, implemented in `parcels.interpolators.XLinear`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "grid = parcels.XGrid.from_dataset(ds, mesh=\"flat\")\n", + "U = parcels.Field(\"U\", ds[\"U\"], grid, interp_method=parcels.interpolators.XLinear)\n", + "V = parcels.Field(\"V\", ds[\"V\"], grid, interp_method=parcels.interpolators.XLinear)\n", + "UV = parcels.VectorField(\"UV\", U, V)\n", + "P = parcels.Field(\"P\", ds[\"P\"], grid, interp_method=parcels.interpolators.XLinear)\n", + "fieldset = parcels.FieldSet([U, V, UV, P])" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We create a Particle class that can sample this field `P`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "SampleParticle = parcels.Particle.add_variable(\n", + " parcels.Variable(\"p\", dtype=np.float32, initial=np.nan)\n", + ")\n", + "\n", + "\n", + "def SampleP(particles, fieldset):\n", + " particles.p = fieldset.P[particles]" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now, we perform four different interpolations on `P`, which we can control by setting `fieldset.P.interp_method`. Note that this can always be done _after_ the `FieldSet` creation. We store the results of each interpolation method in an entry in the dictionary `pset`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-output" + ] + }, + "outputs": [], + "source": [ + "pset = {}\n", + "from parcels.interpolators import (\n", + " CGrid_Tracer,\n", + " XLinear,\n", + " XLinearInvdistLandTracer,\n", + " XNearest,\n", + ")\n", + "\n", + "interp_methods = [XLinear, XLinearInvdistLandTracer, XNearest, CGrid_Tracer]\n", + "for p_interp in interp_methods:\n", + " fieldset.P.interp_method = (\n", + " p_interp # setting the interpolation method for fieldset.P\n", + " )\n", + "\n", + " print(fieldset.P.interp_method.__name__)\n", + " xv, yv = np.meshgrid(np.linspace(0, 1, 8), np.linspace(0, 1, 8))\n", + " pset[p_interp.__name__] = parcels.ParticleSet(\n", + " fieldset, pclass=SampleParticle, lon=xv.flatten(), lat=yv.flatten()\n", + " )\n", + " pset[p_interp.__name__].execute(\n", + " SampleP, runtime=np.timedelta64(1, \"D\"), dt=np.timedelta64(1, \"D\")\n", + " )" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "And then we can show each of the four interpolation methods, by plotting the interpolated values on the `Particles` locations (circles) on top of the `Field` values (background colors)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(1, 4, figsize=(18, 6))\n", + "for i, p in enumerate(pset.keys()):\n", + " data = np.copy(fieldset.P.data[0, 0, :, :])\n", + " data[1, 1] = np.nan\n", + " x = np.linspace(-dx / 2, 1 + dx / 2, len(ds.XG) + 1)\n", + " y = np.linspace(-dy / 2, 1 + dy / 2, len(ds.YG) + 1)\n", + " if p == \"CGrid_Tracer\":\n", + " for lat in fieldset.P.grid.lat:\n", + " ax[i].axhline(lat, color=\"k\", linestyle=\"--\")\n", + " for lon in fieldset.P.grid.lon:\n", + " ax[i].axvline(lon, color=\"k\", linestyle=\"--\")\n", + " pc = ax[i].pcolormesh(x, y, data, vmin=0.1, vmax=1.1)\n", + " ax[i].scatter(\n", + " pset[p].lon, pset[p].lat, c=pset[p].p, edgecolors=\"k\", s=50, vmin=0.1, vmax=1.1\n", + " )\n", + " xp, yp = np.meshgrid(fieldset.P.grid.lon, fieldset.P.grid.lat)\n", + " ax[i].plot(xp, yp, \"kx\")\n", + " ax[i].set_title(f\"Using interp_method='{p}'\")\n", + "plt.colorbar(pc, ax=ax, orientation=\"horizontal\", shrink=0.5)\n", + "plt.show()" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The white box is here the 'land' point where the tracer is set to zero and the crosses are the locations of the grid points. As you see, the interpolated value is always equal to the field value if the particle is exactly on the grid point (circles on crosses).\n", + "\n", + "For `interp_method='XNearest'`, the particle values are the same for all particles in a grid cell. They are also the same for `interp_method='CGrid_Tracer'`, but the grid cells have then shifted. That is because in a C-grid, the tracer is defined on the corners of the velocity grid (black dashed lines in right-most panel).\n", + "\n", + "```{note}\n", + "TODO: check C-grid indexing after implementation in v4\n", + "```\n", + "\n", + "For `interp_method='XLinearInvdistLandTracer'`, we see that values are the same as `interp_method='XLinear'` for grid cells that don't border the land point. For grid cells that do border the land cell, the `XLinearInvdistLandTracer` interpolation method gives higher values, as also shown in the difference plot below.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plt.scatter(\n", + " pset[\"XLinear\"].lon,\n", + " pset[\"XLinear\"].lat,\n", + " c=pset[\"XLinearInvdistLandTracer\"].p - pset[\"XLinear\"].p,\n", + " edgecolors=\"k\",\n", + " s=50,\n", + " cmap=cm.bwr,\n", + " vmin=-0.05,\n", + " vmax=0.05,\n", + ")\n", + "xp, yp = np.meshgrid(fieldset.P.grid.lon, fieldset.P.grid.lat)\n", + "plt.plot(xp, yp, \"kx\")\n", + "plt.colorbar()\n", + "plt.title(\n", + " \"Difference between 'interp_method=XLinear' \"\n", + " \"and 'interp_method=XLinearInvdistLandTracer'\"\n", + ")\n", + "plt.show()" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "So in summary, Parcels has four different interpolation schemes for tracers on structured grids:\n", + "\n", + "1. `interp_method=parcels.interpolators.XLinear`: compute linear interpolation\n", + "2. `interp_method=parcels.interpolators.XLinearInvdistLandTracer`: compute linear interpolation except near land (where field value is zero). In that case, inverse distance weighting interpolation is computed, weighting by squares of the distance.\n", + "3. `interp_method=parcels.interpolators.XNearest`: return nearest field value\n", + "4. `interp_method=parcels.interpolators.CGridTracer`: return nearest field value supposing C cells\n", + "\n", + "In the special case where a `parcels.Field` is constant in time and space, such as implementing constant diffusion on a spherical mesh, we can use:\n", + "\n", + "5. `interp_method=parcels.interpolators.XConstantField`: return single value of a Constant Field" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "```{note}\n", + "TODO: link to reference API with all `parcels.interpolators`\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Interpolation on unstructured grids\n", + "```{note}\n", + "TODO: add example on simple unstructured fieldset\n", + "```\n", + "- UXPiecewiseConstantFace\n", + "- UXPiecewiseLinearNode" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Interpolation at boundaries\n", + "In some cases, we need to implement specific boundary conditions, for example to prevent particles from \n", + "getting \"stuck\" near land. [This guide](../examples_v3/documentation_unstuck_Agrid.ipynb) describes \n", + "how to implement this in parcels using `parcels.interpolators.XFreeslip` and `parcels.interpolators.XPartialslip`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Write your own interpolator\n", + "The interpolators included in Parcels are designed for common interpolation schemes in Parcels simulations. If we want to add a custom interpolation method, we can write a custom function (e.g. spline) using the [interpolator API](./explanation_interpolation.md#interpolator-api)." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "test-notebooks", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/docs/user_guide/examples_v3/tutorial_interpolation.ipynb b/docs/user_guide/examples_v3/tutorial_interpolation.ipynb deleted file mode 100644 index 73fcb2cd7f..0000000000 --- a/docs/user_guide/examples_v3/tutorial_interpolation.ipynb +++ /dev/null @@ -1,218 +0,0 @@ -{ - "cells": [ - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Interpolation\n" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Parcels support a range of different interpolation methods for tracers, such as temperature. Here, we will show how these work, in an idealised example.\n", - "\n", - "We first import the relevant modules\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from matplotlib import cm\n", - "\n", - "import parcels" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We create a small 2D grid where `P` is a tracer that we want to interpolate. In each grid cell, `P` has a random value between 0.1 and 1.1. We then set `P[1,1]` to `0`, which for Parcels specifies that this is a land cell\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dims = [5, 4]\n", - "dx, dy = 1.0 / dims[0], 1.0 / dims[1]\n", - "dimensions = {\n", - " \"lat\": np.linspace(0.0, 1.0, dims[0], dtype=np.float32),\n", - " \"lon\": np.linspace(0.0, 1.0, dims[1], dtype=np.float32),\n", - "}\n", - "data = {\n", - " \"U\": np.zeros(dims, dtype=np.float32),\n", - " \"V\": np.zeros(dims, dtype=np.float32),\n", - " \"P\": np.random.rand(dims[0], dims[1]) + 0.1,\n", - "}\n", - "data[\"P\"][1, 1] = 0.0\n", - "fieldset = parcels.FieldSet.from_data(data, dimensions, mesh=\"flat\")" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We create a Particle class that can sample this field\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "SampleParticle = parcels.Particle.add_variable(\"p\", dtype=np.float32)\n", - "\n", - "\n", - "def SampleP(particle, fieldset, time):\n", - " particle.p = fieldset.P[time, particle.depth, particle.lat, particle.lon]" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now, we perform four different interpolation on `P`, which we can control by setting `fieldset.P.interp_method`. Note that this can always be done _after_ the `FieldSet` creation. We store the results of each interpolation method in an entry in the dictionary `pset`.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "pset = {}\n", - "interp_methods = [\"linear\", \"linear_invdist_land_tracer\", \"nearest\", \"cgrid_tracer\"]\n", - "for p_interp in interp_methods:\n", - " fieldset.P.interp_method = (\n", - " p_interp # setting the interpolation method for fieldset.P\n", - " )\n", - "\n", - " xv, yv = np.meshgrid(np.linspace(0, 1, 8), np.linspace(0, 1, 8))\n", - " pset[p_interp] = parcels.ParticleSet(\n", - " fieldset, pclass=SampleParticle, lon=xv.flatten(), lat=yv.flatten()\n", - " )\n", - " pset[p_interp].execute(SampleP, endtime=1, dt=1)" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "And then we can show each of the four interpolation methods, by plotting the interpolated values on the `Particle` locations (circles) on top of the `Field` values (background colors)\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "fig, ax = plt.subplots(1, 4, figsize=(18, 5))\n", - "for i, p in enumerate(pset.keys()):\n", - " data = fieldset.P.data[0, :, :]\n", - " data[1, 1] = np.nan\n", - " x = np.linspace(-dx / 2, 1 + dx / 2, dims[0] + 1)\n", - " y = np.linspace(-dy / 2, 1 + dy / 2, dims[1] + 1)\n", - " if p == \"cgrid_tracer\":\n", - " for lat in fieldset.P.grid.lat:\n", - " ax[i].axhline(lat, color=\"k\", linestyle=\"--\")\n", - " for lon in fieldset.P.grid.lon:\n", - " ax[i].axvline(lon, color=\"k\", linestyle=\"--\")\n", - " ax[i].pcolormesh(y, x, data, vmin=0.1, vmax=1.1)\n", - " ax[i].scatter(\n", - " pset[p].lon, pset[p].lat, c=pset[p].p, edgecolors=\"k\", s=50, vmin=0.1, vmax=1.1\n", - " )\n", - " xp, yp = np.meshgrid(fieldset.P.lon, fieldset.P.lat)\n", - " ax[i].plot(xp, yp, \"kx\")\n", - " ax[i].set_title(f\"Using interp_method='{p}'\")\n", - "plt.show()" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The white box is here the 'land' point where the tracer is set to zero and the crosses are the locations of the grid points. As you see, the interpolated value is always equal to the field value if the particle is exactly on the grid point (circles on crosses).\n", - "\n", - "For `interp_method='nearest'`, the particle values are the same for all particles in a grid cell. They are also the same for `interp_method='cgrid_tracer'`, but the grid cells have then shifted. That is because in a C-grid, the tracer grid cell is on the top-right corner (black dashed lines in right-most panel).\n", - "\n", - "For `interp_method='linear_invdist_land_tracer'`, we see that values are the same as `interp_method='linear'` for grid cells that don't border the land point. For grid cells that do border the land cell, the `linear_invdist_land_tracer` interpolation method gives higher values, as also shown in the difference plot below\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "plt.scatter(\n", - " pset[\"linear\"].lon,\n", - " pset[\"linear\"].lat,\n", - " c=pset[\"linear_invdist_land_tracer\"].p - pset[\"linear\"].p,\n", - " edgecolors=\"k\",\n", - " s=50,\n", - " cmap=cm.bwr,\n", - " vmin=-0.25,\n", - " vmax=0.25,\n", - ")\n", - "plt.colorbar()\n", - "plt.title(\n", - " \"Difference between 'interp_method=linear' \"\n", - " \"and 'interp_method=linear_invdist_land_tracer'\"\n", - ")\n", - "plt.show()" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "So in summary, Parcels has four different interpolation schemes for tracers:\n", - "\n", - "1. `interp_method=linear`: compute linear interpolation\n", - "2. `interp_method=linear_invdist_land_tracer`: compute linear interpolation except near land (where field value is zero). In that case, inverse distance weighting interpolation is computed, weighting by squares of the distance.\n", - "3. `interp_method=nearest`: return nearest field value\n", - "4. `interp_method=cgrid_tracer`: return nearest field value supposing C cells\n" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.11.6" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/user_guide/index.md b/docs/user_guide/index.md index e1ee35f0bf..a251d4c06a 100644 --- a/docs/user_guide/index.md +++ b/docs/user_guide/index.md @@ -49,6 +49,14 @@ examples/tutorial_gsw_density.ipynb examples/tutorial_Argofloats.ipynb ``` +```{toctree} +:caption: Set interpolation method +:titlesonly: + +examples/explanation_interpolation.md +examples/tutorial_interpolation.ipynb +``` + diff --git a/src/parcels/interpolators.py b/src/parcels/interpolators.py index 85215b7c94..abd707eea2 100644 --- a/src/parcels/interpolators.py +++ b/src/parcels/interpolators.py @@ -604,29 +604,39 @@ def XLinearInvdistLandTracer( corner_data = _get_corner_data_Agrid(field.data, ti, zi, yi, xi, lenT, lenZ, len(xsi), axis_dim) - land_mask = np.isnan(corner_data) + land_mask = np.isclose(corner_data, 0.0) 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 = np.asarray(~all_land_mask, dtype=bool) - if np.any(not_all_land): + some_land = np.logical_and(nb_land > 0, nb_land < 4 * lenZ * lenT) + if np.any(some_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) + dist2 = (eta_b - j_grid) ** 2 + (xsi_b - i_grid) ** 2 valid_mask = ~land_mask + # Normal inverse-distance weighting + inv_dist = 1.0 / dist2 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] + values[some_land] = val[some_land] / w_sum[some_land] + + # If a particle hits exactly one of the 8 corner points, extract it + exact_mask = dist2 == 0 & valid_mask + exact_vals = np.sum(np.where(exact_mask, corner_data, 0.0), axis=(0, 1, 2, 3)) + has_exact = np.any(exact_mask, axis=(0, 1, 2, 3)) + + exact_particles = some_land & has_exact + values[exact_particles] = exact_vals[exact_particles] return values.compute() if is_dask_collection(values) else values diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py index 3bd866ce1d..05dac3d4ce 100644 --- a/tests/test_interpolation.py +++ b/tests/test_interpolation.py @@ -155,7 +155,7 @@ def test_spatial_slip_interpolation(field, func, t, z, y, x, expected): ) 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.data[:, :, 1:3, 1:3] = 0 # Set NaN land value to test inv_dist field.interp_method = func value = field[t, z, y, x]