Skip to content

Removing explicit field chunking #1892

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
erikvansebille merged 5 commits into from
Feb 28, 2025
Merged

Removing explicit field chunking #1892

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
4fbff29
Removing support for custom chunking of Field data
erikvansebille Feb 27, 2025
121aa3c
Removing DaskFileBuffer
erikvansebille Feb 27, 2025
38770d1
Further removals of DaskFielBuffer
erikvansebille Feb 27, 2025
bfb1f62
Review feedback
VeckoTheGecko Feb 27, 2025
ced2559
Merge branch 'v4-dev' into removing_explicit_field_chunking
erikvansebille Feb 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 0 additions & 7 deletions docs/documentation/additional_examples.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,13 +8,6 @@ example_brownian.py
:language: python
:linenos:

example_dask_chunk_OCMs.py
--------------------------

.. literalinclude:: ../examples/example_dask_chunk_OCMs.py
:language: python
:linenos:

example_decaying_moving_eddy.py
-------------------------------

Expand Down
233 changes: 1 addition & 232 deletions docs/examples/documentation_MPI.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"# Parallelisation with MPI and Field chunking with dask\n"
"# Parallelisation with MPI\n"
]
},
{
Expand Down Expand Up @@ -159,237 +159,6 @@
"For small projects, the above instructions are sufficient. If your project is large, then it is helpful to combine the `proc*` directories into a single zarr dataset and to optimise the chunking for your analysis. What is \"large\"? If you find yourself running out of memory while doing your analysis, saving the results, or sorting the dataset, or if reading the data is taking longer than you can tolerate, your problem is \"large.\" Another rule of thumb is if the size of your output directory is 1/3 or more of the memory of your machine, your problem is large. Chunking and combining the `proc*` data in order to speed up analysis is discussed [in the documentation on runs with large output](https://docs.oceanparcels.org/en/latest/examples/documentation_LargeRunsOutput.html).\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Chunking the FieldSet with dask\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"The basic idea of field-chunking in Parcels is that we use the `dask` library to load only these regions of the `Field` that are occupied by `Particles`. The advantage is that if different MPI processors take care of Particles in different parts of the domain, each only needs to load a small section of the full `FieldSet` (although note that this load-balancing functionality is still in [development](#Future-developments:-load-balancing)). Furthermore, the field-chunking in principle makes the `indices` keyword superfluous, as Parcels will determine which part of the domain to load itself.\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"The default behaviour is for `dask` to control the chunking, via a call to `da.from_array(data, chunks='auto')`. The chunksizes are then determined by the layout of the NetCDF files.\n",
"\n",
"However, in tests we have experienced that this may not necessarily be the most efficient chunking. Therefore, Parcels provides control over the chunksize via the `chunksize` keyword in `Field` creation, which requires a dictionary that sets the typical size of chunks for each dimension.\n",
"\n",
"It is strongly encouraged to explore what the best value for chunksize is for your experiment, which will depend on the `FieldSet`, the `ParticleSet` and the type of simulation (2D versus 3D). As a guidance, we have found that chunksizes in the zonal and meridional direction of approximately around 128 to 512 are typically most effective. The binning relates to the size of the model and its data size, so power-of-two values are advantageous but not required.\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"The notebook below shows an example script to explore the scaling of the time taken for `pset.execute()` as a function of zonal and meridional `chunksize` for a dataset from the [Copernicus Marine Service](http://marine.copernicus.eu/) portal.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"pycharm": {
"is_executing": true
}
},
"outputs": [],
"source": [
"%pylab inline\n",
"import os\n",
"import time\n",
"from datetime import timedelta\n",
"from glob import glob\n",
"\n",
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"import psutil\n",
"\n",
"import parcels\n",
"\n",
"\n",
"def set_cms_fieldset(cs):\n",
" data_dir_head = \"/data/oceanparcels/input_data\"\n",
" data_dir = os.path.join(data_dir_head, \"CMS/GLOBAL_REANALYSIS_PHY_001_030/\")\n",
" files = sorted(glob(data_dir + \"mercatorglorys12v1_gl12_mean_201607*.nc\"))\n",
" variables = {\"U\": \"uo\", \"V\": \"vo\"}\n",
" dimensions = {\"lon\": \"longitude\", \"lat\": \"latitude\", \"time\": \"time\"}\n",
"\n",
" if cs not in [\"auto\", False]:\n",
" cs = {\"time\": (\"time\", 1), \"lat\": (\"latitude\", cs), \"lon\": (\"longitude\", cs)}\n",
" return parcels.FieldSet.from_netcdf(files, variables, dimensions, chunksize=cs)\n",
"\n",
"\n",
"func_time = []\n",
"mem_used_GB = []\n",
"chunksize = [128, 256, 512, 768, 1024, 1280, 1536, 1792, 2048, 2610, \"auto\", False]\n",
"for cs in chunksize:\n",
" fieldset = set_cms_fieldset(cs)\n",
" pset = parcels.ParticleSet(\n",
" fieldset=fieldset,\n",
" pclass=parcels.Particle,\n",
" lon=[0],\n",
" lat=[0],\n",
" repeatdt=timedelta(hours=1),\n",
" )\n",
"\n",
" tic = time.time()\n",
" pset.execute(parcels.AdvectionRK4, dt=timedelta(hours=1))\n",
" func_time.append(time.time() - tic)\n",
" process = psutil.Process(os.getpid())\n",
" mem_B_used = process.memory_info().rss\n",
" mem_used_GB.append(mem_B_used / (1024 * 1024))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"pycharm": {
"is_executing": true
}
},
"outputs": [],
"source": [
"fig, ax = plt.subplots(1, 1, figsize=(15, 7))\n",
"\n",
"ax.plot(chunksize[:-2], func_time[:-2], \"o-\")\n",
"ax.plot([0, 2800], [func_time[-2], func_time[-2]], \"--\", label=chunksize[-2])\n",
"ax.plot([0, 2800], [func_time[-1], func_time[-1]], \"--\", label=chunksize[-1])\n",
"plt.xlim([0, 2800])\n",
"plt.legend()\n",
"ax.set_xlabel(\"chunksize\")\n",
"ax.set_ylabel(\"Time spent in pset.execute() [s]\")\n",
"plt.show()"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"The plot above shows that in this case, `chunksize='auto'` and `chunksize=False` (the two dashed lines) are roughly the same speed, but that the fastest run is for `chunksize=(1, 256, 256)`. But note that the actual thresholds and numbers depend on the `FieldSet` used and the specifics of your experiment.\n",
"\n",
"Furthermore, one of the major advantages of field chunking is the efficient utilization of memory. This permits the distribution of the particle advection to many cores, as otherwise the processing unit (e.g. a CPU core; a node in a cluster) would exhaust the memory rapidly. This is shown in the following plot of the memory behaviour.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"pycharm": {
"is_executing": true
}
},
"outputs": [],
"source": [
"fig, ax = plt.subplots(1, 1, figsize=(15, 12))\n",
"ax.plot(chunksize[:-2], mem_used_GB[:-2], \"--\", label=\"memory_blocked [MB]\")\n",
"ax.plot([0, 2800], [mem_used_GB[-2], mem_used_GB[-2]], \"x-\", label=\"auto [MB]\")\n",
"ax.plot([0, 2800], [mem_used_GB[-1], mem_used_GB[-1]], \"--\", label=\"no chunking [MB]\")\n",
"plt.legend()\n",
"ax.set_xlabel(\"chunksize\")\n",
"ax.set_ylabel(\"Memory blocked in pset.execute() [MB]\")\n",
"plt.show()"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"It can clearly be seen that with `chunksize=False` (green line), all Field data are loaded in full directly into memory, which can lead to MPI-run simulations being aborted for memory reasons. Furthermore, one can see that even the automatic method is not optimal, and optimizing the chunksize for a specific hydrodynamic dataset can make a large difference to the memory used.\n",
"\n",
"It may - depending on your simulation goal - be necessary to tweak the chunksize to leave more memory space for additional particles that are being simulated. Since particles and fields share the same memory space, lower memory utilisation by the `FieldSet` means more memory available for a larger `ParticleSet`.\n",
"\n",
"Also note that the above example is for a 2D application. For 3D applications, the `chunksize=False` will almost always be slower than `chunksize='auto'` or any dictionary, and is likely to run into insufficient memory issues, raising a `MemoryError`. The plot below shows the same analysis as above, but this time for a set of simulations using the full 3D Copernicus Marine Service code. In this case, the `chunksize='auto'` is about two orders of magnitude faster than running without chunking, and about 7.5 times faster than with minimal chunk capacity (i.e. `chunksize=(1, 128, 128)`).\n",
"\n",
"Choosing too small chunksizes can make the code slower, again highlighting that it is wise to explore which chunksize is best for your experiment before you perform it.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"pycharm": {
"is_executing": true
}
},
"outputs": [],
"source": [
"from parcels import AdvectionRK4_3D\n",
"\n",
"\n",
"def set_cms_fieldset_3D(cs):\n",
" data_dir_head = \"/data/oceanparcels/input_data\"\n",
" data_dir = os.path.join(data_dir_head, \"CMS/GLOBAL_REANALYSIS_PHY_001_030/\")\n",
" files = sorted(glob(data_dir + \"mercatorglorys12v1_gl12_mean_201607*.nc\"))\n",
" variables = {\"U\": \"uo\", \"V\": \"vo\"}\n",
" dimensions = {\n",
" \"lon\": \"longitude\",\n",
" \"lat\": \"latitude\",\n",
" \"depth\": \"depth\",\n",
" \"time\": \"time\",\n",
" }\n",
"\n",
" if cs not in [\"auto\", False]:\n",
" cs = {\n",
" \"time\": (\"time\", 1),\n",
" \"depth\": {\"depth\", 1},\n",
" \"lat\": (\"latitude\", cs),\n",
" \"lon\": (\"longitude\", cs),\n",
" }\n",
" return parcels.FieldSet.from_netcdf(files, variables, dimensions, chunksize=cs)\n",
"\n",
"\n",
"chunksize_3D = [128, 256, 512, 768, 1024, 1280, 1536, 1792, 2048, 2610, \"auto\", False]\n",
"func_time3D = []\n",
"for cs in chunksize_3D:\n",
" fieldset = set_cms_fieldset_3D(cs)\n",
" pset = parcels.ParticleSet(\n",
" fieldset=fieldset,\n",
" pclass=parcels.Particle,\n",
" lon=[0],\n",
" lat=[0],\n",
" repeatdt=timedelta(hours=1),\n",
" )\n",
"\n",
" tic = time.time()\n",
" pset.execute(AdvectionRK4_3D, dt=timedelta(hours=1))\n",
" func_time3D.append(time.time() - tic)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"pycharm": {
"is_executing": true
}
},
"outputs": [],
"source": [
"fig, ax = plt.subplots(1, 1, figsize=(15, 7))\n",
"\n",
"ax.plot(chunksize[:-2], func_time3D[:-2], \"o-\")\n",
"ax.plot([0, 2800], [func_time3D[-2], func_time3D[-2]], \"--\", label=chunksize_3D[-2])\n",
"plt.xlim([0, 2800])\n",
"plt.legend()\n",
"ax.set_xlabel(\"chunksize\")\n",
"ax.set_ylabel(\"Time spent in pset.execute() [s]\")\n",
"plt.show()"
]
},
{
"attachments": {},
"cell_type": "markdown",
Expand Down
1 change: 0 additions & 1 deletion docs/examples/example_mitgcm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,6 @@ def periodicBC(particle, fieldset, time): # pragma: no cover
str(path),
pset,
outputdt=timedelta(days=1),
chunks=(len(pset), 1),
)
kernels = parcels.AdvectionRK4 + pset.Kernel(periodicBC)
pset.execute(
Expand Down
6 changes: 1 addition & 5 deletions docs/examples/example_nemo_curvilinear.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,11 +28,7 @@ def run_nemo_curvilinear(outfile, advtype="RK4"):
}
variables = {"U": "U", "V": "V"}
dimensions = {"lon": "glamf", "lat": "gphif"}
chunksize = {"lat": ("y", 256), "lon": ("x", 512)}
fieldset = parcels.FieldSet.from_nemo(
filenames, variables, dimensions, chunksize=chunksize
)
assert fieldset.U.chunksize == chunksize
fieldset = parcels.FieldSet.from_nemo(filenames, variables, dimensions)

