Benchmarking Base Branch - DO NOT MERGE - DO NOT CLOSE

environment_py3_win.yml

@@ -4,7 +4,7 @@ channels:
 dependencies:
   - python=3.6
   - cachetools>=1.0.0
-  - cgen
+  - cgen>=2020.1
   - coverage
   - ffmpeg>=3.2.3,<3.2.6
   - flake8>=2.1.0

parcels/__init__.py

@@ -4,6 +4,7 @@
 from parcels.fieldset import *  # noqa
 from parcels.particle import *  # noqa
 from parcels.particleset import *  # noqa
+from parcels.particleset_benchmark import *  # noqa
 from parcels.field import *  # noqa
 from parcels.kernel import *  # noqa
 import parcels.rng as random  # noqa

parcels/codegenerator.py

@@ -1,11 +1,11 @@
 import ast
 import collections
 import math
+import numpy as np
 import random
 from copy import copy
 
 import cgen as c
-import numpy as np
 
 from parcels.field import Field
 from parcels.field import NestedField
@@ -163,7 +163,7 @@ def __getattr__(self, attr):
 
 
 class ErrorCodeNode(IntrinsicNode):
-    symbol_map = {'Success': 'SUCCESS', 'Evaluate': 'EVALUATE', 'Repeat': 'REPEAT', 'Delete': 'DELETE',
+    symbol_map = {'Success': 'SUCCESS', 'Evaluate': 'EVALUATE', 'Repeat': 'REPEAT', 'Delete': 'DELETE', 'StopExecution': 'STOP_EXECUTION',
                   'Error': 'ERROR', 'ErrorInterpolation': 'ERROR_INTERPOLATION',
                   'ErrorOutOfBounds': 'ERROR_OUT_OF_BOUNDS', 'ErrorThroughSurface': 'ERROR_THROUGH_SURFACE'}
 

parcels/examples/example_nemo_curvilinear.py

@@ -1,17 +1,20 @@
+import math
 from argparse import ArgumentParser
 from datetime import timedelta as delta
 from glob import glob
 from os import path
 
 import numpy as np
 import pytest
+import dask
 
 from parcels import AdvectionRK4
 from parcels import FieldSet
 from parcels import JITParticle
 from parcels import ParticleFile
 from parcels import ParticleSet
 from parcels import ScipyParticle
+from parcels import ErrorCode
 
 ptype = {'scipy': ScipyParticle, 'jit': JITParticle}
 
@@ -28,7 +31,7 @@ def run_nemo_curvilinear(mode, outfile):
                        'data': data_path + 'V_purely_zonal-ORCA025_grid_V.nc4'}}
     variables = {'U': 'U', 'V': 'V'}
     dimensions = {'lon': 'glamf', 'lat': 'gphif'}
-    field_chunksize = {'lon': 2, 'lat': 2}
+    field_chunksize = {'y': 2, 'x': 2}
     field_set = FieldSet.from_nemo(filenames, variables, dimensions, field_chunksize=field_chunksize)
     assert field_set.U.field_chunksize == field_chunksize
 
@@ -101,6 +104,102 @@ def test_nemo_3D_samegrid():
     assert fieldset.U.dataFiles is not fieldset.W.dataFiles
 
 
