Particles seem to work withswe

Author: harpolea

advection/simulation.py

@@ -127,8 +127,12 @@ def dovis(self):
         if self.particles is not None:
             particle_positions = self.particles.get_positions()
 
+            # dye particles
+            colors = self.particles.get_init_positions()[:,0]
+
             # plot particles
-            ax.scatter(particle_positions[:, 0], particle_positions[:, 1])
+            ax.scatter(particle_positions[:, 0],
+                particle_positions[:, 1], c=colors, cmap="Greys")
             ax.set_xlim([myg.xmin, myg.xmax])
             ax.set_ylim([myg.ymin, myg.ymax])
 

incompressible/simulation.py

@@ -445,12 +445,15 @@ def dovis(self):
 
             plt.colorbar(img, ax=ax)
 
-        ax = axes.flat[0]
         if self.particles is not None:
+            ax = axes.flat[0]
             particle_positions = self.particles.get_positions()
+            # dye particles
+            colors = self.particles.get_init_positions()[:, 0]
 
             # plot particles
-            ax.scatter(particle_positions[:, 0], particle_positions[:, 1], s=5)
+            ax.scatter(particle_positions[:, 0],
+                particle_positions[:, 1], s=5, c=colors, alpha=0.8, cmap="Greys")
             ax.set_xlim([myg.xmin, myg.xmax])
             ax.set_ylim([myg.ymin, myg.ymax])
 

particles/particles.py

@@ -59,7 +59,7 @@ def __init__(self, sim_data, bc, rp):
 
         self.sim_data = sim_data
         self.bc = bc
-        self.particles = set()
+        self.particles = dict()
         self.rp = rp
 
         # TODO: read something from rp here to determine how to
@@ -94,7 +94,8 @@ def randomly_generate_particles(self, n_particles):
         positions[:, 1] = positions[:, 1] * (myg.ymax - myg.ymin) + \
             myg.ymin
 
-        self.particles = set([Particle(x, y) for (x, y) in positions])
+        for (x, y) in positions:
+            self.particles[(x, y)] = Particle(x, y)
 
     def grid_generate_particles(self, n_particles):
         """
@@ -105,8 +106,6 @@ def grid_generate_particles(self, n_particles):
         """
         sq_n_particles = int(round(np.sqrt(n_particles)))
 
-        print(f'sq_n_particles = {sq_n_particles}')
-
         if sq_n_particles**2 != n_particles:
             msg.warning("WARNING: Changing number of particles from {} to {}".format(n_particles, sq_n_particles**2))
 
@@ -116,10 +115,9 @@ def grid_generate_particles(self, n_particles):
         xs += 0.5 * step
         ys, step = np.linspace(myg.ymin, myg.ymax, num=sq_n_particles, endpoint=False, retstep=True)
         ys += 0.5 * step
-
-        print(f'xs = {xs}')
-
-        self.particles = set([Particle(x, y) for x in xs for y in ys])
+        for x in xs:
+            for y in ys:
+                self.particles[(x, y)] = Particle(x, y)
 
     def update_particles(self, u, v, dt, limiter=0):
         """
@@ -145,7 +143,7 @@ def update_particles(self, u, v, dt, limiter=0):
         ldelta_vx = reconstruction.limit(v, myg, 1, limiter)
         ldelta_vy = reconstruction.limit(v, myg, 2, limiter)
 
-        for p in self.particles:
+        for k, p in self.particles.items():
             # find what cell it lives in
             x_idx = (p.x - myg.xmin) / myg.dx - 0.5
             y_idx = (p.y - myg.ymin) / myg.dy - 0.5
@@ -199,8 +197,8 @@ def enforce_particle_boundaries(self):
         TODO: copying the set and adding everything back again is messy
         - think of a better way to do this?
         """
-        new_particles = self.particles.copy()
-        self.particles = set()
+        old_particles = self.particles.copy()
+        self.particles = dict()
 
         myg = self.sim_data.grid
 
@@ -209,8 +207,8 @@ def enforce_particle_boundaries(self):
         ylb = self.bc.ylb
         yrb = self.bc.yrb
 
-        while new_particles:
-            p = new_particles.pop()
+        while old_particles:
+            k, p = old_particles.popitem()
 
             # -x boundary
             if p.x < myg.xmin:
@@ -256,15 +254,22 @@ def enforce_particle_boundaries(self):
                 else:
                     msg.fail("ERROR: yrb = %s invalid BC" % (yrb))
 
