Repository: OceanParcels/parcels
Pull request #614 (2019-06-26 07:32:39)

Parallel output
Changed in this pull request

- Unique ID at initialization of ParticleSet
- Communication of temp_file location of each process
- Limiting writing to only process 0
- Process 0 gathers output from all temp_files instead of only its own file_list

Issues
- Unique ID not unique when adding particles during run 
- Creation of netCDF file requires everything to be in memory at the same time, does not scale with large runs (in terms of time steps or particles)
- Export depends on to_export variable which is not communicated across processes
- IDs not (necessarily) the same as serial run (because of spreading of particles over processors)

TODO:

- [x] Ensure uniqueness of ID. Suggestions: private/public .add() in ParticleSet (public one only accepts ParticleSet, private also allows particles) & find current max ID on processors using in ParticleSet initialization
- [x] Add flag to enable/disable export of netCDF file in current run (when disabled, temporary files are left untouched, so a follow-up process can collect all output)
- [x] Find alternative for the self.to_export variable in ParticleFile. Suggestion: just assume it is true and then check for existence of any .npy files in export(), throw exception if none exist.
- [x] Change script convert_npydir_to_netcdf.py to include a delete_temp_files flag (when true call .close() instead of .export())
- [x] Review coding style etc.


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/parcels/compiler.py b/parcels/compiler.py
index 045df285..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__("mpicc", 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..485f1edc 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 NEMO 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 2c4ab321..b92331d4 100644
--- a/parcels/kernel.py
+++ b/parcels/kernel.py
@@ -18,7 +18,10 @@ import re
 from hashlib import md5
 import math  # noqa
 import random  # noqa
-from mpi4py import MPI
+try:
+    from mpi4py import MPI
+except:
+    MPI = None
 
 
 __all__ = ['Kernel']
@@ -130,15 +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)
-
-            mpi_comm = MPI.COMM_WORLD
-            mpi_rank = mpi_comm.Get_rank()
-            if mpi_rank == 0:
-                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 = None
-            basename = mpi_comm.bcast(basename, root=0)
-            basename = basename + "_%d" % mpi_rank
+                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
@@ -171,12 +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]]
-            mpi_comm = MPI.COMM_WORLD
-            mpi_rank = mpi_comm.Get_rank()
-            if mpi_rank == 0:
-                basename = path.join(get_cache_dir(), self._cache_key)
-            basename = mpi_comm.bcast(basename, root=0)
-            basename = basename + "_%d" % mpi_rank
+            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 4e6e5f90..65a7ece3 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,20 @@ 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,
+    :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 in MPI mode 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 +94,30 @@ 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)))
+        if MPI:
+            mpi_comm = MPI.COMM_WORLD
+            mpi_rank = mpi_comm.Get_rank()
         else:
-            self.tempwritedir = tempwritedir
+            mpi_rank = 0
 
-        self.delete_tempwritedir()
+        if mpi_rank == 0:
+            if tempwritedir is None:
+                basename = os.path.join(os.path.dirname(str(self.name)), "out-%s" % ''.join(random.choice(string.ascii_uppercase) for _ in range(8)))
+            else:
+                basename = tempwritedir
+        else:
+            basename = None
+
+        if MPI:
+            self.tempwritedir_base = mpi_comm.bcast(basename, root=0)
+            self.tempwritedir = os.path.join(self.tempwritedir_base, "%d" % mpi_rank)
+        else:
+            self.tempwritedir_base = basename
+            self.tempwritedir = os.path.join(self.tempwritedir_base, "0")
+
+        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 +143,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 +199,22 @@ 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):
         """Close the ParticleFile object by exporting and then deleting
         the temporary npy files"""
         self.export()
