Merge pull request #1614 from OceanParcels/code_cleanups

Code cleanups

Author: erikvansebille

parcels/application_kernels/advection.py

@@ -251,21 +251,21 @@ def compute_ds(F0, F1, r, direction, tol):
     s_min = min(abs(ds_x), abs(ds_y), abs(ds_z), abs(ds_t / (dxdy * dz)))
 
     # calculate end position in time s_min
-    def compute_rs(ds, r, B, delta, s_min):
+    def compute_rs(r, B, delta, s_min):
         if abs(B) < tol:
             return -delta * s_min + r
         else:
             return (r + delta / B) * math.exp(-B * s_min) - delta / B
 
-    rs_x = compute_rs(ds_x, xsi, B_x, delta_x, s_min)
-    rs_y = compute_rs(ds_y, eta, B_y, delta_y, s_min)
+    rs_x = compute_rs(xsi, B_x, delta_x, s_min)
+    rs_y = compute_rs(eta, B_y, delta_y, s_min)
 
-    particle_dlon = (1.-rs_x)*(1.-rs_y) * px[0] + rs_x * (1.-rs_y) * px[1] + rs_x * rs_y * px[2] + (1.-rs_x)*rs_y * px[3] - particle.lon  # noqa
-    particle_dlat = (1.-rs_x)*(1.-rs_y) * py[0] + rs_x * (1.-rs_y) * py[1] + rs_x * rs_y * py[2] + (1.-rs_x)*rs_y * py[3] - particle.lat  # noqa
+    particle_dlon += (1.-rs_x)*(1.-rs_y) * px[0] + rs_x * (1.-rs_y) * px[1] + rs_x * rs_y * px[2] + (1.-rs_x)*rs_y * px[3] - particle.lon  # noqa
+    particle_dlat += (1.-rs_x)*(1.-rs_y) * py[0] + rs_x * (1.-rs_y) * py[1] + rs_x * rs_y * py[2] + (1.-rs_x)*rs_y * py[3] - particle.lat  # noqa
 
     if withW:
-        rs_z = compute_rs(ds_z, zeta, B_z, delta_z, s_min)
-        particle.depth = (1.-rs_z) * pz[0] + rs_z * pz[1]
+        rs_z = compute_rs(zeta, B_z, delta_z, s_min)
+        particle_ddepth += (1.-rs_z) * pz[0] + rs_z * pz[1] - particle.depth  # noqa
 
     if particle.dt > 0:
         particle.dt = max(direction * s_min * (dxdy * dz), 1e-7)

parcels/compilation/codegenerator.py

@@ -657,7 +657,8 @@ def visit_Subscript(self, node):
         if isinstance(node.value, FieldNode) or isinstance(node.value, VectorFieldNode):
             node.ccode = node.value.__getitem__(node.slice.ccode).ccode
         elif isinstance(node.value, ParticleXiYiZiTiAttributeNode):
-            node.ccode = f"{node.value.obj}->{node.value.attr}[pnum, {node.slice.ccode}]"
+            ngrid = str(self.fieldset.gridset.size if self.fieldset is not None else 1)
+            node.ccode = f"{node.value.obj}->{node.value.attr}[pnum*{ngrid}+{node.slice.ccode}]"
         elif isinstance(node.value, IntrinsicNode):
             raise NotImplementedError(f"Subscript not implemented for object type {type(node.value).__name__}")
         else:

parcels/kernel.py