-            self.particles.add(p)
+            self.particles[k] = p
 
         self.n_particles = len(self.particles)
 
     def get_positions(self):
         """
         Return an array of particle positions.
         """
-        return np.array([[p.x, p.y] for p in self.particles])
+        return np.array([[p.x, p.y] for p in self.particles.values()])
+
+    def get_init_positions(self):
+        """
+        We defined the particles as a dictionary with their initial positions
+        as the keys, so this just becomes a restructuring operation.
+        """
+        return np.array([[x, y] for (x,y) in self.particles.keys()])
 
     def write_particles(self, filename):
         """

particles/tests/test_particles.py

@@ -115,7 +115,7 @@ def test_particles_advect():
 
     ps.update_particles(u, v, 1)
 
-    positions = set([(p.x, p.y) for p in ps.particles])
+    positions = set([(p.x, p.y) for p in ps.particles.values()])
 
     xs, step = np.linspace(0, 1, num=7, endpoint=False, retstep=True)
     xs += 0.5 * step
@@ -129,7 +129,7 @@ def test_particles_advect():
 
     ps.update_particles(u, v, 1)
 
-    positions = set([(p.x, p.y) for p in ps.particles])
+    positions = set([(p.x, p.y) for p in ps.particles.values()])
 
     correct_positions = set([(x+1, y) for x in xs for y in xs])
 
@@ -142,7 +142,7 @@ def test_particles_advect():
     ps = particles.Particles(sim.cc_data, bc, rp)
     ps.update_particles(u, v, 1)
 
-    positions = set([(p.x, p.y) for p in ps.particles])
+    positions = set([(p.x, p.y) for p in ps.particles.values()])
 
     correct_positions = set([(x, y+1) for x in xs for y in xs])
 
@@ -154,7 +154,7 @@ def test_particles_advect():
     ps = particles.Particles(sim.cc_data, bc, rp)
     ps.update_particles(u, v, 1)
 
-    positions = set([(p.x, p.y) for p in ps.particles])
+    positions = set([(p.x, p.y) for p in ps.particles.values()])
 
     correct_positions = set([(x+1, y+1) for x in xs for y in xs])
 

swe/_defaults

@@ -12,3 +12,8 @@ limiter = 2               ; limiter (0 = none, 1 = 2nd order, 2 = 4th order)
 grav = 1.0                ; gravitational acceleration (in y-direction)
 
 riemann = Roe            ; HLLC or Roe
+
+
+[particles]
+do_particles = 1
+particle_generator = grid

swe/simulation.py

@@ -10,6 +10,7 @@
 import mesh.boundary as bnd
 from simulation_null import NullSimulation, grid_setup, bc_setup
 import util.plot_tools as plot_tools
+import particles.particles as particles
 
 
 class Variables(object):
@@ -129,6 +130,9 @@ def initialize(self, extra_vars=None, ng=4):
 
         self.cc_data = my_data
 
+        if self.rp.get_param("particles.do_particles") == 1:
+            self.particles = particles.Particles(self.cc_data, bc, self.rp)
+
         # some auxillary data that we'll need to fill GC in, but isn't
         # really part of the main solution
         aux_data = self.data_class(my_grid)
@@ -195,6 +199,13 @@ def evolve(self):
                 dtdx*(Flux_x.v(n=n) - Flux_x.ip(1, n=n)) + \
                 dtdy*(Flux_y.v(n=n) - Flux_y.jp(1, n=n))
 
+        if self.particles is not None:
+            limiter = self.rp.get_param("swe.limiter")
+            u, v = self.cc_data.get_var("velocity")
+
+            self.particles.update_particles(u, v, self.dt, limiter)
+            self.particles.enforce_particle_boundaries()
+
         # increment the time
         self.cc_data.t += self.dt
         self.n += 1
@@ -262,6 +273,18 @@ def dovis(self):
             else:
                 ax.set_title(field_names[n])
 
+        if self.particles is not None:
+            ax = axes[0]
+            particle_positions = self.particles.get_positions()
+            # dye particles
+            colors = self.particles.get_init_positions()[:, 0]
+
+            # plot particles
+            ax.scatter(particle_positions[:, 0],
+                particle_positions[:, 1], s=5, c=colors, alpha=0.8, cmap="Greys")
+            ax.set_xlim([myg.xmin, myg.xmax])
+            ax.set_ylim([myg.ymin, myg.ymax])
+
         plt.figtext(0.05, 0.0125, "t = {:10.5g}".format(self.cc_data.t))
 
         plt.pause(0.001)