Repository: OceanParcels/parcels
Pull request #625 (2019-08-16 11:18:52)

First parallel Parcels implementation using MPI
This is the first working implementation of Parallel Parcels using MPI. While it is fully functional, it is not as fast as it could be. Optimisations can particularly be done by using the Field Chunking of #605 and advanced load-balancing. That will be for later PRs


Patch:
diff --git a/.gitignore b/.gitignore
index f77f5e46..9ca60588 100644
--- a/.gitignore
+++ b/.gitignore
@@ -7,6 +7,7 @@ gh-pages/*
 *.nc
 *.nc4
 !*testfields*.nc
+out-*
 *.c
 *.pyc
 *.dSYM/*
diff --git a/.travis.yml b/.travis.yml
index 40f36620..4311855c 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -37,6 +37,11 @@ script:
         sh -e /etc/init.d/xvfb start;
         sleep 3;
       fi
+    - |
+      if [[ "${TRAVIS_OS_NAME}" = "osx" ]]; then
+        export CONDA_BUILD_SYSROOT=/
+        export C_INCLUDE_PATH=$C_INCLUDE_PATH:/Applications/Xcode.app/Contents//Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk/usr/include/
+      fi
     - py.test -v -s tests/
     - |
       # only get examples on linux
diff --git a/environment_py2_linux.yml b/environment_py2_linux.yml
index 5ee8b511..fe4f0ac0 100644
--- a/environment_py2_linux.yml
+++ b/environment_py2_linux.yml
@@ -10,8 +10,10 @@ dependencies:
   - enum34
   - ffmpeg>=3.2.3,<3.2.6
   - flake8>=2.1.0
-  - gcc_linux-64
   - git
+  - gcc_linux-64
+  - mpi4py>=3.0.1
+  - mpich>=3.2.1
   - jupyter
   - matplotlib>=2.0.2
   - netcdf4>=1.1.9
@@ -27,3 +29,4 @@ dependencies:
   - cftime
   - pytest<4.0
   - nbval
+  - scikit-learn
diff --git a/environment_py2_osx.yml b/environment_py2_osx.yml
index a3804059..d5947d2e 100644
--- a/environment_py2_osx.yml
+++ b/environment_py2_osx.yml
@@ -12,6 +12,8 @@ dependencies:
   - ffmpeg>=3.2.3,<3.2.6
   - flake8>=2.1.0
   - git
+  - mpi4py>=3.0.1
+  - mpich>=3.2.1
   - jupyter
   - matplotlib>=2.0.2
   - netcdf4>=1.1.9
@@ -27,3 +29,4 @@ dependencies:
   - cftime
   - pytest<4.0
   - nbval
+  - scikit-learn
diff --git a/environment_py3_linux.yml b/environment_py3_linux.yml
index cb458f9e..f5979b18 100644
--- a/environment_py3_linux.yml
+++ b/environment_py3_linux.yml
@@ -9,8 +9,10 @@ dependencies:
   - coverage
   - ffmpeg>=3.2.3,<3.2.6
   - flake8>=2.1.0
-  - gcc_linux-64
   - git
+  - gcc_linux-64
+  - mpi4py>=3.0.1
+  - mpich>=3.2.1
   - jupyter
   - matplotlib>=2.0.2
   - netcdf4>=1.1.9
@@ -26,3 +28,4 @@ dependencies:
   - dask
   - pytest<4.0
   - nbval
+  - scikit-learn
diff --git a/environment_py3_osx.yml b/environment_py3_osx.yml
index 758ec57f..d1a4b04f 100644
--- a/environment_py3_osx.yml
+++ b/environment_py3_osx.yml
@@ -11,6 +11,8 @@ dependencies:
   - ffmpeg>=3.2.3,<3.2.6
   - flake8>=2.1.0
   - git
+  - mpi4py>=3.0.1
+  - mpich>=3.2.1
   - jupyter
   - matplotlib>=2.0.2
   - netcdf4>=1.1.9
@@ -26,3 +28,4 @@ dependencies:
   - cftime
   - pytest<4.0
   - nbval
+  - scikit-learn
diff --git a/parcels/compiler.py b/parcels/compiler.py
index 844c1c29..7cc98c14 100644
--- a/parcels/compiler.py
+++ b/parcels/compiler.py
@@ -2,6 +2,10 @@ import subprocess
 from os import path, getenv
 from tempfile import gettempdir
 from struct import calcsize
+try:
+    from mpi4py import MPI
+except:
+    MPI = None
 try:
     from pathlib import Path
 except ImportError:
@@ -79,4 +83,5 @@ class GNUCompiler(Compiler):
         cppargs = ['-Wall', '-fPIC', '-I%s' % path.join(get_package_dir(), 'include')] + opt_flags + cppargs
         cppargs += arch_flag
         ldargs = ['-shared'] + ldargs + arch_flag
-        super(GNUCompiler, self).__init__("gcc", cppargs=cppargs, ldargs=ldargs)
+        compiler = "mpicc" if MPI else "gcc"
+        super(GNUCompiler, self).__init__(compiler, cppargs=cppargs, ldargs=ldargs)
diff --git a/parcels/examples/example_stommel.py b/parcels/examples/example_stommel.py
index 0f79eab8..2449462b 100644
--- a/parcels/examples/example_stommel.py
+++ b/parcels/examples/example_stommel.py
@@ -56,7 +56,8 @@ def UpdateP(particle, fieldset, time):
     particle.p = fieldset.P[time, particle.depth, particle.lat, particle.lon]
 
 
-def stommel_example(npart=1, mode='jit', verbose=False, method=AdvectionRK4):
+def stommel_example(npart=1, mode='jit', verbose=False, method=AdvectionRK4,
+                    outfile="StommelParticle.nc", repeatdt=None):
     timer.fieldset = timer.Timer('FieldSet', parent=timer.stommel)
     fieldset = stommel_fieldset()
     filename = 'stommel'
@@ -72,7 +73,7 @@ def stommel_example(npart=1, mode='jit', verbose=False, method=AdvectionRK4):
         p = Variable('p', dtype=np.float32, initial=0.)
         p_start = Variable('p_start', dtype=np.float32, initial=fieldset.P)
 
-    pset = ParticleSet.from_line(fieldset, size=npart, pclass=MyParticle,
+    pset = ParticleSet.from_line(fieldset, size=npart, pclass=MyParticle, repeatdt=repeatdt,
                                  start=(10e3, 5000e3), finish=(100e3, 5000e3), time=0)
 
     if verbose:
@@ -86,7 +87,7 @@ def stommel_example(npart=1, mode='jit', verbose=False, method=AdvectionRK4):
     timer.psetinit.stop()
     timer.psetrun = timer.Timer('Pset_run', parent=timer.pset)
     pset.execute(method + pset.Kernel(UpdateP), runtime=runtime, dt=dt,
-                 moviedt=None, output_file=pset.ParticleFile(name="StommelParticle", outputdt=outputdt))
+                 moviedt=None, output_file=pset.ParticleFile(name=outfile, outputdt=outputdt))
 
     if verbose:
         print("Final particle positions:\n%s" % pset)
@@ -126,12 +127,16 @@ Example of particle advection in the steady-state solution of the Stommel equati
                    help='Print particle information before and after execution')
     p.add_argument('-m', '--method', choices=('RK4', 'EE', 'RK45'), default='RK4',
                    help='Numerical method used for advection')
+    p.add_argument('-o', '--outfile', default='StommelParticle.nc',
+                   help='Name of output file')
+    p.add_argument('-r', '--repeatdt', default=None, type=int,
+                   help='repeatdt of the ParticleSet')
     args = p.parse_args()
 
     timer.args.stop()
     timer.stommel = timer.Timer('Stommel', parent=timer.root)
-    stommel_example(args.particles, mode=args.mode,
-                    verbose=args.verbose, method=method[args.method])
+    stommel_example(args.particles, mode=args.mode, verbose=args.verbose, method=method[args.method],
+                    outfile=args.outfile, repeatdt=args.repeatdt)
     timer.stommel.stop()
     timer.root.stop()
     timer.root.print_tree()
diff --git a/parcels/fieldset.py b/parcels/fieldset.py
index 100de375..c5b08404 100644
--- a/parcels/fieldset.py
+++ b/parcels/fieldset.py
@@ -7,6 +7,10 @@ import numpy as np
 from os import path
 from glob import glob
 from copy import deepcopy
+try:
+    from mpi4py import MPI
+except:
+    MPI = None
 
 
 __all__ = ['FieldSet']
@@ -709,16 +713,18 @@ class FieldSet(object):
         """Write FieldSet to NetCDF file using NEMO convention
 
         :param filename: Basename of the output fileset"""
-        logger.info("Generating NEMO FieldSet output with basename: %s" % filename)
 
-        if hasattr(self, 'U'):
-            self.U.write(filename, varname='vozocrtx')
-        if hasattr(self, 'V'):
-            self.V.write(filename, varname='vomecrty')
+        if MPI is None or MPI.COMM_WORLD.Get_rank() == 0:
+            logger.info("Generating FieldSet output with basename: %s" % filename)
 
-        for v in self.get_fields():
-            if (v.name != 'U') and (v.name != 'V'):
-                v.write(filename)
+            if hasattr(self, 'U'):
+                self.U.write(filename, varname='vozocrtx')
+            if hasattr(self, 'V'):
+                self.V.write(filename, varname='vomecrty')
+
+            for v in self.get_fields():
+                if (v.name != 'U') and (v.name != 'V'):
+                    v.write(filename)
 
     def advancetime(self, fieldset_new):
         """Replace oldest time on FieldSet with new FieldSet
diff --git a/parcels/kernel.py b/parcels/kernel.py
index 8a3131d9..b92331d4 100644
--- a/parcels/kernel.py
+++ b/parcels/kernel.py
@@ -18,6 +18,10 @@ import re
 from hashlib import md5
 import math  # noqa
 import random  # noqa
+try:
+    from mpi4py import MPI
+except:
+    MPI = None
 
 
 __all__ = ['Kernel']
@@ -129,8 +133,15 @@ class Kernel(object):
                 c_include_str = c_include
             self.ccode = loopgen.generate(self.funcname, self.field_args, self.const_args,
                                           kernel_ccode, c_include_str)
+            if MPI:
+                mpi_comm = MPI.COMM_WORLD
+                mpi_rank = mpi_comm.Get_rank()
+                basename = path.join(get_cache_dir(), self._cache_key) if mpi_rank == 0 else None
+                basename = mpi_comm.bcast(basename, root=0)
+                basename = basename + "_%d" % mpi_rank
+            else:
+                basename = path.join(get_cache_dir(), "%s_0" % self._cache_key)
 
-            basename = path.join(get_cache_dir(), self._cache_key)
             self.src_file = "%s.c" % basename
             self.lib_file = "%s.%s" % (basename, 'dll' if platform == 'win32' else 'so')
             self.log_file = "%s.log" % basename
@@ -163,7 +174,15 @@ class Kernel(object):
         # Python's ctype does not deal in any sort of manner well with dynamic linked libraries on this OS.
         if path.isfile(self.lib_file):
             [remove(s) for s in [self.src_file, self.lib_file, self.log_file]]
-            basename = path.join(get_cache_dir(), self._cache_key)
+            if MPI:
+                mpi_comm = MPI.COMM_WORLD
+                mpi_rank = mpi_comm.Get_rank()
+                basename = path.join(get_cache_dir(), self._cache_key) if mpi_rank == 0 else None
+                basename = mpi_comm.bcast(basename, root=0)
+                basename = basename + "_%d" % mpi_rank
+            else:
+                basename = path.join(get_cache_dir(), "%s_0" % self._cache_key)
+
             self.src_file = "%s.c" % basename
             self.lib_file = "%s.%s" % (basename, 'dll' if platform == 'win32' else 'so')
             self.log_file = "%s.log" % basename
diff --git a/parcels/particle.py b/parcels/particle.py
index 9e4b8cd7..00bcc4a0 100644
--- a/parcels/particle.py
+++ b/parcels/particle.py
@@ -8,9 +8,6 @@ from ctypes import c_void_p
 __all__ = ['ScipyParticle', 'JITParticle', 'Variable']
 
 
-lastID = 0  # module-level variable keeping track of last Particle ID used
-
-
 class Variable(object):
     """Descriptor class that delegates data access to particle data
 
@@ -112,6 +109,7 @@ class ParticleType(object):
 
 class _Particle(object):
     """Private base class for all particle types"""
+    lastID = 0  # class-level variable keeping track of last Particle ID used
 
     def __init__(self):
         ptype = self.getPType()
@@ -146,6 +144,10 @@ class _Particle(object):
     def getInitialValue(cls, ptype, name):
         return next((v.initial for v in ptype.variables if v.name is name), None)
 
+    @classmethod
+    def setLastID(cls, offset):
+        _Particle.lastID = offset
+
 
 class ScipyParticle(_Particle):
     """Class encapsulating the basic attributes of a particle,
@@ -169,15 +171,14 @@ class ScipyParticle(_Particle):
     state = Variable('state', dtype=np.int32, initial=ErrorCode.Success, to_write=False)
 
     def __init__(self, lon, lat, fieldset, depth=0., time=0., cptr=None):
-        global lastID
 
         # Enforce default values through Variable descriptor
         type(self).lon.initial = lon
         type(self).lat.initial = lat
         type(self).depth.initial = depth
         type(self).time.initial = time
-        type(self).id.initial = lastID
-        lastID += 1
+        type(self).id.initial = _Particle.lastID
+        _Particle.lastID += 1
         type(self).dt.initial = None
         super(ScipyParticle, self).__init__()
 
diff --git a/parcels/particlefile.py b/parcels/particlefile.py
index 8790f9ee..664382e8 100644
--- a/parcels/particlefile.py
+++ b/parcels/particlefile.py
@@ -8,6 +8,11 @@ import shutil
 import string
 import random
 from parcels.tools.error import ErrorCode
+from glob import glob
+try:
+    from mpi4py import MPI
+except:
+    MPI = None
 try:
     from parcels._version import version as parcels_version
 except:
@@ -47,15 +52,19 @@ class ParticleFile(object):
                      while ParticleFile is given as an argument of ParticleSet.execute()
                      It is either a timedelta object or a positive double.
     :param write_ondelete: Boolean to write particle data only when they are deleted. Default is False
-    :param tempwritedir: directory to write temporary files to during executing.
-                        Default is out-XXXXXX where Xs are random capitals
-    :param pset_info: dictionary of info on the ParticleSet, stored in tempwritedir/pset_info.npy,
-                      used to create NetCDF file from npy-files.
+    :param convert_at_end: Boolean to convert npy files to netcdf at end of run. Default is True
+    :param tempwritedir: directories to write temporary files to during executing.
+                     Default is out-XXXXXX where Xs are random capitals. Files for individual
+                     processors are written to subdirectories 0, 1, 2 etc under tempwritedir
+    :param pset_info: dictionary of info on the ParticleSet, stored in tempwritedir/XX/pset_info.npy,
+                     used to create NetCDF file from npy-files.
     """
 
-    def __init__(self, name, particleset, outputdt=np.infty, write_ondelete=False, tempwritedir=None, pset_info=None):
+    def __init__(self, name, particleset, outputdt=np.infty, write_ondelete=False, convert_at_end=True,
+                 tempwritedir=None, pset_info=None):
 
         self.write_ondelete = write_ondelete
+        self.convert_at_end = convert_at_end
         self.outputdt = outputdt
         self.lasttime_written = None  # variable to check if time has been written already
 
@@ -84,14 +93,21 @@ class ParticleFile(object):
             self.file_list = []
             self.time_written = []
             self.maxid_written = -1
-        self.to_export = False
 
         if tempwritedir is None:
-            self.tempwritedir = os.path.join(os.path.dirname(str(self.name)), "out-%s" % ''.join(random.choice(string.ascii_uppercase) for _ in range(8)))
+            tempwritedir = os.path.join(os.path.dirname(str(self.name)), "out-%s"
+                                        % ''.join(random.choice(string.ascii_uppercase) for _ in range(8)))
+
+        if MPI:
+            mpi_rank = MPI.COMM_WORLD.Get_rank()
+            self.tempwritedir_base = MPI.COMM_WORLD.bcast(tempwritedir, root=0)
         else:
-            self.tempwritedir = tempwritedir
+            self.tempwritedir_base = tempwritedir
+            mpi_rank = 0
+        self.tempwritedir = os.path.join(self.tempwritedir_base, "%d" % mpi_rank)
 
-        self.delete_tempwritedir()
+        if pset_info is None:  # otherwise arrive here from convert_npydir_to_netcdf
+            self.delete_tempwritedir()
 
     def open_netcdf_file(self, data_shape):
         """Initialise NetCDF4.Dataset for trajectory output.
@@ -117,7 +133,8 @@ class ParticleFile(object):
         self.dataset.parcels_mesh = self.parcels_mesh
 
         # Create ID variable according to CF conventions
-        self.id = self.dataset.createVariable("trajectory", "i4", coords)
+        self.id = self.dataset.createVariable("trajectory", "i4", coords,
+                                              fill_value=-2**(31))  # maxint32 fill_value
         self.id.long_name = "Unique identifier for each particle"
         self.id.cf_role = "trajectory_id"
 
@@ -172,16 +189,23 @@ class ParticleFile(object):
             setattr(self.dataset, name, message)
 
     def __del__(self):
-        if self.to_export:
+        # The export can only start when all threads are done.
+        if MPI:
+            MPI.COMM_WORLD.Barrier()
+            rank = MPI.COMM_WORLD.Get_rank()
+        else:
+            rank = 0
+        if self.convert_at_end and rank == 0:  # only export once.
             self.close()
 
-    def close(self):
+    def close(self, delete_tempfiles=True):
         """Close the ParticleFile object by exporting and then deleting
         the temporary npy files"""
         self.export()
-        self.delete_tempwritedir()
+        if delete_tempfiles:
+            self.delete_tempwritedir(tempwritedir=self.tempwritedir_base)
         self.dataset.close()
-        self.to_export = False
+        self.convert_at_end = False
 
     def add_metadata(self, name, message):
         """Add metadata to :class:`parcels.particleset.ParticleSet`
@@ -241,7 +265,7 @@ class ParticleFile(object):
         """Buffer data to set of temporary numpy files, using np.save"""
 
         if not os.path.exists(self.tempwritedir):
-            os.mkdir(self.tempwritedir)
+            os.makedirs(self.tempwritedir)
 
         if len(data_dict) > 0:
             tmpfilename = os.path.join(self.tempwritedir, str(len(self.file_list)) + ".npy")
@@ -255,8 +279,6 @@ class ParticleFile(object):
                 np.save(f, data_dict_once)
             self.file_list_once.append(tmpfilename)
 
-        self.to_export = True
-
     def dump_psetinfo_to_npy(self):
         pset_info = {}
         attrs_to_dump = ['name', 'var_names', 'var_names_once', 'time_origin', 'lonlatdepth_dtype',
@@ -314,8 +336,28 @@ class ParticleFile(object):
     def export(self):
         """Exports outputs in temporary NPY-files to NetCDF file"""
 
+        # Retrieve all temporary writing directories and sort them in numerical order
+        temp_names = sorted(glob(os.path.join("%s" % self.tempwritedir_base, "*")),
+                            key=lambda x: int(os.path.basename(x)))
+
+        if len(temp_names) == 0:
+            raise RuntimeError("No npy files found in %s" % self.tempwritedir_base)
+
+        global_maxid_written = -1
+        global_file_list = []
+        if len(self.var_names_once) > 0:
+            global_file_list_once = []
+        for tempwritedir in temp_names:
+            if os.path.exists(tempwritedir):
+                pset_info_local = np.load(os.path.join(tempwritedir, 'pset_info.npy'), allow_pickle=True).item()
+                global_maxid_written = np.max([global_maxid_written, pset_info_local['maxid_written']])
+                global_file_list += pset_info_local['file_list']
+                if len(self.var_names_once) > 0:
+                    global_file_list_once += pset_info_local['file_list_once']
+        self.maxid_written = global_maxid_written
+
         for var in self.var_names:
-            data = self.read_from_npy(self.file_list, len(self.time_written), var)
+            data = self.read_from_npy(global_file_list, len(self.time_written), var)
             if var == self.var_names[0]:
                 self.open_netcdf_file(data.shape)
             varout = 'z' if var == 'depth' else var
@@ -323,9 +365,13 @@ class ParticleFile(object):
 
         if len(self.var_names_once) > 0:
             for var in self.var_names_once:
-                getattr(self, var)[:] = self.read_from_npy(self.file_list_once, 1, var)
+                getattr(self, var)[:] = self.read_from_npy(global_file_list_once, 1, var)
 
-    def delete_tempwritedir(self):
-        """Deleted all temporary npy files"""
-        if os.path.exists(self.tempwritedir):
-            shutil.rmtree(self.tempwritedir)
+    def delete_tempwritedir(self, tempwritedir=None):
+        """Deleted all temporary npy files
+        :param tempwritedir Optional path of the directory to delete
+        """
+        if tempwritedir is None:
+            tempwritedir = self.tempwritedir
+        if os.path.exists(tempwritedir):
+            shutil.rmtree(tempwritedir)
diff --git a/parcels/particleset.py b/parcels/particleset.py
index 9002abfc..d06e4722 100644
--- a/parcels/particleset.py
+++ b/parcels/particleset.py
@@ -12,6 +12,11 @@ import time as time_module
 import collections
 from datetime import timedelta as delta
 from datetime import datetime, date
+try:
+    from mpi4py import MPI
+    from sklearn.cluster import KMeans
+except:
+    MPI = None
 
 __all__ = ['ParticleSet']
 
@@ -38,43 +43,74 @@ class ParticleSet(object):
     def __init__(self, fieldset, pclass=JITParticle, lon=None, lat=None, depth=None, time=None, repeatdt=None, lonlatdepth_dtype=None, **kwargs):
         self.fieldset = fieldset
         self.fieldset.check_complete()
+        partitions = kwargs.pop('partitions', None)
 
-        def convert_to_list(var):
-            # Convert numpy arrays and single integers/floats to one-dimensional lists
-            if isinstance(var, (int, float)):
-                return [var]
-            elif isinstance(var, np.ndarray):
+        def convert_to_array(var):
+            # Convert lists and single integers/floats to one-dimensional numpy arrays
+            if isinstance(var, np.ndarray):
                 return var.flatten()
-            return var
-        if lon is None:
-            lon = []
-        if lat is None:
-            lat = []
-
-        lon = convert_to_list(lon)
-        lat = convert_to_list(lat)
+            elif isinstance(var, (int, float)):
+                return np.array([var])
+            else:
+                return np.array(var)
+
+        lon = np.empty(shape=0) if lon is None else convert_to_array(lon)
+        lat = np.empty(shape=0) if lat is None else convert_to_array(lat)
+
         if depth is None:
             mindepth, _ = self.fieldset.gridset.dimrange('depth')
-            depth = np.ones(len(lon)) * mindepth
-        depth = convert_to_list(depth)
-        assert len(lon) == len(lat) and len(lon) == len(depth), (
+            depth = np.ones(lon.size) * mindepth
+        else:
+            depth = convert_to_array(depth)
+        assert lon.size == lat.size and lon.size == depth.size, (
             'lon, lat, depth don''t all have the same lenghts')
 
-        time = time.tolist() if isinstance(time, np.ndarray) else time
-        time = [time] * len(lat) if not isinstance(time, list) else time
-        time = [np.datetime64(t) if isinstance(t, datetime) else t for t in time]
-        time = [np.datetime64(t) if isinstance(t, date) else t for t in time]
+        time = convert_to_array(time)
+        time = np.repeat(time, lon.size) if time.size == 1 else time
+        if time.size > 0 and type(time[0]) in [datetime, date]:
+            time = np.array([np.datetime64(t) for t in time])
         self.time_origin = fieldset.time_origin
-        if len(time) > 0 and isinstance(time[0], np.timedelta64) and not self.time_origin:
+        if time.size > 0 and isinstance(time[0], np.timedelta64) and not self.time_origin:
             raise NotImplementedError('If fieldset.time_origin is not a date, time of a particle must be a double')
-        time = [self.time_origin.reltime(t) if isinstance(t, np.datetime64) else t for t in time]
+        time = np.array([self.time_origin.reltime(t) if isinstance(t, np.datetime64) else t for t in time])
+        assert lon.size == time.size, (
+            'time and positions (lon, lat, depth) don''t have the same lengths.')
 
         for kwvar in kwargs:
-            kwargs[kwvar] = convert_to_list(kwargs[kwvar])
+            kwargs[kwvar] = convert_to_array(kwargs[kwvar])
+            assert lon.size == kwargs[kwvar].size, (
+                '%s and positions (lon, lat, depth) don''t have the same lengths.' % kwargs[kwvar])
+
+        if MPI:
+            mpi_comm = MPI.COMM_WORLD
+            mpi_rank = mpi_comm.Get_rank()
+            mpi_size = mpi_comm.Get_size()
+
+            if lon.size < mpi_size and mpi_size > 1:
+                raise RuntimeError('Cannot initialise with fewer particles than MPI processors')
+
+            if mpi_size > 1:
+                if partitions is None:
+                    if mpi_rank == 0:
+                        coords = np.vstack((lon, lat)).transpose()
+                        kmeans = KMeans(n_clusters=mpi_size, random_state=0).fit(coords)
+                        partitions = kmeans.labels_
+                    else:
+                        partitions = None
+                    partitions = mpi_comm.bcast(partitions, root=0)
+                    lon = lon[partitions == mpi_rank]
+                    lat = lat[partitions == mpi_rank]
+                    time = time[partitions == mpi_rank]
+                    depth = depth[partitions == mpi_rank]
+                    for kwvar in kwargs:
+                        kwargs[kwvar] = kwargs[kwvar][partitions == mpi_rank]
+                offset = MPI.COMM_WORLD.allreduce(pclass.lastID, op=MPI.MAX) + sum(i < mpi_rank for i in partitions)
+            else:
+                offset = pclass.lastID
+        else:
+            offset = pclass.lastID
 
-        assert len(lon) == len(time), (
-            'time and positions (lon, lat, depth) don''t have the same '
-            'lengths.')
+        pclass.setLastID(offset)
 
         self.repeatdt = repeatdt.total_seconds() if isinstance(repeatdt, delta) else repeatdt
         if self.repeatdt:
@@ -87,6 +123,7 @@ class ParticleSet(object):
             self.repeatlat = lat
             self.repeatdepth = depth
             self.repeatpclass = pclass
+            self.partitions = partitions
             self.repeatkwargs = kwargs
 
         if lonlatdepth_dtype is None:
@@ -97,14 +134,13 @@ class ParticleSet(object):
             'lon lat depth precision should be set to either np.float32 or np.float64'
         JITParticle.set_lonlatdepth_dtype(self.lonlatdepth_dtype)
 
-        size = len(lon)
-        self.particles = np.empty(size, dtype=pclass)
+        self.particles = np.empty(lon.size, dtype=pclass)
         self.ptype = pclass.getPType()
         self.kernel = None
 
         if self.ptype.uses_jit:
             # Allocate underlying data for C-allocated particles
-            self._particle_data = np.empty(size, dtype=self.ptype.dtype)
+            self._particle_data = np.empty(lon.size, dtype=self.ptype.dtype)
 
             def cptr(i):
                 return self._particle_data[i]
@@ -113,11 +149,11 @@ class ParticleSet(object):
                 return None
 
         if lon is not None and lat is not None:
-            # Initialise from lists of lon/lat coordinates
-            assert size == len(lon) and size == len(lat), (
+            # Initialise from arrays of lon/lat coordinates
+            assert self.particles.size == lon.size and self.particles.size == lat.size, (
                 'Size of ParticleSet does not match lenght of lon and lat.')
 
-            for i in range(size):
+            for i in range(lon.size):
                 self.particles[i] = pclass(lon[i], lat[i], fieldset=fieldset, depth=depth[i], cptr=cptr(i), time=time[i])
                 # Set other Variables if provided
                 for kwvar in kwargs:
@@ -412,9 +448,9 @@ class ParticleSet(object):
             if verbose_progress is None and time_module.time() - walltime_start > 10:
                 # Showing progressbar if runtime > 10 seconds
                 if output_file:
-                    logger.info('Temporary output files are stored in %s.' % output_file.tempwritedir)
+                    logger.info('Temporary output files are stored in %s.' % output_file.tempwritedir_base)
                     logger.info('You can use "parcels_convert_npydir_to_netcdf %s" to convert these '
-                                'to a NetCDF file during the run.' % output_file.tempwritedir)
+                                'to a NetCDF file during the run.' % output_file.tempwritedir_base)
                 pbar = progressbar.ProgressBar(max_value=abs(endtime - _starttime)).start()
                 verbose_progress = True
             if dt > 0:
@@ -426,7 +462,7 @@ class ParticleSet(object):
                 pset_new = ParticleSet(fieldset=self.fieldset, time=time, lon=self.repeatlon,
                                        lat=self.repeatlat, depth=self.repeatdepth,
                                        pclass=self.repeatpclass, lonlatdepth_dtype=self.lonlatdepth_dtype,
-                                       **self.repeatkwargs)
+                                       partitions=self.partitions, **self.repeatkwargs)
                 for p in pset_new:
                     p.dt = dt
                 self.add(pset_new)
diff --git a/parcels/scripts/convert_npydir_to_netcdf.py b/parcels/scripts/convert_npydir_to_netcdf.py
index 0413b1c3..11b77af2 100644
--- a/parcels/scripts/convert_npydir_to_netcdf.py
+++ b/parcels/scripts/convert_npydir_to_netcdf.py
@@ -1,14 +1,18 @@
 from parcels import ParticleFile
 import numpy as np
 from os import path
+from glob import glob
 from argparse import ArgumentParser
 
 
-def convert_npydir_to_netcdf(tempwritedir):
+def convert_npydir_to_netcdf(tempwritedir_base, delete_tempfiles=False):
     """Convert npy files in tempwritedir to a NetCDF file
-    :param tempwritedir: directory where the temporary npy files are stored (can be obtained from ParticleFile.tempwritedir attribute)
+    :param tempwritedir_base: directory where the directories for temporary npy files
+            are stored (can be obtained from ParticleFile.tempwritedir_base attribute)
     """
 
+    tempwritedir = sorted(glob(path.join("%s" % tempwritedir_base, "*")),
+                          key=lambda x: int(path.basename(x)))[0]
     pyset_file = path.join(tempwritedir, 'pset_info.npy')
     if not path.isdir(tempwritedir):
         raise ValueError('Output directory "%s" does not exist' % tempwritedir)
@@ -16,19 +20,23 @@ def convert_npydir_to_netcdf(tempwritedir):
         raise ValueError('Output directory "%s" does not contain a pset_info.npy file' % tempwritedir)
 
     pset_info = np.load(pyset_file, allow_pickle=True).item()
-    pfile = ParticleFile(None, None, pset_info=pset_info)
-    pfile.export()
-    pfile.dataset.close()
+    pfile = ParticleFile(None, None, pset_info=pset_info, tempwritedir=tempwritedir_base, convert_at_end=False)
+    pfile.close(delete_tempfiles)
 
 
-def main(tempwritedir=None):
-    if tempwritedir is None:
+def main(tempwritedir_base=None, delete_tempfiles=False):
+    if tempwritedir_base is None:
         p = ArgumentParser(description="""Script to convert temporary npy output files to NetCDF""")
-        p.add_argument('tempwritedir', help='Name of directory where temporary npy files are stored')
+        p.add_argument('tempwritedir', help='Name of directory where temporary npy files are stored '
+                                            '(not including numbered subdirectories)')
+        p.add_argument('-d', '--delete_tempfiles', default=False,
+                       help='Flag to delete temporary files at end of call (default False)')
         args = p.parse_args()
-        tempwritedir = args.tempwritedir
+        tempwritedir_base = args.tempwritedir
+        if hasattr(args, 'delete_tempfiles'):
+            delete_tempfiles = args.delete_tempfiles
 
-    convert_npydir_to_netcdf(tempwritedir)
+    convert_npydir_to_netcdf(tempwritedir_base, delete_tempfiles)
 
 
 if __name__ == "__main__":


Test patch:
diff --git a/tests/test_mpirun.py b/tests/test_mpirun.py
new file mode 100644
index 00000000..7e538c99
--- /dev/null
+++ b/tests/test_mpirun.py
@@ -0,0 +1,27 @@
+import os
+from netCDF4 import Dataset
+import numpy as np
+try:
+    from mpi4py import MPI
+except:
+    MPI = None
+
+
+def test_mpi_run():
+    if MPI:
+        repeatdt = 200*86400
+        os.system('mpirun -np 2 python parcels/examples/example_stommel.py -p 4 -o StommelMPI.nc -r %d' % repeatdt)
+        os.system('python parcels/examples/example_stommel.py -p 4 -o StommelNoMPI.nc -r %d' % repeatdt)
+
+        ncfile1 = Dataset('StommelMPI.nc', 'r', 'NETCDF4')
+        ncfile2 = Dataset('StommelNoMPI.nc', 'r', 'NETCDF4')
+
+        for v in ncfile2.variables.keys():
+            assert np.allclose(ncfile1.variables[v][:], ncfile2.variables[v][:])
+
+        for a in ncfile2.ncattrs():
+            if a != 'parcels_version':
+                assert getattr(ncfile1, a) == getattr(ncfile2, a)
+
+        ncfile1.close()
+        ncfile2.close()
diff --git a/tests/test_particle_file.py b/tests/test_particle_file.py
index 61f52bc3..44415f56 100644
--- a/tests/test_particle_file.py
+++ b/tests/test_particle_file.py
@@ -25,10 +25,10 @@ def fieldset(xdim=40, ydim=100):
 
 def close_and_compare_netcdffiles(filepath, ofile, assystemcall=False):
     if assystemcall:
-        os.system('parcels_convert_npydir_to_netcdf %s' % ofile.tempwritedir)
+        os.system('parcels_convert_npydir_to_netcdf %s' % ofile.tempwritedir_base)
     else:
         import parcels.scripts.convert_npydir_to_netcdf as convert
-        convert.convert_npydir_to_netcdf(ofile.tempwritedir)
+        convert.convert_npydir_to_netcdf(ofile.tempwritedir_base)
 
     ncfile1 = Dataset(filepath, 'r', 'NETCDF4')
 
@@ -41,7 +41,8 @@ def close_and_compare_netcdffiles(filepath, ofile, assystemcall=False):
         assert np.allclose(ncfile1.variables[v][:], ncfile2.variables[v][:])
 
     for a in ncfile2.ncattrs():
-        assert getattr(ncfile1, a) == getattr(ncfile2, a)
+        if a != 'parcels_version':
+            assert getattr(ncfile1, a) == getattr(ncfile2, a)
 
     ncfile2.close()
     return ncfile1