diff --git a/docs/user_guide/examples/tutorial_nestedgrids.ipynb b/docs/user_guide/examples/tutorial_nestedgrids.ipynb new file mode 100644 index 000000000..7cc465100 --- /dev/null +++ b/docs/user_guide/examples/tutorial_nestedgrids.ipynb @@ -0,0 +1,581 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# 🖥️ Nested Grids\n", + "\n", + "In some applications, you may have access to fields on different grids that each cover only part of the region of interest. Then, you would like to combine them all together. You may for example have one grid covering the entire region and one or more others that cover only part of the region at a finer resolution. The set of those grids form what we call **nested grids**.\n", + "\n", + "In Parcels v4, we can use the new `uxarray` integration to determine in which grid a particle is located. We will demonstrate how to set up a simulation with multiple nested grids, and how to handle particle transitions between these grids, with a very idealised example." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import matplotlib.tri as mtri\n", + "import numpy as np\n", + "import uxarray as ux\n", + "import xarray as xr\n", + "from matplotlib.colors import ListedColormap\n", + "from shapely.geometry import MultiPoint, Point, Polygon\n", + "from triangle import triangulate\n", + "\n", + "import parcels\n", + "from parcels._datasets.structured.generated import simple_UV_dataset" + ] + }, + { + "cell_type": "markdown", + "id": "2", + "metadata": {}, + "source": [ + "## Setting up the individual Grids and Fields\n", + "\n", + "First we define a helper function to quickly set up a nested dataset" + ] + }, + { + "cell_type": "markdown", + "id": "3", + "metadata": {}, + "source": [ + "```{note}\n", + "If your nested data come from hydrodynamic models, you only have to provide the `grid_polygons` list with polygon coordinates for each grid, but don't need to create the `ds_in` like below.\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "ds_in = []\n", + "grid_polygons = []\n", + "\n", + "\n", + "def setup_nested_ds(polygon):\n", + " xdim, ydim = 20, 20\n", + " ds = simple_UV_dataset(dims=(1, 1, ydim, xdim), mesh=\"spherical\")\n", + " ds[\"lon\"][:] = np.linspace(polygon[:, 0].min(), polygon[:, 0].max(), xdim)\n", + " ds[\"lat\"][:] = np.linspace(polygon[:, 1].min(), polygon[:, 1].max(), ydim)\n", + " return ds" + ] + }, + { + "cell_type": "markdown", + "id": "5", + "metadata": {}, + "source": [ + "Now we create a set of zonal and meridional velocities defined on a small grid. Both the zonal (1m/s) and meridional (0 m/s) velocity are uniform." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "polygon = np.array([(10, 15), (25, 10), (25, 25), (17, 36), (10, 32)])\n", + "grid_polygons.append(polygon)\n", + "\n", + "ds_in.append(setup_nested_ds(polygon))\n", + "ds_in[-1][\"U\"][:] = 1.0\n", + "ds_in[-1][\"V\"][:] = 0.0" + ] + }, + { + "cell_type": "markdown", + "id": "7", + "metadata": {}, + "source": [ + "Then, we define another set of zonal and meridional velocities on a slightly larger grid. In this case, both the zonal and meridional velocity are 1 m/s." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "polygon = np.array([(0, -5), (35, 0), (35, 25), (0, 20)])\n", + "grid_polygons.append(polygon)\n", + "\n", + "ds_in.append(setup_nested_ds(polygon))\n", + "ds_in[-1][\"U\"][:] = 1.0\n", + "ds_in[-1][\"V\"][:] = 1.0" + ] + }, + { + "cell_type": "markdown", + "id": "9", + "metadata": {}, + "source": [ + "Finally, we define another set of velocities on an even larger grid. The zonal velocity is again 1 m/s, but the meridional velocity is now a cosine as a function of longitude." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "polygon = np.array([(-10, -20), (60, -20), (60, 40), (-10, 40)])\n", + "grid_polygons.append(polygon)\n", + "\n", + "ds_in.append(setup_nested_ds(polygon))\n", + "lon_g, lat_g = np.meshgrid(ds_in[-1][\"lon\"], ds_in[-1][\"lat\"])\n", + "ds_in[-1][\"U\"][:] = 1.0\n", + "ds_in[-1][\"V\"][:] = np.cos(lon_g / 5 * np.pi / 2)" + ] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": {}, + "source": [ + "We plot the velocity fields on top of each other, indicating the grid boundaries in red." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "n_grids = len(ds_in)\n", + "cmap = ListedColormap([f\"C{i}\" for i in range(n_grids)])\n", + "\n", + "fig, ax = plt.subplots(figsize=(10, 5))\n", + "\n", + "for i in range(n_grids):\n", + " ds = ds_in[i].isel(time=0, depth=0)\n", + " ax.quiver(ds.lon, ds.lat, ds.U, ds.V, color=f\"C{i}\")\n", + " poly = grid_polygons[i]\n", + " ax.plot(\n", + " np.append(poly[:, 0], poly[0, 0]),\n", + " np.append(poly[:, 1], poly[0, 1]),\n", + " \"-r\",\n", + " lw=2,\n", + " )\n", + "\n", + "plt.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "13", + "metadata": {}, + "source": [ + "Note, as seen in the plot above, that the dataset domains in this case are rectangular, but the polygons that will later define the nested Grid boundaries don't have to be. They can have any shape! This means that we can also use this method to subset parts of a Grid, or to define subregions of a domain." + ] + }, + { + "cell_type": "markdown", + "id": "14", + "metadata": {}, + "source": [ + "## Creating a Delaunay triangulation of the nested Grids\n", + "\n", + "Now comes the important part: we need to create a Delaunay triangulation of the nested Grids, so that we can efficiently determine in which Grid a particle is located at any given time. We use the [`triangle` package](https://rufat.be/triangle/delaunay.html) to perform the triangulation, and `shapely` to handle the geometric operations. \n", + "\n", + "Note that we need to make sure that all the edges of the polygons are also edges of the triangulation. We do this by using a [constrained (PSLG) Delaunay triangulation](https://en.wikipedia.org/wiki/Constrained_Delaunay_triangulation).\n", + "\n", + "The result is a set of triangles covering the nested Grids, which we can use to determine in which Grid a particle is located at any given time. It is important that the list of polygons is ordered from smallest to largest Grid, so that triangles in overlapping areas are assigned to the correct Grid." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "from shapely.geometry import Point, Polygon\n", + "from triangle import triangulate\n", + "\n", + "\n", + "def constrained_triangulate_keep_edges(polygons):\n", + " \"\"\"Constrained triangulation while keeping polygon edges.\n", + "\n", + " Args:\n", + " polygons: list of (Ni,2) numpy arrays (polygon boundary points in order)\n", + "\n", + " Returns:\n", + " pts (P,2) array of vertex coordinates,\n", + " tris (M,3) array of triangle vertex indices,\n", + " face_poly (M,) mapping each triangle to the polygon index or -1.\n", + " \"\"\"\n", + " # build vertices + segments for Triangle PSLG\n", + " verts = []\n", + " segments = []\n", + " offset = 0\n", + " for poly in polygons:\n", + " Ni = len(poly)\n", + " verts.extend(poly.tolist())\n", + " segments.extend([[offset + j, offset + ((j + 1) % Ni)] for j in range(Ni)])\n", + " offset += Ni\n", + " verts = np.asarray(verts, dtype=float)\n", + " segments = np.asarray(segments, dtype=int)\n", + "\n", + " mode = \"p\" # \"p\" = PSLG (constrained triangulation)\n", + " B = triangulate({\"vertices\": verts, \"segments\": segments}, mode)\n", + "\n", + " pts = B[\"vertices\"]\n", + " tris = B[\"triangles\"].astype(int)\n", + "\n", + " # assign triangles to polygons using centroid test\n", + " shapely_polys = [Polygon(p) for p in polygons]\n", + " centers = pts[tris].mean(axis=1)\n", + " face_poly = np.full(len(tris), -1, dtype=int)\n", + " for ti, c in enumerate(centers):\n", + " for ip in range(len(shapely_polys)):\n", + " if shapely_polys[ip].contains(Point(c)):\n", + " face_poly[ti] = ip\n", + " break\n", + "\n", + " return pts, tris, face_poly\n", + "\n", + "\n", + "points, face_tris, face_poly = constrained_triangulate_keep_edges(grid_polygons)" + ] + }, + { + "cell_type": "markdown", + "id": "16", + "metadata": {}, + "source": [ + "We can then plot the resulting triangles to verify that they correctly cover the nested Grids." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(figsize=(10, 5))\n", + "for i in range(n_grids)[::-1]:\n", + " tris = face_tris[face_poly == i]\n", + " ax.triplot(points[:, 0], points[:, 1], tris, label=f\"Grid {i}\", color=f\"C{i}\")\n", + "ax.scatter(points[:, 0], points[:, 1], s=10, c=\"k\")\n", + "ax.set_aspect(\"equal\")\n", + "ax.legend()\n", + "\n", + "# reverse legend labels\n", + "handles, labels = ax.get_legend_handles_labels()\n", + "ax.legend(handles[::-1], labels[::-1])\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "18", + "metadata": {}, + "source": [ + "Then, we convert the triangulation into an (unstructured) `parcels.FieldSet`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "19", + "metadata": {}, + "outputs": [], + "source": [ + "# build an xarray dataset compatible with UGRID / uxarray\n", + "n_node = points.shape[0]\n", + "n_face = face_tris.shape[0]\n", + "n_max_face_nodes = face_tris.shape[1]\n", + "\n", + "ds_tri = xr.Dataset(\n", + " {\n", + " \"node_lon\": ((\"n_node\",), points[:, 0]),\n", + " \"node_lat\": ((\"n_node\",), points[:, 1]),\n", + " \"face_node_connectivity\": ((\"n_face\", \"n_max_face_nodes\"), face_tris),\n", + " \"face_polygon\": (\n", + " (\n", + " \"time\",\n", + " \"nz\",\n", + " \"n_face\",\n", + " ),\n", + " face_poly[np.newaxis, np.newaxis, :],\n", + " {\n", + " \"long_name\": \"Grid ID\",\n", + " \"location\": \"face\",\n", + " \"mesh\": \"delaunay\",\n", + " },\n", + " ),\n", + " },\n", + " coords={\n", + " \"time\": np.array([np.timedelta64(0, \"ns\")]),\n", + " \"nz\": np.array([0]),\n", + " \"n_node\": np.arange(n_node),\n", + " \"n_face\": np.arange(n_face),\n", + " },\n", + " attrs={\"Conventions\": \"UGRID-1.0\"},\n", + ")\n", + "\n", + "uxda = ux.UxDataArray(ds_tri[\"face_polygon\"], uxgrid=ux.Grid(ds_tri))\n", + "\n", + "GridID = parcels.Field(\n", + " \"GridID\",\n", + " uxda,\n", + " # TODO note that here we need to use mesh=\"flat\" otherwise the hashing doesn't work. See https://github.com/Parcels-code/Parcels/pull/2439#discussion_r2627664010\n", + " parcels.UxGrid(uxda.uxgrid, z=uxda[\"nz\"], mesh=\"flat\"),\n", + " interp_method=parcels.interpolators.UxPiecewiseConstantFace,\n", + ")\n", + "fieldset = parcels.FieldSet([GridID])" + ] + }, + { + "cell_type": "markdown", + "id": "20", + "metadata": {}, + "source": [ + "We can confirm that the FieldSet has been created correctly by running a Parcels simulation where particles sample the `GridID` field, which indicates in which Grid each particle is located at any given time." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "21", + "metadata": {}, + "outputs": [], + "source": [ + "X, Y = np.meshgrid(\n", + " np.linspace(-8, 58, 25),\n", + " np.linspace(-18, 38, 20),\n", + ")\n", + "\n", + "NestedGridParticle = parcels.Particle.add_variable(\n", + " parcels.Variable(\"gridID\", dtype=np.int32)\n", + ")\n", + "pset = parcels.ParticleSet(\n", + " fieldset, pclass=NestedGridParticle, lon=X.flatten(), lat=Y.flatten()\n", + ")\n", + "\n", + "\n", + "def SampleGridID(particles, fieldset):\n", + " particles.gridID = fieldset.GridID[particles]\n", + "\n", + "\n", + "pset.execute(\n", + " SampleGridID,\n", + " runtime=np.timedelta64(1, \"s\"),\n", + " dt=np.timedelta64(1, \"s\"),\n", + " verbose_progress=False,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "22", + "metadata": {}, + "source": [ + "Indeed, the visualisation below shows that particles correctly identify the grid they are in based on their location." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "23", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(1, 1, figsize=(10, 5))\n", + "plot_args = {\"cmap\": cmap, \"edgecolors\": \"k\", \"linewidth\": 0.5}\n", + "\n", + "triang = mtri.Triangulation(\n", + " fieldset.GridID.grid.uxgrid.node_lon.values,\n", + " fieldset.GridID.grid.uxgrid.node_lat.values,\n", + " triangles=fieldset.GridID.grid.uxgrid.face_node_connectivity.values,\n", + ")\n", + "facecolors = np.squeeze(uxda[0, :].values)\n", + "\n", + "ax.tripcolor(triang, facecolors=facecolors, shading=\"flat\", **plot_args)\n", + "ax.scatter(pset.lon, pset.lat, c=pset.gridID, **plot_args)\n", + "ax.set_aspect(\"equal\")\n", + "ax.set_title(\n", + " \"Nested Grids visualisation (triangulation and interpolated particle values)\"\n", + ")\n", + "plt.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "24", + "metadata": {}, + "source": [ + "## Advecting particles with nested Grid transitions\n", + "\n", + "We can now set up a particle advection simulation using the nested Grids. We first combine all the Fields into a single FieldSet. \n", + "\n", + "We rename the individual Fields by appending the Grid index to their names, so that we can easily identify which Field belongs to which Grid. We also add the `GridID` Field to the FieldSet (note that Parcels v4 supports combining structured and unstructured Fields into one FieldSet, which is very convenient for this usecase)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "25", + "metadata": {}, + "outputs": [], + "source": [ + "fields = [GridID]\n", + "for i, ds in enumerate(ds_in):\n", + " # TODO : use FieldSet.from_sgrid_convetion here once #2437 is merged\n", + " grid = parcels.XGrid.from_dataset(ds, mesh=\"spherical\")\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", + " fset = parcels.FieldSet([U, V, UV])\n", + "\n", + " for fld in fset.fields.values():\n", + " fld.name = f\"{fld.name}{i}\"\n", + " fields.append(fld)\n", + "fieldset = parcels.FieldSet(fields)" + ] + }, + { + "cell_type": "markdown", + "id": "26", + "metadata": {}, + "source": [ + "We then define a custom Advection kernel that advects particles using the appropriate velocity Field based on the `GridID` at the particle's location. Note that for simplicity, we use Eulerian advection here, but in a real-world application you would typically use a higher-order scheme." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "27", + "metadata": {}, + "outputs": [], + "source": [ + "def AdvectEE_NestedGrids(particles, fieldset):\n", + " particles.gridID = fieldset.GridID[particles]\n", + "\n", + " # TODO because of KernelParticle bug (GH #2143), we need to copy lon/lat/time to local variables\n", + " time = particles.time\n", + " z = particles.z\n", + " lat = particles.lat\n", + " lon = particles.lon\n", + " u = np.zeros_like(particles.lon)\n", + " v = np.zeros_like(particles.lat)\n", + "\n", + " unique_ids = np.unique(particles.gridID)\n", + " for gid in unique_ids:\n", + " mask = particles.gridID == gid\n", + " UVField = getattr(fieldset, f\"UV{gid}\")\n", + " (u[mask], v[mask]) = UVField[time[mask], z[mask], lat[mask], lon[mask]]\n", + "\n", + " particles.dlon += u * particles.dt\n", + " particles.dlat += v * particles.dt\n", + "\n", + " # TODO particle states have to be updated manually because UVField is not called with `particles` argument (becaise of GH #2143)\n", + " particles.state = np.where(\n", + " np.isnan(u) | np.isnan(v),\n", + " parcels.StatusCode.ErrorInterpolation,\n", + " particles.state,\n", + " )\n", + "\n", + "\n", + "lat = np.linspace(-17, 35, 10)\n", + "lon = np.full(len(lat), -5)\n", + "\n", + "pset = parcels.ParticleSet(fieldset, pclass=NestedGridParticle, lon=lon, lat=lat)\n", + "ofile = parcels.ParticleFile(\n", + " \"nestedgrid_particles.zarr\", outputdt=np.timedelta64(1, \"D\")\n", + ")\n", + "pset.execute(\n", + " AdvectEE_NestedGrids,\n", + " runtime=np.timedelta64(60, \"D\"),\n", + " dt=np.timedelta64(1, \"D\"),\n", + " output_file=ofile,\n", + " verbose_progress=False,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "28", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(1, 1, figsize=(10, 5))\n", + "\n", + "ds_out = xr.open_zarr(\"nestedgrid_particles.zarr\")\n", + "\n", + "plt.plot(ds_out.lon.T, ds_out.lat.T, \"k\", linewidth=0.5)\n", + "sc = ax.scatter(ds_out.lon, ds_out.lat, c=ds_out.gridID, s=4, cmap=cmap, vmin=0, vmax=2)\n", + "xl, yl = ax.get_xlim(), ax.get_ylim()\n", + "\n", + "for i in range(n_grids - 1):\n", + " poly = grid_polygons[i]\n", + " ax.plot(\n", + " np.append(poly[:, 0], poly[0, 0]),\n", + " np.append(poly[:, 1], poly[0, 1]),\n", + " \"-r\",\n", + " lw=1,\n", + " )\n", + "ax.set_xlim(xl)\n", + "ax.set_ylim(yl)\n", + "ax.set_aspect(\"equal\")\n", + "\n", + "cbar = plt.colorbar(sc, ticks=[0, 1, 2], ax=ax)\n", + "cbar.set_label(\"Grid ID\")\n", + "ax.set_title(\"Particle advection through nested Grids\")\n", + "plt.tight_layout\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "29", + "metadata": {}, + "source": [ + "Indeed, we can see that the particles follow the cosine oscillation pattern in the coarsest grid, move northeast in the finer grid, and move purely zonally in the finest grid.\n", + "\n", + "Note that in the plot above the particles at higher latitudes move farther eastward because of the curvature of the earth, so that is expected." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "docs", + "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.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user_guide/index.md b/docs/user_guide/index.md index a46289455..419594162 100644 --- a/docs/user_guide/index.md +++ b/docs/user_guide/index.md @@ -31,6 +31,7 @@ TODO: Add links to Reference API throughout examples/explanation_grids.md examples/tutorial_nemo_curvilinear.ipynb examples/tutorial_unitconverters.ipynb +examples/tutorial_nestedgrids.ipynb ``` diff --git a/docs/user_guide/v4-migration.md b/docs/user_guide/v4-migration.md index a5685e58e..44633515b 100644 --- a/docs/user_guide/v4-migration.md +++ b/docs/user_guide/v4-migration.md @@ -39,6 +39,7 @@ Version 4 of Parcels is unreleased at the moment. The information in this migrat - `Field.eval()` returns an array of floats instead of a single float (related to the vectorization) - `Field.eval()` does not throw OutOfBounds or other errors +- The `NestedField` class has been removed. See the Nested Grids how-to guide for how to set up Nested Grids in v4. ## GridSet diff --git a/pixi.toml b/pixi.toml index 7b6fafac1..3e28cc342 100644 --- a/pixi.toml +++ b/pixi.toml @@ -30,6 +30,7 @@ xgcm = ">=0.9.0" cf_xarray = ">=0.8.6" cftime = ">=1.6.3" pooch = ">=1.8.0" +py-triangle = ">=20250106,<20250107" [dependencies] parcels = { path = "." } diff --git a/tests-v3/test_fieldset_sampling.py b/tests-v3/test_fieldset_sampling.py index 11dec1988..1a3c68623 100644 --- a/tests-v3/test_fieldset_sampling.py +++ b/tests-v3/test_fieldset_sampling.py @@ -14,7 +14,6 @@ Geographic, Particle, ParticleSet, - StatusCode, Variable, ) from tests.utils import create_fieldset_global @@ -796,82 +795,6 @@ def test_multiple_grid_addlater_error(): assert fail -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="Implementation of NestedFields is being reconsidered in v4.") -def test_nestedfields(): - from parcels import NestedField - - xdim = 10 - ydim = 20 - - U1 = Field( - "U1", - 0.1 * np.ones((ydim, xdim), dtype=np.float32), - lon=np.linspace(0.0, 1.0, xdim, dtype=np.float32), - lat=np.linspace(0.0, 1.0, ydim, dtype=np.float32), - ) - V1 = Field( - "V1", - 0.2 * np.ones((ydim, xdim), dtype=np.float32), - lon=np.linspace(0.0, 1.0, xdim, dtype=np.float32), - lat=np.linspace(0.0, 1.0, ydim, dtype=np.float32), - ) - U2 = Field( - "U2", - 0.3 * np.ones((ydim, xdim), dtype=np.float32), - lon=np.linspace(0.0, 2.0, xdim, dtype=np.float32), - lat=np.linspace(0.0, 2.0, ydim, dtype=np.float32), - ) - V2 = Field( - "V2", - 0.4 * np.ones((ydim, xdim), dtype=np.float32), - lon=np.linspace(0.0, 2.0, xdim, dtype=np.float32), - lat=np.linspace(0.0, 2.0, ydim, dtype=np.float32), - ) - U = NestedField("U", [U1, U2]) - V = NestedField("V", [V1, V2]) - fieldset = FieldSet(U, V) - - P1 = Field( - "P1", - 0.1 * np.ones((ydim, xdim), dtype=np.float32), - lon=np.linspace(0.0, 1.0, xdim, dtype=np.float32), - lat=np.linspace(0.0, 1.0, ydim, dtype=np.float32), - ) - P2 = Field( - "P2", - 0.2 * np.ones((ydim, xdim), dtype=np.float32), - lon=np.linspace(0.0, 2.0, xdim, dtype=np.float32), - lat=np.linspace(0.0, 2.0, ydim, dtype=np.float32), - ) - P = NestedField("P", [P1, P2]) - fieldset.add_field(P) - - def Recover(particle, fieldset, time): # pragma: no cover - if particle.state == StatusCode.ErrorOutOfBounds: - particle.dlon = 0 - particle.dlat = 0 - particle.ddepth = 0 - particle.lon = 0 - particle.lat = 0 - particle.p = 999 - particle.state = StatusCode.Evaluate - - pset = ParticleSet(fieldset, pclass=pclass(), lon=[0], lat=[0.3]) - pset.execute(AdvectionRK4 + pset.Kernel(SampleP), runtime=2, dt=1) - assert np.isclose(pset.lat[0], 0.5) - assert np.isclose(pset.p[0], 0.1) - pset = ParticleSet(fieldset, pclass=pclass(), lon=[0], lat=[1.1]) - pset.execute(AdvectionRK4 + pset.Kernel(SampleP), runtime=2, dt=1) - assert np.isclose(pset.lat[0], 1.5) - assert np.isclose(pset.p[0], 0.2) - pset = ParticleSet(fieldset, pclass=pclass(), lon=[0], lat=[2.3]) - pset.execute(pset.Kernel(AdvectionRK4) + SampleP + Recover, runtime=1, dt=1) - assert np.isclose(pset.lat[0], 0) - assert np.isclose(pset.p[0], 999) - assert np.allclose(fieldset.UV[0][0, 0, 0, 0], [0.1, 0.2]) - - def test_fieldset_sampling_updating_order(tmp_zarrfile): def calc_p(t, y, x): return 10 * t + x + 0.2 * y