Scoping numba performance

environment_py3_linux.yml

@@ -3,7 +3,7 @@ channels:
   - defaults
   - conda-forge
 dependencies:
-  - python>3.4
+  - python>=3.7
   - cachetools>=1.0.0
   - cgen
   - coverage
@@ -30,3 +30,4 @@ dependencies:
   - nbval
   - scikit-learn
   - pykdtree
+  - numba>=0.50.0

environment_py3_osx.yml

@@ -3,7 +3,7 @@ channels:
   - defaults
   - conda-forge
 dependencies:
-  - python>=3.4
+  - python>=3.7
   - cachetools>=1.0.0
   - cgen
   - clang_osx-64
@@ -30,3 +30,4 @@ dependencies:
   - nbval
   - scikit-learn
   - pykdtree
+  - numba>=0.50.0

environment_py3_win.yml

@@ -2,7 +2,7 @@ name: py3_parcels
 channels:
   - conda-forge
 dependencies:
-  - python>3.4
+  - python>=3.7
   - cachetools>=1.0.0
   - cgen>=2020.1
   - coverage
@@ -27,3 +27,4 @@ dependencies:
   - pytest
   - nbval
   - pykdtree
+  - numba>=0.50.0

parcels/__init__.py

@@ -12,8 +12,7 @@
 from parcels.interaction import *  # noqa
 from parcels.application_kernels import *  # noqa
 import parcels.rng as ParcelsRandom  # noqa
-from parcels.compilation import *  # noqa
 from parcels.scripts import *  # noqa
 from parcels.gridset import *  # noqa
-from parcels.grid import *  # noqa
 from parcels.tools import *  # noqa
+from parcels.numba.grid import * # noqa

parcels/application_kernels/advection.py

@@ -1,7 +1,7 @@
 """Collection of pre-built advection kernels"""
 import math
 
