Updating Argo tutorial to v4

docs/community/v4-migration.md

@@ -19,6 +19,10 @@ Version 4 of Parcels is unreleased at the moment. The information in this migrat
 
 - `interp_method` has to be an Interpolation function, instead of a string.
 
+## Particle
+
+- `Particle.add_variables()` has been replaced by `Particle.add_variable()`, which now also takes a list of `Variables`.
+
 ## ParticleSet
 
 - `repeatdt` and `lonlatdepth_dtype` have been removed from the ParticleSet.
@@ -28,3 +32,4 @@ Version 4 of Parcels is unreleased at the moment. The information in this migrat
 ## ParticleFile
 
 - Particlefiles should be created by `ParticleFile(...)` instead of `pset.ParticleFile(...)`
+- The `name` argument in `ParticleFile` has been replaced by `store` and can now be a string, a Path or a zarr store.

docs/examples/tutorial_Argofloats.ipynb

@@ -13,7 +13,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "This tutorial shows how simple it is to construct a Kernel in Parcels that mimics the [vertical movement of Argo floats](https://www.aoml.noaa.gov/phod/argo/images/argo_float_mission.jpg).\n"
+    "This tutorial shows how simple it is to construct a Kernel in Parcels that mimics the [vertical movement of Argo floats](https://www.aoml.noaa.gov/phod/argo/images/argo_float_mission.jpg).\n",
+    "\n",
+    "We first define the kernels for each phase of the Argo cycle."
    ]
   },
   {
@@ -22,52 +24,90 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Define the new Kernel that mimics Argo vertical movement\n",
-    "def ArgoVerticalMovement(particle, fieldset, time):\n",
-    "    driftdepth = 1000  # maximum depth in m\n",
-    "    maxdepth = 2000  # maximum depth in m\n",
-    "    vertical_speed = 0.10  # sink and rise speed in m/s\n",
-    "    cycletime = 10 * 86400  # total time of cycle in seconds\n",
-    "    drifttime = 9 * 86400  # time of deep drift in seconds\n",
-    "\n",
-    "    if particle.cycle_phase == 0:\n",
-    "        # Phase 0: Sinking with vertical_speed until depth is driftdepth\n",
-    "        particle_ddepth += vertical_speed * particle.dt\n",
-    "        if particle.depth + particle_ddepth >= driftdepth:\n",
-    "            particle_ddepth = driftdepth - particle.depth\n",
-    "            particle.cycle_phase = 1\n",
-    "\n",
-    "    elif particle.cycle_phase == 1:\n",
-    "        # Phase 1: Drifting at depth for drifttime seconds\n",
-    "        particle.drift_age += particle.dt\n",
-    "        if particle.drift_age >= drifttime:\n",
-    "            particle.drift_age = 0  # reset drift_age for next cycle\n",
-    "            particle.cycle_phase = 2\n",
-    "\n",
-    "    elif particle.cycle_phase == 2:\n",
-    "        # Phase 2: Sinking further to maxdepth\n",
-    "        particle_ddepth += vertical_speed * particle.dt\n",
-    "        if particle.depth + particle_ddepth >= maxdepth:\n",
-    "            particle_ddepth = maxdepth - particle.depth\n",
-    "            particle.cycle_phase = 3\n",
-    "\n",
-    "    elif particle.cycle_phase == 3:\n",
-    "        # Phase 3: Rising with vertical_speed until at surface\n",
-    "        particle_ddepth -= vertical_speed * particle.dt\n",
-    "        # particle.temp = fieldset.temp[time, particle.depth, particle.lat, particle.lon]  # if fieldset has temperature\n",
-    "        if particle.depth + particle_ddepth <= fieldset.mindepth:\n",
-    "            particle_ddepth = fieldset.mindepth - particle.depth\n",
-    "            # particle.temp = 0./0.  # reset temperature to NaN at end of sampling cycle\n",
-    "            particle.cycle_phase = 4\n",
-    "\n",
-    "    elif particle.cycle_phase == 4:\n",
-    "        # Phase 4: Transmitting at surface until cycletime is reached\n",
-    "        if particle.cycle_age > cycletime:\n",
-    "            particle.cycle_phase = 0\n",
-    "            particle.cycle_age = 0\n",
-    "\n",
-    "    if particle.state == StatusCode.Evaluate:\n",
-    "        particle.cycle_age += particle.dt  # update cycle_age"
+    "import numpy as np\n",
+    "\n",
+    "# Define the new Kernels that mimic Argo vertical movement\n",
+    "driftdepth = 1000  # maximum depth in m\n",
+    "maxdepth = 2000  # maximum depth in m\n",
+    "vertical_speed = 0.10  # sink and rise speed in m/s\n",
+    "cycletime = (\n",
+    "    10 * 86400\n",
+    ")  # total time of cycle in seconds  # TODO update to \"timedelta64[s]\"\n",
+    "drifttime = 9 * 86400  # time of deep drift in seconds\n",
+    "\n",
+    "\n",
+    "def ArgoPhase1(particles, fieldset):\n",
+    "    dt = particles.dt / np.timedelta64(1, \"s\")  # convert dt to seconds\n",
+    "\n",
+    "    def SinkingPhase(p):\n",
+    "        \"\"\"Phase 0: Sinking with vertical_speed until depth is driftdepth\"\"\"\n",
+    "        p.ddepth += vertical_speed * dt\n",
+    "        p.cycle_phase = np.where(p.depth + p.ddepth >= driftdepth, 1, p.cycle_phase)\n",
+    "        p.ddepth = np.where(\n",
+    "            p.depth + p.ddepth >= driftdepth, driftdepth - p.depth, p.ddepth\n",
+    "        )\n",
+    "\n",
+    "    SinkingPhase(particles[particles.cycle_phase == 0])\n",
+    "\n",
+    "\n",
+    "def ArgoPhase2(particles, fieldset):\n",
+    "    dt = particles.dt / np.timedelta64(1, \"s\")  # convert dt to seconds\n",
+    "\n",
+    "    def DriftingPhase(p):\n",
+    "        \"\"\"Phase 1: Drifting at depth for drifttime seconds\"\"\"\n",
+    "        p.drift_age += dt\n",
+    "        p.cycle_phase = np.where(p.drift_age >= drifttime, 2, p.cycle_phase)\n",
+    "        p.drift_age = np.where(p.drift_age >= drifttime, 0, p.drift_age)\n",
+    "\n",
+    "    DriftingPhase(particles[particles.cycle_phase == 1])\n",
+    "\n",
+    "\n",
+    "def ArgoPhase3(particles, fieldset):\n",
+    "    dt = particles.dt / np.timedelta64(1, \"s\")  # convert dt to seconds\n",
+    "\n",
+    "    def SecondSinkingPhase(p):\n",
+    "        \"\"\"Phase 2: Sinking further to maxdepth\"\"\"\n",
+    "        p.ddepth += vertical_speed * dt\n",
+    "        p.cycle_phase = np.where(p.depth + p.ddepth >= maxdepth, 3, p.cycle_phase)\n",
+    "        p.ddepth = np.where(\n",
+    "            p.depth + p.ddepth >= maxdepth, maxdepth - p.depth, p.ddepth\n",
+    "        )\n",
+    "\n",
+    "    SecondSinkingPhase(particles[particles.cycle_phase == 2])\n",
+    "\n",
+    "\n",
+    "def ArgoPhase4(particles, fieldset):\n",
+    "    dt = particles.dt / np.timedelta64(1, \"s\")  # convert dt to seconds\n",
+    "\n",
+    "    def RisingPhase(p):\n",
+    "        \"\"\"Phase 3: Rising with vertical_speed until at surface\"\"\"\n",
+    "        p.ddepth -= vertical_speed * dt\n",
+    "        p.temp = fieldset.thetao[p.time, p.depth, p.lat, p.lon]\n",
+    "        p.cycle_phase = np.where(\n",
+    "            p.depth + p.ddepth <= fieldset.mindepth, 4, p.cycle_phase\n",
+    "        )\n",
+    "        p.ddepth = np.where(\n",
+    "            p.depth + p.ddepth <= fieldset.mindepth,\n",
+    "            fieldset.mindepth - p.depth,\n",
+    "            p.ddepth,\n",
+    "        )\n",
+    "\n",
+    "    RisingPhase(particles[particles.cycle_phase == 3])\n",
+    "\n",
+    "\n",
+    "def ArgoPhase5(particles, fieldset):\n",
+    "    def TransmittingPhase(p):\n",
+    "        \"\"\"Phase 4: Transmitting at surface until cycletime is reached\"\"\"\n",
+    "        p.cycle_phase = np.where(p.cycle_age >= cycletime, 0, p.cycle_phase)\n",
+    "        p.cycle_age = np.where(p.cycle_age >= cycletime, 0, p.cycle_age)\n",
+    "        p.temp = np.nan  # no temperature measurement when at surface\n",
+    "\n",
+    "    TransmittingPhase(particles[particles.cycle_phase == 4])\n",
+    "\n",
+    "\n",
+    "def ArgoPhase6(particles, fieldset):\n",
+    "    dt = particles.dt / np.timedelta64(1, \"s\")  # convert dt to seconds\n",
+    "    particles.cycle_age += dt  # update cycle_age"
    ]
   },
   {
@@ -77,11 +117,7 @@
    "source": [
     "And then we can run Parcels with this 'custom kernel'.\n",
     "\n",
-    "Note that below we use the two-dimensional velocity fields of GlobCurrent, as these are provided as example_data with Parcels.\n",
-    "\n",
-    "We therefore assume that the horizontal velocities are the same throughout the entire water column. However, the `ArgoVerticalMovement` kernel will work on any `FieldSet`, including from full three-dimensional hydrodynamic data.\n",
-    "\n",
-    "If the hydrodynamic data also has a Temperature Field, then uncommenting the lines about temperature will also simulate the sampling of temperature.\n"
+    "Below we use the horizontal velocity fields of CopernicusMarine, which are provided as example_data with Parcels.\n"
    ]
   },
   {
@@ -92,54 +128,58 @@
    "source": [
     "from datetime import timedelta\n",
     "\n",
-    "import numpy as np\n",
+    "import xarray as xr\n",
     "\n",
     "import parcels\n",
     "\n",
-    "# Load the GlobCurrent data in the Agulhas region from the example_data\n",
-    "example_dataset_folder = parcels.download_example_dataset(\"GlobCurrent_example_data\")\n",
-    "filenames = {\n",
-    "    \"U\": f\"{example_dataset_folder}/20*.nc\",\n",
-    "    \"V\": f\"{example_dataset_folder}/20*.nc\",\n",
-    "}\n",
-    "variables = {\n",
-    "    \"U\": \"eastward_eulerian_current_velocity\",\n",
-    "    \"V\": \"northward_eulerian_current_velocity\",\n",
-    "}\n",
-    "dimensions = {\"lat\": \"lat\", \"lon\": \"lon\", \"time\": \"time\"}\n",
-    "fieldset = parcels.FieldSet.from_netcdf(filenames, variables, dimensions)\n",
-    "# uppermost layer in the hydrodynamic data\n",
-    "fieldset.mindepth = fieldset.U.depth[0]\n",
+    "# Load the CopernicusMarine data in the Agulhas region from the example_datasets\n",
+    "example_dataset_folder = parcels.download_example_dataset(\n",
+    "    \"CopernicusMarine_data_for_Argo_tutorial\"\n",
+    ")\n",
+    "\n",
+    "ds = xr.open_mfdataset(f\"{example_dataset_folder}/*.nc\", combine=\"by_coords\")\n",
+    "\n",
+    "# TODO check how we can get good performance without loading full dataset in memory\n",
+    "ds.load()  # load the dataset into memory\n",
     "\n",
+    "fieldset = parcels.FieldSet.from_copernicusmarine(ds)\n",
+    "fieldset.add_constant(\"mindepth\", 1.0)\n",
     "\n",
     "# Define a new Particle type including extra Variables\n",
-    "ArgoParticle = parcels.Particle.add_variables(\n",
+    "ArgoParticle = parcels.Particle.add_variable(\n",
     "    [\n",
-    "        # Phase of cycle:\n",
-    "        # init_descend=0,\n",
-    "        # drift=1,\n",
-    "        # profile_descend=2,\n",
-    "        # profile_ascend=3,\n",
-    "        # transmit=4\n",
     "        parcels.Variable(\"cycle_phase\", dtype=np.int32, initial=0.0),\n",
-    "        parcels.Variable(\"cycle_age\", dtype=np.float32, initial=0.0),\n",
+    "        parcels.Variable(\n",
+    "            \"cycle_age\", dtype=np.float32, initial=0.0\n",
+    "        ),  # TODO update to \"timedelta64[s]\"\n",
     "        parcels.Variable(\"drift_age\", dtype=np.float32, initial=0.0),\n",
-    "        # if fieldset has temperature\n",
-    "        # Variable('temp', dtype=np.float32, initial=np.nan),\n",
+    "        parcels.Variable(\"temp\", dtype=np.float32, initial=np.nan),\n",
     "    ]\n",
     ")\n",
     "\n",
     "# Initiate one Argo float in the Agulhas Current\n",
     "pset = parcels.ParticleSet(\n",
-    "    fieldset=fieldset, pclass=ArgoParticle, lon=[32], lat=[-31], depth=[0]\n",
+    "    fieldset=fieldset,\n",
+    "    pclass=ArgoParticle,\n",
+    "    lon=[32],\n",
+    "    lat=[-31],\n",
+    "    depth=[fieldset.mindepth],\n",
     ")\n",
     "\n",
     "# combine Argo vertical movement kernel with built-in Advection kernel\n",
-    "kernels = [ArgoVerticalMovement, parcels.AdvectionRK4]\n",
+    "kernels = [\n",
+    "    ArgoPhase1,\n",
+    "    ArgoPhase2,\n",
+    "    ArgoPhase3,\n",
+    "    ArgoPhase4,\n",
+    "    ArgoPhase5,\n",
+    "    ArgoPhase6,\n",
+    "    parcels.AdvectionRK4,\n",
+    "]\n",
     "\n",
     "# Create a ParticleFile object to store the output\n",
-    "output_file = pset.ParticleFile(\n",
-    "    name=\"argo_float\",\n",
+    "output_file = parcels.ParticleFile(\n",
+    "    store=\"argo_float.zarr\",\n",
     "    outputdt=timedelta(minutes=15),\n",
     "    chunks=(1, 500),  # setting to write in chunks of 500 observations\n",
     ")\n",
@@ -158,7 +198,25 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Now we can plot the trajectory of the Argo float with some simple calls to netCDF4 and matplotlib\n"
+    "Now we can plot the trajectory of the Argo float with some simple calls to netCDF4 and matplotlib.\n",
+    "\n",
+    "First plot the depth as a function of time, with the temperature as color (only on the upcast)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds_out = xr.open_zarr(\n",
+    "    output_file.store, decode_times=False\n",
+    ")  # TODO fix without using decode_times=False\n",
+    "x = ds_out[\"lon\"][:].squeeze()\n",
+    "y = ds_out[\"lat\"][:].squeeze()\n",
+    "z = ds_out[\"z\"][:].squeeze()\n",
+    "time = ds_out[\"time\"][:].squeeze() / 86400  # convert time to days\n",
+    "temp = ds_out[\"temp\"][:].squeeze()"
    ]
   },
   {
@@ -168,29 +226,49 @@
    "outputs": [],
    "source": [
     "import matplotlib.pyplot as plt\n",
-    "import xarray as xr\n",
-    "from mpl_toolkits.mplot3d import Axes3D\n",
     "\n",
-    "ds = xr.open_zarr(\"argo_float.zarr\")\n",
-    "x = ds[\"lon\"][:].squeeze()\n",
-    "y = ds[\"lat\"][:].squeeze()\n",
-    "z = ds[\"z\"][:].squeeze()\n",
-    "ds.close()\n",
+    "fig = plt.figure(figsize=(13, 6))\n",
+    "ax = plt.axes()\n",
+    "ax.plot(time, z, color=\"gray\")\n",
+    "cb = ax.scatter(time, z, c=temp, s=20, marker=\"o\", zorder=2)\n",
+    "ax.set_xlabel(\"Time [days]\")\n",
+    "ax.set_ylabel(\"Depth (m)\")\n",
+    "ax.invert_yaxis()\n",
+    "fig.colorbar(cb, label=\"Temperature (°C)\")\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can also make a 3D plot of the trajectory colored by temperature."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from mpl_toolkits.mplot3d import Axes3D\n",
     "\n",
-    "fig = plt.figure(figsize=(13, 10))\n",
+    "fig = plt.figure(figsize=(13, 8))\n",
     "ax = plt.axes(projection=\"3d\")\n",
-    "cb = ax.scatter(x, y, z, c=z, s=20, marker=\"o\")\n",
+    "ax.plot3D(x, y, z, color=\"gray\")\n",
+    "cb = ax.scatter(x, y, z, c=temp, s=20, marker=\"o\", zorder=2)\n",
     "ax.set_xlabel(\"Longitude\")\n",
     "ax.set_ylabel(\"Latitude\")\n",
     "ax.set_zlabel(\"Depth (m)\")\n",
     "ax.set_zlim(np.max(z), 0)\n",
+    "fig.colorbar(cb, label=\"Temperature (°C)\")\n",
     "plt.show()"
    ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "parcels",
+   "display_name": "parcels-v4",
    "language": "python",
    "name": "python3"
   },
@@ -204,7 +282,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.3"
+   "version": "3.12.11"
   }
  },
  "nbformat": 4,