+def fieldset_nemo_setup():
+    data_path = path.join(path.dirname(__file__), 'NemoNorthSeaORCA025-N006_data/')
+    ufiles = sorted(glob(data_path + 'ORCA*U.nc'))
+    vfiles = sorted(glob(data_path + 'ORCA*V.nc'))
+    wfiles = sorted(glob(data_path + 'ORCA*W.nc'))
+    mesh_mask = data_path + 'coordinates.nc'
+
+    filenames = {'U': {'lon': mesh_mask, 'lat': mesh_mask, 'depth': wfiles[0], 'data': ufiles},
+                 'V': {'lon': mesh_mask, 'lat': mesh_mask, 'depth': wfiles[0], 'data': vfiles},
+                 'W': {'lon': mesh_mask, 'lat': mesh_mask, 'depth': wfiles[0], 'data': wfiles}}
+    variables = {'U': 'uo',
+                 'V': 'vo',
+                 'W': 'wo'}
+    dimensions = {'U': {'lon': 'glamf', 'lat': 'gphif', 'depth': 'depthw', 'time': 'time_counter'},
+                  'V': {'lon': 'glamf', 'lat': 'gphif', 'depth': 'depthw', 'time': 'time_counter'},
+                  'W': {'lon': 'glamf', 'lat': 'gphif', 'depth': 'depthw', 'time': 'time_counter'}}
+
+    return filenames, variables, dimensions
+
+
+def compute_particle_advection(field_set, mode, lonp, latp):
+
+    def periodicBC(particle, fieldSet, time):
+        if particle.lon > 15.0:
+            particle.lon -= 15.0
+        if particle.lon < 0:
+            particle.lon += 15.0
+        if particle.lat > 60.0:
+            particle.lat -= 11.0
+        if particle.lat < 49.0:
+            particle.lat += 11.0
+
+    def OutOfBounds_reinitialisation(particle, fieldset, time):
+        particle.lat = 2.5
+        particle.lon = 52.0 + (-1e-3 + np.random.rand() * 2.0 * 1e-3)
+
+    pset = ParticleSet.from_list(field_set, ptype[mode], lon=lonp, lat=latp)
+    pfile = ParticleFile("nemo_particles", pset, outputdt=delta(days=1))
+    kernels = pset.Kernel(AdvectionRK4) + periodicBC
+    pset.execute(kernels, runtime=delta(days=4), dt=delta(hours=6),
+                 output_file=pfile, recovery={ErrorCode.ErrorOutOfBounds: OutOfBounds_reinitialisation})
+    return pset
+
+
+@pytest.mark.parametrize('mode', ['jit'])  # Only testing jit as scipy is very slow
+def test_nemo_curvilinear_auto_chunking(mode):
+    dask.config.set({'array.chunk-size': '2MiB'})
+    filenames, variables, dimensions = fieldset_nemo_setup()
+    field_set = FieldSet.from_nemo(filenames, variables, dimensions, field_chunksize='auto')
+    assert field_set.U.dataFiles is not field_set.W.dataFiles
+    npart = 20
+    lonp = 2.5 * np.ones(npart)
+    latp = [i for i in 52.0+(-1e-3+np.random.rand(npart)*2.0*1e-3)]
+    compute_particle_advection(field_set, mode, lonp, latp)
+    # Nemo sample file dimensions: depthu=75, y=201, x=151
+    assert (len(field_set.U.grid.load_chunk) == len(field_set.V.grid.load_chunk))
+    assert (len(field_set.U.grid.load_chunk) == len(field_set.W.grid.load_chunk))
+    assert (len(field_set.U.grid.load_chunk) != 1)
+
+
+@pytest.mark.parametrize('mode', ['jit'])  # Only testing jit as scipy is very slow
+def test_nemo_curvilinear_no_chunking(mode):
+    dask.config.set({'array.chunk-size': '128MiB'})
+    filenames, variables, dimensions = fieldset_nemo_setup()
+    field_set = FieldSet.from_nemo(filenames, variables, dimensions, field_chunksize=False)
+    assert field_set.U.dataFiles is not field_set.W.dataFiles
+    npart = 20
+    lonp = 2.5 * np.ones(npart)
+    latp = [i for i in 52.0+(-1e-3+np.random.rand(npart)*2.0*1e-3)]
+    compute_particle_advection(field_set, mode, lonp, latp)
+    # Nemo sample file dimensions: depthu=75, y=201, x=151
+    assert (len(field_set.U.grid.load_chunk) == len(field_set.V.grid.load_chunk))
+    assert (len(field_set.U.grid.load_chunk) == len(field_set.W.grid.load_chunk))
+    assert (len(field_set.U.grid.load_chunk) == 1)
+
+
+@pytest.mark.parametrize('mode', ['jit'])  # Only testing jit as scipy is very slow
+def test_nemo_curvilinear_specific_chunking(mode):
+    dask.config.set({'array.chunk-size': '128MiB'})
+    filenames, variables, dimensions = fieldset_nemo_setup()
+    chs = {'U': {'depthu': 75, 'y': 16, 'x': 16},
+           'V': {'depthv': 75, 'y': 16, 'x': 16},
+           'W': {'depthw': 75, 'y': 16, 'x': 16}}
+
+    field_set = FieldSet.from_nemo(filenames, variables, dimensions, field_chunksize=chs)
+    assert field_set.U.dataFiles is not field_set.W.dataFiles
+    npart = 20
+    lonp = 2.5 * np.ones(npart)
+    latp = [i for i in 52.0+(-1e-3+np.random.rand(npart)*2.0*1e-3)]
+    compute_particle_advection(field_set, mode, lonp, latp)
+    # Nemo sample file dimensions: depthu=75, y=201, x=151
+    assert (len(field_set.U.grid.load_chunk) == len(field_set.V.grid.load_chunk))
+    assert (len(field_set.U.grid.load_chunk) == len(field_set.W.grid.load_chunk))
+    assert (len(field_set.U.grid.load_chunk) == (1 * int(math.ceil(201.0/16.0)) * int(math.ceil(151.0/16.0))))
+
+
 if __name__ == "__main__":
     p = ArgumentParser(description="""Chose the mode using mode option""")
     p.add_argument('--mode', choices=('scipy', 'jit'), nargs='?', default='jit',

parcels/examples/example_recursive_errorhandling.py

@@ -13,7 +13,7 @@
 
 @pytest.mark.parametrize('mode', ['scipy', 'jit'])
 def test_recursive_errorhandling(mode, xdim=2, ydim=2):
-    """Example script to show how recursaive error handling can work.
+    """Example script to show how recursive error handling can work.
 
     In this example, a set of Particles is started at Longitude 0.5.
     These are run through a Kernel that throws an error if the