Skip to content

Commit 4dd71fe

Browse files
fluidnumericsJoefluidnumericsJoe
authored
Merge pull request #1946 from OceanParcels/feature/uxarray_xarray_fields
Feature/uxarray xarray fields
2 parents 1316f7a + 687c1e1 commit 4dd71fe

19 files changed

Lines changed: 1427 additions & 2947 deletions

‎docs/examples/tutorial_stommel_uxarray.ipynb‎

Lines changed: 364 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,364 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {},
6+
"source": [
7+
"# Stommel Gyre on Unstructured Grid\n",
8+
"This tutorial walks through creating a UXArray dataset using the Stommel Gyre analytical solution for a closed rectangular domain on a beta-plane"
9+
]
10+
},
11+
{
12+
"cell_type": "code",
13+
"execution_count": null,
14+
"metadata": {},
15+
"outputs": [],
16+
"source": [
17+
"def stommel_fieldset_uxarray(xdim=200, ydim=200):\n",
18+
" \"\"\"Simulate a periodic current along a western boundary, with significantly\n",
19+
" larger velocities along the western edge than the rest of the region\n",
20+
"\n",
21+
" The original test description can be found in: N. Fabbroni, 2009,\n",
22+
" Numerical Simulation of Passive tracers dispersion in the sea,\n",
23+
" Ph.D. dissertation, University of Bologna\n",
24+
" http://amsdottorato.unibo.it/1733/1/Fabbroni_Nicoletta_Tesi.pdf\n",
25+
" \"\"\"\n",
26+
" import math\n",
27+
"\n",
28+
" import numpy as np\n",
29+
" import pandas as pd\n",
30+
" import uxarray as ux\n",
31+
"\n",
32+
" a = b = 66666 * 1e3\n",
33+
" scalefac = 0.00025 # to scale for physically meaningful velocities\n",
34+
"\n",
35+
" # Coordinates of the test fieldset\n",
36+
" # Crowd points to the west edge of the domain\n",
37+
" # using a polyonmial map on x-direction\n",
38+
" x = np.linspace(0, 1, xdim, dtype=np.float32)\n",
39+
" lon, lat = np.meshgrid(a * x, np.linspace(0, b, ydim, dtype=np.float32))\n",
40+
" points = (lon.flatten() / 1111111.111111111, lat.flatten() / 1111111.111111111)\n",
41+
"\n",
42+
" # Create the grid\n",
43+
" uxgrid = ux.Grid.from_points(points, method=\"regional_delaunay\")\n",
44+
" uxgrid.construct_face_centers()\n",
45+
"\n",
46+
" # Define arrays U (zonal), V (meridional) and P (sea surface height)\n",
47+
" U = np.zeros((1, 1, lat.size), dtype=np.float32)\n",
48+
" V = np.zeros((1, 1, lat.size), dtype=np.float32)\n",
49+
" P = np.zeros((1, 1, lat.size), dtype=np.float32)\n",
50+
"\n",
51+
" beta = 2e-11\n",
52+
" r = 1 / (11.6 * 86400)\n",
53+
" es = r / (beta * a)\n",
54+
"\n",
55+
" i = 0\n",
56+
" for x, y in zip(lon.flatten(), lat.flatten()):\n",
57+
" xi = x / a\n",
58+
" yi = y / b\n",
59+
" P[0, 0, i] = (\n",
60+
" (1 - math.exp(-xi / es) - xi) * math.pi * np.sin(math.pi * yi) * scalefac\n",
61+
" )\n",
62+
" U[0, 0, i] = (\n",
63+
" -(1 - math.exp(-xi / es) - xi)\n",
64+
" * math.pi**2\n",
65+
" * np.cos(math.pi * yi)\n",
66+
" * scalefac\n",
67+
" )\n",
68+
" V[0, 0, i] = (\n",
69+
" (math.exp(-xi / es) / es - 1) * math.pi * np.sin(math.pi * yi) * scalefac\n",
70+
" )\n",
71+
" i += 1\n",
72+
"\n",
73+
" u = ux.UxDataArray(\n",
74+
" data=U,\n",
75+
" name=\"u\",\n",
76+
" uxgrid=uxgrid,\n",
77+
" dims=[\"time\", \"nz1\", \"n_node\"],\n",
78+
" coords=dict(\n",
79+
" time=([\"time\"], pd.to_datetime([\"2000-01-01\"])),\n",
80+
" nz1=([\"nz1\"], [0]),\n",
81+
" ),\n",
82+
" attrs=dict(\n",
83+
" description=\"zonal velocity\",\n",
84+
" units=\"m/s\",\n",
85+
" location=\"node\",\n",
86+
" mesh=\"delaunay\",\n",
87+
" ),\n",
88+
" )\n",
89+
" v = ux.UxDataArray(\n",
90+
" data=V,\n",
91+
" name=\"v\",\n",
92+
" uxgrid=uxgrid,\n",
93+
" dims=[\"time\", \"nz1\", \"n_node\"],\n",
94+
" coords=dict(\n",
95+
" time=([\"time\"], pd.to_datetime([\"2000-01-01\"])),\n",
96+
" nz1=([\"nz1\"], [0]),\n",
97+
" ),\n",
98+
" attrs=dict(\n",
99+
" description=\"meridional velocity\",\n",
100+
" units=\"m/s\",\n",
101+
" location=\"node\",\n",
102+
" mesh=\"delaunay\",\n",
103+
" ),\n",
104+
" )\n",
105+
" p = ux.UxDataArray(\n",
106+
" data=P,\n",
107+
" name=\"p\",\n",
108+
" uxgrid=uxgrid,\n",
109+
" dims=[\"time\", \"nz1\", \"n_node\"],\n",
110+
" coords=dict(\n",
111+
" time=([\"time\"], pd.to_datetime([\"2000-01-01\"])),\n",
112+
" nz1=([\"nz1\"], [0]),\n",
113+
" ),\n",
114+
" attrs=dict(\n",
115+
" description=\"pressure\",\n",
116+
" units=\"N/m^2\",\n",
117+
" location=\"node\",\n",
118+
" mesh=\"delaunay\",\n",
119+
" ),\n",
120+
" )\n",
121+
"\n",
122+
" return ux.UxDataset({\"u\": u, \"v\": v, \"p\": p}, uxgrid=uxgrid)\n",
123+
"\n",
124+
"\n",
125+
"uxds = stommel_fieldset_uxarray(50, 50)\n",
126+
"\n",
127+
"uxds.uxgrid.plot(\n",
128+
" line_width=0.5,\n",
129+
" height=500,\n",
130+
" width=1000,\n",
131+
" title=\"Regional Delaunay Regions\",\n",
132+
")"
133+
]
134+
},
135+
{
136+
"cell_type": "code",
137+
"execution_count": null,
138+
"metadata": {},
139+
"outputs": [],
140+
"source": [
141+
"def stommel_fieldset_xarray(xdim=200, ydim=200, grid_type=\"A\"):\n",
142+
" \"\"\"Simulate a periodic current along a western boundary, with significantly\n",
143+
" larger velocities along the western edge than the rest of the region\n",
144+
"\n",
145+
" The original test description can be found in: N. Fabbroni, 2009,\n",
146+
" Numerical Simulation of Passive tracers dispersion in the sea,\n",
147+
" Ph.D. dissertation, University of Bologna\n",
148+
" http://amsdottorato.unibo.it/1733/1/Fabbroni_Nicoletta_Tesi.pdf\n",
149+
" \"\"\"\n",
150+
" import math\n",
151+
"\n",
152+
" import numpy as np\n",
153+
" import pandas as pd\n",
154+
" import xarray as xr\n",
155+
"\n",
156+
" a = b = 10000 * 1e3\n",
157+
" scalefac = 0.05 # to scale for physically meaningful velocities\n",
158+
" dx, dy = a / xdim, b / ydim\n",
159+
"\n",
160+
" # Coordinates of the test fieldset (on A-grid in deg)\n",
161+
" lon = np.linspace(0, a, xdim, dtype=np.float32)\n",
162+
" lat = np.linspace(0, b, ydim, dtype=np.float32)\n",
163+
"\n",
164+
" # Define arrays U (zonal), V (meridional) and P (sea surface height)\n",
165+
" U = np.zeros((1, 1, lat.size, lon.size), dtype=np.float32)\n",
166+
" V = np.zeros((1, 1, lat.size, lon.size), dtype=np.float32)\n",
167+
" P = np.zeros((1, 1, lat.size, lon.size), dtype=np.float32)\n",
168+
"\n",
169+
" beta = 2e-11\n",
170+
" r = 1 / (11.6 * 86400)\n",
171+
" es = r / (beta * a)\n",
172+
"\n",
173+
" for j in range(lat.size):\n",
174+
" for i in range(lon.size):\n",
175+
" xi = lon[i] / a\n",
176+
" yi = lat[j] / b\n",
177+
" P[..., j, i] = (\n",
178+
" (1 - math.exp(-xi / es) - xi)\n",
179+
" * math.pi\n",
180+
" * np.sin(math.pi * yi)\n",
181+
" * scalefac\n",
182+
" )\n",
183+
" if grid_type == \"A\":\n",
184+
" U[..., j, i] = (\n",
185+
" -(1 - math.exp(-xi / es) - xi)\n",
186+
" * math.pi**2\n",
187+
" * np.cos(math.pi * yi)\n",
188+
" * scalefac\n",
189+
" )\n",
190+
" V[..., j, i] = (\n",
191+
" (math.exp(-xi / es) / es - 1)\n",
192+
" * math.pi\n",
193+
" * np.sin(math.pi * yi)\n",
194+
" * scalefac\n",
195+
" )\n",
196+
"\n",
197+
" time = pd.to_datetime([\"2000-01-01\"])\n",
198+
" z = [0]\n",
199+
" if grid_type == \"C\":\n",
200+
" V[..., :, 1:] = (P[..., :, 1:] - P[..., :, 0:-1]) / dx * a\n",
201+
" U[..., 1:, :] = -(P[..., 1:, :] - P[..., 0:-1, :]) / dy * b\n",
202+
" u_dims = [\"time\", \"nz1\", \"face_lat\", \"node_lon\"]\n",
203+
" u_lat = lat\n",
204+
" u_lon = lon - dx * 0.5\n",
205+
" u_location = \"x_edge\"\n",
206+
" v_dims = [\"time\", \"nz1\", \"node_lat\", \"face_lon\"]\n",
207+
" v_lat = lat - dy * 0.5\n",
208+
" v_lon = lon\n",
209+
" v_location = \"y_edge\"\n",
210+
" p_dims = [\"time\", \"nz1\", \"face_lat\", \"face_lon\"]\n",
211+
" p_lat = lat\n",
212+
" p_lon = lon\n",
213+
" p_location = \"face\"\n",
214+
"\n",
215+
" else:\n",
216+
" u_dims = [\"time\", \"nz1\", \"node_lat\", \"node_lon\"]\n",
217+
" v_dims = [\"time\", \"nz1\", \"node_lat\", \"node_lon\"]\n",
218+
" p_dims = [\"time\", \"nz1\", \"node_lat\", \"node_lon\"]\n",
219+
" u_lat = lat\n",
220+
" u_lon = lon\n",
221+
" v_lat = lat\n",
222+
" v_lon = lon\n",
223+
" u_location = \"node\"\n",
224+
" v_location = \"node\"\n",
225+
" p_lat = lat\n",
226+
" p_lon = lon\n",
227+
" p_location = \"node\"\n",
228+
"\n",
229+
" u = xr.DataArray(\n",
230+
" data=U,\n",
231+
" name=\"u\",\n",
232+
" dims=u_dims,\n",
233+
" coords=[time, z, u_lat, u_lon],\n",
234+
" attrs=dict(\n",
235+
" description=\"zonal velocity\",\n",
236+
" units=\"m/s\",\n",
237+
" location=u_location,\n",
238+
" mesh=f\"Arakawa-{grid_type}\",\n",
239+
" ),\n",
240+
" )\n",
241+
" v = xr.DataArray(\n",
242+
" data=V,\n",
243+
" name=\"v\",\n",
244+
" dims=v_dims,\n",
245+
" coords=[time, z, v_lat, v_lon],\n",
246+
" attrs=dict(\n",
247+
" description=\"meridional velocity\",\n",
248+
" units=\"m/s\",\n",
249+
" location=v_location,\n",
250+
" mesh=f\"Arakawa-{grid_type}\",\n",
251+
" ),\n",
252+
" )\n",
253+
" p = xr.DataArray(\n",
254+
" data=P,\n",
255+
" name=\"p\",\n",
256+
" dims=p_dims,\n",
257+
" coords=[time, z, p_lat, p_lon],\n",
258+
" attrs=dict(\n",
259+
" description=\"pressure\",\n",
260+
" units=\"N/m^2\",\n",
261+
" location=p_location,\n",
262+
" mesh=f\"Arakawa-{grid_type}\",\n",
263+
" ),\n",
264+
" )\n",
265+
"\n",
266+
" return xr.Dataset({\"u\": u, \"v\": v, \"p\": p})\n",
267+
"\n",
268+
"\n",
269+
"ds_arakawa_a = stommel_fieldset_xarray(50, 50, \"A\")\n",
270+
"ds_arakawa_c = stommel_fieldset_xarray(50, 50, \"C\")"
271+
]
272+
},
273+
{
274+
"cell_type": "code",
275+
"execution_count": null,
276+
"metadata": {},
277+
"outputs": [],
278+
"source": [
279+
"ds_arakawa_a"
280+
]
281+
},
282+
{
283+
"cell_type": "code",
284+
"execution_count": null,
285+
"metadata": {},
286+
"outputs": [],
287+
"source": [
288+
"ds_arakawa_a[\"u\"].attrs"
289+
]
290+
},
291+
{
292+
"cell_type": "code",
293+
"execution_count": null,
294+
"metadata": {},
295+
"outputs": [],
296+
"source": [
297+
"ds_arakawa_c"
298+
]
299+
},
300+
{
301+
"cell_type": "code",
302+
"execution_count": null,
303+
"metadata": {},
304+
"outputs": [],
305+
"source": [
306+
"import numpy as np\n",
307+
"\n",
308+
"min_length_scale = 1111111.111111111 * np.sqrt(np.min(uxds.uxgrid.face_areas))\n",
309+
"print(min_length_scale)\n",
310+
"\n",
311+
"max_v = np.sqrt(uxds[\"u\"] ** 2 + uxds[\"v\"] ** 2).max()\n",
312+
"print(max_v)\n",
313+
"\n",
314+
"cfl = 0.1\n",
315+
"dt = cfl * min_length_scale / max_v\n",
316+
"print(dt)"
317+
]
318+
},
319+
{
320+
"cell_type": "code",
321+
"execution_count": null,
322+
"metadata": {},
323+
"outputs": [],
324+
"source": [
325+
"from datetime import timedelta\n",
326+
"\n",
327+
"import numpy as np\n",
328+
"import uxarray as ux\n",
329+
"\n",
330+
"from parcels import Particle, ParticleSet, UxAdvectionEuler, UXFieldSet\n",
331+
"\n",
332+
"npart = 10\n",
333+
"fieldset = UXFieldSet(uxds)\n",
334+
"# pset = ParticleSet(\n",
335+
"# fieldset,\n",
336+
"# pclass=Particle,\n",
337+
"# lon=np.linspace(1, 59, npart),\n",
338+
"# lat=np.zeros(npart)+30)\n",
339+
"# pset.execute(UxAdvectionEuler, runtime=timedelta(hours=24), dt=timedelta(seconds=dt))"
340+
]
341+
}
342+
],
343+
"metadata": {
344+
"kernelspec": {
345+
"display_name": "parcels",
346+
"language": "python",
347+
"name": "python3"
348+
},
349+
"language_info": {
350+
"codemirror_mode": {
351+
"name": "ipython",
352+
"version": 3
353+
},
354+
"file_extension": ".py",
355+
"mimetype": "text/x-python",
356+
"name": "python",
357+
"nbconvert_exporter": "python",
358+
"pygments_lexer": "ipython3",
359+
"version": "3.13.2"
360+
}
361+
},
362+
"nbformat": 4,
363+
"nbformat_minor": 2
364+
}

‎environment.yml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ dependencies: #! Keep in sync with [tool.pixi.dependencies] in pyproject.toml
1717
- dask>=2.0
1818
- scikit-learn
1919
- zarr>=2.11.0,!=2.18.0,<3
20+
- uxaray>=2025.3.0
2021
- pooch
2122

2223
# Notebooks

0 commit comments

Comments
 (0)