Removing JITParticle Class

Author: erikvansebille

parcels/compilation/codegenerator.py

@@ -9,7 +9,7 @@
 
 from parcels.field import Field, NestedField, VectorField
 from parcels.grid import Grid
-from parcels.particle import JITParticle
+from parcels.particle import ScipyParticle
 from parcels.tools.statuscodes import StatusCode
 from parcels.tools.warnings import KernelWarning
 
@@ -214,7 +214,7 @@ class IntrinsicTransformer(ast.NodeTransformer):
     and propagates attribute access.
     """
 
-    def __init__(self, fieldset=None, ptype=JITParticle):
+    def __init__(self, fieldset=None, ptype=ScipyParticle):
         self.fieldset = fieldset
         self.ptype = ptype
 
@@ -421,7 +421,7 @@ class KernelGenerator(ast.NodeVisitor):
     kernel_vars = ["particle", "fieldset", "time", "output_time", "tol"]
     array_vars: list[str] = []
 
-    def __init__(self, fieldset=None, ptype=JITParticle):
+    def __init__(self, fieldset=None, ptype=ScipyParticle):
         self.fieldset = fieldset
         self.ptype = ptype
         self.field_args = collections.OrderedDict()

parcels/kernel.py

@@ -680,7 +680,7 @@ def evaluate_particle(self, p, endtime):
         Parameters
         ----------
         p :
-            object of (sub-)type (ScipyParticle, JITParticle)
+            object of (sub-)type ScipyParticle
         endtime :
             endtime of this overall kernel evaluation step
         dt :

parcels/particle.py

@@ -1,14 +1,11 @@
-from ctypes import c_void_p
 from operator import attrgetter
 from typing import Literal
 
 import numpy as np
 
 from parcels.tools.statuscodes import StatusCode
 
-__all__ = ["JITParticle", "ScipyInteractionParticle", "ScipyParticle", "Variable"]
-
-indicators_64bit = [np.float64, np.uint64, np.int64, c_void_p]
+__all__ = ["ScipyInteractionParticle", "ScipyParticle", "Variable"]
 
 
 class Variable:
@@ -41,24 +38,14 @@ def name(self):
     def __get__(self, instance, cls):
         if instance is None:
             return self
-        if issubclass(cls, JITParticle):
-            return instance._cptr.__getitem__(self.name)
-        else:
-            return getattr(instance, f"_{self.name}", self.initial)
+        return getattr(instance, f"_{self.name}", self.initial)
 
     def __set__(self, instance, value):
-        if isinstance(instance, JITParticle):
-            instance._cptr.__setitem__(self.name, value)
-        else:
-            setattr(instance, f"_{self.name}", value)
+        setattr(instance, f"_{self.name}", value)
 
     def __repr__(self):
         return f"Variable(name={self._name}, dtype={self.dtype}, initial={self.initial}, to_write={self.to_write})"
 
-    def is64bit(self):
-        """Check whether variable is 64-bit."""
-        return True if self.dtype in indicators_64bit else False
-
 
 class ParticleType:
     """Class encapsulating the type information for custom particles.
@@ -75,7 +62,7 @@ def __init__(self, pclass):
         if not issubclass(pclass, ScipyParticle):
             raise TypeError("Class object does not inherit from parcels.ScipyParticle")
         self.name = pclass.__name__
-        self.uses_jit = issubclass(pclass, JITParticle)
+        self.uses_jit = False  # TODO v4: remove this attribute
         # Pick Variable objects out of __dict__.
         self.variables = [v for v in pclass.__dict__.values() if isinstance(v, Variable)]
         for cls in pclass.__bases__:
@@ -92,8 +79,6 @@ def __init__(self, pclass):
                             "Custom Variable name 'z' is not allowed, as it is used for depth in ParticleFile"
                         )
                 self.variables = ptype.variables + self.variables
-        # Sort variables with all the 64-bit first so that they are aligned for the JIT cptr
-        self.variables = [v for v in self.variables if v.is64bit()] + [v for v in self.variables if not v.is64bit()]
 
     def __repr__(self):
         return f"{type(self).__name__}(pclass={self.name})"
@@ -103,35 +88,6 @@ def __getitem__(self, item):
             if v.name == item:
                 return v
 
