List of Particle Arrays structure for ParticleSet - DO NOT MERGE - DO NOT CLOSE

parcels/kernel.py

@@ -219,9 +219,13 @@ def load_lib(self):
 
     def execute_jit(self, pset, endtime, dt):
         """Invokes JIT engine to perform the core update loop"""
-        if len(pset.particles) > 0:
-            assert pset.fieldset.gridset.size == len(pset.particles[0].xi), \
+        if len(pset.particles_a) > 0 and pset.particles_a[0].shape[0]:
+            assert pset.fieldset.gridset.size == len(pset.particles_a[0][0].xi), \
                 'FieldSet has different amount of grids than Particle.xi. Have you added Fields after creating the ParticleSet?'
+
+        # if len(pset.particles) > 0:
+        #     assert pset.fieldset.gridset.size == len(pset.particles[0].xi), \
+        #         'FieldSet has different amount of grids than Particle.xi. Have you added Fields after creating the ParticleSet?'
         for g in pset.fieldset.gridset.grids:
             g.cstruct = None  # This force to point newly the grids from Python to C
         # Make a copy of the transposed array to enforce
@@ -250,11 +254,18 @@ def execute_jit(self, pset, endtime, dt):
 
         fargs = [byref(f.ctypes_struct) for f in self.field_args.values()]
         fargs += [c_double(f) for f in self.const_args.values()]
-        particle_data = pset._particle_data.ctypes.data_as(c_void_p)
-        self._function(c_int(len(pset)), particle_data,
-                       c_double(endtime),
-                       c_double(dt),
-                       *fargs)
+
+        for subfield_c in pset.particles_c:
+            particle_data = subfield_c.ctypes.data_as(c_void_p)
+            self._function(c_int(subfield_c.shape[0]), particle_data,
+                           c_double(endtime), c_double(dt),
+                           *fargs)
+
+        # particle_data = pset._particle_data.ctypes.data_as(c_void_p)
+        # self._function(c_int(len(pset)), particle_data,
+        #                c_double(endtime),
+        #                c_double(dt),
+        #                *fargs)
 
     def execute_python(self, pset, endtime, dt):
         """Performs the core update loop via Python"""
@@ -268,104 +279,122 @@ def execute_python(self, pset, endtime, dt):
                 continue
             f.data = np.array(f.data)
 
-        for p in pset.particles:
-            ptype = p.getPType()
-            # Don't execute particles that aren't started yet
-            sign_end_part = np.sign(endtime - p.time)
+        # for p in pset.particles:
+        for subfield in pset.particles_a:
+            for p in subfield:
+                ptype = p.getPType()
+                # Don't execute particles that aren't started yet
+                sign_end_part = np.sign(endtime - p.time)
 
-            dt_pos = min(abs(p.dt), abs(endtime - p.time))
+                dt_pos = min(abs(p.dt), abs(endtime - p.time))
 
-            # ==== numerically stable; also making sure that continuously-recovered particles do end successfully,
-            # as they fulfil the condition here on entering at the final calculation here. ==== #
-            if ((sign_end_part != sign_dt) or np.isclose(dt_pos, 0)) and not np.isclose(dt, 0):
-                if abs(p.time) >= abs(endtime):
-                    p.state = ErrorCode.Success
-                continue
+                # ==== numerically stable; also making sure that continuously-recovered particles do end successfully,
+                # as they fulfil the condition here on entering at the final calculation here. ==== #
+                if ((sign_end_part != sign_dt) or np.isclose(dt_pos, 0)) and not np.isclose(dt, 0):
+                    if abs(p.time) >= abs(endtime):
+                        p.state = ErrorCode.Success
+                    continue
 
