Some assorted JIT cleanup

Author: erikvansebille

docs/examples/tutorial_parcels_structure.ipynb

@@ -25,7 +25,7 @@
    "source": [
     "1. [**FieldSet**](#1.-FieldSet). Load and set up the fields. These can be velocity fields that are used to advect the particles, but it can also be e.g. temperature.\n",
     "2. [**ParticleSet**](#2.-ParticleSet). Define the type of particles. Also additional `Variables` can be added to the particles (e.g. temperature, to keep track of the temperature that particles experience).\n",
-    "3. [**Kernels**](#3.-Kernels). Define and compile kernels. Kernels perform some specific operation on the particles every time step (e.g. interpolate the temperature from the temperature field to the particle location).\n",
+    "3. [**Kernels**](#3.-Kernels). Kernels perform some specific operation on the particles every time step (e.g. interpolate the temperature from the temperature field to the particle location).\n",
     "4. [**Execution and output**](#4.-Execution-and-Output). Execute the simulation and write and store the output in a Zarr file.\n",
     "5. [**Optimising and parallelising**](#5.-Optimising-and-parallelising). Optimise and parallelise the code to run faster.\n",
     "\n",

parcels/field.py

@@ -149,7 +149,6 @@ class Field:
         Maximum allowed value on the field. Data above this value are set to zero
     cast_data_dtype : str
         Cast Field data to dtype. Supported dtypes are "float32" (np.float32 (default)) and "float64 (np.float64).
-        Note that dtype can only be "float32" in JIT mode
     time_origin : parcels.tools.converters.TimeConverter
         Time origin of the time axis (only if grid is None)
     interp_method : str

parcels/fieldset.py

@@ -1447,8 +1447,7 @@ def get_fields(self) -> list[Field | VectorField]:
 
     def add_constant(self, name, value):
         """Add a constant to the FieldSet. Note that all constants are
-        stored as 32-bit floats. While constants can be updated during
-        execution in SciPy mode, they can not be updated in JIT mode.
+        stored as 32-bit floats.
 
         Parameters
         ----------

parcels/kernel.py

@@ -94,7 +94,7 @@ class Kernel(BaseKernel):
 
     Notes
     -----
-    A Kernel is either created from a compiled <function ...> object
+    A Kernel is either created from a <function ...> object
     or the necessary information (funcname, funccode, funcvars) is provided.
     The py_ast argument may be derived from the code string, but for
     concatenation, the merged AST plus the new header definition is required.
@@ -159,7 +159,7 @@ def __init__(
                 user_ctx = globals()
             finally:
                 del stack  # Remove cyclic references
-            # Compile and generate Python function from AST
+            # Generate Python function from AST
             py_mod = ast.parse("")
             py_mod.body = [self.py_ast]
             exec(compile(py_mod, "<ast>", "exec"), user_ctx)

parcels/particleset.py

@@ -932,13 +932,13 @@ def execute(
 
         Notes
         -----
-        ``ParticleSet.execute()`` acts as the main entrypoint for simulations, and provides the simulation time-loop. This method encapsulates the logic controlling the switching between kernel execution (where control in handed to C in JIT mode), output file writing, reading in fields for new timesteps, adding new particles to the simulation domain, stopping the simulation, and executing custom functions (``postIterationCallbacks`` provided by the user).
+        ``ParticleSet.execute()`` acts as the main entrypoint for simulations, and provides the simulation time-loop. This method encapsulates the logic controlling the switching between kernel execution, output file writing, reading in fields for new timesteps, adding new particles to the simulation domain, stopping the simulation, and executing custom functions (``postIterationCallbacks`` provided by the user).
         """
         # check if particleset is empty. If so, return immediately
         if len(self) == 0:
             return
 
-        # check if pyfunc has changed since last compile. If so, recompile
+        # check if pyfunc has changed since last generation. If so, regenerate
         if self._kernel is None or (self._kernel.pyfunc is not pyfunc and self._kernel is not pyfunc):
             # Generate and store Kernel
             if isinstance(pyfunc, Kernel):

tests/test_grids.py

@@ -255,7 +255,7 @@ def bath_func(lon):
     for i in range(depth_g0.shape[0]):
         for k in range(depth_g0.shape[2]):
             depth_g0[i, :, k] = bath[i] * k / (depth_g0.shape[2] - 1)
-    depth_g0 = depth_g0.transpose()  # we don't change it on purpose, to check if the transpose op if fixed in jit
+    depth_g0 = depth_g0.transpose()
 
     grid = RectilinearSGrid(lon_g0, lat_g0, depth=depth_g0)
 

tests/test_interpolation.py

@@ -90,10 +90,7 @@ def data_2d():
     ],
 )
 def test_raw_2d_interpolation(data_2d, func, eta, xsi, expected):
-    """Test the 2D interpolation functions on the raw arrays.
-
-    Interpolation via the other interpolation methods are tested in `test_scipy_vs_jit`.
-    """
+    """Test the 2D interpolation functions on the raw arrays."""
     ti = 0
     yi, xi = 1, 1
     ctx = interpolation.InterpolationContext2D(data_2d, eta, xsi, ti, yi, xi)

tests/test_particles.py

@@ -47,7 +47,7 @@ def addOne(particle, fieldset, time):  # pragma: no cover
 
 @pytest.mark.parametrize("type", ["np.int8", "mp.float", "np.int16"])
 def test_variable_unsupported_dtypes(fieldset, type):
-    """Test that checks errors thrown for unsupported dtypes in JIT mode."""
+    """Test that checks errors thrown for unsupported dtypes."""
     TestParticle = ScipyParticle.add_variable("p", dtype=type, initial=10.0)
     with pytest.raises((RuntimeError, TypeError)):
         ParticleSet(fieldset, pclass=TestParticle, lon=[0], lat=[0])