@@ -55,8 +55,6 @@ def __init__(self, fieldset, ptype, pyfunc=None, funcname=None, funccode=None, p
         self._ptype = ptype
         self._lib = None
         self.delete_cfiles = delete_cfiles
-        self._cleanup_files = None
-        self._cleanup_lib = None
         self._c_include = c_include
 
         # Derive meta information from pyfunc, if not given

tests/test_particlefile.py

@@ -284,7 +284,7 @@ def Update_lon(particle, fieldset, time):
 def test_write_xiyi(fieldset, mode, tmpdir):
     outfilepath = tmpdir.join("pfile_xiyi.zarr")
     fieldset.U.data[:] = 1  # set a non-zero zonal velocity
-    fieldset.add_field(Field(name='P', data=np.zeros((2, 20)), lon=np.linspace(0, 1, 20), lat=[0, 2]))
+    fieldset.add_field(Field(name='P', data=np.zeros((3, 20)), lon=np.linspace(0, 1, 20), lat=[-2, 0, 2]))
     dt = 3600
 
     XiYiParticle = ptype[mode].add_variables([
@@ -304,26 +304,29 @@ def Get_XiYi(particle, fieldset, time):
 
     def SampleP(particle, fieldset, time):
         if time > 5*3600:
-            tmp = fieldset.P[particle]  # noqa
+            _ = fieldset.P[particle]  # To trigger sampling of the P field
 
-    pset = ParticleSet(fieldset, pclass=XiYiParticle, lon=[0], lat=[0.2], lonlatdepth_dtype=np.float64)
+    pset = ParticleSet(fieldset, pclass=XiYiParticle, lon=[0, 0.2], lat=[0.2, 1], lonlatdepth_dtype=np.float64)
     pfile = pset.ParticleFile(name=outfilepath, outputdt=dt)
-    pset.execute([Get_XiYi, SampleP, AdvectionRK4], endtime=10*dt, dt=dt, output_file=pfile)
+    pset.execute([SampleP, Get_XiYi, AdvectionRK4], endtime=10*dt, dt=dt, output_file=pfile)
 
     ds = xr.open_zarr(outfilepath)
-    pxi0 = ds['pxi0'][:].values[0].astype(np.int32)
-    pxi1 = ds['pxi1'][:].values[0].astype(np.int32)
-    lons = ds['lon'][:].values[0]
-    pyi = ds['pyi'][:].values[0].astype(np.int32)
-    lats = ds['lat'][:].values[0]
-
-    assert (pxi0[0] == 0) and (pxi0[-1] == 11)  # check that particle has moved
-    assert np.all(pxi1[:7] == 0)  # check that particle has not been sampled on grid 1 until time 6
-    assert np.all(pxi1[7:] > 0)  # check that particle has not been sampled on grid 1 after time 6
-    for xi, lon in zip(pxi0[1:], lons[1:]):
-        assert fieldset.U.grid.lon[xi] <= lon < fieldset.U.grid.lon[xi+1]
-    for yi, lat in zip(pyi[1:], lats[1:]):
-        assert fieldset.U.grid.lat[yi] <= lat < fieldset.U.grid.lat[yi+1]
+    pxi0 = ds['pxi0'][:].values.astype(np.int32)
+    pxi1 = ds['pxi1'][:].values.astype(np.int32)
+    lons = ds['lon'][:].values
+    pyi = ds['pyi'][:].values.astype(np.int32)
+    lats = ds['lat'][:].values
+
+    for p in range(pyi.shape[0]):
+        assert (pxi0[p, 0] == 0) and (pxi0[p, -1] == pset[p].pxi0)  # check that particle has moved
+        assert np.all(pxi1[p, :6] == 0)  # check that particle has not been sampled on grid 1 until time 6
+        assert np.all(pxi1[p, 6:] > 0)  # check that particle has not been sampled on grid 1 after time 6
+        for xi, lon in zip(pxi0[p, 1:], lons[p, 1:]):
+            assert fieldset.U.grid.lon[xi] <= lon < fieldset.U.grid.lon[xi+1]
+        for xi, lon in zip(pxi1[p, 6:], lons[p, 6:]):
+            assert fieldset.P.grid.lon[xi] <= lon < fieldset.P.grid.lon[xi+1]
+        for yi, lat in zip(pyi[p, 1:], lats[p, 1:]):
+            assert fieldset.U.grid.lat[yi] <= lat < fieldset.U.grid.lat[yi+1]
     ds.close()