Simplify Passive Tracers

Author: Romain Beucher

UWGeodynamics/_model.py

@@ -1739,8 +1739,9 @@ def _update(self):
             self._update_stress_history(dt)
 
         if self.passive_tracers:
-            for key in self.passive_tracers:
-                self.passive_tracers[key].integrate(dt)
+            for key, val in self.passive_tracers.items():
+                if val.advector:
+                    val.advector.integrate(dt)
 
         # Do pop control
         self.population_control.repopulate()
@@ -1769,7 +1770,8 @@ def mesh_advector(self, axis):
         self._advector = Mesh_advector(self, axis)
 
     def add_passive_tracers(self, name, vertices=None,
-                            particleEscape=True, centroids=None, zOnly=False):
+                            particleEscape=True, centroids=None,
+                            advect=True):
         """ Add a swarm of passive tracers to the Model
 
         Parameters:
@@ -1832,10 +1834,8 @@ def add_passive_tracers(self, name, vertices=None,
         if not centroids:
 
             tracers = PassiveTracers(self.mesh,
-                                     self.velocityField,
                                      name=name,
-                                     particleEscape=particleEscape,
-                                     zOnly=zOnly)
+                                     particleEscape=particleEscape)
             tracers.add_particles_with_coordinates(vertices)
 
         else:
@@ -1857,15 +1857,16 @@ def add_passive_tracers(self, name, vertices=None,
                 vertices[:, 2] = z
 
             tracers = PassiveTracers(self.mesh,
-                                     self.velocityField,
                                      name=name,
-                                     particleEscape=particleEscape,
-                                     zOnly=zOnly)
+                                     particleEscape=particleEscape)
             tracers.add_particles_with_coordinates(vertices)
 
+        if advect:
+            tracers.advector = uw.systems.SwarmAdvector(
+                self.velocityField, tracers, order=2)
+
         self.passive_tracers[name] = tracers
         setattr(self, name.lower() + "_tracers", tracers)
-
         return tracers
 
     def _get_melt_fraction(self):
@@ -2696,7 +2697,6 @@ def reload_restart_variables(self, step):
                 sys.stdout.flush()
 
     def reload_passive_tracers(self, step):
-        
 
         Model = self.Model
 
@@ -2707,11 +2707,11 @@ def reload_passive_tracers(self, step):
             tracked_fields = tracer.tracked_fields
 
             obj = PassiveTracers(Model.mesh,
-                                 Model.velocityField,
                                  tracer.name,
-                                 zOnly=tracer.zOnly,
                                  particleEscape=tracer.particleEscape)
             obj.load(fpath)
+            if tracer.advector:
+                tracer.advector = uw.systems.SwarmAdvector(Model.velocityField, tracer, order=2)
 
             # Reload global indices
             fpath = os.path.join(self.restartDir, tracer.name + '_global_index-%s.h5' % step)

UWGeodynamics/_utils.py

@@ -69,19 +69,24 @@ def smooth(self):
 
 class PassiveTracers(Swarm):
 
-    def __init__(self, mesh, velocityField, name=None,
-                 particleEscape=True, zOnly=False):
+    def __init__(self, mesh, name=None, particleEscape=True):
 
         super(PassiveTracers, self).__init__(mesh,
                                              particleEscape=particleEscape)
         self.name = name
-        self.velocityField = velocityField
-        self.angular_velocity = 0.5 * (self.velocityField.fn_gradient[1] -
-                                       self.velocityField.fn_gradient[2])
         self.particleEscape = particleEscape
-        self.zOnly = zOnly
-
         self.tracked_fields = {}
+        self._advector = None
+
+    @property
+    def advector(self):
+        return self._advector
+
+    @advector.setter
+    def advector(self, uw_advector):
+        if not isinstance(uw_advector, uw.systems.SwarmAdvector):
+            raise ValueError(""" The advector must be an Underworld SwarmAdvector object""")
+        self._advector = uw_advector
 
     def _global_indices(self):
         indices = np.arange(self.particleLocalCount)
@@ -93,41 +98,12 @@ def _global_indices(self):
         self.global_index.data[:, 0] = pairs.view(np.int64)
 
     def add_particles_with_coordinates(self, vertices, **kwargs):
-
         vals = super(PassiveTracers,
                      self).add_particles_with_coordinates(nd(vertices),
                                                           **kwargs)
-        self.advector = uw.systems.SwarmAdvector(
-            swarm=self,
-            velocityField=self.velocityField, order=2)
-
         self._global_indices()
         return vals
 
-    def integrate(self, dt, **kwargs):
-        """ Integrate swarm velocity in time """
-        if self.zOnly:
-            saved_velocities = np.copy(self.velocityField.data)
-            self.velocityField.data[:, :-1] = 0.
-            self.velocityField.syncronise()
-            self.advector.integrate(dt, **kwargs)
-            self.velocityField.data[...] = saved_velocities
-            self.velocityField.syncronise()
-        else:
-            self.advector.integrate(dt, **kwargs)
-
-        # Integrate tracked field variables over time
-        for name, field in self.tracked_fields.items():
-            if field["timeIntegration"]:
-                obj = getattr(self, name)
-                if self.mesh.dim == 2 and obj.data.shape[-1] > 2:
-                    ang_vel = self.angular_velocity
-                    dtheta = dt * ang_vel.evaluate(self).reshape(1, -1)
-                    rotateTensor2D(obj.data[:, 0:3], dtheta)
-                    obj.data[:, 0:3] += field["value"].evaluate(self)
-                else:
-                    obj.data[...] += field["value"].evaluate(self) * dt
-
     def add_tracked_field(self, value, name, units, dataType, count=1,
                           overwrite=True, timeIntegration=False):
         """ Add a field to be tracked """
@@ -206,7 +182,6 @@ def save(self, outputDir, checkpointID, time):
         comm.Barrier()
 
         # get swarm parameters - serially read from hdf5 file to get size
-
         if rank == 0:
             with h5py.File(name=swarm_fpath, mode="r") as h5f:
                 dset = h5f.get('data')