-            while p.state in [ErrorCode.Evaluate, ErrorCode.Repeat] or np.isclose(dt, 0):
-
-                for var in ptype.variables:
-                    p_var_back[var.name] = getattr(p, var.name)
-                try:
-                    pdt_prekernels = sign_dt * dt_pos
-                    p.dt = pdt_prekernels
-                    state_prev = p.state
-                    res = self.pyfunc(p, pset.fieldset, p.time)
-                    if res is None:
-                        res = ErrorCode.Success
-
-                    if res is ErrorCode.Success and p.state != state_prev:
-                        res = p.state
-
-                    if res == ErrorCode.Success and not np.isclose(p.dt, pdt_prekernels):
-                        res = ErrorCode.Repeat
-
-                except FieldOutOfBoundError as fse_xy:
-                    res = ErrorCode.ErrorOutOfBounds
-                    p.exception = fse_xy
-                except FieldOutOfBoundSurfaceError as fse_z:
-                    res = ErrorCode.ErrorThroughSurface
-                    p.exception = fse_z
-                except TimeExtrapolationError as fse_t:
-                    res = ErrorCode.ErrorTimeExtrapolation
-                    p.exception = fse_t
-                except Exception as e:
-                    res = ErrorCode.Error
-                    p.exception = e
-
-                # Handle particle time and time loop
-                if res in [ErrorCode.Success, ErrorCode.Delete]:
-                    # Update time and repeat
-                    p.time += p.dt
-                    p.update_next_dt()
-                    dt_pos = min(abs(p.dt), abs(endtime - p.time))
-
-                    sign_end_part = np.sign(endtime - p.time)
-                    if res != ErrorCode.Delete and not np.isclose(dt_pos, 0) and (sign_end_part == sign_dt):
-                        res = ErrorCode.Evaluate
-                    if sign_end_part != sign_dt:
-                        dt_pos = 0
-
-                    p.state = res
-                    if np.isclose(dt, 0):
-                        break
-                else:
-                    p.state = res
-                    # Try again without time update
-                    for var in ptype.variables:
-                        if var.name not in ['dt', 'state']:
-                            setattr(p, var.name, p_var_back[var.name])
-                    dt_pos = min(abs(p.dt), abs(endtime - p.time))
+                while p.state in [ErrorCode.Evaluate, ErrorCode.Repeat] or np.isclose(dt, 0):
 
-                    sign_end_part = np.sign(endtime - p.time)
-                    if sign_end_part != sign_dt:
-                        dt_pos = 0
-                    break
+                    for var in ptype.variables:
+                        p_var_back[var.name] = getattr(p, var.name)
+                    try:
+                        pdt_prekernels = sign_dt * dt_pos
+                        p.dt = pdt_prekernels
+                        state_prev = p.state
+                        res = self.pyfunc(p, pset.fieldset, p.time)
+                        if res is None:
+                            res = ErrorCode.Success
+
+                        if res is ErrorCode.Success and p.state != state_prev:
+                            res = p.state
+
+                        if res == ErrorCode.Success and not np.isclose(p.dt, pdt_prekernels):
+                            res = ErrorCode.Repeat
+
+                    except FieldOutOfBoundError as fse_xy:
+                        res = ErrorCode.ErrorOutOfBounds
+                        p.exception = fse_xy
+                    except FieldOutOfBoundSurfaceError as fse_z:
+                        res = ErrorCode.ErrorThroughSurface
+                        p.exception = fse_z
+                    except TimeExtrapolationError as fse_t:
+                        res = ErrorCode.ErrorTimeExtrapolation
+                        p.exception = fse_t
+                    except Exception as e:
+                        res = ErrorCode.Error
+                        p.exception = e
+
+                    # Handle particle time and time loop
+                    if res in [ErrorCode.Success, ErrorCode.Delete]:
+                        # Update time and repeat
+                        p.time += p.dt
+                        p.update_next_dt()
+                        dt_pos = min(abs(p.dt), abs(endtime - p.time))
+
+                        sign_end_part = np.sign(endtime - p.time)
+                        if res != ErrorCode.Delete and not np.isclose(dt_pos, 0) and (sign_end_part == sign_dt):
+                            res = ErrorCode.Evaluate
+                        if sign_end_part != sign_dt:
+                            dt_pos = 0
+
+                        p.state = res
+                        if np.isclose(dt, 0):
+                            break
+                    else:
+                        p.state = res
+                        # Try again without time update
+                        for var in ptype.variables:
+                            if var.name not in ['dt', 'state']:
+                                setattr(p, var.name, p_var_back[var.name])
+                        dt_pos = min(abs(p.dt), abs(endtime - p.time))
+
+                        sign_end_part = np.sign(endtime - p.time)
+                        if sign_end_part != sign_dt:
+                            dt_pos = 0
+                        break
 
     def remove_deleted(self, pset, output_file, endtime):
