diff --git a/docs/community/v4-migration.md b/docs/community/v4-migration.md index 7a5bf03659..21da7fdd78 100644 --- a/docs/community/v4-migration.md +++ b/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. diff --git a/docs/examples/tutorial_Argofloats.ipynb b/docs/examples/tutorial_Argofloats.ipynb index ac9600c8b0..3391156b16 100644 --- a/docs/examples/tutorial_Argofloats.ipynb +++ b/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, diff --git a/parcels/particleset.py b/parcels/particleset.py index caded06f53..03b2246c14 100644 --- a/parcels/particleset.py +++ b/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) diff --git a/parcels/tools/exampledata_utils.py b/parcels/tools/exampledata_utils.py index 6b253b42f4..5513987770 100644 --- a/parcels/tools/exampledata_utils.py +++ b/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", diff --git a/tests/v4/test_particleset_execute.py b/tests/v4/test_particleset_execute.py index c5aef6905c..ee8b781770 100644 --- a/tests/v4/test_particleset_execute.py +++ b/tests/v4/test_particleset_execute.py @@ -364,6 +364,61 @@ def PythonFail(particles, fieldset): # pragma: no cover assert all(pset.time == fieldset.time_interval.left + np.timedelta64(10, "s")) +@pytest.mark.parametrize( + "kernel_names, expected", + [ + ("Lat1", [0, 1]), + ("Lat2", [2, 0]), + pytest.param( + "Lat1and2", + [2, 1], + marks=pytest.mark.xfail( + reason="Will be fixed alongside GH #2143 . Failing due to https://github.com/OceanParcels/Parcels/pull/2199#issuecomment-3285278876." + ), + ), + ("Lat1then2", [2, 1]), + ], +) +def test_execution_update_particle_in_kernel_function(fieldset, kernel_names, expected): + npart = 2 + + pset = ParticleSet(fieldset, lon=np.linspace(0, 1, npart), lat=np.zeros(npart)) + + def Lat1(particles, fieldset): # pragma: no cover + def SetLat1(p): + p.lat = 1 + + SetLat1(particles[(particles.lat == 0) & (particles.lon > 0.5)]) + + def Lat2(particles, fieldset): # pragma: no cover + def SetLat2(p): + p.lat = 2 + + SetLat2(particles[(particles.lat == 0) & (particles.lon < 0.5)]) + + def Lat1and2(particles, fieldset): # pragma: no cover + def SetLat1(p): + p.lat = 1 + + def SetLat2(p): + p.lat = 2 + + SetLat1(particles[(particles.lat == 0) & (particles.lon > 0.5)]) + SetLat2(particles[(particles.lat == 0) & (particles.lon < 0.5)]) + + if kernel_names == "Lat1": + kernels = [Lat1] + elif kernel_names == "Lat2": + kernels = [Lat2] + elif kernel_names == "Lat1and2": + kernels = [Lat1and2] + elif kernel_names == "Lat1then2": + kernels = [Lat1, Lat2] + + pset.execute(kernels, runtime=np.timedelta64(2, "s"), dt=np.timedelta64(1, "s")) + np.testing.assert_allclose(pset.lat, expected, rtol=1e-5) + + def test_uxstommelgyre_pset_execute(): ds = datasets_unstructured["stommel_gyre_delaunay"] grid = UxGrid(grid=ds.uxgrid, z=ds.coords["nz"], mesh="spherical")