Merge pull request #1683 from manuelaghito/dev

Adding gallery example of parallel execution in multiple CPUs

Author: knutfrode

environment.yml

@@ -31,7 +31,7 @@ dependencies:
   - cmocean
   - utm
   - roaring-landmask>=0.7
-  - adios_db>=1.2
+  - adios_db>=1.2,<1.2.7
   - copernicusmarine>=2.0
   - bottleneck
   - pip

examples/example_multiprocessing.py

@@ -0,0 +1,95 @@
+#!/usr/bin/env python
+"""
+Multiprocessing - Parallel execution of multiple OpenDrift instances in distinct CPUs
+======================================================================================
+"""
+
+from datetime import datetime
+import numpy as np
+import xarray as xr
+from opendrift.readers import reader_netCDF_CF_generic
+from opendrift.models.oceandrift import OceanDrift
+from multiprocessing import Pool
+import glob
+import opendrift
+
+def concatenate_outputs(input_files, output_file):
+    """
+    Concatenate OpenDrift NetCDF outputs from multiple workers into one single output file.
+
+    """
+
+    ds_list = []
+    trajectory_offset = 0
+    for f in input_files:
+            dsi = xr.open_dataset(f)
+            traj_vars = [var for var in list(dsi.variables) if "trajectory" in dsi[var].dims]
+            ds_traj = dsi[traj_vars]
+            ds_traj = ds_traj.assign_coords(trajectory=ds_traj.trajectory + trajectory_offset)
+
+            trajectory_offset += ds_traj.dims["trajectory"]
+            ds_list.append(ds_traj)
+
+    ds_conc = xr.concat(ds_list, dim="trajectory")
+    ds_conc.to_netcdf(output_file)
+
+    return ds_conc
+
+def RunOceanDrift(pool_number):
+    """
+    Configure and run one OpenDrift instance
+
+    """
+
+    o = OceanDrift(loglevel=0, logfile="output_W"+str(pool_number)+".log", seed=0)
+
+    reader_topaz4 = reader_netCDF_CF_generic.Reader("https://thredds.met.no/thredds/dodsC/topaz/dataset-topaz4-arc-myoceanv2-be")
+    o.add_reader(reader_topaz4, variables=['x_sea_water_velocity', 'y_sea_water_velocity','sea_water_temperature','sea_water_salinity','sea_floor_depth_below_sea_level'])
+
+    o.set_config('drift:horizontal_diffusivity', 50)
+
+    time = datetime(2023,1,1)
+
+    ntraj=500
+    iniz=np.random.rand(ntraj) * -10. # seeding the chemicals in the upper 10m
+
+    o.seed_elements(lat=positions[pool_number][0], lon=positions[pool_number][1], z=iniz, radius=2000, number=ntraj, time=time, origin_marker=np.ones(ntraj)*(pool_number))
+
+    o.run(steps=7*4, time_step=3600*6, time_step_output=3600*6, outfile = "output_W"+str(pool_number)+".nc")
+
+
+positions=[(58.5,3),(58.2, 2.5),(58,1),(57.8,1.2),(57,1),(56.8,1.8),(56.5,2),(56,2.2)]
+pool_size=len(positions)
+
+#%% Run pool of OpenDrift instances in parallel using distinct CPUs and concatenate results
+with Pool(pool_size) as p:
+    p.starmap(RunOceanDrift, [(i,) for i in range(pool_size)])
+
+concatenate_outputs(input_files=glob.glob("output_W*.nc"), output_file="output_total.nc")
+
+
+#%% Generates animation of the first instance and the concatenated output
+
+o0 = opendrift.open("output_W0.nc")
+
+#%%
+# .. image:: /gallery/animations/example_multiprocessing_0.gif
+
+o0.animation(color='origin_marker',
+            markersize=3,
+            vmin=0,vmax=pool_size-1,
+            colorbar=False,
+            fast = True,
+            lscale = 'l')
+
+o = opendrift.open("output_total.nc")
+
+#%%
+# .. image:: /gallery/animations/example_multiprocessing_1.gif
+
+o.animation(color='origin_marker',
+            markersize=3,
+            vmin=0,vmax=pool_size-1,
+            colorbar=False,
+            fast = True,
+            lscale = 'l')