Feature/uxarray xarray fields

docs/examples/tutorial_stommel_uxarray.ipynb

@@ -0,0 +1,364 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Stommel Gyre on Unstructured Grid\n",
+    "This tutorial walks through creating a UXArray dataset using the Stommel Gyre analytical solution for a closed rectangular domain on a beta-plane"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def stommel_fieldset_uxarray(xdim=200, ydim=200):\n",
+    "    \"\"\"Simulate a periodic current along a western boundary, with significantly\n",
+    "    larger velocities along the western edge than the rest of the region\n",
+    "\n",
+    "    The original test description can be found in: N. Fabbroni, 2009,\n",
+    "    Numerical Simulation of Passive tracers dispersion in the sea,\n",
+    "    Ph.D. dissertation, University of Bologna\n",
+    "    http://amsdottorato.unibo.it/1733/1/Fabbroni_Nicoletta_Tesi.pdf\n",
+    "    \"\"\"\n",
+    "    import math\n",
+    "\n",
+    "    import numpy as np\n",
+    "    import pandas as pd\n",
+    "    import uxarray as ux\n",
+    "\n",
+    "    a = b = 66666 * 1e3\n",
+    "    scalefac = 0.00025  # to scale for physically meaningful velocities\n",
+    "\n",
+    "    # Coordinates of the test fieldset\n",
+    "    # Crowd points to the west edge of the domain\n",
+    "    # using a polyonmial map on x-direction\n",
+    "    x = np.linspace(0, 1, xdim, dtype=np.float32)\n",
+    "    lon, lat = np.meshgrid(a * x, np.linspace(0, b, ydim, dtype=np.float32))\n",
+    "    points = (lon.flatten() / 1111111.111111111, lat.flatten() / 1111111.111111111)\n",
+    "\n",
+    "    # Create the grid\n",
+    "    uxgrid = ux.Grid.from_points(points, method=\"regional_delaunay\")\n",
+    "    uxgrid.construct_face_centers()\n",
+    "\n",
+    "    # Define arrays U (zonal), V (meridional) and P (sea surface height)\n",
+    "    U = np.zeros((1, 1, lat.size), dtype=np.float32)\n",
+    "    V = np.zeros((1, 1, lat.size), dtype=np.float32)\n",
+    "    P = np.zeros((1, 1, lat.size), dtype=np.float32)\n",
+    "\n",
+    "    beta = 2e-11\n",
+    "    r = 1 / (11.6 * 86400)\n",
+    "    es = r / (beta * a)\n",
+    "\n",
+    "    i = 0\n",
+    "    for x, y in zip(lon.flatten(), lat.flatten()):\n",
+    "        xi = x / a\n",
+    "        yi = y / b\n",
+    "        P[0, 0, i] = (\n",
+    "            (1 - math.exp(-xi / es) - xi) * math.pi * np.sin(math.pi * yi) * scalefac\n",
+    "        )\n",
+    "        U[0, 0, i] = (\n",
+    "            -(1 - math.exp(-xi / es) - xi)\n",
+    "            * math.pi**2\n",
+    "            * np.cos(math.pi * yi)\n",
+    "            * scalefac\n",
+    "        )\n",
+    "        V[0, 0, i] = (\n",
+    "            (math.exp(-xi / es) / es - 1) * math.pi * np.sin(math.pi * yi) * scalefac\n",
+    "        )\n",
+    "        i += 1\n",
+    "\n",
+    "    u = ux.UxDataArray(\n",
+    "        data=U,\n",
+    "        name=\"u\",\n",
+    "        uxgrid=uxgrid,\n",
+    "        dims=[\"time\", \"nz1\", \"n_node\"],\n",
+    "        coords=dict(\n",
+    "            time=([\"time\"], pd.to_datetime([\"2000-01-01\"])),\n",
+    "            nz1=([\"nz1\"], [0]),\n",
+    "        ),\n",
+    "        attrs=dict(\n",
+    "            description=\"zonal velocity\",\n",
+    "            units=\"m/s\",\n",
+    "            location=\"node\",\n",
+    "            mesh=\"delaunay\",\n",
+    "        ),\n",
+    "    )\n",
+    "    v = ux.UxDataArray(\n",
+    "        data=V,\n",
+    "        name=\"v\",\n",
+    "        uxgrid=uxgrid,\n",
+    "        dims=[\"time\", \"nz1\", \"n_node\"],\n",
+    "        coords=dict(\n",
+    "            time=([\"time\"], pd.to_datetime([\"2000-01-01\"])),\n",
+    "            nz1=([\"nz1\"], [0]),\n",
+    "        ),\n",
+    "        attrs=dict(\n",
+    "            description=\"meridional velocity\",\n",
+    "            units=\"m/s\",\n",
+    "            location=\"node\",\n",
+    "            mesh=\"delaunay\",\n",
+    "        ),\n",
+    "    )\n",
+    "    p = ux.UxDataArray(\n",
+    "        data=P,\n",
+    "        name=\"p\",\n",
+    "        uxgrid=uxgrid,\n",
+    "        dims=[\"time\", \"nz1\", \"n_node\"],\n",
+    "        coords=dict(\n",
+    "            time=([\"time\"], pd.to_datetime([\"2000-01-01\"])),\n",
+    "            nz1=([\"nz1\"], [0]),\n",
+    "        ),\n",
+    "        attrs=dict(\n",
+    "            description=\"pressure\",\n",
+    "            units=\"N/m^2\",\n",
+    "            location=\"node\",\n",
+    "            mesh=\"delaunay\",\n",
+    "        ),\n",
+    "    )\n",
+    "\n",
+    "    return ux.UxDataset({\"u\": u, \"v\": v, \"p\": p}, uxgrid=uxgrid)\n",
+    "\n",
+    "\n",
+    "uxds = stommel_fieldset_uxarray(50, 50)\n",
+    "\n",
+    "uxds.uxgrid.plot(\n",
+    "    line_width=0.5,\n",
+    "    height=500,\n",
+    "    width=1000,\n",
+    "    title=\"Regional Delaunay Regions\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def stommel_fieldset_xarray(xdim=200, ydim=200, grid_type=\"A\"):\n",
+    "    \"\"\"Simulate a periodic current along a western boundary, with significantly\n",
+    "    larger velocities along the western edge than the rest of the region\n",
+    "\n",
+    "    The original test description can be found in: N. Fabbroni, 2009,\n",
+    "    Numerical Simulation of Passive tracers dispersion in the sea,\n",
+    "    Ph.D. dissertation, University of Bologna\n",
+    "    http://amsdottorato.unibo.it/1733/1/Fabbroni_Nicoletta_Tesi.pdf\n",
+    "    \"\"\"\n",
+    "    import math\n",
+    "\n",
+    "    import numpy as np\n",
+    "    import pandas as pd\n",
+    "    import xarray as xr\n",
+    "\n",
+    "    a = b = 10000 * 1e3\n",
+    "    scalefac = 0.05  # to scale for physically meaningful velocities\n",
+    "    dx, dy = a / xdim, b / ydim\n",
+    "\n",
+    "    # Coordinates of the test fieldset (on A-grid in deg)\n",
+    "    lon = np.linspace(0, a, xdim, dtype=np.float32)\n",
+    "    lat = np.linspace(0, b, ydim, dtype=np.float32)\n",
+    "\n",
+    "    # Define arrays U (zonal), V (meridional) and P (sea surface height)\n",
+    "    U = np.zeros((1, 1, lat.size, lon.size), dtype=np.float32)\n",
+    "    V = np.zeros((1, 1, lat.size, lon.size), dtype=np.float32)\n",
+    "    P = np.zeros((1, 1, lat.size, lon.size), dtype=np.float32)\n",
+    "\n",
+    "    beta = 2e-11\n",
+    "    r = 1 / (11.6 * 86400)\n",
+    "    es = r / (beta * a)\n",
+    "\n",
+    "    for j in range(lat.size):\n",
+    "        for i in range(lon.size):\n",
+    "            xi = lon[i] / a\n",
+    "            yi = lat[j] / b\n",
+    "            P[..., j, i] = (\n",
+    "                (1 - math.exp(-xi / es) - xi)\n",
+    "                * math.pi\n",
+    "                * np.sin(math.pi * yi)\n",
+    "                * scalefac\n",
+    "            )\n",
+    "            if grid_type == \"A\":\n",
+    "                U[..., j, i] = (\n",
+    "                    -(1 - math.exp(-xi / es) - xi)\n",
+    "                    * math.pi**2\n",
+    "                    * np.cos(math.pi * yi)\n",
+    "                    * scalefac\n",
+    "                )\n",
+    "                V[..., j, i] = (\n",
+    "                    (math.exp(-xi / es) / es - 1)\n",
+    "                    * math.pi\n",
+    "                    * np.sin(math.pi * yi)\n",
+    "                    * scalefac\n",
+    "                )\n",
+    "\n",
+    "    time = pd.to_datetime([\"2000-01-01\"])\n",
+    "    z = [0]\n",
+    "    if grid_type == \"C\":\n",
+    "        V[..., :, 1:] = (P[..., :, 1:] - P[..., :, 0:-1]) / dx * a\n",
+    "        U[..., 1:, :] = -(P[..., 1:, :] - P[..., 0:-1, :]) / dy * b\n",
+    "        u_dims = [\"time\", \"nz1\", \"face_lat\", \"node_lon\"]\n",
+    "        u_lat = lat\n",
+    "        u_lon = lon - dx * 0.5\n",
+    "        u_location = \"x_edge\"\n",
+    "        v_dims = [\"time\", \"nz1\", \"node_lat\", \"face_lon\"]\n",
+    "        v_lat = lat - dy * 0.5\n",
+    "        v_lon = lon\n",
+    "        v_location = \"y_edge\"\n",
+    "        p_dims = [\"time\", \"nz1\", \"face_lat\", \"face_lon\"]\n",
+    "        p_lat = lat\n",
+    "        p_lon = lon\n",
+    "        p_location = \"face\"\n",
+    "\n",
+    "    else:\n",
+    "        u_dims = [\"time\", \"nz1\", \"node_lat\", \"node_lon\"]\n",
+    "        v_dims = [\"time\", \"nz1\", \"node_lat\", \"node_lon\"]\n",
+    "        p_dims = [\"time\", \"nz1\", \"node_lat\", \"node_lon\"]\n",
+    "        u_lat = lat\n",
+    "        u_lon = lon\n",
+    "        v_lat = lat\n",
+    "        v_lon = lon\n",
+    "        u_location = \"node\"\n",
+    "        v_location = \"node\"\n",
+    "        p_lat = lat\n",
+    "        p_lon = lon\n",
+    "        p_location = \"node\"\n",
+    "\n",
+    "    u = xr.DataArray(\n",
+    "        data=U,\n",
+    "        name=\"u\",\n",
+    "        dims=u_dims,\n",
+    "        coords=[time, z, u_lat, u_lon],\n",
+    "        attrs=dict(\n",
+    "            description=\"zonal velocity\",\n",
+    "            units=\"m/s\",\n",
+    "            location=u_location,\n",
+    "            mesh=f\"Arakawa-{grid_type}\",\n",
+    "        ),\n",
+    "    )\n",
+    "    v = xr.DataArray(\n",
+    "        data=V,\n",
+    "        name=\"v\",\n",
+    "        dims=v_dims,\n",
+    "        coords=[time, z, v_lat, v_lon],\n",
+    "        attrs=dict(\n",
+    "            description=\"meridional velocity\",\n",
+    "            units=\"m/s\",\n",
+    "            location=v_location,\n",
+    "            mesh=f\"Arakawa-{grid_type}\",\n",
+    "        ),\n",
+    "    )\n",
+    "    p = xr.DataArray(\n",
+    "        data=P,\n",
+    "        name=\"p\",\n",
+    "        dims=p_dims,\n",
+    "        coords=[time, z, p_lat, p_lon],\n",
+    "        attrs=dict(\n",
+    "            description=\"pressure\",\n",
+    "            units=\"N/m^2\",\n",
+    "            location=p_location,\n",
+    "            mesh=f\"Arakawa-{grid_type}\",\n",
+    "        ),\n",
+    "    )\n",
+    "\n",
+    "    return xr.Dataset({\"u\": u, \"v\": v, \"p\": p})\n",
+    "\n",
+    "\n",
+    "ds_arakawa_a = stommel_fieldset_xarray(50, 50, \"A\")\n",
+    "ds_arakawa_c = stommel_fieldset_xarray(50, 50, \"C\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds_arakawa_a"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds_arakawa_a[\"u\"].attrs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds_arakawa_c"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "min_length_scale = 1111111.111111111 * np.sqrt(np.min(uxds.uxgrid.face_areas))\n",
+    "print(min_length_scale)\n",
+    "\n",
+    "max_v = np.sqrt(uxds[\"u\"] ** 2 + uxds[\"v\"] ** 2).max()\n",
+    "print(max_v)\n",
+    "\n",
+    "cfl = 0.1\n",
+    "dt = cfl * min_length_scale / max_v\n",
+    "print(dt)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datetime import timedelta\n",
+    "\n",
+    "import numpy as np\n",
+    "import uxarray as ux\n",
+    "\n",
+    "from parcels import Particle, ParticleSet, UxAdvectionEuler, UXFieldSet\n",
+    "\n",
+    "npart = 10\n",
+    "fieldset = UXFieldSet(uxds)\n",
+    "# pset = ParticleSet(\n",
+    "#         fieldset,\n",
+    "#         pclass=Particle,\n",
+    "#         lon=np.linspace(1, 59, npart),\n",
+    "#         lat=np.zeros(npart)+30)\n",
+    "# pset.execute(UxAdvectionEuler, runtime=timedelta(hours=24), dt=timedelta(seconds=dt))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "parcels",
+   "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.13.2"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

environment.yml

@@ -17,6 +17,7 @@ dependencies: #! Keep in sync with [tool.pixi.dependencies] in pyproject.toml
   - dask>=2.0
   - scikit-learn
   - zarr>=2.11.0,!=2.18.0,<3
+  - uxaray>=2025.3.0
   - pooch
 
   # Notebooks