Wrote tests for the particles class

Author: harpolea

_defaults

@@ -42,3 +42,4 @@ ny = 25                   ; number of zones in the y-direction
 [particles]
 do_particles = 0          ; include particles? (1=yes, 0=no)
 n_particles = 100         ; number of particles
+particle_generator = random ; how do we generate particles? (random, grid)

advection/problems/inputs.smooth

@@ -36,3 +36,4 @@ limiter = 2
 
 [particles]
 do_particles = 1
+particle_generator = grid

advection/problems/inputs.tophat

@@ -35,3 +35,4 @@ limiter = 2
 
 [particles]
 do_particles = 1
+particle_generator = grid

advection/simulation.py

@@ -6,9 +6,9 @@
 import mesh.patch as patch
 from simulation_null import NullSimulation, grid_setup, bc_setup
 import particles.particles as particles
-from util import msg
 import util.plot_tools as plot_tools
 
+
 class Simulation(NullSimulation):
 
     def initialize(self):
@@ -28,12 +28,7 @@ def initialize(self):
         self.cc_data = my_data
 
         if self.rp.get_param("particles.do_particles") == 1:
-            n_particles = self.rp.get_param("particles.n_particles")
-            if n_particles > 0:
-                self.particles = particles.Particles(self.cc_data, bc, n_particles)
-            else:
-                msg.fail("ERROR: n_particles = %s < 0" % (n_particles))
-
+            self.particles = particles.Particles(self.cc_data, bc, self.rp)
 
         # now set the initial conditions for the problem
         problem = importlib.import_module("advection.problems.{}".format(self.problem_name))

advection/tests/test_advection.py

@@ -19,6 +19,7 @@ def setup_method(self):
 
         self.rp.params["mesh.nx"] = 8
         self.rp.params["mesh.ny"] = 8
+        self.rp.params["particles.do_particles"] = 0
 
         self.sim = sn.Simulation("advection", "test", self.rp)
         self.sim.initialize()

particles/particles.py

@@ -9,6 +9,12 @@
 
 
 class Particle(object):
+    """
+    Class to hold properties of a single particle.
+
+    Not sure need velocity (or mass?), but will store it
+    here for now.
+    """
 
     def __init__(self, x, y, u=0, v=0, mass=1):
         self.x = x
@@ -41,7 +47,7 @@ def advect(self, u, v, dt):
 
 class Particles(object):
 
-    def __init__(self, sim_data, bc, n_particles=100):
+    def __init__(self, sim_data, bc, rp):
         """
         Initialize the Particles object.
 
@@ -53,11 +59,26 @@ def __init__(self, sim_data, bc, n_particles=100):
 
         self.sim_data = sim_data
         self.bc = bc
+        self.particles = set()
+        self.rp = rp
 
         # TODO: read something from rp here to determine how to
         # generate the particles - for now, we shall assume random.
 
-        self.randomly_generate_particles(n_particles)
+        particle_generator = self.rp.get_param("particles.particle_generator")
+        n_particles = self.rp.get_param("particles.n_particles")
+        if n_particles <= 0:
+            msg.fail("ERROR: n_particles = %s <= 0" % (n_particles))
+
+        if particle_generator == "random":
+            self.randomly_generate_particles(n_particles)
+        elif particle_generator == "grid":
+            self.grid_generate_particles(n_particles)
+        else:
+            msg.fail("ERROR: do not recognise particle generator %s"
+                     % (particle_generator))
+
+        self.n_particles = len(self.particles)
 
     def randomly_generate_particles(self, n_particles):
         """
@@ -75,6 +96,31 @@ def randomly_generate_particles(self, n_particles):
 
         self.particles = set([Particle(x, y) for (x, y) in positions])
 
+    def grid_generate_particles(self, n_particles):
+        """
+        Generate particles equally spaced across the grid.
+
+        If necessary, shall increase/decrease n_particles
+        in order to
+        """
+        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))
+
+        myg = self.sim_data.grid
+
+        xs, step = np.linspace(myg.xmin, myg.xmax, num=sq_n_particles, endpoint=False, retstep=True)
+        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])
+
     def update_particles(self, u, v, dt, limiter=0):
         """
         Update the particles on the grid. To do this, we need to
@@ -89,7 +135,7 @@ def update_particles(self, u, v, dt, limiter=0):
         differently in different problems.
         """
         myg = self.sim_data.grid