parcels/particleset.py

@@ -7,6 +7,7 @@
 import numpy as np
 import xarray as xr
 from tqdm import tqdm
+from zarr.storage import DirectoryStore
 
 from parcels._core.utils.time import TimeInterval, maybe_convert_python_timedelta_to_numpy
 from parcels._reprs import particleset_repr
@@ -522,7 +523,7 @@ def execute(
 
         # Set up pbar
         if output_file:
-            logger.info(f"Output files are stored in {output_file.store}.")
+            logger.info(f"Output files are stored in {_format_output_location(output_file.store)}")
 
         if verbose_progress:
             pbar = tqdm(total=(end_time - start_time) / np.timedelta64(1, "s"), file=sys.stdout)
@@ -636,3 +637,9 @@ def _get_start_time(first_release_time, time_interval, sign_dt, runtime):
 
     start_time = first_release_time if not np.isnat(first_release_time) else fieldset_start
     return start_time
+
+
+def _format_output_location(zarr_obj):
+    if isinstance(zarr_obj, DirectoryStore):
+        return zarr_obj.path
+    return repr(zarr_obj)

parcels/tools/exampledata_utils.py

@@ -45,6 +45,10 @@
         f"{date.strftime('%Y%m%d')}000000-GLOBCURRENT-L4-CUReul_hs-ALT_SUM-v02.0-fv01.0.nc"
         for date in ([datetime(2002, 1, 1) + timedelta(days=x) for x in range(0, 365)] + [datetime(2003, 1, 1)])
     ],
+    "CopernicusMarine_data_for_Argo_tutorial": [
+        "cmems_mod_glo_phy-cur_anfc_0.083deg_P1D-m_uo-vo_31.00E-33.00E_33.00S-30.00S_0.49-2225.08m_2024-01-01-2024-02-01.nc",
+        "cmems_mod_glo_phy-thetao_anfc_0.083deg_P1D-m_thetao_31.00E-33.00E_33.00S-30.00S_0.49-2225.08m_2024-01-01-2024-02-01.nc",
+    ],
     "DecayingMovingEddy_data": [
         "decaying_moving_eddyU.nc",
         "decaying_moving_eddyV.nc",