-from parcels.tools.statuscodes import OperationCode
+from parcels.tools.statuscodes import OperationCode, StateCode
 
 
 __all__ = ['AdvectionRK4', 'AdvectionEE', 'AdvectionRK45', 'AdvectionRK4_3D',
@@ -12,7 +12,7 @@ def AdvectionRK4(particle, fieldset, time):
     """Advection of particles using fourth-order Runge-Kutta integration.
 
     Function needs to be converted to Kernel object before execution"""
-    (u1, v1) = fieldset.UV[particle]
+    (u1, v1) = fieldset.UV[time, particle.depth, particle.lat, particle.lon, particle]
     lon1, lat1 = (particle.lon + u1*.5*particle.dt, particle.lat + v1*.5*particle.dt)
     (u2, v2) = fieldset.UV[time + .5 * particle.dt, particle.depth, lat1, lon1, particle]
     lon2, lat2 = (particle.lon + u2*.5*particle.dt, particle.lat + v2*.5*particle.dt)
@@ -21,13 +21,14 @@ def AdvectionRK4(particle, fieldset, time):
     (u4, v4) = fieldset.UV[time + particle.dt, particle.depth, lat3, lon3, particle]
     particle.lon += (u1 + 2*u2 + 2*u3 + u4) / 6. * particle.dt
     particle.lat += (v1 + 2*v2 + 2*v3 + v4) / 6. * particle.dt
+    return StateCode.Success
 
 
 def AdvectionRK4_3D(particle, fieldset, time):
     """Advection of particles using fourth-order Runge-Kutta integration including vertical velocity.
 
     Function needs to be converted to Kernel object before execution"""
-    (u1, v1, w1) = fieldset.UVW[particle]
+    (u1, v1, w1) = fieldset.UVW[time + .5 * particle.dt, particle.depth, particle.lat, particle.lon, particle]
     lon1 = particle.lon + u1*.5*particle.dt
     lat1 = particle.lat + v1*.5*particle.dt
     dep1 = particle.depth + w1*.5*particle.dt
@@ -43,6 +44,7 @@ def AdvectionRK4_3D(particle, fieldset, time):
     particle.lon += (u1 + 2*u2 + 2*u3 + u4) / 6. * particle.dt
     particle.lat += (v1 + 2*v2 + 2*v3 + v4) / 6. * particle.dt
     particle.depth += (w1 + 2*w2 + 2*w3 + w4) / 6. * particle.dt
+    return StateCode.Success
 
 
 def AdvectionEE(particle, fieldset, time):
@@ -52,6 +54,7 @@ def AdvectionEE(particle, fieldset, time):
     (u1, v1) = fieldset.UV[particle]
     particle.lon += u1 * particle.dt
     particle.lat += v1 * particle.dt
+    return StateCode.Success
 
 
 def AdvectionRK45(particle, fieldset, time):
@@ -70,7 +73,7 @@ def AdvectionRK45(particle, fieldset, time):
     b4 = [25./216., 0., 1408./2565., 2197./4104., -1./5.]
     b5 = [16./135., 0., 6656./12825., 28561./56430., -9./50., 2./55.]
 
-    (u1, v1) = fieldset.UV[particle]
+    (u1, v1) = fieldset.UV[time, particle.depth, particle.lat, particle.lon, particle]
     lon1, lat1 = (particle.lon + u1 * A[0][0] * particle.dt,
                   particle.lat + v1 * A[0][0] * particle.dt)
     (u2, v2) = fieldset.UV[time + c[0] * particle.dt, particle.depth, lat1, lon1, particle]
@@ -97,10 +100,13 @@ def AdvectionRK45(particle, fieldset, time):
         particle.lon = lon_4th
         particle.lat = lat_4th
         if kappa2 <= math.pow(math.fabs(particle.dt * rk45tol / 10), 2):
-            particle.update_next_dt(particle.dt * 2)
+            particle._next_dt = particle.dt * 2
+#             particle.update_next_dt(particle.dt * 2)
+
     else:
         particle.dt /= 2
         return OperationCode.Repeat
+    return StateCode.Success
 
 
 def AdvectionAnalytical(particle, fieldset, time):

parcels/collection/collectionaos.py

@@ -9,7 +9,7 @@
 from parcels.collection.iterators import BaseParticleAccessor
 from parcels.collection.iterators import BaseParticleCollectionIterator
 from parcels.collection.iterators import BaseParticleCollectionIterable
-from parcels.particle import ScipyParticle, JITParticle  # noqa
+from parcels.particle import ScipyParticle
 from parcels.field import Field
 from parcels.tools.loggers import logger
 from parcels.tools.statuscodes import OperationCode
@@ -449,13 +449,6 @@ def add_single(self, particle_obj):
         assert isinstance(particle_obj, ScipyParticle)
         self._data = np.concatenate([self._data, particle_obj])
         self._ncount = self._data.shape[0]
-        if self._ptype.uses_jit and isinstance(particle_obj, JITParticle):
-            tmp_addr = self._data_c
-            prev_ncount = tmp_addr.shape[0]
-            self._data_c = np.array(self._ncount, dtype=self._ptype.dtype)
-            self._data_c[0:max(prev_ncount-1, 0)] = tmp_addr[:]
-            self._data_c[-1] = particle_obj._cptr
-            # particle_obj._cptr = self._data_c[-1]
 
     def add_same(self, same_class):
         """
@@ -517,11 +510,6 @@ def insert(self, obj, index=None):
             bottom_array = self._data[index:]
             splice_array = np.concatenate([top_array, obj])
             self._data = np.concatenate([splice_array, bottom_array])
-            if self._ptype.uses_jit and isinstance(obj, JITParticle):
-                top_array = self._data_c[0:index-1]
-                bottom_array = self._data_c[index:]
-                splice_array = np.concatenate([top_array, obj._cptr])
-                self._data_c = np.concatenate([splice_array, bottom_array])
             self._ncount = self._data.shape[0]
 
     def push(self, particle_obj):

parcels/collection/collectionsoa.py

@@ -8,8 +8,9 @@
 
 from parcels.collection.collections import ParticleCollection
 from parcels.collection.iterators import BaseParticleAccessor
-from parcels.collection.iterators import BaseParticleCollectionIterator, BaseParticleCollectionIterable
-from parcels.particle import ScipyParticle, JITParticle  # noqa
+from parcels.collection.iterators import BaseParticleCollectionIterator
+from parcels.collection.iterators import BaseParticleCollectionIterable
+from parcels.particle import ScipyParticle
 from parcels.field import Field
 from parcels.tools.loggers import logger
 from parcels.tools.statuscodes import OperationCode
@@ -130,16 +131,12 @@ def __init__(self, pclass, lon, lat, depth, time, lonlatdepth_dtype, pid_orig, p
         self._pclass = pclass
 
         self._ptype = pclass.getPType()
-        self._data = {}
         initialised = set()
 
         self._ncount = len(lon)
-
-        for v in self.ptype.variables:
-            if v.name in ['xi', 'yi', 'zi', 'ti']:
-                self._data[v.name] = np.empty((len(lon), ngrid), dtype=v.dtype)
-            else:
-                self._data[v.name] = np.empty(self._ncount, dtype=v.dtype)
+        self._data = np.empty(self.ncount, dtype=self.ptype.dtype).view(np.recarray)
+        self._pbackup = np.empty(1, dtype=self.ptype.dtype).view(np.recarray)
+        self._data.lat = lat
 
         if lon is not None and lat is not None:
             # Initialise from lists of lon/lat coordinates
@@ -155,7 +152,7 @@ def __init__(self, pclass, lon, lat, depth, time, lonlatdepth_dtype, pid_orig, p
             self._data['fileid'][:] = -1
 
             # special case for exceptions which can only be handled from scipy
-            self._data['exception'] = np.empty(self.ncount, dtype=object)
+#             self._data['exception'] = np.empty(self.ncount, dtype=object)
 
             initialised |= {'lat', 'lon', 'depth', 'time', 'id'}
 
@@ -175,7 +172,7 @@ def __init__(self, pclass, lon, lat, depth, time, lonlatdepth_dtype, pid_orig, p
                     for i in range(self.ncount):
                         if (time[i] is None) or (np.isnan(time[i])):
                             raise RuntimeError('Cannot initialise a Variable with a Field if no time provided (time-type: {} values: {}). Add a "time=" to ParticleSet construction'.format(type(time), time))
-                        v.initial.fieldset.computeTimeChunk(time[i], 0)
+#                         v.initial.fieldset.computeTimeChunk(time[i], 0)
                         self._data[v.name][i] = v.initial[
                             time[i], depth[i], lat[i], lon[i]
                         ]
@@ -234,7 +231,7 @@ def __getattr__(self, name):
         :param name: name of the property
         """
         for v in self.ptype.variables:
-            if v.name == name and name in self._data:
+            if v.name == name and name in list(self._data.dtype.fields):
                 return self._data[name]
         return False
 
@@ -441,8 +438,7 @@ def add_same(self, same_class):
             if not (same_class._sorted
                     and self._data['id'][-1] < same_class._data['id'][0]):
                 self._sorted = False
-            for d in self._data:
-                self._data[d] = np.concatenate((self._data[d], same_class._data[d]))
+            self._data = np.concatenate((self._data, same_class._data))
             self._ncount += same_class.ncount
 
     def __iadd__(self, same_class):
@@ -556,8 +552,7 @@ def remove_single_by_index(self, index):
         """
         super().remove_single_by_index(index)
 
-        for d in self._data:
-            self._data[d] = np.delete(self._data[d], index, axis=0)
+        self._data = np.delete(self._data, index)
 
         self._ncount -= 1
 
@@ -644,8 +639,8 @@ def remove_multi_by_indices(self, indices):
         if type(indices) is dict:
             indices = list(indices.values())
 
-        for d in self._data:
-            self._data[d] = np.delete(self._data[d], indices, axis=0)
+        if len(indices):
+            self._data = np.delete(self._data, indices)
 
         self._ncount -= len(indices)