Skip to content

Commit 070cbea

Browse files
reint-fischerreint-fischer
reint-fischer
authored and
reint-fischer
committed
rewrite nemo curvilinear tutorial and add new periodic BC explanation
1 parent d65eb27 commit 070cbea

3 files changed

Lines changed: 348 additions & 633 deletions

File tree

‎docs/user_guide/examples/tutorial_nemo_curvilinear.ipynb‎

Lines changed: 347 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,347 @@
1+
{
2+
"cells": [
3+
{
4+
"attachments": {},
5+
"cell_type": "markdown",
6+
"metadata": {},
7+
"source": [
8+
"# Curvilinear grids\n"
9+
]
10+
},
11+
{
12+
"attachments": {},
13+
"cell_type": "markdown",
14+
"metadata": {},
15+
"source": [
16+
"Parcels supports [curvilinear grids](https://www.nemo-ocean.eu/doc/node108.html) such as those used in the [NEMO models](https://www.nemo-ocean.eu/).\n",
17+
"\n",
18+
"We will be using the example dataset `NemoCurvilinear_data`. These fields are a purely zonal flow on an aqua-planet (so zonal-velocity is 1 m/s and meridional-velocity is 0 m/s everywhere, and no land). However, because of the curvilinear grid, the `U` and `V` fields vary for the rotated gridcells north of 20N.\n"
19+
]
20+
},
21+
{
22+
"cell_type": "code",
23+
"execution_count": null,
24+
"metadata": {},
25+
"outputs": [],
26+
"source": [
27+
"import matplotlib.pyplot as plt\n",
28+
"import numpy as np\n",
29+
"import xarray as xr\n",
30+
"\n",
31+
"import parcels"
32+
]
33+
},
34+
{
35+
"attachments": {},
36+
"cell_type": "markdown",
37+
"metadata": {},
38+
"source": [
39+
"We can create a `FieldSet` just like we do for normal grids.\n",
40+
"Note that NEMO is discretised on a C-grid. U and V velocities are not located on the same nodes (see https://www.nemo-ocean.eu/doc/node19.html ).\n",
41+
"\n",
42+
"```\n",
43+
" __V1__\n",
44+
"| |\n",
45+
"U0 U1\n",
46+
"|__V0__|\n",
47+
"```\n",
48+
"\n",
49+
"To interpolate U, V velocities on the C-grid, Parcels needs to read the f-nodes, which are located on the corners of the cells (for indexing details: https://www.nemo-ocean.eu/doc/img360.png).\n",
50+
"\n",
51+
"```{note}\n",
52+
"TODO: add link to grid indexing explanation once implemented in v4\n",
53+
"```\n"
54+
]
55+
},
56+
{
57+
"cell_type": "code",
58+
"execution_count": null,
59+
"metadata": {},
60+
"outputs": [],
61+
"source": [
62+
"data_folder = parcels.download_example_dataset(\"NemoCurvilinear_data\")\n",
63+
"files = data_folder.glob(\"*.nc4\")\n",
64+
"ds_fields = xr.open_mfdataset(\n",
65+
" files, combine=\"nested\", data_vars=\"minimal\", coords=\"minimal\", compat=\"override\"\n",
66+
")\n",
67+
"\n",
68+
"\n",
69+
"# TODO: replace manual fieldset creation with FieldSet.from_nemo() once available\n",
70+
"ds_fields = (\n",
71+
" ds_fields.isel(time_counter=0, drop=True)\n",
72+
" .isel(time=0, drop=True)\n",
73+
" .isel(z_a=0, drop=True)\n",
74+
" .rename({\"glamf\": \"lon\", \"gphif\": \"lat\", \"z\": \"depth\"})\n",
75+
")\n",
76+
"\n",
77+
"import xgcm\n",
78+
"\n",
79+
"xgcm_grid = xgcm.Grid(\n",
80+
" ds_fields,\n",
81+
" coords={\n",
82+
" \"X\": {\"left\": \"x\"},\n",
83+
" \"Y\": {\"left\": \"y\"},\n",
84+
" },\n",
85+
" periodic=False,\n",
86+
" autoparse_metadata=False,\n",
87+
")\n",
88+
"grid = parcels.XGrid(xgcm_grid, mesh=\"spherical\")\n",
89+
"\n",
90+
"U = parcels.Field(\n",
91+
" \"U\", ds_fields[\"U\"], grid, interp_method=parcels.interpolators.XLinear\n",
92+
")\n",
93+
"V = parcels.Field(\n",
94+
" \"V\", ds_fields[\"V\"], grid, interp_method=parcels.interpolators.XLinear\n",
95+
")\n",
96+
"U.units = parcels.GeographicPolar()\n",
97+
"V.units = (\n",
98+
" parcels.GeographicPolar()\n",
99+
") # U and V need GeographicPolar for C-Grid interpolation to work correctly\n",
100+
"UV = parcels.VectorField(\n",
101+
" \"UV\", U, V, vector_interp_method=parcels.interpolators.CGrid_Velocity\n",
102+
")\n",
103+
"fieldset = parcels.FieldSet([U, V, UV])"
104+
]
105+
},
106+
{
107+
"attachments": {},
108+
"cell_type": "markdown",
109+
"metadata": {},
110+
"source": [
111+
"And we can plot the `U` field.\n"
112+
]
113+
},
114+
{
115+
"cell_type": "code",
116+
"execution_count": null,
117+
"metadata": {},
118+
"outputs": [],
119+
"source": [
120+
"fig, ax = plt.subplots(1, 2, figsize=(10, 4))\n",
121+
"pc1 = ds_fields.U.plot(cmap=\"viridis\", ax=ax[0], vmin=0)\n",
122+
"pc2 = ax[1].pcolormesh(\n",
123+
" fieldset.U.grid.lon,\n",
124+
" fieldset.U.grid.lat,\n",
125+
" fieldset.U.data[0, 0, :, :],\n",
126+
" vmin=0,\n",
127+
" vmax=1,\n",
128+
")\n",
129+
"ax[1].set_ylabel(\"Latitude [deg N]\")\n",
130+
"ax[1].set_xlabel(\"Longitude [deg E]\")\n",
131+
"plt.colorbar(pc2, label=\"U\", ax=ax[1])\n",
132+
"plt.tight_layout()\n",
133+
"plt.show()"
134+
]
135+
},
136+
{
137+
"attachments": {},
138+
"cell_type": "markdown",
139+
"metadata": {},
140+
"source": [
141+
"As you see above, the `U` field indeed is 1 m/s south of 20N, but varies with longitude and latitude north of that. We can confirm that Parcels will take care to rotate the `U` and `V` fields by doing a field evaluation at (60N, 50E):"
142+
]
143+
},
144+
{
145+
"cell_type": "code",
146+
"execution_count": null,
147+
"metadata": {},
148+
"outputs": [],
149+
"source": [
150+
"u, v = fieldset.UV.eval(\n",
151+
" np.array([0]), np.array([0]), np.array([20]), np.array([50]), applyConversion=False\n",
152+
") # do not convert m/s to deg/s\n",
153+
"print(f\"(u, v) = ({u[0]:.3f}, {v[0]:.3f})\")\n",
154+
"assert np.isclose(u, 1.0, atol=1e-3)"
155+
]
156+
},
157+
{
158+
"cell_type": "markdown",
159+
"metadata": {},
160+
"source": [
161+
"Now we can run particles as normal."
162+
]
163+
},
164+
{
165+
"cell_type": "code",
166+
"execution_count": null,
167+
"metadata": {
168+
"tags": [
169+
"hide-output"
170+
]
171+
},
172+
"outputs": [],
173+
"source": [
174+
"npart = 20\n",
175+
"lonp = 30 * np.ones(npart)\n",
176+
"latp = np.linspace(-70, 88, npart)\n",
177+
"runtime = np.timedelta64(40, \"D\")\n",
178+
"\n",
179+
"pset = parcels.ParticleSet(fieldset, lon=lonp, lat=latp)\n",
180+
"pfile = parcels.ParticleFile(\n",
181+
" store=\"output_curvilinear.zarr\", outputdt=np.timedelta64(1, \"D\")\n",
182+
")\n",
183+
"\n",
184+
"pset.execute(\n",
185+
" [parcels.kernels.AdvectionEE],\n",
186+
" runtime=runtime,\n",
187+
" dt=np.timedelta64(1, \"D\"),\n",
188+
" output_file=pfile,\n",
189+
")\n",
190+
"np.testing.assert_allclose(pset.lat, latp, atol=1e-1)"
191+
]
192+
},
193+
{
194+
"attachments": {},
195+
"cell_type": "markdown",
196+
"metadata": {},
197+
"source": [
198+
"And then we can plot these trajectories. As expected, all trajectories go exactly zonal and due to the curvature of the earth, ones at higher latitude move more degrees eastward (even though the distance in km is equal for all particles). These particles at high latitudes cross the antimeridian (180 deg E) and keep going east.\n"
199+
]
200+
},
201+
{
202+
"cell_type": "code",
203+
"execution_count": null,
204+
"metadata": {},
205+
"outputs": [],
206+
"source": [
207+
"ds = xr.open_zarr(\"output_curvilinear.zarr\")\n",
208+
"\n",
209+
"plt.plot(ds.lon.T, ds.lat.T, \".-\")\n",
210+
"plt.vlines(np.arange(-180, 900, 360), -90, 90, color=\"r\", label=\"antimeridian\")\n",
211+
"plt.ylabel(\"Latitude [deg N]\")\n",
212+
"plt.xlabel(\"Longitude [deg E]\")\n",
213+
"plt.legend(loc=\"lower right\")\n",
214+
"plt.show()"
215+
]
216+
},
217+
{
218+
"cell_type": "markdown",
219+
"metadata": {},
220+
"source": [
221+
"If we want the `particles.lon` to stay within `[-180,180]` (or `[0,360]`), we can either do this in post-processing, or add a periodic boundary Kernel:"
222+
]
223+
},
224+
{
225+
"cell_type": "code",
226+
"execution_count": null,
227+
"metadata": {},
228+
"outputs": [],
229+
"source": [
230+
"# post processing\n",
231+
"ds[\"lon\"] = ds[\"lon\"] % 360\n",
232+
"ds[\"lon\"] = ds[\"lon\"].where(ds[\"lon\"] <= 180, ds[\"lon\"] - 360)"
233+
]
234+
},
235+
{
236+
"cell_type": "code",
237+
"execution_count": null,
238+
"metadata": {
239+
"tags": [
240+
"hide-output"
241+
]
242+
},
243+
"outputs": [],
244+
"source": [
245+
"# with a Kernel\n",
246+
"def periodicBC(particles, fieldset): # pragma: no cover\n",
247+
" particles.dlon = np.where(\n",
248+
" particles.lon + particles.dlon > 180, particles.dlon - 360, particles.dlon\n",
249+
" )\n",
250+
"\n",
251+
"\n",
252+
"pset = parcels.ParticleSet(fieldset, lon=lonp, lat=latp)\n",
253+
"pfile = parcels.ParticleFile(\n",
254+
" store=\"output_curvilinear_periodic.zarr\", outputdt=np.timedelta64(1, \"D\")\n",
255+
")\n",
256+
"\n",
257+
"pset.execute(\n",
258+
" [parcels.kernels.AdvectionEE, periodicBC],\n",
259+
" runtime=runtime,\n",
260+
" dt=np.timedelta64(1, \"D\"),\n",
261+
" output_file=pfile,\n",
262+
")"
263+
]
264+
},
265+
{
266+
"cell_type": "code",
267+
"execution_count": null,
268+
"metadata": {},
269+
"outputs": [],
270+
"source": [
271+
"ds_periodic = xr.open_zarr(\"output_curvilinear_periodic.zarr\")\n",
272+
"\n",
273+
"fig, ax = plt.subplots(1, 2, figsize=(10, 5))\n",
274+
"ax[0].plot(ds.lon.T, ds.lat.T, \".-\")\n",
275+
"ax[0].vlines(np.arange(-180, 360, 360), -90, 90, color=\"r\", label=\"antimeridian\")\n",
276+
"ax[0].set_ylabel(\"Latitude [deg N]\")\n",
277+
"ax[0].set_xlabel(\"Longitude [deg E]\")\n",
278+
"ax[0].set_title(\"in post processing\")\n",
279+
"ax[0].legend(loc=\"lower center\")\n",
280+
"\n",
281+
"ax[1].plot(ds_periodic.lon.T, ds_periodic.lat.T, \".-\")\n",
282+
"ax[1].vlines(np.arange(-180, 360, 360), -90, 90, color=\"r\", label=\"antimeridian\")\n",
283+
"ax[1].set_ylabel(\"Latitude [deg N]\")\n",
284+
"ax[1].set_xlabel(\"Longitude [deg E]\")\n",
285+
"ax[1].set_title(\"with periodic Kernel\")\n",
286+
"ax[1].legend(loc=\"lower center\")\n",
287+
"\n",
288+
"plt.show()"
289+
]
290+
},
291+
{
292+
"attachments": {},
293+
"cell_type": "markdown",
294+
"metadata": {},
295+
"source": [
296+
"## Speeding up `ParticleSet` initialisation by efficiently finding particle start-locations on the `Grid`\n",
297+
"\n",
298+
"On a Curvilinear grid, determining the location of all `particles` on the grid is more complicated and therefore takes longer than on a Rectilinear grid. Since Parcels version 2.2.2, a function is available on the `ParticleSet` class, that speeds up the look-up. After creating the `ParticleSet`, but before running the `ParticleSet.execute()`, simply call the function `ParticleSet.populate_indices()`.\n"
299+
]
300+
},
301+
{
302+
"cell_type": "code",
303+
"execution_count": null,
304+
"metadata": {},
305+
"outputs": [],
306+
"source": [
307+
"pset = parcels.ParticleSet(fieldset, lon=lonp, lat=latp)\n",
308+
"pset.populate_indices()"
309+
]
310+
},
311+
{
312+
"cell_type": "markdown",
313+
"metadata": {},
314+
"source": [
315+
"For this speed-up to remain in effect while evaluating a `parcels.Field` in a `Kernel`, we need to remember to pass the `particles` object in the square brackets, by using one of:\n",
316+
"```\n",
317+
"particles.P = fieldset.P[particles]\n",
318+
"```\n",
319+
"or\n",
320+
"```\n",
321+
"particles.P = fieldset.P[particles.time, particles.z, particles.lat, particles.lon, particles]\n",
322+
"```"
323+
]
324+
}
325+
],
326+
"metadata": {
327+
"kernelspec": {
328+
"display_name": "test-notebooks",
329+
"language": "python",
330+
"name": "python3"
331+
},
332+
"language_info": {
333+
"codemirror_mode": {
334+
"name": "ipython",
335+
"version": 3
336+
},
337+
"file_extension": ".py",
338+
"mimetype": "text/x-python",
339+
"name": "python",
340+
"nbconvert_exporter": "python",
341+
"pygments_lexer": "ipython3",
342+
"version": "3.11.0"
343+
}
344+
},
345+
"nbformat": 4,
346+
"nbformat_minor": 4
347+
}

0 commit comments

Comments
 (0)