fix field creation from code review suggestions

Author: reint-fischer

docs/user_guide/examples/tutorial_diffusion.ipynb

@@ -114,8 +114,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import random\n",
-    "\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "import trajan as ta\n",
@@ -197,7 +195,7 @@
     "ds = simple_UV_dataset(dims=(1, 1, Ny, 1), mesh=\"flat\").isel(time=0, depth=0)\n",
     "ds[\"lat\"][:] = np.linspace(-0.01, 1.01, Ny)\n",
     "ds[\"lon\"][:] = np.ones(len(ds.XG))\n",
-    "ds[\"Kh_meridional\"] = ds[\"U\"] + Kh_meridional[:, None]\n",
+    "ds[\"Kh_meridional\"] = ([\"YG\",\"XG\"], Kh_meridional[:, None])\n",
     "ds"
    ]
   },
@@ -345,7 +343,7 @@
     "    chunks=(len(testParticles), 50),\n",
     "    outputdt=np.timedelta64(1, \"ms\"),\n",
     ")\n",
-    "random.seed(1636)  # Random seed for reproducibility\n",
+    "np.random.seed(1636)  # Random seed for reproducibility\n",
     "testParticles.execute(\n",
     "    parcels.kernels.AdvectionDiffusionEM,\n",
     "    runtime=np.timedelta64(300, \"ms\"),\n",
@@ -499,45 +497,32 @@
    "source": [
     "fieldset = parcels.FieldSet.from_copernicusmarine(ds_fields)\n",
     "\n",
+    "def degree_lat_to_meter(d):\n",
+    "    return d * 1000.0 * 1.852 * 60.0\n",
     "\n",
-    "def calc_cell_edge_sizes(grid):\n",
-    "    \"\"\"Calculate cell sizes based on numpy.gradient method.\n",
     "\n",
-    "    Currently only works for Rectilinear Grids. Operates in place adding a `cell_edge_sizes`\n",
-    "    attribute to the grid.\n",
-    "    \"\"\"\n",
-    "    # TODO: check for gridtypes once added in v4\n",
-    "    # assert grid._gtype in (GridType.RectilinearZGrid, GridType.RectilinearSGrid), (\n",
-    "    #     f\"_cell_edge_sizes() not implemented for {grid._gtype} grids. \"\n",
-    "    #     \"You can provide cell_edge_sizes yourself by in, e.g., \"\n",
-    "    #     \"NEMO using the e1u fields etc from the mesh_mask.nc file.\"\n",
-    "    # )\n",
+    "def degree_lon_to_meter(d, lat):\n",
+    "    return d * 1000.0 * 1.852 * 60.0 * np.cos(lat * np.pi / 180)\n",
     "\n",
-    "    cell_edge_sizes_x = np.zeros((grid.ydim + 1, grid.xdim + 1), dtype=np.float32)\n",
-    "    cell_edge_sizes_y = np.zeros((grid.ydim + 1, grid.xdim + 1), dtype=np.float32)\n",
     "\n",
-    "    x_conv = (\n",
-    "        parcels.GeographicPolar()\n",
-    "    )  # if grid._mesh == \"spherical\" else parcels.UnitConverter()\n",
-    "    y_conv = (\n",
-    "        parcels.Geographic()\n",
-    "    )  # if grid._mesh == \"spherical\" else parcels.UnitConverter()\n",
-    "    for y, (lat, dy) in enumerate(zip(grid.lat, np.gradient(grid.lat), strict=False)):\n",
-    "        for x, (lon, dx) in enumerate(\n",
-    "            zip(grid.lon, np.gradient(grid.lon), strict=False)\n",
-    "        ):\n",
-    "            cell_edge_sizes_x[y, x] = x_conv.to_source(dx, grid.depth[0], lat, lon)\n",
-    "            cell_edge_sizes_y[y, x] = y_conv.to_source(dy, grid.depth[0], lat, lon)\n",
-    "    return cell_edge_sizes_x, cell_edge_sizes_y\n",
+    "def calc_cell_areas(ds):\n",
+    "    \"\"\"calculate cell areas for rectilinear grids\"\"\"\n",
+    "    lon, lat = ds[\"longitude\"], ds[\"latitude\"]\n",
+    "    assert \"degrees\" in lon.attrs[\"units\"]\n",
+    "    assert \"degrees\" in lat.attrs[\"units\"]\n",
     "\n",
+    "    LON, LAT = np.meshgrid(lon, lat)\n",
+    "    X, Y = degree_lon_to_meter(LON, LAT), degree_lat_to_meter(LAT)\n",
     "\n",
-    "def calc_cell_areas(grid):\n",
-    "    cell_edge_sizes_x, cell_edge_sizes_y = calc_cell_edge_sizes(grid)\n",
-    "    return cell_edge_sizes_x * cell_edge_sizes_y\n",
+    "    dX = np.gradient(X, axis=1)\n",
+    "    dY = np.gradient(Y, axis=0)\n",
+    "    cell_areas = dX * dY\n",
+    "\n",
+    "    return cell_areas\n",
     "\n",
     "\n",
     "da_cell_areas = xr.DataArray(\n",
-    "    data=calc_cell_areas(fieldset.U.grid),\n",
+    "    data=calc_cell_areas(ds_fields),\n",
     "    coords=dict(\n",
     "        latitude=([\"lat\"], ds_fields.latitude.values),\n",
     "        longitude=([\"lon\"], ds_fields.longitude.values),\n",
@@ -612,11 +597,12 @@
    "outputs": [],
    "source": [
     "ds_particles = xr.open_zarr(\"smagdiff.zarr\")\n",
+    "\n",
     "temperature = ds_fields.isel(time=0, depth=0).thetao.plot(cmap=\"magma\")\n",
     "velocity = ds_fields.isel(time=0, depth=0).plot.quiver(\n",
     "    x=\"longitude\", y=\"latitude\", u=\"uo\", v=\"vo\"\n",
     ")\n",
-    "ds_particles.traj.plot(color=\"blue\")\n",
+    "particles = ds_particles.traj.plot(color=\"blue\")\n",
     "plt.ylim(-31, -30)\n",
     "plt.xlim(31, 32.1)\n",
     "plt.show()"