diff --git a/docs/examples/tutorial_stommel_uxarray.ipynb b/docs/examples/tutorial_stommel_uxarray.ipynb new file mode 100644 index 0000000000..d5bc4317cf --- /dev/null +++ b/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 +} diff --git a/environment.yml b/environment.yml index 81af0d5a60..49b139f5f3 100644 --- a/environment.yml +++ b/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 diff --git a/parcels/_index_search.py b/parcels/_index_search.py index ddfbadd173..8af53bee27 100644 --- a/parcels/_index_search.py +++ b/parcels/_index_search.py @@ -1,5 +1,6 @@ from __future__ import annotations +from datetime import datetime from typing import TYPE_CHECKING import numpy as np @@ -21,78 +22,91 @@ if TYPE_CHECKING: from .field import Field - from .grid import Grid + # from .grid import Grid -def _search_time_index(grid: Grid, time: float, allow_time_extrapolation=True): +def _search_time_index(field: Field, time: datetime, allow_time_extrapolation=True): """Find and return the index and relative coordinate in the time array associated with a given time. + Parameters + ---------- + field: Field + + time: datetime + This is the amount of time, in seconds (time_delta), in unix epoch Note that we normalize to either the first or the last index if the sampled value is outside the time value range. """ - if not allow_time_extrapolation and (time < grid.time[0] or time > grid.time[-1]): + if not allow_time_extrapolation and (time < field.data.time[0] or time > field.data.time[-1]): _raise_time_extrapolation_error(time, field=None) - time_index = grid.time <= time + time_index = field.data.time <= time if time_index.all(): # If given time > last known field time, use # the last field frame without interpolation - ti = len(grid.time) - 1 + ti = len(field.data.time) - 1 + elif np.logical_not(time_index).all(): # If given time < any time in the field, use # the first field frame without interpolation ti = 0 else: ti = int(time_index.argmin() - 1) if time_index.any() else 0 - if grid.tdim == 1: + if len(field.data.time) == 1: tau = 0 - elif ti == len(grid.time) - 1: + elif ti == len(field.data.time) - 1: tau = 1 else: - tau = (time - grid.time[ti]) / (grid.time[ti + 1] - grid.time[ti]) if grid.time[ti] != grid.time[ti + 1] else 0 + tau = ( + (time - field.data.time[ti]).total_seconds() + / (field.data.time[ti + 1] - field.data.time[ti]).total_seconds() + if field.data.time[ti] != field.data.time[ti + 1] + else 0 + ) return tau, ti -def search_indices_vertical_z(grid: Grid, gridindexingtype: GridIndexingType, z: float): - if grid.depth[-1] > grid.depth[0]: - if z < grid.depth[0]: +def search_indices_vertical_z(depth, gridindexingtype: GridIndexingType, z: float): + if depth[-1] > depth[0]: + if z < depth[0]: # Since MOM5 is indexed at cell bottom, allow z at depth[0] - dz where dz = (depth[1] - depth[0]) - if gridindexingtype == "mom5" and z > 2 * grid.depth[0] - grid.depth[1]: - return (-1, z / grid.depth[0]) + if gridindexingtype == "mom5" and z > 2 * depth[0] - depth[1]: + return (-1, z / depth[0]) else: _raise_field_out_of_bound_surface_error(z, None, None) - elif z > grid.depth[-1]: + elif z > depth[-1]: # In case of CROCO, allow particles in last (uppermost) layer using depth[-1] if gridindexingtype in ["croco"] and z < 0: return (-2, 1) _raise_field_out_of_bound_error(z, None, None) - depth_indices = grid.depth < z - if z >= grid.depth[-1]: - zi = len(grid.depth) - 2 + depth_indices = depth < z + if z >= depth[-1]: + zi = len(depth) - 2 else: - zi = depth_indices.argmin() - 1 if z > grid.depth[0] else 0 + zi = depth_indices.argmin() - 1 if z > depth[0] else 0 else: - if z > grid.depth[0]: + if z > depth[0]: _raise_field_out_of_bound_surface_error(z, None, None) - elif z < grid.depth[-1]: + elif z < depth[-1]: _raise_field_out_of_bound_error(z, None, None) - depth_indices = grid.depth > z - if z <= grid.depth[-1]: - zi = len(grid.depth) - 2 + depth_indices = depth > z + if z <= depth[-1]: + zi = len(depth) - 2 else: - zi = depth_indices.argmin() - 1 if z < grid.depth[0] else 0 - zeta = (z - grid.depth[zi]) / (grid.depth[zi + 1] - grid.depth[zi]) + zi = depth_indices.argmin() - 1 if z < depth[0] else 0 + zeta = (z - depth[zi]) / (depth[zi + 1] - depth[zi]) while zeta > 1: zi += 1 - zeta = (z - grid.depth[zi]) / (grid.depth[zi + 1] - grid.depth[zi]) + zeta = (z - depth[zi]) / (depth[zi + 1] - depth[zi]) while zeta < 0: zi -= 1 - zeta = (z - grid.depth[zi]) / (grid.depth[zi + 1] - grid.depth[zi]) + zeta = (z - depth[zi]) / (depth[zi + 1] - depth[zi]) return (zi, zeta) +## TODO : Still need to implement the search_indices_vertical_s function def search_indices_vertical_s( - grid: Grid, + field: Field, interp_method: InterpMethodOption, time: float, z: float, @@ -107,32 +121,32 @@ def search_indices_vertical_s( if interp_method in ["bgrid_velocity", "bgrid_w_velocity", "bgrid_tracer"]: xsi = 1 eta = 1 - if time < grid.time[ti]: + if time < field.time[ti]: ti -= 1 - if grid._z4d: # type: ignore[attr-defined] - if ti == len(grid.time) - 1: + if field._z4d: # type: ignore[attr-defined] + if ti == len(field.time) - 1: depth_vector = ( - (1 - xsi) * (1 - eta) * grid.depth[-1, :, yi, xi] - + xsi * (1 - eta) * grid.depth[-1, :, yi, xi + 1] - + xsi * eta * grid.depth[-1, :, yi + 1, xi + 1] - + (1 - xsi) * eta * grid.depth[-1, :, yi + 1, xi] + (1 - xsi) * (1 - eta) * field.depth[-1, :, yi, xi] + + xsi * (1 - eta) * field.depth[-1, :, yi, xi + 1] + + xsi * eta * field.depth[-1, :, yi + 1, xi + 1] + + (1 - xsi) * eta * field.depth[-1, :, yi + 1, xi] ) else: dv2 = ( - (1 - xsi) * (1 - eta) * grid.depth[ti : ti + 2, :, yi, xi] - + xsi * (1 - eta) * grid.depth[ti : ti + 2, :, yi, xi + 1] - + xsi * eta * grid.depth[ti : ti + 2, :, yi + 1, xi + 1] - + (1 - xsi) * eta * grid.depth[ti : ti + 2, :, yi + 1, xi] + (1 - xsi) * (1 - eta) * field.depth[ti : ti + 2, :, yi, xi] + + xsi * (1 - eta) * field.depth[ti : ti + 2, :, yi, xi + 1] + + xsi * eta * field.depth[ti : ti + 2, :, yi + 1, xi + 1] + + (1 - xsi) * eta * field.depth[ti : ti + 2, :, yi + 1, xi] ) - tt = (time - grid.time[ti]) / (grid.time[ti + 1] - grid.time[ti]) + tt = (time - field.time[ti]) / (field.time[ti + 1] - field.time[ti]) assert tt >= 0 and tt <= 1, "Vertical s grid is being wrongly interpolated in time" depth_vector = dv2[0, :] * (1 - tt) + dv2[1, :] * tt else: depth_vector = ( - (1 - xsi) * (1 - eta) * grid.depth[:, yi, xi] - + xsi * (1 - eta) * grid.depth[:, yi, xi + 1] - + xsi * eta * grid.depth[:, yi + 1, xi + 1] - + (1 - xsi) * eta * grid.depth[:, yi + 1, xi] + (1 - xsi) * (1 - eta) * field.depth[:, yi, xi] + + xsi * (1 - eta) * field.depth[:, yi, xi + 1] + + xsi * eta * field.depth[:, yi + 1, xi + 1] + + (1 - xsi) * eta * field.depth[:, yi + 1, xi] ) z = np.float32(z) # type: ignore # TODO: remove type ignore once we migrate to float64 @@ -167,32 +181,31 @@ def search_indices_vertical_s( def _search_indices_rectilinear( - field: Field, time: float, z: float, y: float, x: float, ti: int, particle=None, search2D=False + field: Field, time: datetime, z: float, y: float, x: float, ti: int, ei: int | None = None, search2D=False ): - grid = field.grid - - if grid.xdim > 1 and (not grid.zonal_periodic): - if x < grid.lonlat_minmax[0] or x > grid.lonlat_minmax[1]: + # TODO : If ei is provided, check if particle is in the same cell + if field.xdim > 1 and (not field.zonal_periodic): + if x < field.lonlat_minmax[0] or x > field.lonlat_minmax[1]: _raise_field_out_of_bound_error(z, y, x) - if grid.ydim > 1 and (y < grid.lonlat_minmax[2] or y > grid.lonlat_minmax[3]): + if field.ydim > 1 and (y < field.lonlat_minmax[2] or y > field.lonlat_minmax[3]): _raise_field_out_of_bound_error(z, y, x) - if grid.xdim > 1: - if grid.mesh != "spherical": - lon_index = grid.lon < x + if field.xdim > 1: + if field._mesh_type != "spherical": + lon_index = field.lon < x if lon_index.all(): - xi = len(grid.lon) - 2 + xi = len(field.lon) - 2 else: xi = lon_index.argmin() - 1 if lon_index.any() else 0 - xsi = (x - grid.lon[xi]) / (grid.lon[xi + 1] - grid.lon[xi]) + xsi = (x - field.lon[xi]) / (field.lon[xi + 1] - field.lon[xi]) if xsi < 0: xi -= 1 - xsi = (x - grid.lon[xi]) / (grid.lon[xi + 1] - grid.lon[xi]) + xsi = (x - field.lon[xi]) / (field.lon[xi + 1] - field.lon[xi]) elif xsi > 1: xi += 1 - xsi = (x - grid.lon[xi]) / (grid.lon[xi + 1] - grid.lon[xi]) + xsi = (x - field.lon[xi]) / (field.lon[xi + 1] - field.lon[xi]) else: - lon_fixed = grid.lon.copy() + lon_fixed = field.lon.copy() indices = lon_fixed >= lon_fixed[0] if not indices.all(): lon_fixed[indices.argmin() :] += 360 @@ -214,32 +227,33 @@ def _search_indices_rectilinear( else: xi, xsi = -1, 0 - if grid.ydim > 1: - lat_index = grid.lat < y + if field.ydim > 1: + lat_index = field.lat < y if lat_index.all(): - yi = len(grid.lat) - 2 + yi = len(field.lat) - 2 else: yi = lat_index.argmin() - 1 if lat_index.any() else 0 - eta = (y - grid.lat[yi]) / (grid.lat[yi + 1] - grid.lat[yi]) + eta = (y - field.lat[yi]) / (field.lat[yi + 1] - field.lat[yi]) if eta < 0: yi -= 1 - eta = (y - grid.lat[yi]) / (grid.lat[yi + 1] - grid.lat[yi]) + eta = (y - field.lat[yi]) / (field.lat[yi + 1] - field.lat[yi]) elif eta > 1: yi += 1 - eta = (y - grid.lat[yi]) / (grid.lat[yi + 1] - grid.lat[yi]) + eta = (y - field.lat[yi]) / (field.lat[yi + 1] - field.lat[yi]) else: yi, eta = -1, 0 - if grid.zdim > 1 and not search2D: - if grid._gtype == GridType.RectilinearZGrid: + if field.zdim > 1 and not search2D: + if field._gtype == GridType.RectilinearZGrid: try: (zi, zeta) = search_indices_vertical_z(field.grid, field.gridindexingtype, z) except FieldOutOfBoundError: _raise_field_out_of_bound_error(z, y, x) except FieldOutOfBoundSurfaceError: _raise_field_out_of_bound_surface_error(z, y, x) - elif grid._gtype == GridType.RectilinearSGrid: + elif field._gtype == GridType.RectilinearSGrid: + ## TODO : Still need to implement the search_indices_vertical_s function (zi, zeta) = search_indices_vertical_s(field.grid, field.interp_method, time, z, y, x, ti, yi, xi, eta, xsi) else: zi, zeta = -1, 0 @@ -247,18 +261,18 @@ def _search_indices_rectilinear( if not ((0 <= xsi <= 1) and (0 <= eta <= 1) and (0 <= zeta <= 1)): _raise_field_sampling_error(z, y, x) - if particle: - particle.ei[field.igrid] = field.ravel_index(zi, yi, xi) + _ei = field.ravel_index(zi, yi, xi) - return (zeta, eta, xsi, zi, yi, xi) + return (zeta, eta, xsi, _ei) +## TODO : Still need to implement the search_indices_curvilinear def _search_indices_curvilinear(field: Field, time, z, y, x, ti, particle=None, search2D=False): if particle: zi, yi, xi = field.unravel_index(particle.ei) else: - xi = int(field.grid.xdim / 2) - 1 - yi = int(field.grid.ydim / 2) - 1 + xi = int(field.xdim / 2) - 1 + yi = int(field.ydim / 2) - 1 xsi = eta = -1.0 grid = field.grid invA = np.array([[1, 0, 0, 0], [-1, 1, 0, 0], [-1, 0, 0, 1], [1, -1, 1, -1]]) @@ -266,12 +280,12 @@ def _search_indices_curvilinear(field: Field, time, z, y, x, ti, particle=None, it = 0 tol = 1.0e-10 if not grid.zonal_periodic: - if x < grid.lonlat_minmax[0] or x > grid.lonlat_minmax[1]: + if x < field.lonlat_minmax[0] or x > field.lonlat_minmax[1]: if grid.lon[0, 0] < grid.lon[0, -1]: _raise_field_out_of_bound_error(z, y, x) elif x < grid.lon[0, 0] and x > grid.lon[0, -1]: # This prevents from crashing in [160, -160] _raise_field_out_of_bound_error(z, y, x) - if y < grid.lonlat_minmax[2] or y > grid.lonlat_minmax[3]: + if y < field.lonlat_minmax[2] or y > field.lonlat_minmax[3]: _raise_field_out_of_bound_error(z, y, x) while xsi < -tol or xsi > 1 + tol or eta < -tol or eta > 1 + tol: diff --git a/parcels/application_kernels/__init__.py b/parcels/application_kernels/__init__.py index d308bcb1ec..cd933324ed 100644 --- a/parcels/application_kernels/__init__.py +++ b/parcels/application_kernels/__init__.py @@ -1,3 +1,4 @@ from .advection import * from .advectiondiffusion import * from .interaction import * +from .interpolation import * diff --git a/parcels/application_kernels/interpolation.py b/parcels/application_kernels/interpolation.py new file mode 100644 index 0000000000..a1abf14064 --- /dev/null +++ b/parcels/application_kernels/interpolation.py @@ -0,0 +1,53 @@ +"""Collection of pre-built interpolation kernels.""" + +import numpy as np + +from parcels.field import Field + +__all__ = [ + "UXPiecewiseConstantFace", + "UXPiecewiseLinearNode", +] + + +def UXPiecewiseConstantFace( + field: Field, + ti: int, + ei: int, + bcoords: np.ndarray, + tau: np.float32 | np.float64, + t: np.float32 | np.float64, + z: np.float32 | np.float64, + y: np.float32 | np.float64, + x: np.float32 | np.float64, +): + """ + Piecewise constant interpolation kernel for face registered data. + This interpolation method is appropriate for fields that are + face registered, such as u,v in FESOM. + """ + # TODO joe : handle vertical interpolation + zi, fi = field.unravel_index(ei) + return field.data[ti, zi, fi] + + +def UXPiecewiseLinearNode( + field: Field, + ti: int, + ei: int, + bcoords: np.ndarray, + tau: np.float32 | np.float64, + t: np.float32 | np.float64, + z: np.float32 | np.float64, + y: np.float32 | np.float64, + x: np.float32 | np.float64, +): + """ + Piecewise linear interpolation kernel for node registered data. This + interpolation method is appropriate for fields that are node registered + such as the vertical velocity w in FESOM. + """ + # TODO joe : handle vertical interpolation + zi, fi = field.unravel_index(ei) + node_ids = field.data.uxgrid.face_node_connectivity[fi, :] + return np.dot(field.data[ti, zi, node_ids], bcoords) diff --git a/parcels/field.py b/parcels/field.py index 1b40893482..9a78fc578d 100644 --- a/parcels/field.py +++ b/parcels/field.py @@ -1,32 +1,22 @@ -import collections -import math +import inspect import warnings -from typing import TYPE_CHECKING, cast +from collections.abc import Callable +from datetime import datetime +from enum import IntEnum +from typing import TYPE_CHECKING -import dask.array as da import numpy as np +import uxarray as ux import xarray as xr +from uxarray.grid.neighbors import _barycentric_coordinates -import parcels.tools.interpolation_utils as i_u from parcels._compat import add_note -from parcels._interpolation import ( - InterpolationContext2D, - InterpolationContext3D, - get_2d_interpolator_registry, - get_3d_interpolator_registry, -) from parcels._typing import ( - GridIndexingType, - InterpMethod, - InterpMethodOption, Mesh, VectorType, - assert_valid_gridindexingtype, - assert_valid_interp_method, ) -from parcels.tools._helpers import default_repr, field_repr, should_calculate_next_ti +from parcels.tools._helpers import default_repr, field_repr from parcels.tools.converters import ( - TimeConverter, UnitConverter, unitconverters_map, ) @@ -37,20 +27,20 @@ FieldSamplingError, _raise_field_out_of_bound_error, ) -from parcels.tools.warnings import FieldSetWarning -from ._index_search import _search_indices_curvilinear, _search_indices_rectilinear, _search_time_index -from .fieldfilebuffer import ( - NetcdfFileBuffer, -) -from .grid import Grid, GridType +from ._index_search import _search_indices_rectilinear, _search_time_index if TYPE_CHECKING: - import numpy.typing as npt + pass + +__all__ = ["Field", "GridType", "VectorField"] - from parcels.fieldset import FieldSet -__all__ = ["Field", "VectorField"] +class GridType(IntEnum): + RectilinearZGrid = 0 + RectilinearSGrid = 1 + CurvilinearZGrid = 2 + CurvilinearSGrid = 3 def _isParticle(key): @@ -76,142 +66,169 @@ def _deal_with_errors(error, key, vector_type: VectorType): return 0 -def _croco_from_z_to_sigma_scipy(fieldset, time, z, y, x, particle): - """Calculate local sigma level of the particle, by linearly interpolating the - scaling function that maps sigma to depth (using local ocean depth H, - sea-surface Zeta and stretching parameters Cs_w and hc). - See also https://croco-ocean.gitlabpages.inria.fr/croco_doc/model/model.grid.html#vertical-grid-parameters +class Field: + """The Field class that holds scalar field data. + The `Field` object is a wrapper around a xarray.DataArray or uxarray.UxDataArray object. + Additionally, it holds a dynamic Callable procedure that is used to interpolate the field data. + During initialization, the user can supply a custom interpolation method that is used to interpolate the field data, + so long as the interpolation method has the correct signature. + + Notes + ----- + The xarray.DataArray or uxarray.UxDataArray object contains the field data and metadata. + * dims: (time, [nz1 | nz], [face_lat | node_lat | edge_lat], [face_lon | node_lon | edge_lon]) + * attrs: (location, mesh, mesh_type) + + When using a xarray.DataArray object, + * The xarray.DataArray object must have the "location" and "mesh" attributes set. + * The "location" attribute must be set to one of the following to define which pairing of points a field is associated with. + * "node" + * "face" + * "x_edge" + * "y_edge" + * For an A-Grid, the "location" attribute must be set to / is assumed to be "node" (node_lat,node_lon). + * For a C-Grid, the "location" setting for a field has the following interpretation: + * "node" ~> the field is associated with the vorticity points (node_lat, node_lon) + * "face" ~> the field is associated with the tracer points (face_lat, face_lon) + * "x_edge" ~> the field is associated with the u-velocity points (face_lat, node_lon) + * "y_edge" ~> the field is associated with the v-velocity points (node_lat, face_lon) + + When using a uxarray.UxDataArray object, + * The uxarray.UxDataArray.UxGrid object must have the "Conventions" attribute set to "UGRID-1.0" + and the uxarray.UxDataArray object must comply with the UGRID conventions. + See https://ugrid-conventions.github.io/ugrid-conventions/ for more information. + """ - h = fieldset.H.eval(time, 0, y, x, particle=particle, applyConversion=False) - zeta = fieldset.Zeta.eval(time, 0, y, x, particle=particle, applyConversion=False) - sigma_levels = fieldset.U.grid.depth - z0 = fieldset.hc * sigma_levels + (h - fieldset.hc) * fieldset.Cs_w.data[0, :, 0, 0] - zvec = z0 + zeta * (1 + (z0 / h)) - zinds = zvec <= z - if z >= zvec[-1]: - zi = len(zvec) - 2 - else: - zi = zinds.argmin() - 1 if z >= zvec[0] else 0 - return sigma_levels[zi] + (z - zvec[zi]) * (sigma_levels[zi + 1] - sigma_levels[zi]) / (zvec[zi + 1] - zvec[zi]) + @staticmethod + def _interp_template( + self, + ti: int, + ei: int, + bcoords: np.ndarray, + tau: np.float32 | np.float64, + t: np.float32 | np.float64, + z: np.float32 | np.float64, + y: np.float32 | np.float64, + x: np.float32 | np.float64, + ) -> np.float32 | np.float64: + """Template function used for the signature check of the lateral interpolation methods.""" + return 0.0 + + def _validate_interp_function(self, func: Callable) -> bool: + """Ensures that the function has the correct signature.""" + template_sig = inspect.signature(self._interp_template) + func_sig = inspect.signature(func) + + if len(template_sig.parameters) != len(func_sig.parameters): + return False + + for (_name1, param1), (_name2, param2) in zip( + template_sig.parameters.items(), func_sig.parameters.items(), strict=False + ): + if param1.kind != param2.kind: + return False + if param1.annotation != param2.annotation: + return False + return_annotation = func_sig.return_annotation + template_return = template_sig.return_annotation -class Field: - """Class that encapsulates access to field data. - - Parameters - ---------- - name : str - Name of the field - data : np.ndarray - 2D, 3D or 4D numpy array of field data with shape [ydim, xdim], [zdim, ydim, xdim], [tdim, ydim, xdim] or [tdim, zdim, ydim, xdim], - lon : np.ndarray or list - Longitude coordinates (numpy vector or array) of the field (only if grid is None) - lat : np.ndarray or list - Latitude coordinates (numpy vector or array) of the field (only if grid is None) - depth : np.ndarray or list - Depth coordinates (numpy vector or array) of the field (only if grid is None) - time : np.ndarray - Time coordinates (numpy vector) of the field (only if grid is None) - mesh : str - String indicating the type of mesh coordinates and - units used during velocity interpolation: (only if grid is None) - - 1. spherical: Lat and lon in degree, with a - correction for zonal velocity U near the poles. - 2. flat (default): No conversion, lat/lon are assumed to be in m. - grid : parcels.grid.Grid - :class:`parcels.grid.Grid` object containing all the lon, lat depth, time - mesh and time_origin information. Can be constructed from any of the Grid objects - time_origin : parcels.tools.converters.TimeConverter - Time origin of the time axis (only if grid is None) - interp_method : str - Method for interpolation. Options are 'linear' (default), 'nearest', - 'linear_invdist_land_tracer', 'cgrid_velocity', 'cgrid_tracer' and 'bgrid_velocity' - allow_time_extrapolation : bool - boolean whether to allow for extrapolation in time - (i.e. beyond the last available time snapshot) + if return_annotation != template_return: + return False - """ - - allow_time_extrapolation: bool + return True def __init__( self, - name: str | tuple[str, str], - data, - lon=None, - lat=None, - depth=None, - time=None, - grid=None, - mesh: Mesh = "flat", - time_origin: TimeConverter | None = None, - interp_method: InterpMethod = "linear", + name: str, + data: xr.DataArray | ux.UxDataArray, + grid: ux.UxGrid | None = None, # TODO Nick : Once parcels.Grid class is added, allow for it to be passed here + mesh_type: Mesh = "flat", + interp_method: Callable | None = None, allow_time_extrapolation: bool | None = None, - gridindexingtype: GridIndexingType = "nemo", - data_full_zdim=None, ): - if not isinstance(name, tuple): - self.name = name - else: - self.name = name[0] + self.name = name self.data = data - if grid: - self._grid = grid + self.grid = grid + + _validate_dataarray(data, name) + + self._parent_mesh = data.attrs["mesh"] + self._mesh_type = mesh_type + self._location = data.attrs["location"] + self._vertical_location = None + + # Setting the interpolation method dynamically + if interp_method is None: + self._interp_method = self._interp_template # Default to method that returns 0 always else: - if (time is not None) and isinstance(time[0], np.datetime64): - time_origin = TimeConverter(time[0]) - time = np.array([time_origin.reltime(t) for t in time]) - else: - time_origin = TimeConverter(0) - self._grid = Grid.create_grid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh) - self.igrid = -1 - self.units = UnitConverter() - if self.grid.mesh == "spherical": - try: - self.units = unitconverters_map[self.name] - except KeyError: - pass - if isinstance(interp_method, dict): - if self.name in interp_method: - self.interp_method = interp_method[self.name] - else: - raise RuntimeError(f"interp_method is a dictionary but {name} is not in it") + self._validate_interp_function(interp_method) + self._interp_method = interp_method + + self.igrid = -1 # Default the grid index to -1 + + if self._mesh_type == "flat" or (self.name not in unitconverters_map.keys()): + self.units = UnitConverter() + elif self._mesh_type == "spherical": + self.units = unitconverters_map[self.name] else: - self.interp_method = interp_method - assert_valid_gridindexingtype(gridindexingtype) - self._gridindexingtype = gridindexingtype - if self.interp_method in ["bgrid_velocity", "bgrid_w_velocity", "bgrid_tracer"] and self.grid._gtype in [ - GridType.RectilinearSGrid, - GridType.CurvilinearSGrid, - ]: - warnings.warn( - "General s-levels are not supported in B-grid. RectilinearSGrid and CurvilinearSGrid can still be used to deal with shaved cells, but the levels must be horizontal.", - FieldSetWarning, - stacklevel=2, - ) + raise ValueError("Unsupported mesh type in data array attributes. Choose either: 'spherical' or 'flat'") - self.fieldset: FieldSet | None = None if allow_time_extrapolation is None: - self.allow_time_extrapolation = True if len(self.grid.time) == 1 else False + self.allow_time_extrapolation = True if len(getattr(self.data, "time", [])) == 1 else False else: self.allow_time_extrapolation = allow_time_extrapolation - self.data = self._reshape(self.data) + if type(self.data) is ux.UxDataArray: + self._spatialhash = self.grid.get_spatial_hash() + self._gtype = None + # Set the vertical location + if "nz1" in data.dims: + self._vertical_location = "center" + elif "nz" in data.dims: + self._vertical_location = "face" + else: # TODO Nick : This bit probably needs an overhaul once the parcels.Grid class is integrated. + self._spatialhash = None + # Set the grid type + if "x_g" in self.data.coords: + lon = self.data.x_g + elif "x_c" in self.data.coords: + lon = self.data.x_c + else: + lon = self.data.lon + + if "nz1" in self.data.coords: + depth = self.data.nz1 + elif "nz" in self.data.coords: + depth = self.data.nz + elif "depth" in self.data.coords: + depth = self.data.depth + else: + depth = None - # Hack around the fact that NaN and ridiculously large values - # propagate in SciPy's interpolators - self.data[np.isnan(self.data)] = 0.0 + if len(lon.shape) <= 1: + if depth is None or len(depth.shape) <= 1: + self._gtype = GridType.RectilinearZGrid + else: + self._gtype = GridType.RectilinearSGrid + else: + if depth is None or len(depth.shape) <= 1: + self._gtype = GridType.CurvilinearZGrid + else: + self._gtype = GridType.CurvilinearSGrid - # data_full_zdim is the vertical dimension of the complete field data, ignoring the indices. - # (data_full_zdim = grid.zdim if no indices are used, for A- and C-grids and for some B-grids). It is used for the B-grid, - # since some datasets do not provide the deeper level of data (which is ignored by the interpolation). - self.data_full_zdim = data_full_zdim + self._lonlat_minmax = np.array( + [np.nanmin(self.lon), np.nanmax(self.lon), np.nanmin(self.lat), np.nanmax(self.lat)], dtype=np.float32 + ) - def __repr__(self) -> str: + def __repr__(self): return field_repr(self) + @property + def lonlat_minmax(self): + return self._lonlat_minmax + @property def units(self): return self._units @@ -223,368 +240,167 @@ def units(self, value): self._units = value @property - def grid(self): - return self._grid + def lat(self): + if type(self.data) is ux.UxDataArray: + if self._location == "node": + return self.grid.node_lat + elif self._location == "face": + return self.grid.face_lat + elif self._location == "edge": + return self.grid.edge_lat + else: + return self.data.lat @property def lon(self): - """Lon defined on the Grid object""" - return self.grid.lon - - @property - def lat(self): - """Lat defined on the Grid object""" - return self.grid.lat + if type(self.data) is ux.UxDataArray: + if self._location == "node": + return self.grid.node_lon + elif self._location == "face": + return self.grid.face_lon + elif self._location == "edge": + return self.grid.edge_lon + else: + return self.data.lon @property def depth(self): - """Depth defined on the Grid object""" - return self.grid.depth + if type(self.data) is ux.UxDataArray: + if self._vertical_location == "center": + return self.grid.nz1 + elif self._vertical_location == "face": + return self.grid.nz + else: + return self.data.depth @property - def interp_method(self): - return self._interp_method - - @interp_method.setter - def interp_method(self, value): - assert_valid_interp_method(value) - self._interp_method = value + def xdim(self): + if type(self.data) is xr.DataArray: + if "face_lon" in self.data.dims: + return self.data.sizes["face_lon"] + elif "node_lon" in self.data.dims: + return self.data.sizes["node_lon"] + else: + return self.data.sizes["lon"] + else: + return 0 # TODO : Discuss what we want to return as xdim for uxdataarray obj @property - def gridindexingtype(self): - return self._gridindexingtype - - @classmethod - def _get_dim_filenames(cls, filenames, dim): - if isinstance(filenames, str) or not isinstance(filenames, collections.abc.Iterable): - return [filenames] - elif isinstance(filenames, dict): - assert dim in filenames.keys(), "filename dimension keys must be lon, lat, depth or data" - filename = filenames[dim] - if isinstance(filename, str): - return [filename] + def ydim(self): + if type(self.data) is xr.DataArray: + if "face_lat" in self.data.dims: + return self.data.sizes["face_lat"] + elif "node_lat" in self.data.dims: + return self.data.sizes["node_lat"] else: - return filename + return self.data.sizes["lat"] else: - return filenames - - @staticmethod - def _collect_time(data_filenames, dimensions, indices): - time = [] - for fname in data_filenames: - with NetcdfFileBuffer(fname, dimensions, indices) as filebuffer: - ftime = filebuffer.time - time.append(ftime) - time = np.concatenate(time).ravel() - if time.size == 1 and time[0] is None: - time[0] = 0 - time_origin = TimeConverter(time[0]) - time = time_origin.reltime(time) - - return time, time_origin - - @classmethod - def from_netcdf( - cls, - filenames, - variable, - dimensions, - grid=None, - mesh: Mesh = "spherical", - allow_time_extrapolation: bool | None = None, - **kwargs, - ) -> "Field": - """Create field from netCDF file. - - Parameters - ---------- - filenames : list of str or dict - list of filenames to read for the field. filenames can be a list ``[files]`` or - a dictionary ``{dim:[files]}`` (if lon, lat, depth and/or data not stored in same files as data) - In the latter case, time values are in filenames[data] - variable : dict, tuple of str or str - Dict or tuple mapping field name to variable name in the NetCDF file. - dimensions : dict - Dictionary mapping variable names for the relevant dimensions in the NetCDF file - mesh : - String indicating the type of mesh coordinates and - units used during velocity interpolation: - - 1. spherical (default): Lat and lon in degree, with a - correction for zonal velocity U near the poles. - 2. flat: No conversion, lat/lon are assumed to be in m. - allow_time_extrapolation : bool - boolean whether to allow for extrapolation in time - (i.e. beyond the last available time snapshot) - Default is False if dimensions includes time, else True - gridindexingtype : str - The type of gridindexing. Either 'nemo' (default), 'mitgcm', 'mom5', 'pop', or 'croco' are supported. - See also the Grid indexing documentation on oceanparcels.org - grid : - (Default value = None) - **kwargs : - Keyword arguments passed to the :class:`Field` constructor. - """ - if isinstance(variable, str): # for backward compatibility with Parcels < 2.0.0 - variable = (variable, variable) - elif isinstance(variable, dict): - assert ( - len(variable) == 1 - ), "Field.from_netcdf() supports only one variable at a time. Use FieldSet.from_netcdf() for multiple variables." - variable = tuple(variable.items())[0] - assert ( - len(variable) == 2 - ), "The variable tuple must have length 2. Use FieldSet.from_netcdf() for multiple variables" - - data_filenames = cls._get_dim_filenames(filenames, "data") - lonlat_filename = cls._get_dim_filenames(filenames, "lon") - if isinstance(filenames, dict): - assert len(lonlat_filename) == 1 - if lonlat_filename != cls._get_dim_filenames(filenames, "lat"): - raise NotImplementedError( - "longitude and latitude dimensions are currently processed together from one single file" - ) - lonlat_filename = lonlat_filename[0] - if "depth" in dimensions: - depth_filename = cls._get_dim_filenames(filenames, "depth") - if isinstance(filenames, dict) and len(depth_filename) != 1: - raise NotImplementedError("Vertically adaptive meshes not implemented for from_netcdf()") - depth_filename = depth_filename[0] - - gridindexingtype = kwargs.get("gridindexingtype", "nemo") + return 0 # TODO : Discuss what we want to return as ydim for uxdataarray obj - indices: dict[str, npt.NDArray] = {} - - interp_method: InterpMethod = kwargs.pop("interp_method", "linear") - if type(interp_method) is dict: - if variable[0] in interp_method: - interp_method = interp_method[variable[0]] - else: - raise RuntimeError(f"interp_method is a dictionary but {variable[0]} is not in it") - interp_method = cast(InterpMethodOption, interp_method) - - if "lon" in dimensions and "lat" in dimensions: - with NetcdfFileBuffer( - lonlat_filename, - dimensions, - indices, - gridindexingtype=gridindexingtype, - ) as filebuffer: - lat, lon = filebuffer.latlon - indices = filebuffer.indices - # Check if parcels_mesh has been explicitly set in file - if "parcels_mesh" in filebuffer.dataset.attrs: - mesh = filebuffer.dataset.attrs["parcels_mesh"] + @property + def zdim(self): + if "nz1" in self.data.dims: + return self.data.sizes["nz1"] + elif "nz" in self.data.dims: + return self.data.sizes["nz"] else: - lon = 0 - lat = 0 - mesh = "flat" - - if "depth" in dimensions: - with NetcdfFileBuffer( - depth_filename, - dimensions, - indices, - interp_method=interp_method, - gridindexingtype=gridindexingtype, - ) as filebuffer: - filebuffer.name = variable[1] - depth = filebuffer.depth + return 0 + + @property + def n_face(self): + if type(self.data) is ux.uxDataArray: + return self.grid.n_face else: - indices["depth"] = np.array([0]) - depth = np.zeros(1) - - if len(data_filenames) > 1 and "time" not in dimensions: - raise RuntimeError("Multiple files given but no time dimension specified") - - if grid is None: - # Concatenate time variable to determine overall dimension - # across multiple files - if "time" in dimensions: - time, time_origin = cls._collect_time(data_filenames, dimensions, indices) - grid = Grid.create_grid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh) - else: # e.g. for the CROCO CS_w field, see https://github.com/OceanParcels/Parcels/issues/1831 - grid = Grid.create_grid(lon, lat, depth, np.array([0.0]), time_origin=TimeConverter(0.0), mesh=mesh) - data_filenames = [data_filenames[0]] - - if "time" in indices: - warnings.warn( - "time dimension in indices is not necessary anymore. It is then ignored.", FieldSetWarning, stacklevel=2 - ) + return 0 # TODO : Discuss what we want to return as n_face for dataarray obj - with NetcdfFileBuffer( # type: ignore[operator] - data_filenames, - dimensions, - indices, - interp_method=interp_method, - ) as filebuffer: - # If Field.from_netcdf is called directly, it may not have a 'data' dimension - # In that case, assume that 'name' is the data dimension - filebuffer.name = variable[1] - buffer_data = filebuffer.data - if len(buffer_data.shape) == 4: - errormessage = ( - f"Field {filebuffer.name} expecting a data shape of [tdim={grid.tdim}, zdim={grid.zdim}, " - f"ydim={grid.ydim}, xdim={grid.xdim }] " - f"but got shape {buffer_data.shape}." - ) - assert buffer_data.shape[0] == grid.tdim, errormessage - assert buffer_data.shape[2] == grid.ydim, errormessage - assert buffer_data.shape[3] == grid.xdim, errormessage - - data = buffer_data + @property + def interp_method(self): + return self._interp_method - if allow_time_extrapolation is None: - allow_time_extrapolation = False if "time" in dimensions else True - - return cls( - variable, - data, - grid=grid, - allow_time_extrapolation=allow_time_extrapolation, - interp_method=interp_method, - **kwargs, + @interp_method.setter + def interp_method(self, method: Callable): + self._validate_interp_function(method) + self._interp_method = method + + def _get_ux_barycentric_coordinates(self, y, x, fi): + """Checks if a point is inside a given face id. Used for unstructured grids.""" + # Check if particle is in the same face, otherwise search again. + n_nodes = self.grid.n_nodes_per_face[fi].to_numpy() + node_ids = self.grid.face_node_connectivity[fi, 0:n_nodes] + nodes = np.column_stack( + ( + np.deg2rad(self.grid.node_lon[node_ids].to_numpy()), + np.deg2rad(self.grid.node_lat[node_ids].to_numpy()), + ) ) - @classmethod - def from_xarray( - cls, - da: xr.DataArray, - name: str, - dimensions, - mesh: Mesh = "spherical", - allow_time_extrapolation: bool | None = None, - **kwargs, - ): - """Create field from xarray Variable. - - Parameters - ---------- - da : xr.DataArray - Xarray DataArray - name : str - Name of the Field - dimensions : dict - Dictionary mapping variable names for the relevant dimensions in the DataArray - mesh : str - String indicating the type of mesh coordinates and - units used during velocity interpolation: - - 1. spherical (default): Lat and lon in degree, with a - correction for zonal velocity U near the poles. - 2. flat: No conversion, lat/lon are assumed to be in m. - allow_time_extrapolation : bool - boolean whether to allow for extrapolation in time - (i.e. beyond the last available time snapshot) - Default is False if dimensions includes time, else True - **kwargs : - Keyword arguments passed to the :class:`Field` constructor. - """ - data = da.data - interp_method = kwargs.pop("interp_method", "linear") - - time = da[dimensions["time"]].values if "time" in dimensions else np.array([0.0]) - depth = da[dimensions["depth"]].values if "depth" in dimensions else np.array([0]) - lon = da[dimensions["lon"]].values - lat = da[dimensions["lat"]].values - - time_origin = TimeConverter(time[0]) - time = time_origin.reltime(time) # type: ignore[assignment] - - grid = Grid.create_grid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh) - return cls( - name, - data, - grid=grid, - allow_time_extrapolation=allow_time_extrapolation, - interp_method=interp_method, - **kwargs, - ) + coord = np.deg2rad([x, y]) + bcoord = np.asarray(_barycentric_coordinates(nodes, coord)) + err = abs(np.dot(bcoord, nodes[:, 0]) - coord[0]) + abs(np.dot(bcoord, nodes[:, 1]) - coord[1]) + return bcoord, err + + def _search_indices_unstructured(self, z, y, x, ei=None, search2D=False): + tol = 1e-10 + if ei is None: + # Search using global search + fi, bcoords = self._spatialhash.query([[x, y]]) # Get the face id for the particle + if fi == -1: + raise FieldOutOfBoundError(z, y, x) + # TODO Joe : Do the vertical grid search + # zi = self._vertical_search(z) + zi = 0 # For now + return bcoords, self.ravel_index(zi, 0, fi) + else: + zi, fi = self.unravel_index(ei[self.igrid]) # Get the z, and face index of the particle + # Search using nearest neighbors + bcoords, err = self._get_ux_barycentric_coordinates(y, x, fi) - def _reshape(self, data): - # Ensure that field data is the right data type - if not isinstance(data, (np.ndarray)): - data = np.array(data) - - if self.grid.xdim == 1 or self.grid.ydim == 1: - data = np.squeeze(data) # First remove all length-1 dimensions in data, so that we can add them below - if self.grid.xdim == 1 and len(data.shape) < 4: - data = np.expand_dims(data, axis=-1) - if self.grid.ydim == 1 and len(data.shape) < 4: - data = np.expand_dims(data, axis=-2) - if self.grid.tdim == 1: - if len(data.shape) < 4: - data = data.reshape(sum(((1,), data.shape), ())) - if self.grid.zdim == 1: - if len(data.shape) == 4: - data = data.reshape(sum(((data.shape[0],), data.shape[2:]), ())) - if len(data.shape) == 4: - errormessage = f"Field {self.name} expecting a data shape of [tdim, zdim, ydim, xdim]. " - assert data.shape[0] == self.grid.tdim, errormessage - assert data.shape[2] == self.grid.ydim, errormessage - assert data.shape[3] == self.grid.xdim, errormessage - if self.gridindexingtype == "pop": - assert data.shape[1] == self.grid.zdim or data.shape[1] == self.grid.zdim - 1, errormessage + if ((bcoords >= 0).all()) and ((bcoords <= 1.0).all()) and err < tol: + # TODO Joe : Do the vertical grid search + return bcoords, ei else: - assert data.shape[1] == self.grid.zdim, errormessage + # In this case we need to search the neighbors + for neighbor in self.grid.face_face_connectivity[fi, :]: + bcoords, err = self._get_ux_barycentric_coordinates(y, x, neighbor) + if ((bcoords >= 0).all()) and ((bcoords <= 1.0).all()) and err < tol: + # TODO Joe: Do the vertical grid search + return bcoords, self.ravel_index(zi, 0, neighbor) + + # If we reach this point, we do a global search as a last ditch effort the particle is out of bounds + fi, bcoords = self._spatialhash.query([[x, y]]) # Get the face id for the particle + if fi == -1: + raise FieldOutOfBoundError(z, y, x) + + def _search_indices_structured(self, z, y, x, ei=None, search2D=False): + if self._gtype in [GridType.RectilinearSGrid, GridType.RectilinearZGrid]: + (zeta, eta, xsi, zi, yi, xi) = _search_indices_rectilinear(self, z, y, x, ei=ei, search2D=search2D) else: - assert data.shape == ( - self.grid.tdim, - self.grid.ydim, - self.grid.xdim, - ), f"Field {self.name} expecting a data shape of [tdim, ydim, xdim]. " + ## TODO : Still need to implement the search_indices_curvilinear + # (zeta, eta, xsi, zi, yi, xi) = _search_indices_curvilinear( + # self, z, y, x, ei=ei, search2D=search2D + # ) + raise NotImplementedError("Curvilinear grid search not implemented yet") - return data + return (zeta, eta, xsi, zi, yi, xi) - def _search_indices(self, time, z, y, x, particle=None, search2D=False): - tau, ti = _search_time_index(self.grid, time, self.allow_time_extrapolation) + def _search_indices(self, time: datetime, z, y, x, ei=None, search2D=False): + tau, ti = _search_time_index(self, time, self.allow_time_extrapolation) - if self.grid._gtype in [GridType.RectilinearSGrid, GridType.RectilinearZGrid]: - (zeta, eta, xsi, zi, yi, xi) = _search_indices_rectilinear( - self, time, z, y, x, ti, particle=particle, search2D=search2D - ) + if ei is None: + _ei = None else: - (zeta, eta, xsi, zi, yi, xi) = _search_indices_curvilinear( - self, time, z, y, x, ti, particle=particle, search2D=search2D - ) - return (tau, zeta, eta, xsi, ti, zi, yi, xi) - - def _interpolator2D(self, time, z, y, x, particle=None): - """Impelement 2D interpolation with coordinate transformations as seen in Delandmeter and Van Sebille (2019), 10.5194/gmd-12-3571-2019..""" - try: - f = get_2d_interpolator_registry()[self.interp_method] - except KeyError: - if self.interp_method == "cgrid_velocity": - raise RuntimeError( - f"{self.name} is a scalar field. cgrid_velocity interpolation method should be used for vector fields (e.g. FieldSet.UV)" - ) - else: - raise RuntimeError(self.interp_method + " is not implemented for 2D grids") - - (tau, _, eta, xsi, ti, _, yi, xi) = self._search_indices(time, z, y, x, particle=particle) - - ctx = InterpolationContext2D(self.data, tau, eta, xsi, ti, yi, xi) - return f(ctx) - - def _interpolator3D(self, time, z, y, x, particle=None): - """Impelement 3D interpolation with coordinate transformations as seen in Delandmeter and Van Sebille (2019), 10.5194/gmd-12-3571-2019..""" - try: - f = get_3d_interpolator_registry()[self.interp_method] - except KeyError: - raise RuntimeError(self.interp_method + " is not implemented for 3D grids") + _ei = ei[self.igrid] - (tau, zeta, eta, xsi, ti, zi, yi, xi) = self._search_indices(time, z, y, x, particle=particle) - - ctx = InterpolationContext3D(self.data, tau, zeta, eta, xsi, ti, zi, yi, xi, self.gridindexingtype) - return f(ctx) + if type(self.data) is ux.UxDataArray: + bcoords, ei = self._search_indices_unstructured(z, y, x, ei=_ei, search2D=search2D) + else: + bcoords, ei = self._search_indices_structured(z, y, x, ei=_ei, search2D=search2D) + return bcoords, ei, tau, ti - def _interpolate(self, time, z, y, x, particle=None): - """Interpolate spatial field values.""" + def _interpolate(self, time: datetime, z, y, x, ei): try: - if self.grid.zdim == 1: - val = self._interpolator2D(time, z, y, x, particle=particle) - else: - val = self._interpolator3D(time, z, y, x, particle=particle) + bcoords, _ei, tau, ti = self._search_indices(time, z, y, x, ei=ei) + val = self._interp_method(ti, _ei, bcoords, tau, time, z, y, x) if np.isnan(val): # Detect Out-of-bounds sampling and raise exception @@ -614,17 +430,19 @@ def __getitem__(self, key): except tuple(AllParcelsErrorCodes.keys()) as error: return _deal_with_errors(error, key, vector_type=None) - def eval(self, time, z, y, x, particle=None, applyConversion=True): + def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True): """Interpolate field values in space and time. We interpolate linearly in time and apply implicit unit conversion to the result. Note that we defer to scipy.interpolate to perform spatial interpolation. """ - if self.gridindexingtype == "croco" and self not in [self.fieldset.H, self.fieldset.Zeta]: - z = _croco_from_z_to_sigma_scipy(self.fieldset, time, z, y, x, particle=particle) + if particle is None: + _ei = None + else: + _ei = particle.ei[self.igrid] - value = self._interpolate(time, z, y, x, particle=particle) + value = self._interpolate(time, z, y, x, ei=_ei) if applyConversion: return self.units.to_target(value, z, y, x) @@ -645,17 +463,21 @@ def ravel_index(self, zi, yi, xi): yi : int y index xi : int - x index + x index. When using an unstructured grid, this is the face index (fi) Returns ------- int flat index """ - return xi + self.grid.xdim * (yi + self.grid.ydim * zi) + if type(self.data) is xr.DataArray: + return xi + self.xdim * (yi + self.ydim * zi) + else: + return xi + self.n_face * zi def unravel_index(self, ei): """Return the zi, yi, xi indices for a given flat index. + Only used when working with fields on a structured grid. Parameters ---------- @@ -671,48 +493,79 @@ def unravel_index(self, ei): xi : int The x index. """ - _ei = ei[self.igrid] - zi = _ei // (self.grid.xdim * self.grid.ydim) - _ei = _ei % (self.grid.xdim * self.grid.ydim) - yi = _ei // self.grid.xdim - xi = _ei % self.grid.xdim - return zi, yi, xi + if type(self.data) is xr.DataArray: + _ei = ei[self.igrid] + zi = _ei // (self.xdim * self.ydim) + _ei = _ei % (self.xdim * self.ydim) + yi = _ei // self.xdim + xi = _ei % self.xdim + return zi, yi, xi + else: + _ei = ei[self.igrid] + zi = _ei // self.n_face + fi = _ei % self.n_face + return zi, fi + + def __getattr__(self, key: str): + return getattr(self.data, key) + + def __contains__(self, key: str): + return key in self.data class VectorField: - """Class VectorField stores 2 or 3 fields which defines together a vector field. - This enables to interpolate them as one single vector field in the kernels. - - Parameters - ---------- - name : str - Name of the vector field - U : parcels.field.Field - field defining the zonal component - V : parcels.field.Field - field defining the meridional component - W : parcels.field.Field - field defining the vertical component (default: None) - """ + """VectorField class that holds vector field data needed to execute particles.""" - def __init__(self, name: str, U: Field, V: Field, W: Field | None = None): + @staticmethod + def _vector_interp_template( + self, + ti: int, + ei: int, + bcoords: np.ndarray, + t: np.float32 | np.float64, + z: np.float32 | np.float64, + y: np.float32 | np.float64, + x: np.float32 | np.float64, + ) -> np.float32 | np.float64: + """Template function used for the signature check of the lateral interpolation methods.""" + return 0.0 + + def _validate_vector_interp_function(self, func: Callable): + """Ensures that the function has the correct signature.""" + expected_params = ["ti", "ei", "bcoords", "t", "z", "y", "x"] + expected_return_types = (np.float32, np.float64) + + sig = inspect.signature(func) + params = list(sig.parameters.keys()) + + # Check the parameter names and count + if params != expected_params: + raise TypeError(f"Function must have parameters {expected_params}, but got {params}") + + # Check return annotation if present + return_annotation = sig.return_annotation + if return_annotation not in (inspect.Signature.empty, *expected_return_types): + raise TypeError(f"Function must return a float, but got {return_annotation}") + + def __init__( + self, name: str, U: Field, V: Field, W: Field | None = None, vector_interp_method: Callable | None = None + ): self.name = name self.U = U self.V = V self.W = W - if self.U.gridindexingtype == "croco" and self.W: - self.vector_type: VectorType = "3DSigma" - elif self.W: + + if self.W: self.vector_type = "3D" else: self.vector_type = "2D" - self.gridindexingtype = U.gridindexingtype - if self.U.interp_method == "cgrid_velocity": - assert self.V.interp_method == "cgrid_velocity", "Interpolation methods of U and V are not the same." - assert self._check_grid_dimensions(U.grid, V.grid), "Dimensions of U and V are not the same." - if W is not None and self.U.gridindexingtype != "croco": - assert W.interp_method == "cgrid_velocity", "Interpolation methods of U and W are not the same." - assert self._check_grid_dimensions(U.grid, W.grid), "Dimensions of U and W are not the same." + + # Setting the interpolation method dynamically + if vector_interp_method is None: + self._vector_interp_method = None + else: + self._validate_vector_interp_function(vector_interp_method) + self._interp_method = vector_interp_method def __repr__(self): return f"""<{type(self).__name__}> @@ -721,472 +574,105 @@ def __repr__(self): V: {default_repr(self.V)} W: {default_repr(self.W)}""" - @staticmethod - def _check_grid_dimensions(grid1, grid2): - return ( - np.allclose(grid1.lon, grid2.lon) - and np.allclose(grid1.lat, grid2.lat) - and np.allclose(grid1.depth, grid2.depth) - and np.allclose(grid1.time, grid2.time) - ) - - def c_grid_interpolation2D(self, time, z, y, x, particle=None, applyConversion=True): - grid = self.U.grid - (tau, _, eta, xsi, ti, zi, yi, xi) = self.U._search_indices(time, z, y, x, particle=particle) - - if grid._gtype in [GridType.RectilinearSGrid, GridType.RectilinearZGrid]: - px = np.array([grid.lon[xi], grid.lon[xi + 1], grid.lon[xi + 1], grid.lon[xi]]) - py = np.array([grid.lat[yi], grid.lat[yi], grid.lat[yi + 1], grid.lat[yi + 1]]) - else: - px = np.array([grid.lon[yi, xi], grid.lon[yi, xi + 1], grid.lon[yi + 1, xi + 1], grid.lon[yi + 1, xi]]) - py = np.array([grid.lat[yi, xi], grid.lat[yi, xi + 1], grid.lat[yi + 1, xi + 1], grid.lat[yi + 1, xi]]) - - if grid.mesh == "spherical": - px[0] = px[0] + 360 if px[0] < x - 225 else px[0] - px[0] = px[0] - 360 if px[0] > x + 225 else px[0] - px[1:] = np.where(px[1:] - px[0] > 180, px[1:] - 360, px[1:]) - px[1:] = np.where(-px[1:] + px[0] > 180, px[1:] + 360, px[1:]) - xx = (1 - xsi) * (1 - eta) * px[0] + xsi * (1 - eta) * px[1] + xsi * eta * px[2] + (1 - xsi) * eta * px[3] - assert abs(xx - x) < 1e-4 - c1 = i_u._geodetic_distance(py[0], py[1], px[0], px[1], grid.mesh, np.dot(i_u.phi2D_lin(0.0, xsi), py)) - c2 = i_u._geodetic_distance(py[1], py[2], px[1], px[2], grid.mesh, np.dot(i_u.phi2D_lin(eta, 1.0), py)) - c3 = i_u._geodetic_distance(py[2], py[3], px[2], px[3], grid.mesh, np.dot(i_u.phi2D_lin(1.0, xsi), py)) - c4 = i_u._geodetic_distance(py[3], py[0], px[3], px[0], grid.mesh, np.dot(i_u.phi2D_lin(eta, 0.0), py)) - - def _calc_UV(ti, yi, xi): - if grid.zdim == 1: - if self.gridindexingtype == "nemo": - U0 = self.U.data[ti, yi + 1, xi] * c4 - U1 = self.U.data[ti, yi + 1, xi + 1] * c2 - V0 = self.V.data[ti, yi, xi + 1] * c1 - V1 = self.V.data[ti, yi + 1, xi + 1] * c3 - elif self.gridindexingtype in ["mitgcm", "croco"]: - U0 = self.U.data[ti, yi, xi] * c4 - U1 = self.U.data[ti, yi, xi + 1] * c2 - V0 = self.V.data[ti, yi, xi] * c1 - V1 = self.V.data[ti, yi + 1, xi] * c3 - else: - if self.gridindexingtype == "nemo": - U0 = self.U.data[ti, zi, yi + 1, xi] * c4 - U1 = self.U.data[ti, zi, yi + 1, xi + 1] * c2 - V0 = self.V.data[ti, zi, yi, xi + 1] * c1 - V1 = self.V.data[ti, zi, yi + 1, xi + 1] * c3 - elif self.gridindexingtype in ["mitgcm", "croco"]: - U0 = self.U.data[ti, zi, yi, xi] * c4 - U1 = self.U.data[ti, zi, yi, xi + 1] * c2 - V0 = self.V.data[ti, zi, yi, xi] * c1 - V1 = self.V.data[ti, zi, yi + 1, xi] * c3 - U = (1 - xsi) * U0 + xsi * U1 - V = (1 - eta) * V0 + eta * V1 - rad = np.pi / 180.0 - deg2m = 1852 * 60.0 - if applyConversion: - meshJac = (deg2m * deg2m * math.cos(rad * y)) if grid.mesh == "spherical" else 1 + @property + def vector_interp_method(self): + return self._vector_interp_method + + @vector_interp_method.setter + def vector_interp_method(self, method: Callable): + self._validate_vector_interp_function(method) + self._vector_interp_method = method + + # @staticmethod + # TODO : def _check_grid_dimensions(grid1, grid2): + # return ( + # np.allclose(grid1.lon, grid2.lon) + # and np.allclose(grid1.lat, grid2.lat) + # and np.allclose(grid1.depth, grid2.depth) + # and np.allclose(grid1.time, grid2.time) + # ) + def _interpolate(self, time, z, y, x, ei): + bcoords, _ei, ti = self._search_indices(time, z, y, x, ei=ei) + + if self._vector_interp_method is None: + u = self.U.eval(time, z, y, x, _ei, applyConversion=False) + v = self.V.eval(time, z, y, x, _ei, applyConversion=False) + if "3D" in self.vector_type: + w = self.W.eval(time, z, y, x, _ei, applyConversion=False) + return (u, v, w) else: - meshJac = deg2m if grid.mesh == "spherical" else 1 - - jac = i_u._compute_jacobian_determinant(py, px, eta, xsi) * meshJac - - u = ( - (-(1 - eta) * U - (1 - xsi) * V) * px[0] - + ((1 - eta) * U - xsi * V) * px[1] - + (eta * U + xsi * V) * px[2] - + (-eta * U + (1 - xsi) * V) * px[3] - ) / jac - v = ( - (-(1 - eta) * U - (1 - xsi) * V) * py[0] - + ((1 - eta) * U - xsi * V) * py[1] - + (eta * U + xsi * V) * py[2] - + (-eta * U + (1 - xsi) * V) * py[3] - ) / jac - if isinstance(u, da.core.Array): - u = u.compute() - v = v.compute() - return (u, v) - - u, v = _calc_UV(ti, yi, xi) - if should_calculate_next_ti(ti, tau, self.U.grid.tdim): - ut1, vt1 = _calc_UV(ti + 1, yi, xi) - u = (1 - tau) * u + tau * ut1 - v = (1 - tau) * v + tau * vt1 - return (u, v) - - def c_grid_interpolation3D_full(self, time, z, y, x, particle=None): - grid = self.U.grid - (tau, zeta, eta, xsi, ti, zi, yi, xi) = self.U._search_indices(time, z, y, x, particle=particle) - - if grid._gtype in [GridType.RectilinearSGrid, GridType.RectilinearZGrid]: - px = np.array([grid.lon[xi], grid.lon[xi + 1], grid.lon[xi + 1], grid.lon[xi]]) - py = np.array([grid.lat[yi], grid.lat[yi], grid.lat[yi + 1], grid.lat[yi + 1]]) + return (u, v, 0) else: - px = np.array([grid.lon[yi, xi], grid.lon[yi, xi + 1], grid.lon[yi + 1, xi + 1], grid.lon[yi + 1, xi]]) - py = np.array([grid.lat[yi, xi], grid.lat[yi, xi + 1], grid.lat[yi + 1, xi + 1], grid.lat[yi + 1, xi]]) - - if grid.mesh == "spherical": - px[0] = px[0] + 360 if px[0] < x - 225 else px[0] - px[0] = px[0] - 360 if px[0] > x + 225 else px[0] - px[1:] = np.where(px[1:] - px[0] > 180, px[1:] - 360, px[1:]) - px[1:] = np.where(-px[1:] + px[0] > 180, px[1:] + 360, px[1:]) - xx = (1 - xsi) * (1 - eta) * px[0] + xsi * (1 - eta) * px[1] + xsi * eta * px[2] + (1 - xsi) * eta * px[3] - assert abs(xx - x) < 1e-4 - - px = np.concatenate((px, px)) - py = np.concatenate((py, py)) - if grid._z4d: - pz = np.array( - [ - grid.depth[0, zi, yi, xi], - grid.depth[0, zi, yi, xi + 1], - grid.depth[0, zi, yi + 1, xi + 1], - grid.depth[0, zi, yi + 1, xi], - grid.depth[0, zi + 1, yi, xi], - grid.depth[0, zi + 1, yi, xi + 1], - grid.depth[0, zi + 1, yi + 1, xi + 1], - grid.depth[0, zi + 1, yi + 1, xi], - ] - ) + (u, v, w) = self._vector_interp_method(ti, _ei, bcoords, time, z, y, x) + return (u, v, w) + + def eval(self, time, z, y, x, ei=None, applyConversion=True): + if ei is None: + _ei = 0 else: - pz = np.array( - [ - grid.depth[zi, yi, xi], - grid.depth[zi, yi, xi + 1], - grid.depth[zi, yi + 1, xi + 1], - grid.depth[zi, yi + 1, xi], - grid.depth[zi + 1, yi, xi], - grid.depth[zi + 1, yi, xi + 1], - grid.depth[zi + 1, yi + 1, xi + 1], - grid.depth[zi + 1, yi + 1, xi], - ] - ) + _ei = ei[self.igrid] - u0 = self.U.data[ti, zi, yi + 1, xi] - u1 = self.U.data[ti, zi, yi + 1, xi + 1] - v0 = self.V.data[ti, zi, yi, xi + 1] - v1 = self.V.data[ti, zi, yi + 1, xi + 1] - w0 = self.W.data[ti, zi, yi + 1, xi + 1] - w1 = self.W.data[ti, zi + 1, yi + 1, xi + 1] - - if should_calculate_next_ti(ti, tau, self.U.grid.tdim): - u0 = (1 - tau) * u0 + tau * self.U.data[ti + 1, zi, yi + 1, xi] - u1 = (1 - tau) * u1 + tau * self.U.data[ti + 1, zi, yi + 1, xi + 1] - v0 = (1 - tau) * v0 + tau * self.V.data[ti + 1, zi, yi, xi + 1] - v1 = (1 - tau) * v1 + tau * self.V.data[ti + 1, zi, yi + 1, xi + 1] - w0 = (1 - tau) * w0 + tau * self.W.data[ti + 1, zi, yi + 1, xi + 1] - w1 = (1 - tau) * w1 + tau * self.W.data[ti + 1, zi + 1, yi + 1, xi + 1] - - U0 = u0 * i_u.jacobian3D_lin_face(pz, py, px, zeta, eta, 0, "zonal", grid.mesh) - U1 = u1 * i_u.jacobian3D_lin_face(pz, py, px, zeta, eta, 1, "zonal", grid.mesh) - V0 = v0 * i_u.jacobian3D_lin_face(pz, py, px, zeta, 0, xsi, "meridional", grid.mesh) - V1 = v1 * i_u.jacobian3D_lin_face(pz, py, px, zeta, 1, xsi, "meridional", grid.mesh) - W0 = w0 * i_u.jacobian3D_lin_face(pz, py, px, 0, eta, xsi, "vertical", grid.mesh) - W1 = w1 * i_u.jacobian3D_lin_face(pz, py, px, 1, eta, xsi, "vertical", grid.mesh) - - # Computing fluxes in half left hexahedron -> flux_u05 - xx = [ - px[0], - (px[0] + px[1]) / 2, - (px[2] + px[3]) / 2, - px[3], - px[4], - (px[4] + px[5]) / 2, - (px[6] + px[7]) / 2, - px[7], - ] - yy = [ - py[0], - (py[0] + py[1]) / 2, - (py[2] + py[3]) / 2, - py[3], - py[4], - (py[4] + py[5]) / 2, - (py[6] + py[7]) / 2, - py[7], - ] - zz = [ - pz[0], - (pz[0] + pz[1]) / 2, - (pz[2] + pz[3]) / 2, - pz[3], - pz[4], - (pz[4] + pz[5]) / 2, - (pz[6] + pz[7]) / 2, - pz[7], - ] - flux_u0 = u0 * i_u.jacobian3D_lin_face(zz, yy, xx, 0.5, 0.5, 0, "zonal", grid.mesh) - flux_v0_halfx = v0 * i_u.jacobian3D_lin_face(zz, yy, xx, 0.5, 0, 0.5, "meridional", grid.mesh) - flux_v1_halfx = v1 * i_u.jacobian3D_lin_face(zz, yy, xx, 0.5, 1, 0.5, "meridional", grid.mesh) - flux_w0_halfx = w0 * i_u.jacobian3D_lin_face(zz, yy, xx, 0, 0.5, 0.5, "vertical", grid.mesh) - flux_w1_halfx = w1 * i_u.jacobian3D_lin_face(zz, yy, xx, 1, 0.5, 0.5, "vertical", grid.mesh) - flux_u05 = flux_u0 + flux_v0_halfx - flux_v1_halfx + flux_w0_halfx - flux_w1_halfx - - # Computing fluxes in half front hexahedron -> flux_v05 - xx = [ - px[0], - px[1], - (px[1] + px[2]) / 2, - (px[0] + px[3]) / 2, - px[4], - px[5], - (px[5] + px[6]) / 2, - (px[4] + px[7]) / 2, - ] - yy = [ - py[0], - py[1], - (py[1] + py[2]) / 2, - (py[0] + py[3]) / 2, - py[4], - py[5], - (py[5] + py[6]) / 2, - (py[4] + py[7]) / 2, - ] - zz = [ - pz[0], - pz[1], - (pz[1] + pz[2]) / 2, - (pz[0] + pz[3]) / 2, - pz[4], - pz[5], - (pz[5] + pz[6]) / 2, - (pz[4] + pz[7]) / 2, - ] - flux_u0_halfy = u0 * i_u.jacobian3D_lin_face(zz, yy, xx, 0.5, 0.5, 0, "zonal", grid.mesh) - flux_u1_halfy = u1 * i_u.jacobian3D_lin_face(zz, yy, xx, 0.5, 0.5, 1, "zonal", grid.mesh) - flux_v0 = v0 * i_u.jacobian3D_lin_face(zz, yy, xx, 0.5, 0, 0.5, "meridional", grid.mesh) - flux_w0_halfy = w0 * i_u.jacobian3D_lin_face(zz, yy, xx, 0, 0.5, 0.5, "vertical", grid.mesh) - flux_w1_halfy = w1 * i_u.jacobian3D_lin_face(zz, yy, xx, 1, 0.5, 0.5, "vertical", grid.mesh) - flux_v05 = flux_u0_halfy - flux_u1_halfy + flux_v0 + flux_w0_halfy - flux_w1_halfy - - # Computing fluxes in half lower hexahedron -> flux_w05 - xx = [ - px[0], - px[1], - px[2], - px[3], - (px[0] + px[4]) / 2, - (px[1] + px[5]) / 2, - (px[2] + px[6]) / 2, - (px[3] + px[7]) / 2, - ] - yy = [ - py[0], - py[1], - py[2], - py[3], - (py[0] + py[4]) / 2, - (py[1] + py[5]) / 2, - (py[2] + py[6]) / 2, - (py[3] + py[7]) / 2, - ] - zz = [ - pz[0], - pz[1], - pz[2], - pz[3], - (pz[0] + pz[4]) / 2, - (pz[1] + pz[5]) / 2, - (pz[2] + pz[6]) / 2, - (pz[3] + pz[7]) / 2, - ] - flux_u0_halfz = u0 * i_u.jacobian3D_lin_face(zz, yy, xx, 0.5, 0.5, 0, "zonal", grid.mesh) - flux_u1_halfz = u1 * i_u.jacobian3D_lin_face(zz, yy, xx, 0.5, 0.5, 1, "zonal", grid.mesh) - flux_v0_halfz = v0 * i_u.jacobian3D_lin_face(zz, yy, xx, 0.5, 0, 0.5, "meridional", grid.mesh) - flux_v1_halfz = v1 * i_u.jacobian3D_lin_face(zz, yy, xx, 0.5, 1, 0.5, "meridional", grid.mesh) - flux_w0 = w0 * i_u.jacobian3D_lin_face(zz, yy, xx, 0, 0.5, 0.5, "vertical", grid.mesh) - flux_w05 = flux_u0_halfz - flux_u1_halfz + flux_v0_halfz - flux_v1_halfz + flux_w0 - - surf_u05 = i_u.jacobian3D_lin_face(pz, py, px, 0.5, 0.5, 0.5, "zonal", grid.mesh) - jac_u05 = i_u.jacobian3D_lin_face(pz, py, px, zeta, eta, 0.5, "zonal", grid.mesh) - U05 = flux_u05 / surf_u05 * jac_u05 - - surf_v05 = i_u.jacobian3D_lin_face(pz, py, px, 0.5, 0.5, 0.5, "meridional", grid.mesh) - jac_v05 = i_u.jacobian3D_lin_face(pz, py, px, zeta, 0.5, xsi, "meridional", grid.mesh) - V05 = flux_v05 / surf_v05 * jac_v05 - - surf_w05 = i_u.jacobian3D_lin_face(pz, py, px, 0.5, 0.5, 0.5, "vertical", grid.mesh) - jac_w05 = i_u.jacobian3D_lin_face(pz, py, px, 0.5, eta, xsi, "vertical", grid.mesh) - W05 = flux_w05 / surf_w05 * jac_w05 - - jac = i_u.jacobian3D_lin(pz, py, px, zeta, eta, xsi, grid.mesh) - dxsidt = i_u.interpolate(i_u.phi1D_quad, [U0, U05, U1], xsi) / jac - detadt = i_u.interpolate(i_u.phi1D_quad, [V0, V05, V1], eta) / jac - dzetdt = i_u.interpolate(i_u.phi1D_quad, [W0, W05, W1], zeta) / jac - - dphidxsi, dphideta, dphidzet = i_u.dphidxsi3D_lin(zeta, eta, xsi) - - u = np.dot(dphidxsi, px) * dxsidt + np.dot(dphideta, px) * detadt + np.dot(dphidzet, px) * dzetdt - v = np.dot(dphidxsi, py) * dxsidt + np.dot(dphideta, py) * detadt + np.dot(dphidzet, py) * dzetdt - w = np.dot(dphidxsi, pz) * dxsidt + np.dot(dphideta, pz) * detadt + np.dot(dphidzet, pz) * dzetdt - - if isinstance(u, da.core.Array): - u = u.compute() - v = v.compute() - w = w.compute() - return (u, v, w) + (u, v, w) = self._interpolate(time, z, y, x, _ei) - def c_grid_interpolation3D(self, ti, z, y, x, time, particle=None, applyConversion=True): - """Perform C grid interpolation in 3D. :: - - +---+---+---+ - | |V1 | | - +---+---+---+ - |U0 | |U1 | - +---+---+---+ - | |V0 | | - +---+---+---+ - - The interpolation is done in the following by - interpolating linearly U depending on the longitude coordinate and - interpolating linearly V depending on the latitude coordinate. - Curvilinear grids are treated properly, since the element is projected to a rectilinear parent element. - """ - if self.U.grid._gtype in [GridType.RectilinearSGrid, GridType.CurvilinearSGrid]: - (u, v, w) = self.c_grid_interpolation3D_full(time, z, y, x, particle=particle) - else: - if self.gridindexingtype == "croco": - z = _croco_from_z_to_sigma_scipy(self.fieldset, time, z, y, x, particle=particle) - (u, v) = self.c_grid_interpolation2D(time, z, y, x, particle=particle) - w = self.W.eval(time, z, y, x, particle=particle, applyConversion=False) - if applyConversion: + if applyConversion: + u = self.U.units.to_target(u, z, y, x) + v = self.V.units.to_target(v, z, y, x) + if "3D" in self.vector_type: w = self.W.units.to_target(w, z, y, x) - return (u, v, w) - def _is_land2D(self, di, yi, xi): - if self.U.data.ndim == 3: - if di < np.shape(self.U.data)[0]: - return np.isclose(self.U.data[di, yi, xi], 0.0) and np.isclose(self.V.data[di, yi, xi], 0.0) - else: - return True - else: - if di < self.U.grid.zdim and yi < np.shape(self.U.data)[-2] and xi < np.shape(self.U.data)[-1]: - return np.isclose(self.U.data[0, di, yi, xi], 0.0) and np.isclose(self.V.data[0, di, yi, xi], 0.0) - else: - return True - - def slip_interpolation(self, time, z, y, x, particle=None, applyConversion=True): - (_, zeta, eta, xsi, ti, zi, yi, xi) = self.U._search_indices(time, z, y, x, particle=particle) - di = ti if self.U.grid.zdim == 1 else zi # general third dimension - - f_u, f_v, f_w = 1, 1, 1 - if ( - self._is_land2D(di, yi, xi) - and self._is_land2D(di, yi, xi + 1) - and self._is_land2D(di + 1, yi, xi) - and self._is_land2D(di + 1, yi, xi + 1) - and eta > 0 - ): - if self.U.interp_method == "partialslip": - f_u = f_u * (0.5 + 0.5 * eta) / eta - if self.vector_type == "3D": - f_w = f_w * (0.5 + 0.5 * eta) / eta - elif self.U.interp_method == "freeslip": - f_u = f_u / eta - if self.vector_type == "3D": - f_w = f_w / eta - if ( - self._is_land2D(di, yi + 1, xi) - and self._is_land2D(di, yi + 1, xi + 1) - and self._is_land2D(di + 1, yi + 1, xi) - and self._is_land2D(di + 1, yi + 1, xi + 1) - and eta < 1 - ): - if self.U.interp_method == "partialslip": - f_u = f_u * (1 - 0.5 * eta) / (1 - eta) - if self.vector_type == "3D": - f_w = f_w * (1 - 0.5 * eta) / (1 - eta) - elif self.U.interp_method == "freeslip": - f_u = f_u / (1 - eta) - if self.vector_type == "3D": - f_w = f_w / (1 - eta) - if ( - self._is_land2D(di, yi, xi) - and self._is_land2D(di, yi + 1, xi) - and self._is_land2D(di + 1, yi, xi) - and self._is_land2D(di + 1, yi + 1, xi) - and xsi > 0 - ): - if self.U.interp_method == "partialslip": - f_v = f_v * (0.5 + 0.5 * xsi) / xsi - if self.vector_type == "3D": - f_w = f_w * (0.5 + 0.5 * xsi) / xsi - elif self.U.interp_method == "freeslip": - f_v = f_v / xsi - if self.vector_type == "3D": - f_w = f_w / xsi - if ( - self._is_land2D(di, yi, xi + 1) - and self._is_land2D(di, yi + 1, xi + 1) - and self._is_land2D(di + 1, yi, xi + 1) - and self._is_land2D(di + 1, yi + 1, xi + 1) - and xsi < 1 - ): - if self.U.interp_method == "partialslip": - f_v = f_v * (1 - 0.5 * xsi) / (1 - xsi) - if self.vector_type == "3D": - f_w = f_w * (1 - 0.5 * xsi) / (1 - xsi) - elif self.U.interp_method == "freeslip": - f_v = f_v / (1 - xsi) - if self.vector_type == "3D": - f_w = f_w / (1 - xsi) - if self.U.grid.zdim > 1: - if ( - self._is_land2D(di, yi, xi) - and self._is_land2D(di, yi, xi + 1) - and self._is_land2D(di, yi + 1, xi) - and self._is_land2D(di, yi + 1, xi + 1) - and zeta > 0 - ): - if self.U.interp_method == "partialslip": - f_u = f_u * (0.5 + 0.5 * zeta) / zeta - f_v = f_v * (0.5 + 0.5 * zeta) / zeta - elif self.U.interp_method == "freeslip": - f_u = f_u / zeta - f_v = f_v / zeta - if ( - self._is_land2D(di + 1, yi, xi) - and self._is_land2D(di + 1, yi, xi + 1) - and self._is_land2D(di + 1, yi + 1, xi) - and self._is_land2D(di + 1, yi + 1, xi + 1) - and zeta < 1 - ): - if self.U.interp_method == "partialslip": - f_u = f_u * (1 - 0.5 * zeta) / (1 - zeta) - f_v = f_v * (1 - 0.5 * zeta) / (1 - zeta) - elif self.U.interp_method == "freeslip": - f_u = f_u / (1 - zeta) - f_v = f_v / (1 - zeta) - - u = f_u * self.U.eval(time, z, y, x, particle, applyConversion=applyConversion) - v = f_v * self.V.eval(time, z, y, x, particle, applyConversion=applyConversion) - if self.vector_type == "3D": - w = f_w * self.W.eval(time, z, y, x, particle, applyConversion=applyConversion) - return u, v, w - else: - return u, v - - def eval(self, time, z, y, x, particle=None, applyConversion=True): - if self.U.interp_method in ["partialslip", "freeslip"]: - return self.slip_interpolation(time, z, y, x, particle=particle, applyConversion=applyConversion) - - if self.U.interp_method not in ["cgrid_velocity"]: - u = self.U.eval(time, z, y, x, particle=particle, applyConversion=False) - v = self.V.eval(time, z, y, x, particle=particle, applyConversion=False) - if applyConversion: - u = self.U.units.to_target(u, z, y, x) - v = self.V.units.to_target(v, z, y, x) - elif self.U.interp_method == "cgrid_velocity": - (u, v) = self.c_grid_interpolation2D(time, z, y, x, particle=particle, applyConversion=applyConversion) - if "3D" in self.vector_type: - w = self.W.eval(time, z, y, x, particle=particle, applyConversion=applyConversion) - return (u, v, w) - else: - return (u, v) + return (u, v, w) def __getitem__(self, key): try: if _isParticle(key): - return self.eval(key.time, key.depth, key.lat, key.lon, key) + return self.eval(key.time, key.depth, key.lat, key.lon, key.ei) else: return self.eval(*key) except tuple(AllParcelsErrorCodes.keys()) as error: return _deal_with_errors(error, key, vector_type=self.vector_type) + + +def _validate_dataarray(data, name): + """Verifies that all the required attributes are present in the xarray.DataArray or + uxarray.UxDataArray object. + """ + if isinstance(data, ux.UxDataArray): + # Validate dimensions + if not ("nz1" in data.dims or "nz" in data.dims): + raise ValueError( + f"Field {name} is missing a 'nz1' or 'nz' dimension in the field's metadata. " + "This attribute is required for xarray.DataArray objects." + ) + + if "time" not in data.dims: + raise ValueError( + f"Field {name} is missing a 'time' dimension in the field's metadata. " + "This attribute is required for xarray.DataArray objects." + ) + + # Validate attributes + required_keys = ["location", "mesh"] + for key in required_keys: + if key not in data.attrs.keys(): + raise ValueError( + f"Field {name} is missing a '{key}' attribute in the field's metadata. " + "This attribute is required for xarray.DataArray objects." + ) + + if type(data) is ux.UxDataArray: + _validate_uxgrid(data.uxgrid, name) + + +def _validate_uxgrid(grid, name): + """Verifies that all the required attributes are present in the uxarray.UxDataArray.UxGrid object.""" + if "Conventions" not in grid.attrs.keys(): + raise ValueError( + f"Field {name} is missing a 'Conventions' attribute in the field's metadata. " + "This attribute is required for uxarray.UxDataArray objects." + ) + if grid.attrs["Conventions"] != "UGRID-1.0": + raise ValueError( + f"Field {name} has a 'Conventions' attribute that is not 'UGRID-1.0'. " + "This attribute is required for uxarray.UxDataArray objects." + "See https://ugrid-conventions.github.io/ugrid-conventions/ for more information." + ) diff --git a/parcels/fieldset.py b/parcels/fieldset.py index b6a42bb282..9c8ee9315f 100644 --- a/parcels/fieldset.py +++ b/parcels/fieldset.py @@ -1,19 +1,10 @@ -import importlib.util -import os -import sys -import warnings -from glob import glob - import numpy as np +import uxarray as ux +import xarray as xr -from parcels._typing import GridIndexingType, InterpMethodOption, Mesh +from parcels._typing import Mesh from parcels.field import Field, VectorField -from parcels.grid import Grid -from parcels.gridset import GridSet -from parcels.particlefile import ParticleFile from parcels.tools._helpers import fieldset_repr -from parcels.tools.converters import TimeConverter -from parcels.tools.warnings import FieldSetWarning __all__ = ["FieldSet"] @@ -23,39 +14,91 @@ class FieldSet: Parameters ---------- - U : parcels.field.Field - Field object for zonal velocity component - V : parcels.field.Field - Field object for meridional velocity component - fields : dict mapping str to Field - Additional fields to include in the FieldSet. These fields can be used - in custom kernels. + ds : xarray.Dataset | uxarray.UxDataset) + xarray.Dataset and/or uxarray.UxDataset objects containing the field data. + + Notes + ----- + The `ds` object is a xarray.Dataset or uxarray.UxDataset object. + In XArray terminology, the (Ux)Dataset holds multiple (Ux)DataArray objects. + Each (Ux)DataArray object is a single "field" that is associated with their own + dimensions and coordinates within the (Ux)Dataset. + + A (Ux)Dataset object is associated with a single mesh, which can have multiple + types of "points" (multiple "grids") (e.g. for UxDataSets, these are "face_lon", + "face_lat", "node_lon", "node_lat", "edge_lon", "edge_lat"). Each (Ux)DataArray is + registered to a specific set of points on the mesh. + + For UxDataset objects, each `UXDataArray.attributes` field dictionary contains + the necessary metadata to help determine which set of points a field is registered + to and what parent model the field is associated with. Parcels uses this metadata + during execution for interpolation. Each `UXDataArray.attributes` field dictionary + must have: + * "location" key set to "face", "node", or "edge" to define which pairing of points a field is associated with. + * "mesh" key to define which parent model the fields are associated with (e.g. "fesom_mesh", "icon_mesh") + """ - def __init__(self, U: Field | None, V: Field | None, fields=None): - self.gridset = GridSet() + def __init__(self, datasets: list[xr.Dataset | ux.UxDataset]): + self.datasets = datasets + self._completed: bool = False - self._particlefile: ParticleFile | None = None - if U: - self.add_field(U, "U") - # see #1663 for type-ignore reason - self.time_origin = self.U.grid.time_origin if isinstance(self.U, Field) else self.U[0].grid.time_origin # type: ignore - if V: - self.add_field(V, "V") - - # Add additional fields as attributes - if fields: - for name, field in fields.items(): - self.add_field(field, name) + self._fieldnames = [] + time_origin = None + # Create pointers to each (Ux)DataArray + for ds in datasets: + for field in ds.data_vars: + if type(ds[field]) is ux.UxDataArray: + self.add_field(Field(field, ds[field], grid=ds[field].uxgrid), field) + else: + self.add_field(Field(field, ds[field]), field) + self._fieldnames.append(field) + + if "time" in ds.coords: + if time_origin is None: + time_origin = ds.time.min().data + else: + time_origin = min(time_origin, ds.time.min().data) + else: + time_origin = 0.0 + self.time_origin = time_origin self._add_UVfield() def __repr__(self): return fieldset_repr(self) - @property - def particlefile(self): - return self._particlefile + def dimrange(self, dim): + """Returns maximum value of a dimension (lon, lat, depth or time) + on 'left' side and minimum value on 'right' side for all grids + in a gridset. Useful for finding e.g. longitude range that + overlaps on all grids in a gridset. + """ + maxleft, minright = (-np.inf, np.inf) + dim2ds = { + "depth": ["nz1", "nz"], + "lat": ["node_lat", "face_lat", "edge_lat"], + "lon": ["node_lon", "face_lon", "edge_lon"], + "time": ["time"], + } + for ds in self.datasets: + for field in ds.data_vars: + for d in dim2ds[dim]: # check all possible dimensions + if d in ds[field].dims: + if dim == "depth": + maxleft = max(maxleft, ds[field][d].min().data) + minright = min(minright, ds[field][d].max().data) + else: + maxleft = max(maxleft, ds[field][d].data[0]) + minright = min(minright, ds[field][d].data[-1]) + maxleft = 0 if maxleft == -np.inf else maxleft # if all len(dim) == 1 + minright = 0 if minright == np.inf else minright # if all len(dim) == 1 + + return maxleft, minright + + # @property + # def particlefile(self): + # return self._particlefile @staticmethod def checkvaliddimensionsdict(dims): @@ -63,87 +106,9 @@ def checkvaliddimensionsdict(dims): if d not in ["lon", "lat", "depth", "time"]: raise NameError(f"{d} is not a valid key in the dimensions dictionary") - @classmethod - def from_data( - cls, - data, - dimensions, - mesh: Mesh = "spherical", - allow_time_extrapolation: bool | None = None, - **kwargs, - ): - """Initialise FieldSet object from raw data. - - Parameters - ---------- - data : - Dictionary mapping field names to numpy arrays. - Note that at least a 'U' and 'V' numpy array need to be given, and that - the built-in Advection kernels assume that U and V are in m/s. - Data shape is either [ydim, xdim], [zdim, ydim, xdim], [tdim, ydim, xdim] or [tdim, zdim, ydim, xdim], - dimensions : dict - Dictionary mapping field dimensions (lon, - lat, depth, time) to numpy arrays. - Note that dimensions can also be a dictionary of dictionaries if - dimension names are different for each variable - (e.g. dimensions['U'], dimensions['V'], etc). - mesh : str - String indicating the type of mesh coordinates and - units used during velocity interpolation, see also `this tutorial <../examples/tutorial_unitconverters.ipynb>`__: - - 1. spherical (default): Lat and lon in degree, with a - correction for zonal velocity U near the poles. - 2. flat: No conversion, lat/lon are assumed to be in m. - allow_time_extrapolation : bool - boolean whether to allow for extrapolation - (i.e. beyond the last available time snapshot) - Default is False if dimensions includes time, else True - **kwargs : - Keyword arguments passed to the :class:`Field` constructor. - - Examples - -------- - For usage examples see the following tutorials: - - * `Analytical advection <../examples/tutorial_analyticaladvection.ipynb>`__ - - * `Diffusion <../examples/tutorial_diffusion.ipynb>`__ - - * `Interpolation <../examples/tutorial_interpolation.ipynb>`__ - - * `Unit converters <../examples/tutorial_unitconverters.ipynb>`__ - """ - fields = {} - for name, datafld in data.items(): - # Use dimensions[name] if dimensions is a dict of dicts - dims = dimensions[name] if name in dimensions else dimensions - cls.checkvaliddimensionsdict(dims) - - if allow_time_extrapolation is None: - allow_time_extrapolation = False if "time" in dims else True - - lon = dims["lon"] - lat = dims["lat"] - depth = np.zeros(1, dtype=np.float32) if "depth" not in dims else dims["depth"] - time = np.zeros(1, dtype=np.float64) if "time" not in dims else dims["time"] - time = np.array(time) - if isinstance(time[0], np.datetime64): - time_origin = TimeConverter(time[0]) - time = np.array([time_origin.reltime(t) for t in time]) - else: - time_origin = kwargs.pop("time_origin", TimeConverter(0)) - grid = Grid.create_grid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh) - - fields[name] = Field( - name, - datafld, - grid=grid, - allow_time_extrapolation=allow_time_extrapolation, - **kwargs, - ) - u = fields.pop("U", None) - v = fields.pop("V", None) - return cls(u, v, fields=fields) + @property + def gridset_size(self): + return len(self._fieldnames) def add_field(self, field: Field, name: str | None = None): """Add a :class:`parcels.field.Field` object to the FieldSet. @@ -174,9 +139,10 @@ def add_field(self, field: Field, name: str | None = None): raise RuntimeError(f"FieldSet already has a Field with name '{name}'") else: setattr(self, name, field) - self.gridset.add_grid(field) + self._gridset_size += 1 + self._fieldnames.append(name) - def add_constant_field(self, name: str, value: float, mesh: Mesh = "flat"): + def add_constant_field(self, name: str, value, mesh: Mesh = "flat"): """Wrapper function to add a Field that is constant in space, useful e.g. when using constant horizontal diffusivity @@ -184,8 +150,8 @@ def add_constant_field(self, name: str, value: float, mesh: Mesh = "flat"): ---------- name : str Name of the :class:`parcels.field.Field` object to be added - value : float - Value of the constant field (stored as 32-bit float) + value : + Value of the constant field mesh : str String indicating the type of mesh coordinates and units used during velocity interpolation, see also `this tutorial <../examples/tutorial_unitconverters.ipynb>`__: @@ -194,7 +160,23 @@ def add_constant_field(self, name: str, value: float, mesh: Mesh = "flat"): correction for zonal velocity U near the poles. 2. flat: No conversion, lat/lon are assumed to be in m. """ - self.add_field(Field(name, value, lon=0, lat=0, mesh=mesh)) + time = 0.0 + values = np.full((1, 1, 1, 1), value) + data = xr.DataArray( + data=values, + name=name, + dims="null", + coords=[time, [0], [0], [0]], + attrs=dict(description="null", units="null", location="node", mesh="constant", mesh_type=mesh), + ) + self.add_field( + Field( + name, + data, + interp_method=None, # TODO : Need to define an interpolation method for constants + allow_time_extrapolation=True, + ) + ) def add_vector_field(self, vfield): """Add a :class:`parcels.field.VectorField` object to the FieldSet. @@ -202,13 +184,24 @@ def add_vector_field(self, vfield): Parameters ---------- vfield : parcels.VectorField - class:`parcels.field.VectorField` object to be added + class:`parcels.FieldSet.VectorField` object to be added """ setattr(self, vfield.name, vfield) for v in vfield.__dict__.values(): if isinstance(v, Field) and (v not in self.get_fields()): self.add_field(v) + def get_fields(self) -> list[Field | VectorField]: + """Returns a list of all the :class:`parcels.field.Field` and :class:`parcels.field.VectorField` + objects associated with this FieldSet. + """ + fields = [] + for v in self.__dict__.values(): + if type(v) in [Field, VectorField]: + if v not in fields: + fields.append(v) + return fields + def _add_UVfield(self): if not hasattr(self, "UV") and hasattr(self, "U") and hasattr(self, "V"): self.add_vector_field(VectorField("UV", self.U, self.V)) @@ -222,742 +215,86 @@ def _check_complete(self): if type(value) is Field: assert value.name == attr, f"Field {value.name}.name ({attr}) is not consistent" - def check_velocityfields(U, V, W): - if (U.interp_method == "cgrid_velocity" and V.interp_method != "cgrid_velocity") or ( - U.interp_method != "cgrid_velocity" and V.interp_method == "cgrid_velocity" - ): - raise ValueError("If one of U,V.interp_method='cgrid_velocity', the other should be too") - - if "linear_invdist_land_tracer" in [U.interp_method, V.interp_method]: - raise NotImplementedError( - "interp_method='linear_invdist_land_tracer' is not implemented for U and V Fields" - ) - - if U.interp_method == "cgrid_velocity": - if U.grid.xdim == 1 or U.grid.ydim == 1 or V.grid.xdim == 1 or V.grid.ydim == 1: - raise NotImplementedError( - "C-grid velocities require longitude and latitude dimensions at least length 2" - ) - - if U.gridindexingtype not in ["nemo", "mitgcm", "mom5", "pop", "croco"]: - raise ValueError("Field.gridindexing has to be one of 'nemo', 'mitgcm', 'mom5', 'pop' or 'croco'") - - if V.gridindexingtype != U.gridindexingtype or (W and W.gridindexingtype != U.gridindexingtype): - raise ValueError("Not all velocity Fields have the same gridindexingtype") - - W = self.W if hasattr(self, "W") else None - check_velocityfields(self.U, self.V, W) - - for g in self.gridset.grids: - g._check_zonal_periodic() - if len(g.time) == 1: - continue - assert isinstance( - g.time_origin.time_origin, type(self.time_origin.time_origin) - ), "time origins of different grids must be have the same type" - g.time = g.time + self.time_origin.reltime(g.time_origin) - g._time_origin = self.time_origin self._add_UVfield() self._completed = True - @classmethod - def _parse_wildcards(cls, paths, filenames, var): - if not isinstance(paths, list): - paths = sorted(glob(str(paths))) - if len(paths) == 0: - notfound_paths = filenames[var] if isinstance(filenames, dict) and var in filenames else filenames - raise OSError(f"FieldSet files not found for variable {var}: {notfound_paths}") - for fp in paths: - if not os.path.exists(fp): - raise OSError(f"FieldSet file not found: {fp}") - return paths - - @classmethod - def from_netcdf( - cls, - filenames, - variables, - dimensions, - mesh: Mesh = "spherical", - allow_time_extrapolation: bool | None = None, - **kwargs, - ): - """Initialises FieldSet object from NetCDF files. - - Parameters - ---------- - filenames : - Dictionary mapping variables to file(s). The - filepath may contain wildcards to indicate multiple files - or be a list of file. - filenames can be a list ``[files]``, a dictionary ``{var:[files]}``, - a dictionary ``{dim:[files]}`` (if lon, lat, depth and/or data not stored in same files as data), - or a dictionary of dictionaries ``{var:{dim:[files]}}``. - time values are in ``filenames[data]`` - variables : dict - Dictionary mapping variables to variable names in the netCDF file(s). - Note that the built-in Advection kernels assume that U and V are in m/s - dimensions : dict - Dictionary mapping data dimensions (lon, - lat, depth, time, data) to dimensions in the netCF file(s). - Note that dimensions can also be a dictionary of dictionaries if - dimension names are different for each variable - (e.g. dimensions['U'], dimensions['V'], etc). - mesh : str - String indicating the type of mesh coordinates and - units used during velocity interpolation, see also `this tutorial <../examples/tutorial_unitconverters.ipynb>`__: - - 1. spherical (default): Lat and lon in degree, with a - correction for zonal velocity U near the poles. - 2. flat: No conversion, lat/lon are assumed to be in m. - allow_time_extrapolation : bool - boolean whether to allow for extrapolation - (i.e. beyond the last available time snapshot) - Default is False if dimensions includes time, else True - interp_method : str - Method for interpolation. Options are 'linear' (default), 'nearest', - 'linear_invdist_land_tracer', 'cgrid_velocity', 'cgrid_tracer' and 'bgrid_velocity' - gridindexingtype : str - The type of gridindexing. Either 'nemo' (default), 'mitgcm', 'mom5', 'pop', or 'croco' are supported. - See also the Grid indexing documentation on oceanparcels.org - **kwargs : - Keyword arguments passed to the :class:`parcels.Field` constructor. - - - Examples - -------- - For usage examples see the following tutorials: - - * `Basic Parcels setup <../examples/parcels_tutorial.ipynb>`__ - - * `Argo floats <../examples/tutorial_Argofloats.ipynb>`__ - - * `Time-evolving depth dimensions <../examples/tutorial_timevaryingdepthdimensions.ipynb>`__ - - """ - fields: dict[str, Field] = {} - for var, name in variables.items(): - # Resolve all matching paths for the current variable - paths = filenames[var] if type(filenames) is dict and var in filenames else filenames - if type(paths) is not dict: - paths = cls._parse_wildcards(paths, filenames, var) - else: - for dim, p in paths.items(): - paths[dim] = cls._parse_wildcards(p, filenames, var) - - # Use dimensions[var] if it's a dict of dicts - dims = dimensions[var] if var in dimensions else dimensions - cls.checkvaliddimensionsdict(dims) - - grid = None - - fields[var] = Field.from_netcdf( - paths, - (var, name), - dims, - grid=grid, - mesh=mesh, - allow_time_extrapolation=allow_time_extrapolation, - **kwargs, - ) - - u = fields.pop("U", None) - v = fields.pop("V", None) - return cls(u, v, fields=fields) - - @classmethod - def from_nemo( - cls, - filenames, - variables, - dimensions, - mesh: Mesh = "spherical", - allow_time_extrapolation: bool | None = None, - tracer_interp_method: InterpMethodOption = "cgrid_tracer", - **kwargs, - ): - """Initialises FieldSet object from NetCDF files of Curvilinear NEMO fields. - - See `here <../examples/tutorial_nemo_curvilinear.ipynb>`__ - for a detailed tutorial on the setup for 2D NEMO fields and `here <../examples/tutorial_nemo_3D.ipynb>`__ - for the tutorial on the setup for 3D NEMO fields. - - See `here <../examples/documentation_indexing.ipynb>`__ - for a more detailed explanation of the different methods that can be used for c-grid datasets. - - Parameters - ---------- - filenames : - Dictionary mapping variables to file(s). The - filepath may contain wildcards to indicate multiple files, - or be a list of file. - filenames can be a list ``[files]``, a dictionary ``{var:[files]}``, - a dictionary ``{dim:[files]}`` (if lon, lat, depth and/or data not stored in same files as data), - or a dictionary of dictionaries ``{var:{dim:[files]}}`` - time values are in ``filenames[data]`` - variables : dict - Dictionary mapping variables to variable names in the netCDF file(s). - Note that the built-in Advection kernels assume that U and V are in m/s - dimensions : dict - Dictionary mapping data dimensions (lon, - lat, depth, time, data) to dimensions in the netCF file(s). - Note that dimensions can also be a dictionary of dictionaries if - dimension names are different for each variable. - Watch out: 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). :: - - +-----------------------------+-----------------------------+-----------------------------+ - | | V[k,j+1,i+1] | | - +-----------------------------+-----------------------------+-----------------------------+ - |U[k,j+1,i] |W[k:k+2,j+1,i+1],T[k,j+1,i+1]|U[k,j+1,i+1] | - +-----------------------------+-----------------------------+-----------------------------+ - | | V[k,j,i+1] | | - +-----------------------------+-----------------------------+-----------------------------+ - - 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 ) - In 3D, the depth is the one corresponding to W nodes - The gridindexingtype is set to 'nemo'. See also the Grid indexing documentation on oceanparcels.org - mesh : str - String indicating the type of mesh coordinates and - units used during velocity interpolation, see also `this tutorial <../examples/tutorial_unitconverters.ipynb>`__: - - 1. spherical (default): Lat and lon in degree, with a - correction for zonal velocity U near the poles. - 2. flat: No conversion, lat/lon are assumed to be in m. - allow_time_extrapolation : bool - boolean whether to allow for extrapolation - (i.e. beyond the last available time snapshot) - Default is False if dimensions includes time, else True - tracer_interp_method : str - Method for interpolation of tracer fields. It is recommended to use 'cgrid_tracer' (default) - Note that in the case of from_nemo() and from_c_grid_dataset(), the velocity fields are default to 'cgrid_velocity' - **kwargs : - Keyword arguments passed to the :func:`Fieldset.from_c_grid_dataset` constructor. - - """ - if kwargs.pop("gridindexingtype", "nemo") != "nemo": - raise ValueError( - "gridindexingtype must be 'nemo' in FieldSet.from_nemo(). Use FieldSet.from_c_grid_dataset otherwise" - ) - fieldset = cls.from_c_grid_dataset( - filenames, - variables, - dimensions, - mesh=mesh, - allow_time_extrapolation=allow_time_extrapolation, - tracer_interp_method=tracer_interp_method, - gridindexingtype="nemo", - **kwargs, - ) - return fieldset - - @classmethod - def from_mitgcm( - cls, - filenames, - variables, - dimensions, - mesh: Mesh = "spherical", - allow_time_extrapolation: bool | None = None, - tracer_interp_method: InterpMethodOption = "cgrid_tracer", - **kwargs, - ): - """Initialises FieldSet object from NetCDF files of MITgcm fields. - All parameters and keywords are exactly the same as for FieldSet.from_nemo(), except that - gridindexing is set to 'mitgcm' for grids that have the shape:: - - +-----------------------------+-----------------------------+-----------------------------+ - | | V[k,j+1,i] | | - +-----------------------------+-----------------------------+-----------------------------+ - |U[k,j,i] | W[k-1:k,j,i], T[k,j,i] |U[k,j,i+1] | - +-----------------------------+-----------------------------+-----------------------------+ - | | V[k,j,i] | | - +-----------------------------+-----------------------------+-----------------------------+ - - For indexing details: https://mitgcm.readthedocs.io/en/latest/algorithm/algorithm.html#spatial-discretization-of-the-dynamical-equations - Note that vertical velocity (W) is assumed positive in the positive z direction (which is upward in MITgcm) - """ - if kwargs.pop("gridindexingtype", "mitgcm") != "mitgcm": - raise ValueError( - "gridindexingtype must be 'mitgcm' in FieldSet.from_mitgcm(). Use FieldSet.from_c_grid_dataset otherwise" - ) - fieldset = cls.from_c_grid_dataset( - filenames, - variables, - dimensions, - mesh=mesh, - allow_time_extrapolation=allow_time_extrapolation, - tracer_interp_method=tracer_interp_method, - gridindexingtype="mitgcm", - **kwargs, - ) - return fieldset - - @classmethod - def from_croco( - cls, - filenames, - variables, - dimensions, - hc: float | None = None, - mesh="spherical", - allow_time_extrapolation=None, - tracer_interp_method="cgrid_tracer", - **kwargs, - ): - """Initialises FieldSet object from NetCDF files of CROCO fields. - All parameters and keywords are exactly the same as for FieldSet.from_nemo(), except that - in order to scale the vertical coordinate in CROCO, the following fields are required: - the bathymetry (``h``), the sea-surface height (``zeta``), the S-coordinate stretching curves - at W-points (``Cs_w``), and the stretching parameter (``hc``). - The horizontal interpolation uses the MITgcm grid indexing as described in FieldSet.from_mitgcm(). - - In 3D, when there is a ``depth`` dimension, the sigma grid scaling means that FieldSet.from_croco() - requires variables ``H: h`` and ``Zeta: zeta``, ``Cs_w: Cs_w``, as well as the stretching parameter ``hc`` - (as an extra input) parameter to work. - - See `the CROCO 3D tutorial <../examples/tutorial_croco_3D.ipynb>`__ for more infomation. - """ - if kwargs.pop("gridindexingtype", "croco") != "croco": - raise ValueError( - "gridindexingtype must be 'croco' in FieldSet.from_croco(). Use FieldSet.from_c_grid_dataset otherwise" - ) - - dimsU = dimensions["U"] if "U" in dimensions else dimensions - croco3D = True if "depth" in dimsU else False - - if croco3D: - if "W" in variables and variables["W"] == "omega": - warnings.warn( - "Note that Parcels expects 'w' for vertical velicites in 3D CROCO fields.\nSee https://docs.oceanparcels.org/en/latest/examples/tutorial_croco_3D.html for more information", - FieldSetWarning, - stacklevel=2, - ) - if "H" not in variables: - raise ValueError("FieldSet.from_croco() requires a bathymetry field 'H' for 3D CROCO fields") - if "Zeta" not in variables: - raise ValueError("FieldSet.from_croco() requires a free-surface field 'Zeta' for 3D CROCO fields") - if "Cs_w" not in variables: - raise ValueError( - "FieldSet.from_croco() requires the S-coordinate stretching curves at W-points 'Cs_w' for 3D CROCO fields" - ) - - interp_method = {} - for v in variables: - if v in ["U", "V"]: - interp_method[v] = "cgrid_velocity" - elif v in ["W", "H"]: - interp_method[v] = "linear" - else: - interp_method[v] = tracer_interp_method - - # Suppress the warning about the velocity interpolation since it is ok for CROCO - warnings.filterwarnings( - "ignore", - "Sampling of velocities should normally be done using fieldset.UV or fieldset.UVW object; tread carefully", - ) - - fieldset = cls.from_netcdf( - filenames, - variables, - dimensions, - mesh=mesh, - allow_time_extrapolation=allow_time_extrapolation, - interp_method=interp_method, - gridindexingtype="croco", - **kwargs, - ) - if croco3D: - if hc is None: - raise ValueError("FieldSet.from_croco() requires the hc parameter for 3D CROCO fields") - fieldset.add_constant("hc", hc) - return fieldset - - @classmethod - def from_c_grid_dataset( - cls, - filenames, - variables, - dimensions, - mesh: Mesh = "spherical", - allow_time_extrapolation: bool | None = None, - tracer_interp_method: InterpMethodOption = "cgrid_tracer", - gridindexingtype: GridIndexingType = "nemo", - **kwargs, - ): - """Initialises FieldSet object from NetCDF files of Curvilinear NEMO fields. - - See `here <../examples/documentation_indexing.ipynb>`__ - for a more detailed explanation of the different methods that can be used for c-grid datasets. - - Parameters - ---------- - filenames : - Dictionary mapping variables to file(s). The - filepath may contain wildcards to indicate multiple files, - or be a list of file. - filenames can be a list ``[files]``, a dictionary ``{var:[files]}``, - a dictionary ``{dim:[files]}`` (if lon, lat, depth and/or data not stored in same files as data), - or a dictionary of dictionaries ``{var:{dim:[files]}}`` - time values are in ``filenames[data]`` - variables : dict - Dictionary mapping variables to variable - names in the netCDF file(s). - dimensions : dict - Dictionary mapping data dimensions (lon, - lat, depth, time, data) to dimensions in the netCF file(s). - Note that dimensions can also be a dictionary of dictionaries if - dimension names are different for each variable. - Watch out: 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 ). :: - - +-----------------------------+-----------------------------+-----------------------------+ - | | V[k,j+1,i+1] | | - +-----------------------------+-----------------------------+-----------------------------+ - |U[k,j+1,i] |W[k:k+2,j+1,i+1],T[k,j+1,i+1]|U[k,j+1,i+1] | - +-----------------------------+-----------------------------+-----------------------------+ - | | V[k,j,i+1] | | - +-----------------------------+-----------------------------+-----------------------------+ - - 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 ) - In 3D, the depth is the one corresponding to W nodes. - mesh : str - String indicating the type of mesh coordinates and - units used during velocity interpolation, see also `this tutorial <../examples/tutorial_unitconverters.ipynb>`__: - - 1. spherical (default): Lat and lon in degree, with a - correction for zonal velocity U near the poles. - 2. flat: No conversion, lat/lon are assumed to be in m. - allow_time_extrapolation : bool - boolean whether to allow for extrapolation - (i.e. beyond the last available time snapshot) - Default is False if dimensions includes time, else True - tracer_interp_method : str - Method for interpolation of tracer fields. It is recommended to use 'cgrid_tracer' (default) - Note that in the case of from_nemo() and from_c_grid_dataset(), the velocity fields are default to 'cgrid_velocity' - gridindexingtype : str - The type of gridindexing. Set to 'nemo' in FieldSet.from_nemo(), 'mitgcm' in FieldSet.from_mitgcm() or 'croco' in FieldSet.from_croco(). - See also the Grid indexing documentation on oceanparcels.org (Default value = 'nemo') - **kwargs : - Keyword arguments passed to the :func:`Fieldset.from_netcdf` constructor. - """ - if "U" in dimensions and "V" in dimensions and dimensions["U"] != dimensions["V"]: - raise ValueError( - "On a C-grid, the dimensions of velocities should be the corners (f-points) of the cells, so the same for U and V. " - "See also https://docs.oceanparcels.org/en/latest/examples/documentation_indexing.html" - ) - if "U" in dimensions and "W" in dimensions and dimensions["U"] != dimensions["W"]: - raise ValueError( - "On a C-grid, the dimensions of velocities should be the corners (f-points) of the cells, so the same for U, V and W. " - "See also https://docs.oceanparcels.org/en/latest/examples/documentation_indexing.html" - ) - if "interp_method" in kwargs.keys(): - raise TypeError("On a C-grid, the interpolation method for velocities should not be overridden") - - interp_method = {} - for v in variables: - if v in ["U", "V", "W"]: - interp_method[v] = "cgrid_velocity" - else: - interp_method[v] = tracer_interp_method - - return cls.from_netcdf( - filenames, - variables, - dimensions, - mesh=mesh, - allow_time_extrapolation=allow_time_extrapolation, - interp_method=interp_method, - gridindexingtype=gridindexingtype, - **kwargs, - ) - - @classmethod - def from_mom5( - cls, - filenames, - variables, - dimensions, - mesh: Mesh = "spherical", - allow_time_extrapolation: bool | None = None, - tracer_interp_method: InterpMethodOption = "bgrid_tracer", - **kwargs, - ): - """Initialises FieldSet object from NetCDF files of MOM5 fields. - - Parameters - ---------- - filenames : - Dictionary mapping variables to file(s). The - filepath may contain wildcards to indicate multiple files, - or be a list of file. - filenames can be a list ``[files]``, a dictionary ``{var:[files]}``, - a dictionary ``{dim:[files]}`` (if lon, lat, depth and/or data not stored in same files as data), - or a dictionary of dictionaries ``{var:{dim:[files]}}`` - time values are in ``filenames[data]`` - variables : dict - Dictionary mapping variables to variable names in the netCDF file(s). - Note that the built-in Advection kernels assume that U and V are in m/s - dimensions : dict - Dictionary mapping data dimensions (lon, - lat, depth, time, data) to dimensions in the netCF file(s). - Note that dimensions can also be a dictionary of dictionaries if - dimension names are different for each variable. :: - - +-------------------------------+-------------------------------+-------------------------------+ - |U[k,j+1,i],V[k,j+1,i] | |U[k,j+1,i+1],V[k,j+1,i+1] | - +-------------------------------+-------------------------------+-------------------------------+ - | |W[k-1:k+1,j+1,i+1],T[k,j+1,i+1]| | - +-------------------------------+-------------------------------+-------------------------------+ - |U[k,j,i],V[k,j,i] | |U[k,j,i+1],V[k,j,i+1] | - +-------------------------------+-------------------------------+-------------------------------+ - - In 2D: U and V nodes are on the cell vertices and interpolated bilinearly as a A-grid. - T node is at the cell centre and interpolated constant per cell as a C-grid. - In 3D: U and V nodes are at the middle of the cell vertical edges, - They are interpolated bilinearly (independently of z) in the cell. - W nodes are at the centre of the horizontal interfaces, but below the U and V. - They are interpolated linearly (as a function of z) in the cell. - Note that W is normally directed upward in MOM5, but Parcels requires W - in the positive z-direction (downward) so W is multiplied by -1. - T node is at the cell centre, and constant per cell. - mesh : str - String indicating the type of mesh coordinates and - units used during velocity interpolation, see also the `Unit converters tutorial <../examples/tutorial_unitconverters.ipynb>`__: - - 1. spherical (default): Lat and lon in degree, with a - correction for zonal velocity U near the poles. - 2. flat: No conversion, lat/lon are assumed to be in m. - allow_time_extrapolation : bool - boolean whether to allow for extrapolation - (i.e. beyond the last available time snapshot) - Default is False if dimensions includes time, else True - tracer_interp_method : str - Method for interpolation of tracer fields. It is recommended to use 'bgrid_tracer' (default) - Note that in the case of from_mom5() and from_b_grid_dataset(), the velocity fields are default to 'bgrid_velocity' - **kwargs : - Keyword arguments passed to the :func:`Fieldset.from_b_grid_dataset` constructor. - """ - fieldset = cls.from_b_grid_dataset( - filenames, - variables, - dimensions, - mesh=mesh, - allow_time_extrapolation=allow_time_extrapolation, - tracer_interp_method=tracer_interp_method, - gridindexingtype="mom5", - **kwargs, - ) - return fieldset - - @classmethod - def from_a_grid_dataset(cls, filenames, variables, dimensions, **kwargs): - """ - Load a FieldSet from an A-grid dataset, which is the default grid type. - - Parameters - ---------- - filenames : - Path(s) to the input files. - variables : - Dictionary of the variables in the NetCDF file. - dimensions : - Dictionary of the dimensions in the NetCDF file. - **kwargs : - Additional keyword arguments for `from_netcdf()`. - - Returns - ------- - FieldSet - A FieldSet object. - """ - return cls.from_netcdf(filenames, variables, dimensions, **kwargs) - - @classmethod - def from_b_grid_dataset( - cls, - filenames, - variables, - dimensions, - mesh: Mesh = "spherical", - allow_time_extrapolation: bool | None = None, - tracer_interp_method: InterpMethodOption = "bgrid_tracer", - **kwargs, - ): - """Initialises FieldSet object from NetCDF files of Bgrid fields. - - Parameters - ---------- - filenames : - Dictionary mapping variables to file(s). The - filepath may contain wildcards to indicate multiple files, - or be a list of file. - filenames can be a list ``[files]``, a dictionary ``{var:[files]}``, - a dictionary ``{dim:[files]}`` (if lon, lat, depth and/or data not stored in same files as data), - or a dictionary of dictionaries ``{var:{dim:[files]}}`` - time values are in ``filenames[data]`` - variables : dict - Dictionary mapping variables to variable - names in the netCDF file(s). - dimensions : dict - Dictionary mapping data dimensions (lon, - lat, depth, time, data) to dimensions in the netCF file(s). - Note that dimensions can also be a dictionary of dictionaries if - dimension names are different for each variable. - U and V velocity nodes are not located as W velocity and T tracer nodes (see http://www2.cesm.ucar.edu/models/cesm1.0/pop2/doc/sci/POPRefManual.pdf ). :: - - +-----------------------------+-----------------------------+-----------------------------+ - |U[k,j+1,i],V[k,j+1,i] | |U[k,j+1,i+1],V[k,j+1,i+1] | - +-----------------------------+-----------------------------+-----------------------------+ - | |W[k:k+2,j+1,i+1],T[k,j+1,i+1]| | - +-----------------------------+-----------------------------+-----------------------------+ - |U[k,j,i],V[k,j,i] | |U[k,j,i+1],V[k,j,i+1] | - +-----------------------------+-----------------------------+-----------------------------+ - - In 2D: U and V nodes are on the cell vertices and interpolated bilinearly as a A-grid. - T node is at the cell centre and interpolated constant per cell as a C-grid. - In 3D: U and V nodes are at the midlle of the cell vertical edges, - They are interpolated bilinearly (independently of z) in the cell. - W nodes are at the centre of the horizontal interfaces. - They are interpolated linearly (as a function of z) in the cell. - T node is at the cell centre, and constant per cell. - mesh : str - String indicating the type of mesh coordinates and - units used during velocity interpolation, see also `this tutorial <../examples/tutorial_unitconverters.ipynb>`__: - - 1. spherical (default): Lat and lon in degree, with a - correction for zonal velocity U near the poles. - 2. flat: No conversion, lat/lon are assumed to be in m. - allow_time_extrapolation : bool - boolean whether to allow for extrapolation - (i.e. beyond the last available time snapshot) - Default is False if dimensions includes time, else True - tracer_interp_method : str - Method for interpolation of tracer fields. It is recommended to use 'bgrid_tracer' (default) - Note that in the case of from_b_grid_dataset(), the velocity fields are default to 'bgrid_velocity' - **kwargs : - Keyword arguments passed to the :func:`Fieldset.from_netcdf` constructor. - """ - if "U" in dimensions and "V" in dimensions and dimensions["U"] != dimensions["V"]: - raise ValueError( - "On a B-grid, the dimensions of velocities should be the (top) corners of the grid cells, so the same for U and V. " - "See also https://docs.oceanparcels.org/en/latest/examples/documentation_indexing.html" - ) - if "U" in dimensions and "W" in dimensions and dimensions["U"] != dimensions["W"]: - raise ValueError( - "On a B-grid, the dimensions of velocities should be the (top) corners of the grid cells, so the same for U, V and W. " - "See also https://docs.oceanparcels.org/en/latest/examples/documentation_indexing.html" - ) - - interp_method = {} - for v in variables: - if v in ["U", "V"]: - interp_method[v] = "bgrid_velocity" - elif v in ["W"]: - interp_method[v] = "bgrid_w_velocity" - else: - interp_method[v] = tracer_interp_method - - return cls.from_netcdf( - filenames, - variables, - dimensions, - mesh=mesh, - allow_time_extrapolation=allow_time_extrapolation, - interp_method=interp_method, - **kwargs, - ) - - @classmethod - def from_xarray_dataset(cls, ds, variables, dimensions, mesh="spherical", allow_time_extrapolation=None, **kwargs): - """Initialises FieldSet data from xarray Datasets. - - Parameters - ---------- - ds : xr.Dataset - xarray Dataset. - Note that the built-in Advection kernels assume that U and V are in m/s - variables : dict - Dictionary mapping parcels variable names to data variables in the xarray Dataset. - dimensions : dict - Dictionary mapping data dimensions (lon, - lat, depth, time, data) to dimensions in the xarray Dataset. - Note that dimensions can also be a dictionary of dictionaries if - dimension names are different for each variable - (e.g. dimensions['U'], dimensions['V'], etc). - mesh : str - String indicating the type of mesh coordinates and - units used during velocity interpolation, see also `this tutorial <../examples/tutorial_unitconverters.ipynb>`__: - - 1. spherical (default): Lat and lon in degree, with a - correction for zonal velocity U near the poles. - 2. flat: No conversion, lat/lon are assumed to be in m. - allow_time_extrapolation : bool - boolean whether to allow for extrapolation - (i.e. beyond the last available time snapshot) - Default is False if dimensions includes time, else True - **kwargs : - Keyword arguments passed to the :func:`Field.from_xarray` constructor. - """ - fields = {} - - for var, name in variables.items(): - dims = dimensions[var] if var in dimensions else dimensions - cls.checkvaliddimensionsdict(dims) - - fields[var] = Field.from_xarray( - ds[name], - var, - dims, - mesh=mesh, - allow_time_extrapolation=allow_time_extrapolation, - **kwargs, - ) - u = fields.pop("U", None) - v = fields.pop("V", None) - return cls(u, v, fields=fields) - - @classmethod - def from_modulefile(cls, filename, modulename="create_fieldset", **kwargs): - """Initialises FieldSet data from a file containing a python module file with a create_fieldset() function. - - Parameters - ---------- - filename: path to a python file containing at least a function which returns a FieldSet object. - modulename: name of the function in the python file that returns a FieldSet object. Default is "create_fieldset". - """ - # check if filename exists - if not os.path.exists(filename): - raise OSError(f"FieldSet module file {filename} does not exist") - - # Importing the source file directly (following https://docs.python.org/3/library/importlib.html#importing-a-source-file-directly) - spec = importlib.util.spec_from_file_location(modulename, filename) - fieldset_module = importlib.util.module_from_spec(spec) - sys.modules[modulename] = fieldset_module - spec.loader.exec_module(fieldset_module) - - if not hasattr(fieldset_module, modulename): - raise OSError(f"{filename} does not contain a {modulename} function") - fieldset = getattr(fieldset_module, modulename)(**kwargs) - if not isinstance(fieldset, FieldSet): - raise OSError(f"Module {filename}.{modulename} does not return a FieldSet object") - return fieldset - - def get_fields(self) -> list[Field | VectorField]: - """Returns a list of all the :class:`parcels.field.Field` and :class:`parcels.field.VectorField` - objects associated with this FieldSet. - """ - fields = [] - for v in self.__dict__.values(): - if type(v) in [Field, VectorField]: - if v not in fields: - fields.append(v) - return fields + # @classmethod + # def from_nemo( + # cls, + # filenames, + # variables, + # dimensions, + # mesh: Mesh = "spherical", + # allow_time_extrapolation: bool | None = None, + # tracer_interp_method: InterpMethodOption = "cgrid_tracer", + # **kwargs, + # ): + + # @classmethod + # def from_mitgcm( + # cls, + # filenames, + # variables, + # dimensions, + # mesh: Mesh = "spherical", + # allow_time_extrapolation: bool | None = None, + # tracer_interp_method: InterpMethodOption = "cgrid_tracer", + # **kwargs, + # ): + + # @classmethod + # def from_croco( + # cls, + # filenames, + # variables, + # dimensions, + # hc: float | None = None, + # mesh="spherical", + # allow_time_extrapolation=None, + # tracer_interp_method="cgrid_tracer", + # **kwargs, + # ): + + # @classmethod + # def from_c_grid_dataset( + # cls, + # filenames, + # variables, + # dimensions, + # mesh: Mesh = "spherical", + # allow_time_extrapolation: bool | None = None, + # tracer_interp_method: InterpMethodOption = "cgrid_tracer", + # gridindexingtype: GridIndexingType = "nemo", + # **kwargs, + # ): + + # @classmethod + # def from_mom5( + # cls, + # filenames, + # variables, + # dimensions, + # mesh: Mesh = "spherical", + # allow_time_extrapolation: bool | None = None, + # tracer_interp_method: InterpMethodOption = "bgrid_tracer", + # **kwargs, + # ): + + # @classmethod + # def from_a_grid_dataset(cls, filenames, variables, dimensions, **kwargs): + + # @classmethod + # def from_b_grid_dataset( + # cls, + # filenames, + # variables, + # dimensions, + # mesh: Mesh = "spherical", + # allow_time_extrapolation: bool | None = None, + # tracer_interp_method: InterpMethodOption = "bgrid_tracer", + # **kwargs, + # ): def add_constant(self, name, value): """Add a constant to the FieldSet. Note that all constants are @@ -980,29 +317,29 @@ def add_constant(self, name, value): """ setattr(self, name, value) - def computeTimeChunk(self, time=0.0, dt=1): - """Load a chunk of three data time steps into the FieldSet. - This is used when FieldSet uses data imported from netcdf, - with default option deferred_load. The loaded time steps are at or immediatly before time - and the two time steps immediately following time if dt is positive (and inversely for negative dt) - - Parameters - ---------- - time : - Time around which the FieldSet data are to be loaded. - Time is provided as a double, relatively to Fieldset.time_origin. - Default is 0. - dt : - time step of the integration scheme, needed to set the direction of time chunk loading. - Default is 1. - """ - nextTime = np.inf if dt > 0 else -np.inf - - if abs(nextTime) == np.inf or np.isnan(nextTime): # Second happens when dt=0 - return nextTime - else: - nSteps = int((nextTime - time) / dt) - if nSteps == 0: - return nextTime - else: - return time + nSteps * dt + # def computeTimeChunk(self, time=0.0, dt=1): + # """Load a chunk of three data time steps into the FieldSet. + # This is used when FieldSet uses data imported from netcdf, + # with default option deferred_load. The loaded time steps are at or immediatly before time + # and the two time steps immediately following time if dt is positive (and inversely for negative dt) + + # Parameters + # ---------- + # time : + # Time around which the FieldSet data are to be loaded. + # Time is provided as a double, relatively to Fieldset.time_origin. + # Default is 0. + # dt : + # time step of the integration scheme, needed to set the direction of time chunk loading. + # Default is 1. + # """ + # nextTime = np.inf if dt > 0 else -np.inf + + # if abs(nextTime) == np.inf or np.isnan(nextTime): # Second happens when dt=0 + # return nextTime + # else: + # nSteps = int((nextTime - time) / dt) + # if nSteps == 0: + # return nextTime + # else: + # return time + nSteps * dt diff --git a/parcels/particleset.py b/parcels/particleset.py index 2a8d625ca0..98ce3d0487 100644 --- a/parcels/particleset.py +++ b/parcels/particleset.py @@ -122,7 +122,7 @@ def ArrayClass_init(self, *args, **kwargs): if type(self).ngrids.initial < 0: numgrids = ngrids if numgrids is None and fieldset is not None: - numgrids = fieldset.gridset.size + numgrids = fieldset.gridset_size assert numgrids is not None, "Neither fieldsets nor number of grids are specified - exiting." type(self).ngrids.initial = numgrids self.ngrids = type(self).ngrids.initial @@ -148,7 +148,7 @@ def ArrayClass_init(self, *args, **kwargs): pid_orig = np.arange(lon.size) if depth is None: - mindepth = self.fieldset.gridset.dimrange("depth")[0] + mindepth = self.fieldset.dimrange("depth")[0] depth = np.ones(lon.size) * mindepth else: depth = convert_to_flat_array(depth) @@ -191,7 +191,7 @@ def ArrayClass_init(self, *args, **kwargs): self._repeatkwargs = kwargs self._repeatkwargs.pop("partition_function", None) - ngrids = fieldset.gridset.size + ngrids = fieldset.gridset_size # Variables used for interaction kernels. inter_dist_horiz = None @@ -962,7 +962,7 @@ def execute( if runtime is not None and endtime is not None: raise RuntimeError("Only one of (endtime, runtime) can be specified") - mintime, maxtime = self.fieldset.gridset.dimrange("time") + mintime, maxtime = self.fieldset.dimrange("time") default_release_time = mintime if dt >= 0 else maxtime if np.any(np.isnan(self.particledata.data["time"])): @@ -980,7 +980,7 @@ def execute( if runtime is not None: endtime = starttime + runtime * np.sign(dt) elif endtime is None: - mintime, maxtime = self.fieldset.gridset.dimrange("time") + mintime, maxtime = self.fieldset.dimrange("time") endtime = maxtime if dt >= 0 else mintime if (abs(endtime - starttime) < 1e-5 or runtime == 0) and dt == 0: diff --git a/parcels/tools/_helpers.py b/parcels/tools/_helpers.py index 7f26273bb6..74f4bdfb33 100644 --- a/parcels/tools/_helpers.py +++ b/parcels/tools/_helpers.py @@ -72,9 +72,8 @@ def field_repr(field: Field) -> str: """Return a pretty repr for Field""" out = f"""<{type(field).__name__}> name : {field.name!r} - grid : {field.grid!r} + data : {field.data!r} extrapolate time: {field.allow_time_extrapolation!r} - gridindexingtype: {field.gridindexingtype!r} """ return textwrap.dedent(out).strip() diff --git a/parcels/tools/exampledata_utils.py b/parcels/tools/exampledata_utils.py index 52a4c64e2d..6b253b42f4 100644 --- a/parcels/tools/exampledata_utils.py +++ b/parcels/tools/exampledata_utils.py @@ -49,6 +49,12 @@ "decaying_moving_eddyU.nc", "decaying_moving_eddyV.nc", ], + "FESOM_periodic_channel": [ + "fesom_channel.nc", + "u.fesom_channel.nc", + "v.fesom_channel.nc", + "w.fesom_channel.nc", + ], "NemoCurvilinear_data": [ "U_purely_zonal-ORCA025_grid_U.nc4", "V_purely_zonal-ORCA025_grid_V.nc4", diff --git a/tests/test_advection.py b/tests/test_advection.py index a62fe878dd..f0f710637e 100644 --- a/tests/test_advection.py +++ b/tests/test_advection.py @@ -57,6 +57,8 @@ def depth(): return np.linspace(0, 30, zdim, dtype=np.float32) +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_advection_zonal(lon, lat, depth): """Particles at high latitude move geographically faster due to the pole correction in `GeographicPolar`.""" npart = 10 @@ -88,6 +90,8 @@ def test_advection_zonal(lon, lat, depth): assert (np.diff(pset3D.lon) > 1.0e-4).all() +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_advection_meridional(lon, lat): """Particles at high latitude move geographically faster due to the pole correction in `GeographicPolar`.""" npart = 10 @@ -101,6 +105,8 @@ def test_advection_meridional(lon, lat): assert np.allclose(np.diff(pset.lat), delta_lat, rtol=1.0e-4) +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_advection_3D(): """Flat 2D zonal flow that increases linearly with depth from 0 m/s to 1 m/s.""" xdim = ydim = zdim = 2 @@ -122,6 +128,8 @@ def test_advection_3D(): assert np.allclose(pset.depth * pset.time, pset.lon, atol=1.0e-1) +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("direction", ["up", "down"]) @pytest.mark.parametrize("wErrorThroughSurface", [True, False]) def test_advection_3D_outofbounds(direction, wErrorThroughSurface): @@ -167,6 +175,8 @@ def SubmergeParticle(particle, fieldset, time): # pragma: no cover assert len(pset) == 0 +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("rk45_tol", [10, 100]) def test_advection_RK45(lon, lat, rk45_tol): npart = 10 @@ -324,6 +334,8 @@ def test_advection_periodic_zonal_meridional(): assert abs(pset.lat[0] - 0.15) < 0.1 +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("u", [-0.3, np.array(0.2)]) @pytest.mark.parametrize("v", [0.2, np.array(1)]) @pytest.mark.parametrize("w", [None, -0.2, np.array(0.7)]) @@ -401,6 +413,8 @@ def fieldset_stationary(): return create_fieldset_stationary() +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize( "method, rtol, diffField", [ @@ -435,6 +449,8 @@ def test_stationary_eddy(fieldset_stationary, method, rtol, diffField): assert np.allclose(pset.lat, exp_lat, rtol=rtol) +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_stationary_eddy_vertical(): npart = 1 lon = np.linspace(12000, 21000, npart) @@ -505,6 +521,8 @@ def fieldset_moving(): return create_fieldset_moving() +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize( "method, rtol, diffField", [ @@ -582,6 +600,8 @@ def fieldset_decaying(): return create_fieldset_decaying() +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize( "method, rtol, diffField", [ @@ -622,6 +642,8 @@ def test_decaying_eddy(fieldset_decaying, method, rtol, diffField): assert np.allclose(pset.lat, exp_lat, rtol=rtol) +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_analyticalAgrid(): lon = np.arange(0, 15, dtype=np.float32) lat = np.arange(0, 15, dtype=np.float32) diff --git a/tests/test_diffusion.py b/tests/test_diffusion.py index 389174c759..1cfb71e018 100644 --- a/tests/test_diffusion.py +++ b/tests/test_diffusion.py @@ -17,6 +17,8 @@ from tests.utils import create_fieldset_zeros_conversion +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("mesh", ["spherical", "flat"]) def test_fieldKh_Brownian(mesh): xdim = 200 @@ -50,6 +52,8 @@ def test_fieldKh_Brownian(mesh): assert np.allclose(np.mean(lats), 0, atol=tol) +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("mesh", ["spherical", "flat"]) @pytest.mark.parametrize("kernel", [AdvectionDiffusionM1, AdvectionDiffusionEM]) def test_fieldKh_SpatiallyVaryingDiffusion(mesh, kernel): @@ -83,6 +87,8 @@ def test_fieldKh_SpatiallyVaryingDiffusion(mesh, kernel): assert stats.skew(lons) > stats.skew(lats) +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("lambd", [1, 5]) def test_randomexponential(lambd): fieldset = create_fieldset_zeros_conversion() @@ -109,6 +115,8 @@ def vertical_randomexponential(particle, fieldset, time): # pragma: no cover assert np.allclose(np.mean(depth), expected_mean, rtol=0.1) +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("mu", [0.8 * np.pi, np.pi]) @pytest.mark.parametrize("kappa", [2, 4]) def test_randomvonmises(mu, kappa): diff --git a/tests/test_field.py b/tests/test_field.py index 7d7ec1ce36..2a29cdf690 100644 --- a/tests/test_field.py +++ b/tests/test_field.py @@ -11,6 +11,8 @@ from tests.utils import TEST_DATA +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_field_from_netcdf_variables(): filename = str(TEST_DATA / "perlinfieldsU.nc") dims = {"lon": "x", "lat": "y"} @@ -30,6 +32,8 @@ def test_field_from_netcdf_variables(): f3 = Field.from_netcdf(filename, variable, dims) +@pytest.mark.v4remove +@pytest.mark.xfail(reason="GH1946") def test_field_from_netcdf(): filenames = { "lon": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), @@ -41,6 +45,8 @@ def test_field_from_netcdf(): Field.from_netcdf(filenames, variable, dimensions, interp_method="cgrid_velocity") +@pytest.mark.v4remove +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize( "calendar, cftime_datetime", zip(_get_cftime_calendars(), _get_cftime_datetimes(), strict=True) ) diff --git a/tests/test_fieldset.py b/tests/test_fieldset.py index f81165bfca..ed70d46d89 100644 --- a/tests/test_fieldset.py +++ b/tests/test_fieldset.py @@ -13,6 +13,7 @@ Variable, ) from parcels.field import Field, VectorField +from parcels.tools.converters import GeographicPolar, UnitConverter from tests.utils import TEST_DATA @@ -54,6 +55,8 @@ def to_xarray_dataset(data: dict[str, np.array], dimensions: dict[str, np.array] ) +@pytest.mark.v4remove +@pytest.mark.xfail(reason="GH1946") @pytest.fixture def multifile_fieldset(tmp_path): stem = "test_subsets" @@ -76,16 +79,18 @@ def multifile_fieldset(tmp_path): return FieldSet.from_netcdf(files, variables, dimensions) +@pytest.mark.v4remove +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("xdim", [100, 200]) @pytest.mark.parametrize("ydim", [100, 200]) def test_fieldset_from_data(xdim, ydim): """Simple test for fieldset initialisation from data.""" data, dimensions = generate_fieldset_data(xdim, ydim) fieldset = FieldSet.from_data(data, dimensions) - assert len(fieldset.U.data.shape) == 3 - assert len(fieldset.V.data.shape) == 3 - assert np.allclose(fieldset.U.data[0, :], data["U"], rtol=1e-12) - assert np.allclose(fieldset.V.data[0, :], data["V"], rtol=1e-12) + assert len(fieldset.U.data.shape) == 2 + assert len(fieldset.V.data.shape) == 2 + assert np.allclose(fieldset.U.data, data["U"], rtol=1e-12) + assert np.allclose(fieldset.V.data, data["V"], rtol=1e-12) @pytest.mark.v4remove @@ -97,6 +102,8 @@ def test_fieldset_vmin_vmax(): assert np.isclose(np.amax(fieldset.U.data), 7) +@pytest.mark.v4remove +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("ttype", ["float", "datetime64"]) @pytest.mark.parametrize("tdim", [1, 20]) def test_fieldset_from_data_timedims(ttype, tdim): @@ -107,9 +114,12 @@ def test_fieldset_from_data_timedims(ttype, tdim): dimensions["time"] = [np.datetime64("2018-01-01") + np.timedelta64(t, "D") for t in range(tdim)] fieldset = FieldSet.from_data(data, dimensions) for i, dtime in enumerate(dimensions["time"]): - assert fieldset.U.grid.time_origin.fulltime(fieldset.U.grid.time[i]) == dtime + print(fieldset.U, dtime) + assert fieldset.U.time[i].data == dtime +@pytest.mark.v4remove +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("xdim", [100, 200]) @pytest.mark.parametrize("ydim", [100, 50]) def test_fieldset_from_data_different_dimensions(xdim, ydim): @@ -130,8 +140,8 @@ def test_fieldset_from_data_different_dimensions(xdim, ydim): } fieldset = FieldSet.from_data(data, dimensions) - assert len(fieldset.U.data.shape) == 3 - assert len(fieldset.V.data.shape) == 3 + assert len(fieldset.U.data.shape) == 4 + assert len(fieldset.V.data.shape) == 4 assert len(fieldset.P.data.shape) == 4 assert fieldset.P.data.shape == (tdim, zdim, ydim / 2, xdim / 2) assert np.allclose(fieldset.U.data, 0.0, rtol=1e-12) @@ -139,6 +149,8 @@ def test_fieldset_from_data_different_dimensions(xdim, ydim): assert np.allclose(fieldset.P.data, 2.0, rtol=1e-12) +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_fieldset_from_modulefile(): nemo_fname = str(TEST_DATA / "fieldset_nemo.py") nemo_error_fname = str(TEST_DATA / "fieldset_nemo_error.py") @@ -157,6 +169,35 @@ def test_fieldset_from_modulefile(): FieldSet.from_modulefile(nemo_error_fname, modulename="none_returning_function") +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") +def test_field_from_netcdf_fieldtypes(): + filenames = { + "varU": { + "lon": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), + "lat": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), + "data": str(TEST_DATA / "Uu_eastward_nemo_cross_180lon.nc"), + }, + "varV": { + "lon": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), + "lat": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), + "data": str(TEST_DATA / "Vv_eastward_nemo_cross_180lon.nc"), + }, + } + variables = {"varU": "U", "varV": "V"} + dimensions = {"lon": "glamf", "lat": "gphif"} + + # first try without setting fieldtype + fset = FieldSet.from_nemo(filenames, variables, dimensions) + assert isinstance(fset.varU.units, UnitConverter) + + # now try with setting fieldtype + fset = FieldSet.from_nemo(filenames, variables, dimensions, fieldtype={"varU": "U", "varV": "V"}) + assert isinstance(fset.varU.units, GeographicPolar) + + +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_fieldset_from_agrid_dataset(): filenames = { "lon": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), @@ -168,6 +209,8 @@ def test_fieldset_from_agrid_dataset(): FieldSet.from_a_grid_dataset(filenames, variable, dimensions) +@pytest.mark.v4remove +@pytest.mark.xfail(reason="GH1946") def test_fieldset_from_cgrid_interpmethod(): filenames = { "lon": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), @@ -182,7 +225,9 @@ def test_fieldset_from_cgrid_interpmethod(): FieldSet.from_c_grid_dataset(filenames, variable, dimensions, interp_method="partialslip") -@pytest.mark.parametrize("calltype", ["from_data", "from_nemo"]) +@pytest.mark.v4future +@pytest.mark.xfail(reason="GH1946") +@pytest.mark.parametrize("calltype", ["from_nemo"]) def test_illegal_dimensionsdict(calltype): with pytest.raises(NameError): if calltype == "from_data": @@ -197,20 +242,24 @@ def test_illegal_dimensionsdict(calltype): FieldSet.from_nemo(filenames, variables, dimensions) +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("xdim", [100, 200]) @pytest.mark.parametrize("ydim", [100, 200]) def test_add_field(xdim, ydim, tmpdir): data, dimensions = generate_fieldset_data(xdim, ydim) - fieldset = FieldSet.from_data(data, dimensions) + fieldset = FieldSet.from_data(data, dimensions) # TODO : Remove from_data field = Field("newfld", fieldset.U.data, lon=fieldset.U.lon, lat=fieldset.U.lat) fieldset.add_field(field) assert fieldset.newfld.data.shape == fieldset.U.data.shape +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("dupobject", ["same", "new"]) def test_add_duplicate_field(dupobject): data, dimensions = generate_fieldset_data(100, 100) - fieldset = FieldSet.from_data(data, dimensions) + fieldset = FieldSet.from_data(data, dimensions) # TODO : Remove from_data field = Field("newfld", fieldset.U.data, lon=fieldset.U.lon, lat=fieldset.U.lat) fieldset.add_field(field) with pytest.raises(RuntimeError): @@ -221,18 +270,20 @@ def test_add_duplicate_field(dupobject): fieldset.add_field(field2) -@pytest.mark.parametrize("field_type", ["normal", "vector"]) -def test_add_field_after_pset(field_type): +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") +@pytest.mark.parametrize("fieldtype", ["normal", "vector"]) +def test_add_field_after_pset(fieldtype): data, dimensions = generate_fieldset_data(100, 100) - fieldset = FieldSet.from_data(data, dimensions) + fieldset = FieldSet.from_data(data, dimensions) # TODO : Remove from_data pset = ParticleSet(fieldset, Particle, lon=0, lat=0) # noqa ; to trigger fieldset._check_complete field1 = Field("field1", fieldset.U.data, lon=fieldset.U.lon, lat=fieldset.U.lat) field2 = Field("field2", fieldset.U.data, lon=fieldset.U.lon, lat=fieldset.U.lat) vfield = VectorField("vfield", field1, field2) with pytest.raises(RuntimeError): - if field_type == "normal": + if fieldtype == "normal": fieldset.add_field(field1) - elif field_type == "vector": + elif fieldtype == "vector": fieldset.add_vector_field(vfield) @@ -242,9 +293,11 @@ def test_fieldset_samegrids_from_file(multifile_fieldset): assert multifile_fieldset.U.grid == multifile_fieldset.V.grid +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("gridtype", ["A", "C"]) def test_fieldset_dimlength1_cgrid(gridtype): - fieldset = FieldSet.from_data({"U": 0, "V": 0}, {"lon": 0, "lat": 0}) + fieldset = FieldSet.from_data({"U": 0, "V": 0}, {"lon": 0, "lat": 0}) # TODO : Remove from_data if gridtype == "C": fieldset.U.interp_method = "cgrid_velocity" fieldset.V.interp_method = "cgrid_velocity" @@ -263,6 +316,8 @@ def assign_dataset_timestamp_dim(ds, timestamp): return ds +@pytest.mark.v4remove +@pytest.mark.xfail(reason="GH1946") def test_fieldset_diffgrids_from_file(tmp_path): """Test for subsetting fieldset from file using indices dict.""" stem = "test_subsets" @@ -293,10 +348,12 @@ def test_fieldset_diffgrids_from_file(tmp_path): assert fieldset.U.grid != fieldset.V.grid +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_fieldset_diffgrids_from_file_data(multifile_fieldset): """Test for subsetting fieldset from file using indices dict.""" data, dimensions = generate_fieldset_data(100, 100) - field_U = FieldSet.from_data(data, dimensions).U + field_U = FieldSet.from_data(data, dimensions).U # TODO : Remove from_data field_U.name = "B" multifile_fieldset.add_field(field_U, "B") @@ -306,10 +363,12 @@ def test_fieldset_diffgrids_from_file_data(multifile_fieldset): assert multifile_fieldset.U.grid != multifile_fieldset.B.grid +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_fieldset_samegrids_from_data(): """Test for subsetting fieldset from file using indices dict.""" data, dimensions = generate_fieldset_data(100, 100) - fieldset1 = FieldSet.from_data(data, dimensions) + fieldset1 = FieldSet.from_data(data, dimensions) # TODO : Remove from_data field_data = fieldset1.U field_data.name = "B" fieldset1.add_field(field_data, "B") @@ -321,9 +380,11 @@ def addConst(particle, fieldset, time): # pragma: no cover particle.lon = particle.lon + fieldset.movewest + fieldset.moveeast +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_fieldset_constant(): data, dimensions = generate_fieldset_data(100, 100) - fieldset = FieldSet.from_data(data, dimensions) + fieldset = FieldSet.from_data(data, dimensions) # TODO : Remove from_data westval = -0.2 eastval = 0.3 fieldset.add_constant("movewest", westval) @@ -335,6 +396,8 @@ def test_fieldset_constant(): assert abs(pset.lon[0] - (0.5 + westval + eastval)) < 1e-4 +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("swapUV", [False, True]) def test_vector_fields(swapUV): lon = np.linspace(0.0, 10.0, 12, dtype=np.float32) @@ -343,7 +406,7 @@ def test_vector_fields(swapUV): V = np.zeros((10, 12), dtype=np.float32) data = {"U": U, "V": V} dimensions = {"U": {"lat": lat, "lon": lon}, "V": {"lat": lat, "lon": lon}} - fieldset = FieldSet.from_data(data, dimensions, mesh="flat") + fieldset = FieldSet.from_data(data, dimensions, mesh="flat") # TODO : Remove from_data if swapUV: # we test that we can freely edit whatever UV field UV = VectorField("UV", fieldset.V, fieldset.U) fieldset.add_vector_field(UV) @@ -358,6 +421,8 @@ def test_vector_fields(swapUV): assert abs(pset.lat[0] - 0.5) < 1e-9 +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_add_second_vector_field(): lon = np.linspace(0.0, 10.0, 12, dtype=np.float32) lat = np.linspace(0.0, 10.0, 10, dtype=np.float32) @@ -365,11 +430,11 @@ def test_add_second_vector_field(): V = np.zeros((10, 12), dtype=np.float32) data = {"U": U, "V": V} dimensions = {"U": {"lat": lat, "lon": lon}, "V": {"lat": lat, "lon": lon}} - fieldset = FieldSet.from_data(data, dimensions, mesh="flat") + fieldset = FieldSet.from_data(data, dimensions, mesh="flat") # TODO : Remove from_data data2 = {"U2": U, "V2": V} dimensions2 = {"lon": [ln + 0.1 for ln in lon], "lat": [lt - 0.1 for lt in lat]} - fieldset2 = FieldSet.from_data(data2, dimensions2, mesh="flat") + fieldset2 = FieldSet.from_data(data2, dimensions2, mesh="flat") # TODO : Remove from_data UV2 = VectorField("UV2", fieldset2.U2, fieldset2.V2) fieldset.add_vector_field(UV2) @@ -399,11 +464,11 @@ def test_timestamps(datetype, tmpdir): dims1["time"] = np.arange("2005-02-01", "2005-02-11", dtype="datetime64[D]") dims2["time"] = np.arange("2005-02-11", "2005-02-15", dtype="datetime64[D]") - fieldset1 = FieldSet.from_data(data1, dims1) + fieldset1 = FieldSet.from_data(data1, dims1) # TODO : Remove from_data fieldset1.U.data[0, :, :] = 2.0 fieldset1.write(tmpdir.join("file1")) - fieldset2 = FieldSet.from_data(data2, dims2) + fieldset2 = FieldSet.from_data(data2, dims2) # TODO : Remove from_data fieldset2.U.data[0, :, :] = 0.0 fieldset2.write(tmpdir.join("file2")) @@ -470,7 +535,7 @@ def temp_func(time): data = {"U": U, "V": V, "W": W, "temp": temp, "D": D} fieldset = FieldSet.from_data( data, dimensions, mesh="flat", time_periodic=time_periodic, allow_time_extrapolation=True - ) + ) # TODO : Remove from_data def sampleTemp(particle, fieldset, time): # pragma: no cover particle.temp = fieldset.temp[time, particle.depth, particle.lat, particle.lon] @@ -509,6 +574,8 @@ def sampleTemp(particle, fieldset, time): # pragma: no cover assert np.allclose(pset.d[0], 1.0) +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("tdim", [10, None]) def test_fieldset_from_xarray(tdim): def generate_dataset(xdim, ydim, zdim=1, tdim=1): @@ -563,6 +630,8 @@ def test_fieldset_frompop(): pset.execute(AdvectionRK4, runtime=3, dt=1) +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_fieldset_from_data_gridtypes(): """Simple test for fieldset initialisation from data.""" xdim, ydim, zdim = 20, 10, 4 @@ -581,7 +650,7 @@ def test_fieldset_from_data_gridtypes(): depth_s[k, :, :] = depth[k] # Rectilinear Z grid - fieldset = FieldSet.from_data(data, dimensions, mesh="flat") + fieldset = FieldSet.from_data(data, dimensions, mesh="flat") # TODO : Remove from_data pset = ParticleSet(fieldset, Particle, [0, 0], [0, 0], [0, 0.4]) pset.execute(AdvectionRK4, runtime=1.5, dt=0.5) plon = pset.lon @@ -592,7 +661,7 @@ def test_fieldset_from_data_gridtypes(): # Rectilinear S grid dimensions["depth"] = depth_s - fieldset = FieldSet.from_data(data, dimensions, mesh="flat") + fieldset = FieldSet.from_data(data, dimensions, mesh="flat") # TODO : Remove from_data pset = ParticleSet(fieldset, Particle, [0, 0], [0, 0], [0, 0.4]) pset.execute(AdvectionRK4, runtime=1.5, dt=0.5) assert np.allclose(plon, pset.lon) @@ -602,7 +671,7 @@ def test_fieldset_from_data_gridtypes(): dimensions["lon"] = lonm dimensions["lat"] = latm dimensions["depth"] = depth - fieldset = FieldSet.from_data(data, dimensions, mesh="flat") + fieldset = FieldSet.from_data(data, dimensions, mesh="flat") # TODO : Remove from_data pset = ParticleSet(fieldset, Particle, [0, 0], [0, 0], [0, 0.4]) pset.execute(AdvectionRK4, runtime=1.5, dt=0.5) assert np.allclose(plon, pset.lon) @@ -610,7 +679,7 @@ def test_fieldset_from_data_gridtypes(): # Curvilinear S grid dimensions["depth"] = depth_s - fieldset = FieldSet.from_data(data, dimensions, mesh="flat") + fieldset = FieldSet.from_data(data, dimensions, mesh="flat") # TODO : Remove from_data pset = ParticleSet(fieldset, Particle, [0, 0], [0, 0], [0, 0.4]) pset.execute(AdvectionRK4, runtime=1.5, dt=0.5) assert np.allclose(plon, pset.lon) diff --git a/tests/test_fieldset_sampling.py b/tests/test_fieldset_sampling.py index c8033a658d..ab4e5af4e7 100644 --- a/tests/test_fieldset_sampling.py +++ b/tests/test_fieldset_sampling.py @@ -85,6 +85,8 @@ def fieldset_geometric_polar(): return create_fieldset_geometric_polar() +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_fieldset_sample(fieldset): """Sample the fieldset using indexing notation.""" xdim, ydim = 120, 80 @@ -98,6 +100,8 @@ def test_fieldset_sample(fieldset): assert np.allclose(u_s, lat, rtol=1e-5) +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_fieldset_sample_eval(fieldset): """Sample the fieldset using the explicit eval function.""" xdim, ydim = 60, 60 @@ -120,6 +124,8 @@ def test_fieldset_polar_with_halo(fieldset_geometric_polar): assert pset.lon[0] == 0.0 +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("zdir", [-1, 1]) def test_verticalsampling(zdir): dims = (4, 2, 2) @@ -171,6 +177,7 @@ def test_pset_from_field(): assert np.allclose(pdens / sum(pdens.flatten()), startfield / sum(startfield.flatten()), atol=1e-2) +@pytest.mark.v4alpha def test_nearest_neighbor_interpolation2D(): npart = 81 dims = (2, 2) @@ -193,6 +200,8 @@ def test_nearest_neighbor_interpolation2D(): assert np.allclose(pset.p[(pset.lon > 0.5) | (pset.lat < 0.5)], 0.0, rtol=1e-5) +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_nearest_neighbor_interpolation3D(): npart = 81 dims = (2, 2, 2) @@ -219,6 +228,8 @@ def test_nearest_neighbor_interpolation3D(): assert np.allclose(pset.p[(pset.lon > 0.5) | (pset.lat < 0.5) & (pset.depth < 0.5)], 0.0, rtol=1e-5) +@pytest.mark.v4future +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("withDepth", [True, False]) @pytest.mark.parametrize("arrtype", ["ones", "rand"]) def test_inversedistance_nearland(withDepth, arrtype): @@ -260,6 +271,8 @@ def test_inversedistance_nearland(withDepth, arrtype): assert success +@pytest.mark.v4future +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("boundaryslip", ["freeslip", "partialslip"]) @pytest.mark.parametrize("withW", [False, True]) @pytest.mark.parametrize("withT", [False, True]) @@ -311,6 +324,8 @@ def test_partialslip_nearland_zonal(boundaryslip, withW, withT): assert np.allclose([p.depth for p in pset], 0.1) +@pytest.mark.v4future +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("boundaryslip", ["freeslip", "partialslip"]) @pytest.mark.parametrize("withW", [False, True]) def test_partialslip_nearland_meridional(boundaryslip, withW): @@ -353,6 +368,8 @@ def test_partialslip_nearland_meridional(boundaryslip, withW): assert np.allclose([p.depth for p in pset], 0.1) +@pytest.mark.v4future +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("boundaryslip", ["freeslip", "partialslip"]) def test_partialslip_nearland_vertical(boundaryslip): npart = 20 @@ -381,6 +398,8 @@ def test_partialslip_nearland_vertical(boundaryslip): assert np.allclose([p.lat for p in pset], 0.1) +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_fieldset_sample_particle(): """Sample the fieldset using an array of particles.""" npart = 120 @@ -403,6 +422,8 @@ def test_fieldset_sample_particle(): assert np.allclose(pset.u, lat, rtol=1e-6) +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_fieldset_sample_geographic(fieldset_geometric): """Sample a fieldset with conversion to geographic units (degrees).""" npart = 120 @@ -419,6 +440,8 @@ def test_fieldset_sample_geographic(fieldset_geometric): assert np.allclose(pset.u, lat, rtol=1e-6) +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_fieldset_sample_geographic_noconvert(fieldset_geometric): """Sample a fieldset without conversion to geographic units.""" npart = 120 @@ -435,6 +458,8 @@ def test_fieldset_sample_geographic_noconvert(fieldset_geometric): assert np.allclose(pset.u, lat * 1000 * 1.852 * 60, rtol=1e-6) +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_fieldset_sample_geographic_polar(fieldset_geometric_polar): """Sample a fieldset with conversion to geographic units and a pole correction.""" npart = 120 @@ -451,6 +476,8 @@ def test_fieldset_sample_geographic_polar(fieldset_geometric_polar): assert np.allclose(pset.u, lat, rtol=1e-2) +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_meridionalflow_spherical(): """Create uniform NORTHWARD flow on spherical earth and advect particles. @@ -480,6 +507,8 @@ def test_meridionalflow_spherical(): assert pset.lon[1] - lonstart[1] < 1e-4 +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_zonalflow_spherical(): """Create uniform EASTWARD flow on spherical earth and advect particles. @@ -516,6 +545,8 @@ def test_zonalflow_spherical(): assert abs(pset.p[1] - p_fld) < 1e-4 +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") def test_random_field(): """Sampling test that tests for overshoots by sampling a field of random numbers between 0 and 1.""" xdim, ydim = 20, 20 @@ -540,6 +571,8 @@ def test_random_field(): assert (sampled >= 0.0).all() +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("allow_time_extrapolation", [True, False]) def test_sampling_out_of_bounds_time(allow_time_extrapolation): xdim, ydim, tdim = 10, 10, 10 @@ -658,6 +691,8 @@ def test_sample(particle, fieldset, time): # pragma: no cover assert np.allclose(pset.sample_var, 10.0) +@pytest.mark.v4future +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("npart", [1, 10]) def test_sampling_multigrids_non_vectorfield(npart): xdim, ydim = 100, 200 @@ -697,6 +732,8 @@ def test_sample(particle, fieldset, time): # pragma: no cover assert np.allclose(pset.sample_var, 10.0) +@pytest.mark.v4future +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize("ugridfactor", [1, 10]) def test_sampling_multiple_grid_sizes(ugridfactor): xdim, ydim = 10, 20 @@ -724,6 +761,8 @@ def test_sampling_multiple_grid_sizes(ugridfactor): assert np.all((0 <= pset.xi) & (pset.xi < xdim * ugridfactor)) +@pytest.mark.v4future +@pytest.mark.xfail(reason="GH1946") def test_multiple_grid_addlater_error(): xdim, ydim = 10, 20 U = Field( diff --git a/tests/test_grids.py b/tests/test_grids.py deleted file mode 100644 index 59b6537dca..0000000000 --- a/tests/test_grids.py +++ /dev/null @@ -1,992 +0,0 @@ -import math -from datetime import timedelta - -import numpy as np -import pytest -import xarray as xr - -from parcels import ( - AdvectionRK4, - AdvectionRK4_3D, - CurvilinearZGrid, - Field, - FieldSet, - Particle, - ParticleSet, - RectilinearSGrid, - RectilinearZGrid, - StatusCode, - UnitConverter, - Variable, -) -from tests.utils import TEST_DATA - - -def test_multi_structured_grids(): - def temp_func(lon, lat): - return 20 + lat / 1000.0 + 2 * np.sin(lon * 2 * np.pi / 5000.0) - - a = 10000 - b = 10000 - - # Grid 0 - xdim_g0 = 201 - ydim_g0 = 201 - # Coordinates of the test fieldset (on A-grid in deg) - lon_g0 = np.linspace(0, a, xdim_g0, dtype=np.float32) - lat_g0 = np.linspace(0, b, ydim_g0, dtype=np.float32) - time_g0 = np.linspace(0.0, 1000.0, 2, dtype=np.float64) - grid_0 = RectilinearZGrid(lon_g0, lat_g0, time=time_g0) - - # Grid 1 - xdim_g1 = 51 - ydim_g1 = 51 - # Coordinates of the test fieldset (on A-grid in deg) - lon_g1 = np.linspace(0, a, xdim_g1, dtype=np.float32) - lat_g1 = np.linspace(0, b, ydim_g1, dtype=np.float32) - time_g1 = np.linspace(0.0, 1000.0, 2, dtype=np.float64) - grid_1 = RectilinearZGrid(lon_g1, lat_g1, time=time_g1) - - u_data = np.ones((time_g0.size, lat_g0.size, lon_g0.size), dtype=np.float32) - u_data = 2 * u_data - u_field = Field("U", u_data, grid=grid_0) - - temp0_data = np.empty((time_g0.size, lat_g0.size, lon_g0.size), dtype=np.float32) - for i in range(lon_g0.size): - for j in range(lat_g0.size): - temp0_data[:, j, i] = temp_func(lon_g0[i], lat_g0[j]) - temp0_field = Field("temp0", temp0_data, grid=grid_0) - - v_data = np.zeros((time_g1.size, lat_g1.size, lon_g1.size), dtype=np.float32) - v_field = Field("V", v_data, grid=grid_1) - - temp1_data = np.empty((time_g1.size, lat_g1.size, lon_g1.size), dtype=np.float32) - for i in range(lon_g1.size): - for j in range(lat_g1.size): - temp1_data[:, j, i] = temp_func(lon_g1[i], lat_g1[j]) - temp1_field = Field("temp1", temp1_data, grid=grid_1) - - other_fields = {} - other_fields["temp0"] = temp0_field - other_fields["temp1"] = temp1_field - - fieldset = FieldSet(u_field, v_field, fields=other_fields) - - def sampleTemp(particle, fieldset, time): # pragma: no cover - # Note that fieldset.temp is interpolated at time=time+dt. - # Indeed, sampleTemp is called at time=time, but the result is written - # at time=time+dt, after the Kernel update - particle.temp0 = fieldset.temp0[time + particle.dt, particle.depth, particle.lat, particle.lon] - particle.temp1 = fieldset.temp1[time + particle.dt, particle.depth, particle.lat, particle.lon] - - MyParticle = Particle.add_variables( - [Variable("temp0", dtype=np.float32, initial=20.0), Variable("temp1", dtype=np.float32, initial=20.0)] - ) - - pset = ParticleSet.from_list(fieldset, MyParticle, lon=[3001], lat=[5001], repeatdt=1) - - pset.execute(AdvectionRK4 + pset.Kernel(sampleTemp), runtime=3, dt=1) - - # check if particle xi and yi are different for the two grids - # xi check from unraveled index - assert np.all( - [fieldset.U.unravel_index(pset[i].ei)[2] != fieldset.V.unravel_index(pset[i].ei)[2] for i in range(3)] - ) - # yi check from unraveled index - assert np.all( - [fieldset.U.unravel_index(pset[i].ei)[1] != fieldset.V.unravel_index(pset[i].ei)[1] for i in range(3)] - ) - # advect without updating temperature to test particle deletion - pset.remove_indices(np.array([1])) - pset.execute(AdvectionRK4, runtime=1, dt=1) - - assert np.all([np.isclose(p.temp0, p.temp1, atol=1e-3) for p in pset]) - - -def test_time_format_in_grid(): - lon = np.linspace(0, 1, 2, dtype=np.float32) - lat = np.linspace(0, 1, 2, dtype=np.float32) - time = np.array([np.datetime64("2000-01-01")] * 2) - with pytest.raises(AssertionError, match="Time vector"): - RectilinearZGrid(lon, lat, time=time) - - -@pytest.mark.v4remove -@pytest.mark.xfail(reason="negate_depth removed in v4") -def test_negate_depth(): - depth = np.linspace(0, 5, 10, dtype=np.float32) - fieldset = FieldSet.from_data( - {"U": np.zeros((10, 1, 1)), "V": np.zeros((10, 1, 1))}, {"lon": [0], "lat": [0], "depth": depth} - ) - assert np.all(fieldset.gridset.grids[0].depth == depth) - fieldset.U.grid.negate_depth() - assert np.all(fieldset.gridset.grids[0].depth == -depth) - - -def test_avoid_repeated_grids(): - lon_g0 = np.linspace(0, 1000, 11, dtype=np.float32) - lat_g0 = np.linspace(0, 1000, 11, dtype=np.float32) - time_g0 = np.linspace(0, 1000, 2, dtype=np.float64) - grid_0 = RectilinearZGrid(lon_g0, lat_g0, time=time_g0) - - lon_g1 = np.linspace(0, 1000, 21, dtype=np.float32) - lat_g1 = np.linspace(0, 1000, 21, dtype=np.float32) - time_g1 = np.linspace(0, 1000, 2, dtype=np.float64) - grid_1 = RectilinearZGrid(lon_g1, lat_g1, time=time_g1) - - u_data = np.zeros((time_g0.size, lat_g0.size, lon_g0.size), dtype=np.float32) - u_field = Field("U", u_data, grid=grid_0) - - v_data = np.zeros((time_g1.size, lat_g1.size, lon_g1.size), dtype=np.float32) - v_field = Field("V", v_data, grid=grid_1) - - temp0_field = Field("temp", u_data, lon=lon_g0, lat=lat_g0, time=time_g0) - - other_fields = {} - other_fields["temp"] = temp0_field - - fieldset = FieldSet(u_field, v_field, fields=other_fields) - assert fieldset.gridset.size == 2 - assert fieldset.U.grid is fieldset.temp.grid - assert fieldset.V.grid is not fieldset.U.grid - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="Calls fieldset.add_periodic_halo(). Should adapt this test case.") -def test_multigrids_pointer(): - lon_g0 = np.linspace(0, 1e4, 21, dtype=np.float32) - lat_g0 = np.linspace(0, 1000, 2, dtype=np.float32) - depth_g0 = np.zeros((5, lat_g0.size, lon_g0.size), dtype=np.float32) - - def bath_func(lon): - return lon / 1000.0 + 10 - - bath = bath_func(lon_g0) - - zdim = depth_g0.shape[0] - for i in range(lon_g0.size): - for k in range(zdim): - depth_g0[k, :, i] = bath[i] * k / (zdim - 1) - - grid_0 = RectilinearSGrid(lon_g0, lat_g0, depth=depth_g0) - grid_1 = RectilinearSGrid(lon_g0, lat_g0, depth=depth_g0) - - u_data = np.zeros((zdim, lat_g0.size, lon_g0.size), dtype=np.float32) - v_data = np.zeros((zdim, lat_g0.size, lon_g0.size), dtype=np.float32) - w_data = np.zeros((zdim, lat_g0.size, lon_g0.size), dtype=np.float32) - - u_field = Field("U", u_data, grid=grid_0) - v_field = Field("V", v_data, grid=grid_0) - w_field = Field("W", w_data, grid=grid_1) - - fieldset = FieldSet(u_field, v_field, fields={"W": w_field}) - fieldset.add_periodic_halo(zonal=3, meridional=2) # unit test of halo for SGrid - - assert u_field.grid == v_field.grid - assert u_field.grid == w_field.grid # w_field.grid is now supposed to be grid_1 - - pset = ParticleSet.from_list(fieldset, Particle, lon=[0], lat=[0], depth=[1]) - - for _ in range(10): - pset.execute(AdvectionRK4_3D, runtime=1000, dt=500) - - -@pytest.mark.parametrize("z4d", ["True", "False"]) -def test_rectilinear_s_grid_sampling(z4d): - lon_g0 = np.linspace(-3e4, 3e4, 61, dtype=np.float32) - lat_g0 = np.linspace(0, 1000, 2, dtype=np.float32) - time_g0 = np.linspace(0, 1000, 2, dtype=np.float64) - if z4d: - depth_g0 = np.zeros((time_g0.size, 5, lat_g0.size, lon_g0.size), dtype=np.float32) - else: - depth_g0 = np.zeros((5, lat_g0.size, lon_g0.size), dtype=np.float32) - - def bath_func(lon): - bath = (lon <= -2e4) * 20.0 - bath += (lon > -2e4) * (lon < 2e4) * (110.0 + 90 * np.sin(lon / 2e4 * np.pi / 2.0)) - bath += (lon >= 2e4) * 200.0 - return bath - - bath = bath_func(lon_g0) - - zdim = depth_g0.shape[-3] - for i in range(depth_g0.shape[-1]): - for k in range(zdim): - if z4d: - depth_g0[:, k, :, i] = bath[i] * k / (zdim - 1) - else: - depth_g0[k, :, i] = bath[i] * k / (zdim - 1) - - grid = RectilinearSGrid(lon_g0, lat_g0, depth=depth_g0, time=time_g0) - - u_data = np.zeros((grid.tdim, grid.zdim, grid.ydim, grid.xdim), dtype=np.float32) - v_data = np.zeros((grid.tdim, grid.zdim, grid.ydim, grid.xdim), dtype=np.float32) - temp_data = np.zeros((grid.tdim, grid.zdim, grid.ydim, grid.xdim), dtype=np.float32) - for k in range(1, zdim): - temp_data[:, k, :, :] = k / (zdim - 1.0) - u_field = Field("U", u_data, grid=grid) - v_field = Field("V", v_data, grid=grid) - temp_field = Field("temp", temp_data, grid=grid) - - other_fields = {} - other_fields["temp"] = temp_field - fieldset = FieldSet(u_field, v_field, fields=other_fields) - - def sampleTemp(particle, fieldset, time): # pragma: no cover - particle.temp = fieldset.temp[time, particle.depth, particle.lat, particle.lon] - - MyParticle = Particle.add_variable("temp", dtype=np.float32, initial=20.0) - - lon = 400 - lat = 0 - ratio = 0.3 - pset = ParticleSet.from_list(fieldset, MyParticle, lon=[lon], lat=[lat], depth=[bath_func(lon) * ratio]) - - pset.execute(pset.Kernel(sampleTemp), runtime=1) - assert np.allclose(pset.temp[0], ratio, atol=1e-4) - - -def test_rectilinear_s_grids_advect1(): - # Constant water transport towards the east. check that the particle stays at the same relative depth (z/bath) - lon_g0 = np.linspace(0, 1e4, 21, dtype=np.float32) - lat_g0 = np.linspace(0, 1000, 2, dtype=np.float32) - depth_g0 = np.zeros((lon_g0.size, lat_g0.size, 5), dtype=np.float32) - - def bath_func(lon): - return lon / 1000.0 + 10 - - bath = bath_func(lon_g0) - - for i in range(depth_g0.shape[0]): - for k in range(depth_g0.shape[2]): - depth_g0[i, :, k] = bath[i] * k / (depth_g0.shape[2] - 1) - depth_g0 = depth_g0.transpose() - - grid = RectilinearSGrid(lon_g0, lat_g0, depth=depth_g0) - - zdim = depth_g0.shape[0] - u_data = np.zeros((zdim, lat_g0.size, lon_g0.size), dtype=np.float32) - v_data = np.zeros((zdim, lat_g0.size, lon_g0.size), dtype=np.float32) - w_data = np.zeros((zdim, lat_g0.size, lon_g0.size), dtype=np.float32) - for i in range(lon_g0.size): - u_data[:, :, i] = 1 * 10 / bath[i] - for k in range(zdim): - w_data[k, :, i] = u_data[k, :, i] * depth_g0[k, :, i] / bath[i] * 1e-3 - - u_field = Field("U", u_data, grid=grid) - v_field = Field("V", v_data, grid=grid) - w_field = Field("W", w_data, grid=grid) - - fieldset = FieldSet(u_field, v_field, fields={"W": w_field}) - - lon = np.zeros(11) - lat = np.zeros(11) - ratio = [min(i / 10.0, 0.99) for i in range(11)] - depth = bath_func(lon) * ratio - pset = ParticleSet.from_list(fieldset, Particle, lon=lon, lat=lat, depth=depth) - - pset.execute(AdvectionRK4_3D, runtime=10000, dt=500) - assert np.allclose(pset.depth / bath_func(pset.lon), ratio) - - -def test_rectilinear_s_grids_advect2(): - # Move particle towards the east, check relative depth evolution - lon_g0 = np.linspace(0, 1e4, 21, dtype=np.float32) - lat_g0 = np.linspace(0, 1000, 2, dtype=np.float32) - depth_g0 = np.zeros((5, lat_g0.size, lon_g0.size), dtype=np.float32) - - def bath_func(lon): - return lon / 1000.0 + 10 - - bath = bath_func(lon_g0) - - zdim = depth_g0.shape[0] - for i in range(lon_g0.size): - for k in range(zdim): - depth_g0[k, :, i] = bath[i] * k / (zdim - 1) - - grid = RectilinearSGrid(lon_g0, lat_g0, depth=depth_g0) - - u_data = np.zeros((zdim, lat_g0.size, lon_g0.size), dtype=np.float32) - v_data = np.zeros((zdim, lat_g0.size, lon_g0.size), dtype=np.float32) - rel_depth_data = np.zeros((zdim, lat_g0.size, lon_g0.size), dtype=np.float32) - for k in range(1, zdim): - rel_depth_data[k, :, :] = k / (zdim - 1.0) - - u_field = Field("U", u_data, grid=grid) - v_field = Field("V", v_data, grid=grid) - rel_depth_field = Field("relDepth", rel_depth_data, grid=grid) - fieldset = FieldSet(u_field, v_field, fields={"relDepth": rel_depth_field}) - - MyParticle = Particle.add_variable("relDepth", dtype=np.float32, initial=20.0) - - def moveEast(particle, fieldset, time): # pragma: no cover - particle_dlon += 5 * particle.dt # noqa - particle.relDepth = fieldset.relDepth[time, particle.depth, particle.lat, particle.lon] - - depth = 0.9 - pset = ParticleSet.from_list(fieldset, MyParticle, lon=[0], lat=[0], depth=[depth]) - - kernel = pset.Kernel(moveEast) - for _ in range(10): - pset.execute(kernel, runtime=100, dt=50) - assert np.allclose(pset.relDepth[0], depth / bath_func(pset.lon[0])) - - -def test_curvilinear_grids(): - x = np.linspace(0, 1e3, 7, dtype=np.float32) - y = np.linspace(0, 1e3, 5, dtype=np.float32) - (xx, yy) = np.meshgrid(x, y) - - r = np.sqrt(xx * xx + yy * yy) - theta = np.arctan2(yy, xx) - theta = theta + np.pi / 6.0 - - lon = r * np.cos(theta) - lat = r * np.sin(theta) - time = np.array([0, 86400], dtype=np.float64) - grid = CurvilinearZGrid(lon, lat, time=time) - - u_data = np.ones((2, y.size, x.size), dtype=np.float32) - v_data = np.zeros((2, y.size, x.size), dtype=np.float32) - u_data[0, :, :] = lon[:, :] + lat[:, :] - u_field = Field("U", u_data, grid=grid) - v_field = Field("V", v_data, grid=grid) - fieldset = FieldSet(u_field, v_field) - - def sampleSpeed(particle, fieldset, time): # pragma: no cover - u, v = fieldset.UV[time, particle.depth, particle.lat, particle.lon] - particle.speed = math.sqrt(u * u + v * v) - - MyParticle = Particle.add_variable("speed", dtype=np.float32, initial=0.0) - - pset = ParticleSet.from_list(fieldset, MyParticle, lon=[400, -200], lat=[600, 600]) - pset.execute(pset.Kernel(sampleSpeed), runtime=1) - assert np.allclose(pset.speed[0], 1000) - - -def test_nemo_grid(): - filenames = { - "U": { - "lon": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), - "lat": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), - "data": str(TEST_DATA / "Uu_eastward_nemo_cross_180lon.nc"), - }, - "V": { - "lon": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), - "lat": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), - "data": str(TEST_DATA / "Vv_eastward_nemo_cross_180lon.nc"), - }, - } - variables = {"U": "U", "V": "V"} - dimensions = {"lon": "glamf", "lat": "gphif"} - fieldset = FieldSet.from_nemo(filenames, variables, dimensions) - - # test ParticleSet.from_field on curvilinear grids - ParticleSet.from_field(fieldset, Particle, start_field=fieldset.U, size=5) - - def sampleVel(particle, fieldset, time): # pragma: no cover - (particle.zonal, particle.meridional) = fieldset.UV[time, particle.depth, particle.lat, particle.lon] - - MyParticle = Particle.add_variables( - [Variable("zonal", dtype=np.float32, initial=0.0), Variable("meridional", dtype=np.float32, initial=0.0)] - ) - - lonp = 175.5 - latp = 81.5 - pset = ParticleSet.from_list(fieldset, MyParticle, lon=[lonp], lat=[latp]) - pset.execute(pset.Kernel(sampleVel), runtime=1) - u = fieldset.U.units.to_source(pset.zonal[0], 0, latp, lonp) - v = fieldset.V.units.to_source(pset.meridional[0], 0, latp, lonp) - assert abs(u - 1) < 1e-4 - assert abs(v) < 1e-4 - - -def test_advect_nemo(): - filenames = { - "U": { - "lon": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), - "lat": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), - "data": str(TEST_DATA / "Uu_eastward_nemo_cross_180lon.nc"), - }, - "V": { - "lon": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), - "lat": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), - "data": str(TEST_DATA / "Vv_eastward_nemo_cross_180lon.nc"), - }, - } - variables = {"U": "U", "V": "V"} - dimensions = {"lon": "glamf", "lat": "gphif"} - fieldset = FieldSet.from_nemo(filenames, variables, dimensions) - - lonp = 175.5 - latp = 81.5 - pset = ParticleSet.from_list(fieldset, Particle, lon=[lonp], lat=[latp]) - pset.execute(AdvectionRK4, runtime=timedelta(days=2), dt=timedelta(hours=6)) - assert abs(pset.lat[0] - latp) < 1e-3 - - -@pytest.mark.parametrize("time", [True, False]) -def test_cgrid_uniform_2dvel(time): - lon = np.array([[0, 2], [0.4, 1.5]]) - lat = np.array([[0, -0.5], [0.8, 0.5]]) - U = np.array([[-99, -99], [4.4721359549995793e-01, 1.3416407864998738e00]]) - V = np.array([[-99, 1.2126781251816650e00], [-99, 1.2278812270298409e00]]) - - if time: - U = np.stack((U, U)) - V = np.stack((V, V)) - dimensions = {"lat": lat, "lon": lon, "time": np.array([0, 10])} - else: - dimensions = {"lat": lat, "lon": lon} - data = {"U": np.array(U, dtype=np.float32), "V": np.array(V, dtype=np.float32)} - fieldset = FieldSet.from_data(data, dimensions, mesh="flat") - fieldset.U.interp_method = "cgrid_velocity" - fieldset.V.interp_method = "cgrid_velocity" - - def sampleVel(particle, fieldset, time): # pragma: no cover - (particle.zonal, particle.meridional) = fieldset.UV[time, particle.depth, particle.lat, particle.lon] - - MyParticle = Particle.add_variables( - [Variable("zonal", dtype=np.float32, initial=0.0), Variable("meridional", dtype=np.float32, initial=0.0)] - ) - - pset = ParticleSet.from_list(fieldset, MyParticle, lon=0.7, lat=0.3) - pset.execute(pset.Kernel(sampleVel), runtime=1) - assert (pset[0].zonal - 1) < 1e-6 - assert (pset[0].meridional - 1) < 1e-6 - - -@pytest.mark.v4alpha -@pytest.mark.parametrize("vert_mode", ["zlev"]) # , "slev1", "slev2"]) # TODO v4: time-varying depth not supported yet -@pytest.mark.parametrize("time", [True, False]) -def test_cgrid_uniform_3dvel(vert_mode, time): - lon = np.array([[0, 2], [0.4, 1.5]]) - lat = np.array([[0, -0.5], [0.8, 0.5]]) - - u0 = 4.4721359549995793e-01 - u1 = 1.3416407864998738e00 - v0 = 1.2126781251816650e00 - v1 = 1.2278812270298409e00 - w0 = 1 - w1 = 1 - - if vert_mode == "zlev": - depth = np.array([0, 1]) - elif vert_mode == "slev1": - depth = np.array([[[0, 0], [0, 0]], [[1, 1], [1, 1]]]) - elif vert_mode == "slev2": - depth = np.array([[[-1, -0.6], [-1.1257142857142859, -0.9]], [[1, 1.5], [0.50857142857142845, 0.8]]]) - w0 = 1.0483007922296661e00 - w1 = 1.3098951476312375e00 - - U = np.array([[[-99, -99], [u0, u1]], [[-99, -99], [-99, -99]]]) - V = np.array([[[-99, v0], [-99, v1]], [[-99, -99], [-99, -99]]]) - W = np.array([[[-99, -99], [-99, w0]], [[-99, -99], [-99, w1]]]) - - if time: - U = np.stack((U, U)) - V = np.stack((V, V)) - W = np.stack((W, W)) - dimensions = {"lat": lat, "lon": lon, "depth": depth, "time": np.array([0, 10])} - else: - dimensions = {"lat": lat, "lon": lon, "depth": depth} - data = {"U": np.array(U, dtype=np.float32), "V": np.array(V, dtype=np.float32), "W": np.array(W, dtype=np.float32)} - fieldset = FieldSet.from_data(data, dimensions, mesh="flat") - fieldset.U.interp_method = "cgrid_velocity" - fieldset.V.interp_method = "cgrid_velocity" - fieldset.W.interp_method = "cgrid_velocity" - - def sampleVel(particle, fieldset, time): # pragma: no cover - (particle.zonal, particle.meridional, particle.vertical) = fieldset.UVW[ - time, particle.depth, particle.lat, particle.lon - ] - - MyParticle = Particle.add_variables( - [ - Variable("zonal", dtype=np.float32, initial=0.0), - Variable("meridional", dtype=np.float32, initial=0.0), - Variable("vertical", dtype=np.float32, initial=0.0), - ] - ) - - pset = ParticleSet.from_list(fieldset, MyParticle, lon=0.7, lat=0.3, depth=0.2) - pset.execute(pset.Kernel(sampleVel), runtime=1) - assert abs(pset[0].zonal - 1) < 1e-6 - assert abs(pset[0].meridional - 1) < 1e-6 - assert abs(pset[0].vertical - 1) < 1e-6 - - -@pytest.mark.parametrize("vert_mode", ["zlev", "slev1"]) -@pytest.mark.parametrize("time", [True, False]) -def test_cgrid_uniform_3dvel_spherical(vert_mode, time): - dim_file = xr.open_dataset(TEST_DATA / "mask_nemo_cross_180lon.nc") - u_file = xr.open_dataset(TEST_DATA / "Uu_eastward_nemo_cross_180lon.nc") - v_file = xr.open_dataset(TEST_DATA / "Vv_eastward_nemo_cross_180lon.nc") - j = 4 - i = 11 - lon = np.array(dim_file.glamf[0, j : j + 2, i : i + 2]) - lat = np.array(dim_file.gphif[0, j : j + 2, i : i + 2]) - U = np.array(u_file.U[0, j : j + 2, i : i + 2]) - V = np.array(v_file.V[0, j : j + 2, i : i + 2]) - trash = np.zeros((2, 2)) - U = np.stack((U, trash)) - V = np.stack((V, trash)) - w0 = 1 - w1 = 1 - W = np.array([[[-99, -99], [-99, w0]], [[-99, -99], [-99, w1]]]) - - if vert_mode == "zlev": - depth = np.array([0, 1]) - elif vert_mode == "slev1": - depth = np.array([[[0, 0], [0, 0]], [[1, 1], [1, 1]]]) - - if time: - U = np.stack((U, U)) - V = np.stack((V, V)) - W = np.stack((W, W)) - dimensions = {"lat": lat, "lon": lon, "depth": depth, "time": np.array([0, 10])} - else: - dimensions = {"lat": lat, "lon": lon, "depth": depth} - data = {"U": np.array(U, dtype=np.float32), "V": np.array(V, dtype=np.float32), "W": np.array(W, dtype=np.float32)} - fieldset = FieldSet.from_data(data, dimensions, mesh="spherical") - fieldset.U.interp_method = "cgrid_velocity" - fieldset.V.interp_method = "cgrid_velocity" - fieldset.W.interp_method = "cgrid_velocity" - - def sampleVel(particle, fieldset, time): # pragma: no cover - (particle.zonal, particle.meridional, particle.vertical) = fieldset.UVW[ - time, particle.depth, particle.lat, particle.lon - ] - - MyParticle = Particle.add_variables( - [ - Variable("zonal", dtype=np.float32, initial=0.0), - Variable("meridional", dtype=np.float32, initial=0.0), - Variable("vertical", dtype=np.float32, initial=0.0), - ] - ) - - lonp = 179.8 - latp = 81.35 - pset = ParticleSet.from_list(fieldset, MyParticle, lon=lonp, lat=latp, depth=0.2) - pset.execute(pset.Kernel(sampleVel), runtime=1) - pset.zonal[0] = fieldset.U.units.to_source(pset.zonal[0], 0, latp, lonp) - pset.meridional[0] = fieldset.V.units.to_source(pset.meridional[0], 0, latp, lonp) - assert abs(pset[0].zonal - 1) < 1e-3 - assert abs(pset[0].meridional) < 1e-3 - assert abs(pset[0].vertical - 1) < 1e-3 - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="From_pop is not supported during v4-alpha development. This will be reconsidered in v4.") -@pytest.mark.parametrize("vert_discretisation", ["zlevel", "slevel", "slevel2"]) -def test_popgrid(vert_discretisation): - if vert_discretisation == "zlevel": - w_dep = "w_dep" - elif vert_discretisation == "slevel": - w_dep = "w_deps" # same as zlevel, but defined as slevel - elif vert_discretisation == "slevel2": - w_dep = "w_deps2" # contains shaved cells - - filenames = str(TEST_DATA / "POPtestdata_time.nc") - variables = {"U": "U", "V": "V", "W": "W", "T": "T"} - dimensions = {"lon": "lon", "lat": "lat", "depth": w_dep, "time": "time"} - - fieldset = FieldSet.from_pop(filenames, variables, dimensions, mesh="flat") - - def sampleVel(particle, fieldset, time): # pragma: no cover - (particle.zonal, particle.meridional, particle.vert) = fieldset.UVW[particle] - particle.tracer = fieldset.T[particle] - - def OutBoundsError(particle, fieldset, time): # pragma: no cover - if particle.state == StatusCode.ErrorOutOfBounds: - particle.out_of_bounds = 1 - particle_ddepth -= 3 # noqa - particle.state = StatusCode.Success - - MyParticle = Particle.add_variables( - [ - Variable("zonal", dtype=np.float32, initial=0.0), - Variable("meridional", dtype=np.float32, initial=0.0), - Variable("vert", dtype=np.float32, initial=0.0), - Variable("tracer", dtype=np.float32, initial=0.0), - Variable("out_of_bounds", dtype=np.float32, initial=0.0), - ] - ) - - pset = ParticleSet.from_list(fieldset, MyParticle, lon=[3, 5, 1], lat=[3, 5, 1], depth=[3, 7, 11]) - pset.execute(pset.Kernel(sampleVel) + OutBoundsError, runtime=1) - if vert_discretisation == "slevel2": - assert np.isclose(pset.vert[0], 0.0) - assert np.isclose(pset.zonal[0], 0.0) - assert np.isclose(pset.tracer[0], 99.0) - assert np.isclose(pset.vert[1], -0.0066666666) - assert np.isclose(pset.zonal[1], 0.015) - assert np.isclose(pset.tracer[1], 1.0) - assert pset.out_of_bounds[0] == 0 - assert pset.out_of_bounds[1] == 0 - assert pset.out_of_bounds[2] == 1 - else: - assert np.allclose(pset.zonal, 0.015) - assert np.allclose(pset.meridional, 0.01) - assert np.allclose(pset.vert, -0.01) - assert np.allclose(pset.tracer, 1) - - -@pytest.mark.parametrize("gridindexingtype", ["mitgcm", "nemo"]) -@pytest.mark.parametrize("coordtype", ["rectilinear", "curvilinear"]) -def test_cgrid_indexing(gridindexingtype, coordtype): - xdim, ydim = 151, 201 - a = b = 20000 # domain size - lon = np.linspace(-a / 2, a / 2, xdim, dtype=np.float32) - lat = np.linspace(-b / 2, b / 2, ydim, dtype=np.float32) - dx, dy = lon[2] - lon[1], lat[2] - lat[1] - omega = 2 * np.pi / timedelta(days=1).total_seconds() - - index_signs = {"nemo": -1, "mitgcm": 1} - isign = index_signs[gridindexingtype] - - def rotate_coords(lon, lat, alpha=0): - rotmat = np.array([[np.cos(alpha), np.sin(alpha)], [-np.sin(alpha), np.cos(alpha)]]) - lons, lats = np.meshgrid(lon, lat) - rotated = np.einsum("ji, mni -> jmn", rotmat, np.dstack([lons, lats])) - return rotated[0], rotated[1] - - if coordtype == "rectilinear": - alpha = 0 - elif coordtype == "curvilinear": - alpha = 15 * np.pi / 180 - lon, lat = rotate_coords(lon, lat, alpha) - - def calc_r_phi(ln, lt): - return np.sqrt(ln**2 + lt**2), np.arctan2(ln, lt) - - if coordtype == "rectilinear": - - def calculate_UVR(lat, lon, dx, dy, omega, alpha): - U = np.zeros((lat.size, lon.size), dtype=np.float32) - V = np.zeros((lat.size, lon.size), dtype=np.float32) - R = np.zeros((lat.size, lon.size), dtype=np.float32) - for i in range(lon.size): - for j in range(lat.size): - r, phi = calc_r_phi(lon[i], lat[j]) - R[j, i] = r - r, phi = calc_r_phi(lon[i] + isign * dx / 2, lat[j]) - V[j, i] = -omega * r * np.sin(phi) - r, phi = calc_r_phi(lon[i], lat[j] + isign * dy / 2) - U[j, i] = omega * r * np.cos(phi) - return U, V, R - elif coordtype == "curvilinear": - - def calculate_UVR(lat, lon, dx, dy, omega, alpha): - U = np.zeros(lat.shape, dtype=np.float32) - V = np.zeros(lat.shape, dtype=np.float32) - R = np.zeros(lat.shape, dtype=np.float32) - for i in range(lat.shape[1]): - for j in range(lat.shape[0]): - r, phi = calc_r_phi(lon[j, i], lat[j, i]) - R[j, i] = r - r, phi = calc_r_phi( - lon[j, i] + isign * (dx / 2) * np.cos(alpha), lat[j, i] - isign * (dx / 2) * np.sin(alpha) - ) - V[j, i] = np.sin(alpha) * (omega * r * np.cos(phi)) + np.cos(alpha) * (-omega * r * np.sin(phi)) - r, phi = calc_r_phi( - lon[j, i] + isign * (dy / 2) * np.sin(alpha), lat[j, i] + isign * (dy / 2) * np.cos(alpha) - ) - U[j, i] = np.cos(alpha) * (omega * r * np.cos(phi)) - np.sin(alpha) * (-omega * r * np.sin(phi)) - return U, V, R - - U, V, R = calculate_UVR(lat, lon, dx, dy, omega, alpha) - - data = {"U": U, "V": V, "R": R} - dimensions = {"lon": lon, "lat": lat} - fieldset = FieldSet.from_data(data, dimensions, mesh="flat", gridindexingtype=gridindexingtype) - fieldset.U.interp_method = "cgrid_velocity" - fieldset.V.interp_method = "cgrid_velocity" - - def UpdateR(particle, fieldset, time): # pragma: no cover - if time == 0: - particle.radius_start = fieldset.R[time, particle.depth, particle.lat, particle.lon] - particle.radius = fieldset.R[time, particle.depth, particle.lat, particle.lon] - - MyParticle = Particle.add_variables( - [Variable("radius", dtype=np.float32, initial=0.0), Variable("radius_start", dtype=np.float32, initial=0.0)] - ) - - pset = ParticleSet(fieldset, pclass=MyParticle, lon=0, lat=4e3, time=0) - - pset.execute(pset.Kernel(UpdateR) + AdvectionRK4, runtime=timedelta(hours=14), dt=timedelta(minutes=5)) - assert np.allclose(pset.radius, pset.radius_start, atol=10) - - -@pytest.mark.parametrize("gridindexingtype", ["mitgcm", "nemo"]) -@pytest.mark.parametrize("withtime", [False, True]) -def test_cgrid_indexing_3D(gridindexingtype, withtime): - xdim = zdim = 201 - ydim = 2 - a = c = 20000 # domain size - b = 2 - lon = np.linspace(-a / 2, a / 2, xdim, dtype=np.float32) - lat = np.linspace(-b / 2, b / 2, ydim, dtype=np.float32) - depth = np.linspace(-c / 2, c / 2, zdim, dtype=np.float32) - dx, dz = lon[1] - lon[0], depth[1] - depth[0] - omega = 2 * np.pi / timedelta(days=1).total_seconds() - if withtime: - time = np.linspace(0, 24 * 60 * 60, 10) - dimensions = {"lon": lon, "lat": lat, "depth": depth, "time": time} - dsize = (time.size, depth.size, lat.size, lon.size) - else: - dimensions = {"lon": lon, "lat": lat, "depth": depth} - dsize = (depth.size, lat.size, lon.size) - - hindex_signs = {"nemo": -1, "mitgcm": 1} - hsign = hindex_signs[gridindexingtype] - - def calc_r_phi(ln, dp): - # r = np.sqrt(ln ** 2 + dp ** 2) - # phi = np.arcsin(dp/r) if r > 0 else 0 - return np.sqrt(ln**2 + dp**2), np.arctan2(ln, dp) - - def populate_UVWR(lat, lon, depth, dx, dz, omega): - U = np.zeros(dsize, dtype=np.float32) - V = np.zeros(dsize, dtype=np.float32) - W = np.zeros(dsize, dtype=np.float32) - R = np.zeros(dsize, dtype=np.float32) - - for i in range(lon.size): - for j in range(lat.size): - for k in range(depth.size): - r, phi = calc_r_phi(lon[i], depth[k]) - if withtime: - R[:, k, j, i] = r - else: - R[k, j, i] = r - r, phi = calc_r_phi(lon[i] + hsign * dx / 2, depth[k]) - if withtime: - W[:, k, j, i] = -omega * r * np.sin(phi) - else: - W[k, j, i] = -omega * r * np.sin(phi) - r, phi = calc_r_phi(lon[i], depth[k] + dz / 2) - if withtime: - U[:, k, j, i] = omega * r * np.cos(phi) - else: - U[k, j, i] = omega * r * np.cos(phi) - return U, V, W, R - - U, V, W, R = populate_UVWR(lat, lon, depth, dx, dz, omega) - data = {"U": U, "V": V, "W": W, "R": R} - fieldset = FieldSet.from_data(data, dimensions, mesh="flat", gridindexingtype=gridindexingtype) - fieldset.U.interp_method = "cgrid_velocity" - fieldset.V.interp_method = "cgrid_velocity" - fieldset.W.interp_method = "cgrid_velocity" - - def UpdateR(particle, fieldset, time): # pragma: no cover - if time == 0: - particle.radius_start = fieldset.R[time, particle.depth, particle.lat, particle.lon] - particle.radius = fieldset.R[time, particle.depth, particle.lat, particle.lon] - - MyParticle = Particle.add_variables( - [Variable("radius", dtype=np.float32, initial=0.0), Variable("radius_start", dtype=np.float32, initial=0.0)] - ) - - pset = ParticleSet(fieldset, pclass=MyParticle, depth=4e3, lon=0, lat=0, time=0) - - pset.execute(pset.Kernel(UpdateR) + AdvectionRK4_3D, runtime=timedelta(hours=14), dt=timedelta(minutes=5)) - assert np.allclose(pset.radius, pset.radius_start, atol=10) - - -@pytest.mark.parametrize("gridindexingtype", ["pop", "mom5"]) -@pytest.mark.parametrize("withtime", [False, True]) -def test_bgrid_indexing_3D(gridindexingtype, withtime): - xdim = zdim = 201 - ydim = 2 - a = c = 20000 # domain size - b = 2 - lon = np.linspace(-a / 2, a / 2, xdim, dtype=np.float32) - lat = np.linspace(-b / 2, b / 2, ydim, dtype=np.float32) - depth = np.linspace(-c / 2, c / 2, zdim, dtype=np.float32) - dx, dz = lon[1] - lon[0], depth[1] - depth[0] - omega = 2 * np.pi / timedelta(days=1).total_seconds() - if withtime: - time = np.linspace(0, 24 * 60 * 60, 10) - dimensions = {"lon": lon, "lat": lat, "depth": depth, "time": time} - dsize = (time.size, depth.size, lat.size, lon.size) - else: - dimensions = {"lon": lon, "lat": lat, "depth": depth} - dsize = (depth.size, lat.size, lon.size) - - vindex_signs = {"pop": 1, "mom5": -1} - vsign = vindex_signs[gridindexingtype] - - def calc_r_phi(ln, dp): - return np.sqrt(ln**2 + dp**2), np.arctan2(ln, dp) - - def populate_UVWR(lat, lon, depth, dx, dz, omega): - U = np.zeros(dsize, dtype=np.float32) - V = np.zeros(dsize, dtype=np.float32) - W = np.zeros(dsize, dtype=np.float32) - R = np.zeros(dsize, dtype=np.float32) - - for i in range(lon.size): - for j in range(lat.size): - for k in range(depth.size): - r, phi = calc_r_phi(lon[i], depth[k]) - if withtime: - R[:, k, j, i] = r - else: - R[k, j, i] = r - r, phi = calc_r_phi(lon[i] - dx / 2, depth[k]) - if withtime: - W[:, k, j, i] = -omega * r * np.sin(phi) - else: - W[k, j, i] = -omega * r * np.sin(phi) - # Since Parcels loads as dimensions only the depth of W-points - # and lon/lat of UV-points, W-points are similarly interpolated - # in MOM5 and POP. Indexing is shifted for UV-points. - r, phi = calc_r_phi(lon[i], depth[k] + vsign * dz / 2) - if withtime: - U[:, k, j, i] = omega * r * np.cos(phi) - else: - U[k, j, i] = omega * r * np.cos(phi) - return U, V, W, R - - U, V, W, R = populate_UVWR(lat, lon, depth, dx, dz, omega) - data = {"U": U, "V": V, "W": W, "R": R} - fieldset = FieldSet.from_data(data, dimensions, mesh="flat", gridindexingtype=gridindexingtype) - fieldset.U.interp_method = "bgrid_velocity" - fieldset.V.interp_method = "bgrid_velocity" - fieldset.W.interp_method = "bgrid_w_velocity" - - def UpdateR(particle, fieldset, time): # pragma: no cover - if time == 0: - particle.radius_start = fieldset.R[time, particle.depth, particle.lat, particle.lon] - particle.radius = fieldset.R[time, particle.depth, particle.lat, particle.lon] - - MyParticle = Particle.add_variables( - [Variable("radius", dtype=np.float32, initial=0.0), Variable("radius_start", dtype=np.float32, initial=0.0)] - ) - - pset = ParticleSet(fieldset, pclass=MyParticle, depth=-9.995e3, lon=0, lat=0, time=0) - - pset.execute(pset.Kernel(UpdateR) + AdvectionRK4_3D, runtime=timedelta(hours=14), dt=timedelta(minutes=5)) - assert np.allclose(pset.radius, pset.radius_start, atol=10) - - -@pytest.mark.parametrize( - "gridindexingtype", - [ - pytest.param( - "mom5", - marks=[ - pytest.mark.v4alpha, - pytest.mark.xfail(reason="https://github.com/OceanParcels/Parcels/pull/1936#issuecomment-2717408483"), - ], - ) - ], -) # TODO v4: add pop in params? -@pytest.mark.parametrize("extrapolation", [True, False]) -def test_bgrid_interpolation(gridindexingtype, extrapolation): - xi, yi = 3, 2 - if extrapolation: - zi = 0 if gridindexingtype == "mom5" else -1 - else: - zi = 2 - if gridindexingtype == "mom5": - ufile = str(TEST_DATA / "access-om2-01_u.nc") - vfile = str(TEST_DATA / "access-om2-01_v.nc") - wfile = str(TEST_DATA / "access-om2-01_wt.nc") - - filenames = { - "U": {"lon": ufile, "lat": ufile, "depth": wfile, "data": ufile}, - "V": {"lon": ufile, "lat": ufile, "depth": wfile, "data": vfile}, - "W": {"lon": ufile, "lat": ufile, "depth": wfile, "data": wfile}, - } - - variables = {"U": "u", "V": "v", "W": "wt"} - - dimensions = { - "U": {"lon": "xu_ocean", "lat": "yu_ocean", "depth": "sw_ocean", "time": "time"}, - "V": {"lon": "xu_ocean", "lat": "yu_ocean", "depth": "sw_ocean", "time": "time"}, - "W": {"lon": "xu_ocean", "lat": "yu_ocean", "depth": "sw_ocean", "time": "time"}, - } - - fieldset = FieldSet.from_mom5(filenames, variables, dimensions) - ds_u = xr.open_dataset(ufile) - ds_v = xr.open_dataset(vfile) - ds_w = xr.open_dataset(wfile) - u = ds_u.u.isel(time=0, st_ocean=zi, yu_ocean=yi, xu_ocean=xi) - v = ds_v.v.isel(time=0, st_ocean=zi, yu_ocean=yi, xu_ocean=xi) - w = ds_w.wt.isel(time=0, sw_ocean=zi, yt_ocean=yi, xt_ocean=xi) - - elif gridindexingtype == "pop": - datafname = str(TEST_DATA / "popdata.nc") - coordfname = str(TEST_DATA / "popcoordinates.nc") - filenames = { - "U": {"lon": coordfname, "lat": coordfname, "depth": coordfname, "data": datafname}, - "V": {"lon": coordfname, "lat": coordfname, "depth": coordfname, "data": datafname}, - "W": {"lon": coordfname, "lat": coordfname, "depth": coordfname, "data": datafname}, - } - - variables = {"U": "UVEL", "V": "VVEL", "W": "WVEL"} - dimensions = {"lon": "U_LON_2D", "lat": "U_LAT_2D", "depth": "w_dep"} - - fieldset = FieldSet.from_pop(filenames, variables, dimensions) - dsc = xr.open_dataset(coordfname) - dsd = xr.open_dataset(datafname) - u = dsd.UVEL.isel(k=zi, j=yi, i=xi) - v = dsd.VVEL.isel(k=zi, j=yi, i=xi) - w = dsd.WVEL.isel(k=zi, j=yi, i=xi) - - fieldset.U.units = UnitConverter() - fieldset.V.units = UnitConverter() - - def VelocityInterpolator(particle, fieldset, time): # pragma: no cover - particle.Uvel = fieldset.U[time, particle.depth, particle.lat, particle.lon] - particle.Vvel = fieldset.V[time, particle.depth, particle.lat, particle.lon] - particle.Wvel = fieldset.W[time, particle.depth, particle.lat, particle.lon] - - myParticle = Particle.add_variables( - [ - Variable("Uvel", dtype=np.float32, initial=0.0), - Variable("Vvel", dtype=np.float32, initial=0.0), - Variable("Wvel", dtype=np.float32, initial=0.0), - ] - ) - - for pointtype in ["U", "V", "W"]: - if gridindexingtype == "pop": - if pointtype in ["U", "V"]: - lons = dsc.U_LON_2D[yi, xi].values - lats = dsc.U_LAT_2D[yi, xi].values - deps = dsc.depth_t[zi].values - elif pointtype == "W": - lons = dsc.T_LON_2D[yi, xi].values - lats = dsc.T_LAT_2D[yi, xi].values - deps = dsc.w_dep[zi].values - if extrapolation: - deps = 5499.0 - elif gridindexingtype == "mom5": - if pointtype in ["U", "V"]: - lons = u.xu_ocean.data.reshape(1) - lats = u.yu_ocean.data.reshape(1) - deps = u.st_ocean.data.reshape(1) - elif pointtype == "W": - lons = w.xt_ocean.data.reshape(1) - lats = w.yt_ocean.data.reshape(1) - deps = w.sw_ocean.data.reshape(1) - if extrapolation: - deps = 0 - - pset = ParticleSet.from_list(fieldset=fieldset, pclass=myParticle, lon=lons, lat=lats, depth=deps) - pset.execute(VelocityInterpolator, runtime=1) - - convfactor = 0.01 if gridindexingtype == "pop" else 1.0 - if pointtype in ["U", "V"]: - assert np.allclose(pset.Uvel[0], u * convfactor) - assert np.allclose(pset.Vvel[0], v * convfactor) - elif pointtype == "W": - if extrapolation: - assert np.allclose(pset.Wvel[0], 0, atol=1e-9) - else: - assert np.allclose(pset.Wvel[0], w * convfactor) diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py index 9b76481598..cac7abc893 100644 --- a/tests/test_interpolation.py +++ b/tests/test_interpolation.py @@ -56,6 +56,8 @@ def data_2d(): return create_interpolation_data().isel(depth=0).values +@pytest.mark.v4remove +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize( "func, eta, xsi, expected", [ @@ -74,6 +76,8 @@ def test_raw_2d_interpolation(data_2d, func, eta, xsi, expected): assert func(ctx) == expected +@pytest.mark.v4remove +@pytest.mark.xfail(reason="GH1946") @pytest.mark.usefixtures("tmp_interpolator_registry") def test_interpolator_override(): fieldset = create_fieldset_zeros_3d() @@ -86,6 +90,8 @@ def test_interpolator(ctx: interpolation.InterpolationContext3D): fieldset.U[0, 0.5, 0.5, 0.5] +@pytest.mark.v4remove +@pytest.mark.xfail(reason="GH1946") @pytest.mark.usefixtures("tmp_interpolator_registry") def test_full_depth_provided_to_interpolators(): """The full depth needs to be provided to the interpolation schemes as some interpolators @@ -105,6 +111,8 @@ def test_interpolator2(ctx: interpolation.InterpolationContext3D): fieldset.U[0.5, 0.5, 0.5, 0.5] +@pytest.mark.v4remove +@pytest.mark.xfail(reason="GH1946") @pytest.mark.parametrize( "interp_method", [ diff --git a/tests/v4/test_uxarray_fieldset.py b/tests/v4/test_uxarray_fieldset.py new file mode 100644 index 0000000000..0e168fd098 --- /dev/null +++ b/tests/v4/test_uxarray_fieldset.py @@ -0,0 +1,59 @@ +from datetime import timedelta + +import pytest +import uxarray as ux + +from parcels import ( + FieldSet, + Particle, + ParticleSet, + UXPiecewiseConstantFace, + UXPiecewiseLinearNode, + download_example_dataset, +) + + +@pytest.fixture +def ds_fesom_channel() -> ux.UxDataset: + fesom_path = download_example_dataset("FESOM_periodic_channel") + grid_path = f"{fesom_path}/fesom_channel.nc" + data_path = [ + f"{fesom_path}/u.fesom_channel.nc", + f"{fesom_path}/v.fesom_channel.nc", + f"{fesom_path}/w.fesom_channel.nc", + ] + ds = ux.open_mfdataset(grid_path, data_path).rename_vars({"u": "U", "v": "V", "w": "W"}) + return ds + + +def test_fesom_fieldset(ds_fesom_channel): + fieldset = FieldSet([ds_fesom_channel]) + fieldset._check_complete() + # Check that the fieldset has the expected properties + assert fieldset.datasets[0] == ds_fesom_channel + + +def test_fesom_in_particleset(ds_fesom_channel): + fieldset = FieldSet([ds_fesom_channel]) + # Check that the fieldset has the expected properties + assert fieldset.datasets[0] == ds_fesom_channel + pset = ParticleSet(fieldset, pclass=Particle) + assert pset.fieldset == fieldset + + +def test_set_interp_methods(ds_fesom_channel): + fieldset = FieldSet([ds_fesom_channel]) + # Set the interpolation method for each field + fieldset.U.interp_method = UXPiecewiseConstantFace + fieldset.V.interp_method = UXPiecewiseConstantFace + fieldset.W.interp_method = UXPiecewiseLinearNode + + +def test_fesom_channel(ds_fesom_channel): + fieldset = FieldSet([ds_fesom_channel]) + # Set the interpolation method for each field + fieldset.U.interp_method = UXPiecewiseConstantFace + fieldset.V.interp_method = UXPiecewiseConstantFace + fieldset.W.interp_method = UXPiecewiseLinearNode + pset = ParticleSet(fieldset, pclass=Particle) + pset.execute(endtime=timedelta(days=1), dt=timedelta(hours=1)) diff --git a/tests/v4/test_xarray_fieldset.py b/tests/v4/test_xarray_fieldset.py new file mode 100644 index 0000000000..e69de29bb2