-    @property
-    def _cache_key(self):
-        return "-".join([f"{v.name}:{v.dtype}" for v in self.variables])
-
-    @property
-    def dtype(self):
-        """Numpy.dtype object that defines the C struct."""
-        type_list = [(v.name, v.dtype) for v in self.variables]
-        for v in self.variables:
-            if v.dtype not in self.supported_dtypes:
-                raise RuntimeError(str(v.dtype) + " variables are not implemented in JIT mode")
-        if self.size % 8 > 0:
-            # Add padding to be 64-bit aligned
-            type_list += [("pad", np.float32)]
-        return np.dtype(type_list)
-
-    @property
-    def size(self):
-        """Size of the underlying particle struct in bytes."""
-        return sum([8 if v.is64bit() else 4 for v in self.variables])
-
-    @property
-    def supported_dtypes(self):
-        """List of all supported numpy dtypes. All others are not supported."""
-        # Developer note: other dtypes (mostly 2-byte ones) are not supported now
-        # because implementing and aligning them in cgen.GenerableStruct is a
-        # major headache. Perhaps in a later stage
-        return [np.int32, np.uint32, np.int64, np.uint64, np.float32, np.double, np.float64, c_void_p]
-
 
 class ScipyParticle:
     """Class encapsulating the basic attributes of a particle, to be executed in SciPy mode.
@@ -192,9 +148,7 @@ def __init__(self, lon, lat, pid, fieldset=None, ngrids=None, depth=0.0, time=0.
                 initial = v.initial(self)
             else:
                 initial = v.initial
-            # Enforce type of initial value
-            if v.dtype != c_void_p:
-                setattr(self, v.name, v.dtype(initial))
+            setattr(self, v.name, v.dtype(initial))
 
     def __del__(self):
         pass  # superclass is 'object', and object itself has no destructor, hence 'pass'
@@ -267,7 +221,7 @@ def set_lonlatdepth_dtype(cls, dtype):
         cls.depth_nextloop.dtype = dtype
 
     @classmethod
-    def setLastID(cls, offset):
+    def setLastID(cls, offset):  # TODO v4: check if we can implement this in another way
         ScipyParticle.lastID = offset
 
 
@@ -277,36 +231,5 @@ def setLastID(cls, offset):
 
 
 class JITParticle(ScipyParticle):
-    """Particle class for JIT-based (Just-In-Time) Particle objects.
-
-    Parameters
-    ----------
-    lon : float
-        Initial longitude of particle
-    lat : float
-        Initial latitude of particle
-    fieldset : parcels.fieldset.FieldSet
-        mod:`parcels.fieldset.FieldSet` object to track this particle on
-    dt :
-        Execution timestep for this particle
-    time :
-        Current time of the particle
-
-
-    Notes
-    -----
-    Additional Variables can be added via the :Class Variable: objects
-
-    Users should use JITParticles for faster advection computation.
-    """
-
     def __init__(self, *args, **kwargs):
-        self._cptr = kwargs.pop("cptr", None)
-        if self._cptr is None:
-            # Allocate data for a single particle
-            ptype = self.getPType()
-            self._cptr = np.empty(1, dtype=ptype.dtype)[0]
-        super().__init__(*args, **kwargs)
-
-    def __del__(self):
-        super().__del__()
+        raise NotImplementedError("JITParticle has been deprecated in Parcels v4. Use ScipyParticle instead.")

parcels/particleset.py

@@ -24,7 +24,7 @@
     KDTreeFlatNeighborSearch,
 )
 from parcels.kernel import Kernel
-from parcels.particle import JITParticle, Variable
+from parcels.particle import ScipyParticle, Variable
 from parcels.particledata import ParticleData, ParticleDataIterator
 from parcels.particlefile import ParticleFile
 from parcels.tools._helpers import particleset_repr, timedelta_to_float
@@ -56,9 +56,8 @@ class ParticleSet:
     ----------
     fieldset :
         mod:`parcels.fieldset.FieldSet` object from which to sample velocity.
-    pclass : parcels.particle.JITParticle or parcels.particle.ScipyParticle
-        Optional :mod:`parcels.particle.JITParticle` or
-        :mod:`parcels.particle.ScipyParticle` object that defines custom particle
+    pclass : parcels.particle.ScipyParticle
+        Optional object that inherits from :mod:`parcels.particle.ScipyParticle` object that defines custom particle
     lon :
         List of initial longitude values for particles
     lat :
@@ -87,7 +86,7 @@ class ParticleSet:
     def __init__(
         self,
         fieldset,
-        pclass=JITParticle,
+        pclass=ScipyParticle,
         lon=None,
         lat=None,
         depth=None,
@@ -481,8 +480,8 @@ def from_list(
         ----------
         fieldset :
             mod:`parcels.fieldset.FieldSet` object from which to sample velocity
-        pclass : parcels.particle.JITParticle or parcels.particle.ScipyParticle
-            Particle class. May be a particle class as defined in parcels, or a subclass defining a custom particle.
+        pclass :
+            Particle class. May be a parcels.particle.ScipyParticle class as defined in parcels, or a subclass defining a custom particle.
         lon :
             List of initial longitude values for particles
         lat :
@@ -537,8 +536,8 @@ def from_line(
         ----------
         fieldset :
             mod:`parcels.fieldset.FieldSet` object from which to sample velocity
-        pclass : parcels.particle.JITParticle or parcels.particle.ScipyParticle
-            Particle class. May be a particle class as defined in parcels, or a subclass defining a custom particle.
+        pclass :
+            Particle class. May be a parcels.particle.ScipyParticle as defined in parcels, or a subclass defining a custom particle.
         start :
             Start point (longitude, latitude) for initialisation of particles on a straight line.
         finish :
@@ -653,8 +652,8 @@ def from_field(
         ----------
         fieldset : parcels.fieldset.FieldSet
             mod:`parcels.fieldset.FieldSet` object from which to sample velocity
-        pclass : parcels.particle.JITParticle or parcels.particle.ScipyParticle
-            Particle class. May be a particle class as defined in parcels, or a subclass defining a custom particle.
+        pclass :
+            Particle class. May be a parcels.particle.ScipyParticle class as defined in parcels, or a subclass defining a custom particle.
         start_field : parcels.field.Field
             Field for initialising particles stochastically (horizontally)  according to the presented density field.
         size :
@@ -697,8 +696,8 @@ def from_particlefile(
         ----------
         fieldset : parcels.fieldset.FieldSet
             mod:`parcels.fieldset.FieldSet` object from which to sample velocity
-        pclass : parcels.particle.JITParticle or parcels.particle.ScipyParticle
-            Particle class. May be a particle class as defined in parcels, or a subclass defining a custom particle.
+        pclass :
+            Particle class. May be a parcels.particle.ScipyParticle class as defined in parcels, or a subclass defining a custom particle.
         filename : str
             Name of the particlefile from which to read initial conditions
         restart : bool