-        self.delete_tempwritedir()
+        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`
@@ -254,7 +287,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")
@@ -268,8 +301,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',
@@ -320,8 +351,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(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
@@ -329,9 +380,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 9fd9e624..3f5874ae 100644
--- a/parcels/particleset.py
+++ b/parcels/particleset.py
@@ -12,8 +12,12 @@ import time as time_module
 import collections
 from datetime import timedelta as delta
 from datetime import datetime, date
-from sklearn.cluster import KMeans
-from mpi4py import MPI
+try:
+    from mpi4py import MPI
+except:
+    MPI = None
+if MPI:
+    from sklearn.cluster import KMeans
 
 __all__ = ['ParticleSet']
 
@@ -40,6 +44,7 @@ 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()
+        partition = kwargs.pop('partition', None)
 
         def convert_to_list(var):
             # Convert numpy arrays and single integers/floats to one-dimensional lists
@@ -67,27 +72,40 @@ class ParticleSet(object):
         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]
 
-
-        mpi_comm = MPI.COMM_WORLD
-        mpi_rank = mpi_comm.Get_rank()
-        mpi_size = mpi_comm.Get_size()
-        if mpi_rank == 0:
-            coords = np.vstack((lon, lat)).transpose()
-            kmeans = KMeans(n_clusters=mpi_size, random_state=0).fit(coords)
-            partition = kmeans.labels_
+        if MPI:
+            mpi_comm = MPI.COMM_WORLD
+            mpi_rank = mpi_comm.Get_rank()
+            mpi_size = mpi_comm.Get_size()
+
+            if len(lon) < mpi_size and mpi_size > 1:
+                raise RuntimeError('Cannot initialise with fewer particles than MPI processors')
+
+            if mpi_size > 1:
+                if partition is None:
+                    if mpi_rank == 0:
+                        coords = np.vstack((lon, lat)).transpose()
+                        kmeans = KMeans(n_clusters=mpi_size, random_state=0).fit(coords)
+                        partition = kmeans.labels_
+                    else:
+                        partition = None
+                    partition = mpi_comm.bcast(partition, root=0)
+                    lon = np.array(lon)[partition == mpi_rank]
+                    lat = np.array(lat)[partition == mpi_rank]
+                    time = np.array(time)[partition == mpi_rank]
+                    depth = np.array(depth)[partition == mpi_rank]
+                offset = MPI.COMM_WORLD.allreduce(pclass.lastID, op=MPI.MAX) + sum(i < mpi_rank for i in partition)
+            else:
+                offset = MPI.COMM_WORLD.allreduce(pclass.lastID, op=MPI.MAX)
         else:
-            partition = None
-        partition = mpi_comm.bcast(partition, root=0)
-        lon = np.array(lon)[partition == mpi_rank]
-        lat = np.array(lat)[partition == mpi_rank]
-        time = np.array(time)[partition == mpi_rank]
-        depth = np.array(depth)[partition == mpi_rank]
+            offset = pclass.lastID
+
+        pclass.setLastID(offset)
+
         lon = convert_to_list(lon)
         lat = convert_to_list(lat)
         time = convert_to_list(time)
         depth = convert_to_list(depth)
 
-
         self.time_origin = fieldset.time_origin
         if len(time) > 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')
@@ -111,6 +129,7 @@ class ParticleSet(object):
             self.repeatlat = lat
             self.repeatdepth = depth
             self.repeatpclass = pclass
+            self.partition = partition
             self.repeatkwargs = kwargs
 
         if lonlatdepth_dtype is None:
@@ -436,9 +455,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:
@@ -450,7 +469,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)
+                                       partition=self.partition, **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..9533177b 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,27 @@ 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)
+    if delete_tempfiles:
+        pfile.close()
+    else:
+        pfile.export()
+        pfile.dataset.close()
 
 
-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..e1da2cf9
--- /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:
+        releasedt = 200*86400
+        os.system('mpirun -np 2 python parcels/examples/example_stommel.py -p 4 -o StommelMPI.nc -r %d' % releasedt)
+        os.system('python parcels/examples/example_stommel.py -p 4 -o StommelNoMPI.nc -r %d' % releasedt)
+
+        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