# Now run particles as normal
npart = 20
Expand Down
1 change: 0 additions & 1 deletion docs/examples/example_ofam.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,6 @@ def set_ofam_fieldset(deferred_load=True, use_xarray=False):
dimensions,
allow_time_extrapolation=True,
deferred_load=deferred_load,
chunksize=False,
)


Expand Down
20 changes: 0 additions & 20 deletions docs/examples/tutorial_parcels_structure.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -418,26 +418,6 @@
"- [Optimising the partitioning of the particles with a user-defined partition function](https://docs.oceanparcels.org/en/latest/examples/documentation_MPI.html#Optimising-the-partitioning-of-the-particles-with-a-user-defined-partition_function)\n",
"- [Future developments: load balancing](https://docs.oceanparcels.org/en/latest/examples/documentation_MPI.html#Future-developments:-load-balancing)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Another good way to optimise Parcels and speed-up execution is to chunk the `FieldSet` with `dask`, using the `chunksize` argument in the `FieldSet` creation. This will allow Parcels to load the `FieldSet` in chunks. \n",
"\n",
"Using chunking can be especially useful when working with large datasets _and_ when the particles only occupy a small region of the domain.\n",
"\n",
"Note that the **default** is `chunksize=None`, which means that the `FieldSet` is loaded in its entirety. This is generally the most efficient way to load the `FieldSet` when the particles are spread out over the entire domain.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### For more tutorials chunking and dask:\n",
"\n",
"- [Chunking the FieldSet with dask](https://docs.oceanparcels.org/en/latest/examples/documentation_MPI.html#Chunking-the-FieldSet-with-dask)"
]
}
],
"metadata": {
Expand Down
Loading
Loading