-        myd = self.sim_data.data
+        # myd = self.sim_data.data
 
         # limit the velocity
 
@@ -120,7 +166,7 @@ def update_particles(self, u, v, dt, limiter=0):
             if x_idx >= myg.nx:
                 x_frac += (x_idx - myg.nx) + 1
                 x_idx = myg.nx - 1
-            if y_idx >=  myg.ny:
+            if y_idx >= myg.ny:
                 y_frac += (y_idx - myg.ny) + 1
                 y_idx = myg.ny - 1
 
@@ -167,47 +213,53 @@ def enforce_particle_boundaries(self):
             p = new_particles.pop()
 
             # -x boundary
-            if xlb == "outflow":
-                if p.x < myg.xmin:
+            if p.x < myg.xmin:
+                if xlb in ["outflow", "neumann"]:
                     continue
-            elif xlb == "periodic":
-                if p.x < myg.xmin:
+                elif xlb == "periodic":
                     p.x = myg.xmax + p.x - myg.xmin
-            else:
-                msg.fail("ERROR: xlb = %s invalid BC" % (xlb))
+                elif xlb in ["reflect-even", "reflect-odd"]:
+                    p.x = 2 * myg.xmin - p.x
+                else:
+                    msg.fail("ERROR: xlb = %s invalid BC" % (xlb))
 
             # +x boundary
-            if xrb == "outflow":
-                if p.x > myg.xmax:
+            if p.x > myg.xmax:
+                if xrb in ["outflow", "neumann"]:
                     continue
-            elif xrb == "periodic":
-                if p.x > myg.xmax:
+                elif xrb == "periodic":
                     p.x = myg.xmin + p.x - myg.xmax
-            else:
-                msg.fail("ERROR: xrb = %s invalid BC" % (xrb))
+                elif xrb in ["reflect-even", "reflect-odd"]:
+                    p.x = 2 * myg.xmax - p.x
+                else:
+                    msg.fail("ERROR: xrb = %s invalid BC" % (xrb))
 
             # -y boundary
-            if ylb == "outflow":
-                if p.y < myg.ymin:
+            if p.y < myg.ymin:
+                if ylb in ["outflow", "neumann"]:
                     continue
-            elif ylb == "periodic":
-                if p.y < myg.ymin:
+                elif ylb == "periodic":
                     p.y = myg.ymax + p.y - myg.ymin
-            else:
-                msg.fail("ERROR: ylb = %s invalid BC" % (ylb))
-
-            # +x boundary
-            if yrb == "outflow":
-                if p.y > myg.ymax:
+                elif ylb in ["reflect-even", "reflect-odd"]:
+                    p.y = 2 * myg.ymin - p.y
+                else:
+                    msg.fail("ERROR: ylb = %s invalid BC" % (ylb))
+
+            # +y boundary
+            if p.y > myg.ymax:
+                if yrb in ["outflow", "neumann"]:
                     continue
-            elif yrb == "periodic":
-                if p.y > myg.ymax:
+                elif yrb == "periodic":
                     p.y = myg.ymin + p.y - myg.ymax
-            else:
-                msg.fail("ERROR: yrb = %s invalid BC" % (yrb))
+                elif yrb in ["reflect-even", "reflect-odd"]:
+                    p.y = 2 * myg.ymax - p.y
+                else:
+                    msg.fail("ERROR: yrb = %s invalid BC" % (yrb))
 
             self.particles.add(p)
 
