{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Multiprocessing - Parallel execution of multiple OpenDrift instances in distinct CPUs\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\nimport numpy as np\nimport xarray as xr\nfrom opendrift.readers import reader_netCDF_CF_generic\nfrom opendrift.models.oceandrift import OceanDrift\nfrom multiprocessing import Pool\nimport glob\nimport opendrift\n\ndef concatenate_outputs(input_files, output_file):\n    \"\"\"\n    Concatenate OpenDrift NetCDF outputs from multiple workers into one single output file.\n\n    \"\"\"\n\n    ds_list = []\n    trajectory_offset = 0\n    for f in input_files:\n            dsi = xr.open_dataset(f)\n            traj_vars = [var for var in list(dsi.variables) if \"trajectory\" in dsi[var].dims]\n            ds_traj = dsi[traj_vars]\n            ds_traj = ds_traj.assign_coords(trajectory=ds_traj.trajectory + trajectory_offset)\n\n            trajectory_offset += ds_traj.dims[\"trajectory\"]\n            ds_list.append(ds_traj)\n\n    ds_conc = xr.concat(ds_list, dim=\"trajectory\")\n    ds_conc.to_netcdf(output_file)\n\n    return ds_conc\n\ndef RunOceanDrift(pool_number):\n    \"\"\"\n    Configure and run one OpenDrift instance\n\n    \"\"\"\n\n    o = OceanDrift(loglevel=0, logfile=\"output_W\"+str(pool_number)+\".log\", seed=0)\n\n    reader_topaz4 = reader_netCDF_CF_generic.Reader(\"https://thredds.met.no/thredds/dodsC/topaz/dataset-topaz4-arc-myoceanv2-be\")\n    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'])\n\n    o.set_config('drift:horizontal_diffusivity', 50)\n\n    time = datetime(2023,1,1)\n\n    ntraj=500\n    iniz=np.random.rand(ntraj) * -10. # seeding the chemicals in the upper 10m\n\n    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))\n\n    o.run(steps=7*4, time_step=3600*6, time_step_output=3600*6, outfile = \"output_W\"+str(pool_number)+\".nc\")\n\n\npositions=[(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)]\npool_size=len(positions)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "with Pool(pool_size) as p:\n    p.starmap(RunOceanDrift, [(i,) for i in range(pool_size)])\n\nconcatenate_outputs(input_files=glob.glob(\"output_W*.nc\"), output_file=\"output_total.nc\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "o0 = opendrift.open(\"output_W0.nc\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<img src=\"file://gallery/animations/example_multiprocessing_0.gif\">\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "o0.animation(color='origin_marker',\n            markersize=3,\n            vmin=0,vmax=pool_size-1,\n            colorbar=False,\n            fast = True,\n            lscale = 'l')\n\no = opendrift.open(\"output_total.nc\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<img src=\"file://gallery/animations/example_multiprocessing_1.gif\">\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "o.animation(color='origin_marker',\n            markersize=3,\n            vmin=0,vmax=pool_size-1,\n            colorbar=False,\n            fast = True,\n            lscale = 'l')"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "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.14.2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}