-        """Utility to remove all particles that signalled deletion"""
-        indices = [i for i, p in enumerate(pset.particles)
-                   if p.state in [ErrorCode.Delete]]
-        if len(indices) > 0 and output_file is not None:
-            output_file.write(pset[indices], endtime, deleted_only=True)
-        pset.remove(indices)
+            """Utility to remove all particles that signalled deletion"""
+            indices = []
+            dparticles = []
+            for subfield in pset.particles_a:
+                # del_items = [(i, p) for i, p in enumerate(subfield) if p.state in [ErrorCode.Delete]]     # we could unzip that, but in Python <3.6, more than 255 elem would make it crash
+                local_indices = []
+                local_particles = []
+                for i, p in enumerate(subfield):
+                    if p.state == ErrorCode.Delete:
+                        local_indices.append(i)
+                        local_particles.append(p)
+                if len(local_particles) > 0:
+                    local_particles = np.array(local_particles)
+                    local_indices = np.array(local_indices)
+                else:
+                    local_particles = None
+                    local_indices = None
+                indices.append(local_indices)
+                dparticles.append(local_particles)
+            # ==== do / fix ParticleFile TODO ==== #
+            if output_file is not None:
+                output_file.write(dparticles, endtime, deleted_only=True)    # => individual entries being None -> handled in ParticleFile
+            # indices = [i for i, p in enumerate(pset.particles) if p.state in [ErrorCode.Delete]]
+            # if len(indices) > 0 and output_file is not None:
+            #     output_file.write(pset[indices], endtime, deleted_only=True)
+            pset.remove(indices)
 
     def execute(self, pset, endtime, dt, recovery=None, output_file=None, execute_once=False):
         """Execute this Kernel over a ParticleSet for several timesteps"""
-        for p in pset.particles:
-            p.reset_state()
+        for subfield in pset.particles_a:
+            for p in subfield:
+                p.reset_state()
+
+        # for p in pset.particles:
+        #     p.reset_state()
 
         if abs(dt) < 1e-6 and not execute_once:
             logger.warning_once("'dt' is too small, causing numerical accuracy limit problems. Please chose a higher 'dt' and rather scale the 'time' axis of the field accordingly. (related issue #762)")
 
-        # def remove_deleted(pset):
-        #     """Utility to remove all particles that signalled deletion"""
-        #     indices = [i for i, p in enumerate(pset.particles)
-        #                if p.state in [ErrorCode.Delete]]
-        #     if len(indices) > 0 and output_file is not None:
-        #         output_file.write(pset[indices], endtime, deleted_only=True)
-        #     pset.remove(indices)
-
         if recovery is None:
             recovery = {}
         elif ErrorCode.ErrorOutOfBounds in recovery and ErrorCode.ErrorThroughSurface not in recovery:
@@ -386,8 +415,8 @@ def execute(self, pset, endtime, dt, recovery=None, output_file=None, execute_on
         # Remove all particles that signalled deletion
         self.remove_deleted(pset, output_file=output_file, endtime=endtime)
 
-        # Identify particles that threw errors
-        error_particles = [p for p in pset.particles if p.state not in [ErrorCode.Success, ErrorCode.Evaluate]]
+        # == Identify particles that threw errors == #
+        error_particles = [p for subfield in pset.particles_a for p in subfield if p.state not in [ErrorCode.Success, ErrorCode.Evaluate]]
 
         while len(error_particles) > 0:
             # Apply recovery kernel
@@ -397,7 +426,6 @@ def execute(self, pset, endtime, dt, recovery=None, output_file=None, execute_on
                 if p.state == ErrorCode.Repeat:
                     p.reset_state()
                 elif p.state == ErrorCode.Delete:
-                    # p.delete()
                     pass
                 elif p.state in recovery_map:
                     recovery_kernel = recovery_map[p.state]
@@ -419,7 +447,7 @@ def execute(self, pset, endtime, dt, recovery=None, output_file=None, execute_on
             else:
                 self.execute_python(pset, endtime, dt)
 
-            error_particles = [p for p in pset.particles if p.state not in [ErrorCode.Success, ErrorCode.Evaluate]]
+            error_particles = [p for subfield in pset.particles_a for p in subfield if p.state not in [ErrorCode.Success, ErrorCode.Evaluate]]
 
     def merge(self, kernel):
         funcname = self.funcname + kernel.funcname