+        self.n_particles = len(self.particles)
+
     def get_positions(self):
         """
         Return an array of particle positions.

particles/tests/test_particles.py

@@ -0,0 +1,161 @@
+import particles.particles as particles
+import mesh.patch as patch
+from util import runparams
+import numpy as np
+from numpy.testing import assert_array_equal
+from simulation_null import NullSimulation, grid_setup, bc_setup
+
+
+def test_particle():
+    n_particles = 5
+
+    for n in range(n_particles):
+        x, y = np.random.rand(2)
+        p = particles.Particle(x, y)
+
+        assert_array_equal(p.pos(), [x, y])
+        assert_array_equal(p.velocity(), [0, 0])
+
+
+def test_particles_random_gen():
+    rp = runparams.RuntimeParameters()
+
+    rp.params["mesh.nx"] = 8
+    rp.params["mesh.ny"] = 8
+    rp.params["particles.do_particles"] = 0
+    rp.params["particles.n_particles"] = 50
+    rp.params["particles.particle_generator"] = "random"
+
+    # set up sim
+    sim = NullSimulation("", "", rp)
+
+    # set up grid
+    my_grid = grid_setup(rp)
+    my_data = patch.CellCenterData2d(my_grid)
+    bc = bc_setup(rp)[0]
+    my_data.create()
+    sim.cc_data = my_data
+
+    # seed random number generator
+    np.random.seed(3287469)
+
+    ps = particles.Particles(sim.cc_data, bc, rp)
+    positions = set([(p.x, p.y) for p in ps.particles])
+
+    assert len(ps.particles) == 50, "There should be 50 particles"
+
+    # reseed random number generator
+    np.random.seed(3287469)
+
+    correct_positions = np.random.rand(50, 2)
+    correct_positions = set([(x, y) for (x, y) in correct_positions])
+
+    assert positions == correct_positions, "sets are not the same"
+
+
+def test_particles_grid_gen():
+    rp = runparams.RuntimeParameters()
+
+    rp.params["mesh.nx"] = 8
+    rp.params["mesh.ny"] = 8
+    rp.params["particles.do_particles"] = 0
+    rp.params["particles.n_particles"] = 50
+    rp.params["particles.particle_generator"] = "grid"
+
+    # set up sim
+    sim = NullSimulation("", "", rp)
+
+    # set up grid
+    my_grid = grid_setup(rp)
+    my_data = patch.CellCenterData2d(my_grid)
+    bc = bc_setup(rp)[0]
+    my_data.create()
+    sim.cc_data = my_data
+
+    ps = particles.Particles(sim.cc_data, bc, rp)
+
+    assert len(ps.particles) == 49, "There should be 49 particles"
+
+    positions = set([(p.x, p.y) for p in ps.particles])
+
+    xs, step = np.linspace(0, 1, num=7, endpoint=False, retstep=True)
+    xs += 0.5 * step
+
+    correct_positions = set([(x, y) for x in xs for y in xs])
+
+    assert positions == correct_positions, "sets are not the same"
+
+
+def test_particles_advect():
+    rp = runparams.RuntimeParameters()
+
+    rp.params["mesh.nx"] = 8
+    rp.params["mesh.ny"] = 8
+    rp.params["particles.do_particles"] = 0
+    rp.params["particles.n_particles"] = 50
+    rp.params["particles.particle_generator"] = "grid"
+
+    # set up sim
+    sim = NullSimulation("", "", rp)
+
+    # set up grid
+    my_grid = grid_setup(rp, ng=4)
+    my_data = patch.CellCenterData2d(my_grid)
+    bc = bc_setup(rp)[0]
+    my_data.create()
+    sim.cc_data = my_data
+
+    ps = particles.Particles(sim.cc_data, bc, rp)
+
+    # advect with constant velocity
+
+    # first try 0 velocity
+    u = my_grid.scratch_array()
+    v = my_grid.scratch_array()
+
+    ps.update_particles(u, v, 1)
+
+    positions = set([(p.x, p.y) for p in ps.particles])
+
+    xs, step = np.linspace(0, 1, num=7, endpoint=False, retstep=True)
+    xs += 0.5 * step
+
+    correct_positions = set([(x, y) for x in xs for y in xs])
+
+    assert positions == correct_positions, "sets are not the same"
+
+    # now move constant speed to the right
+    u[:, :] = 1
+
+    ps.update_particles(u, v, 1)
+
+    positions = set([(p.x, p.y) for p in ps.particles])
+
+    correct_positions = set([(x+1, y) for x in xs for y in xs])
+
+    assert positions == correct_positions, "sets are not the same"
+
+    # constant speed up
+    u[:, :] = 0
+    v[:, :] = 1
+
+    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])
+
+    correct_positions = set([(x, y+1) for x in xs for y in xs])
+
+    assert positions == correct_positions, "sets are not the same"
+
+    # constant speed right + up
+    u[:, :] = 1
+
+    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])
+
+    correct_positions = set([(x+1, y+1) for x in xs for y in xs])
+
+    assert positions == correct_positions, "sets are not the same"