Added io for particles

harpolea · harpolea · commit 56239b959822 · 2018-07-05T14:25:05.000+01:00
diff --git a/advection/_defaults b/advection/_defaults
@@ -9,5 +9,6 @@ v = 1.0      ; advective velocity in y-direction
 limiter = 2  ; limiter (0 = none, 1 = 2nd order, 2 = 4th order)
 
 
-
-
+[particles]
+do_particles = 0
+particle_generator = grid
diff --git a/advection/problems/inputs.tophat b/advection/problems/inputs.tophat
@@ -32,7 +32,3 @@ u = 1.0
 v = 1.0
 
 limiter = 2
-
-[particles]
-do_particles = 1
-particle_generator = grid
diff --git a/advection/simulation.py b/advection/simulation.py
@@ -28,7 +28,9 @@ def initialize(self):
         self.cc_data = my_data
 
         if self.rp.get_param("particles.do_particles") == 1:
-            self.particles = particles.Particles(self.cc_data, bc, self.rp)
+            n_particles = self.rp.get_param("particles.n_particles")
+            particle_generator = self.rp.get_param("particles.particle_generator")
+            self.particles = particles.Particles(self.cc_data, bc, n_particles, particle_generator)
 
         # now set the initial conditions for the problem
         problem = importlib.import_module("advection.problems.{}".format(self.problem_name))
@@ -86,9 +88,7 @@ def evolve(self):
             u2d = myg.scratch_array() + u
             v2d = myg.scratch_array() + v
 
-            self.particles.update_particles(u2d, v2d, self.dt)
-
-            self.particles.enforce_particle_boundaries()
+            self.particles.update_particles(self.dt, u2d, v2d)
 
         # increment the time
         self.cc_data.t += self.dt
diff --git a/advection_fv4/simulation.py b/advection_fv4/simulation.py
@@ -6,6 +6,7 @@
 import mesh.array_indexer as ai
 import mesh.fv as fv
 from simulation_null import grid_setup, bc_setup
+import particles.particles as particles
 
 
 class Simulation(advection_rk.Simulation):
@@ -26,6 +27,11 @@ 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")
+            particle_generator = self.rp.get_param("particles.particle_generator")
+            self.particles = particles.Particles(self.cc_data, bc, n_particles, particle_generator)
+
         # now set the initial conditions for the problem
         problem = importlib.import_module("advection_fv4.problems.{}".format(self.problem_name))
         problem.init_data(self.cc_data, self.rp)
diff --git a/advection_rk/simulation.py b/advection_rk/simulation.py
@@ -72,6 +72,15 @@ def evolve(self):
 
         rk.compute_final_update()
 
+        if self.particles is not None:
+            myg = self.cc_data.grid
+            u = self.rp.get_param("advection.u")
+            v = self.rp.get_param("advection.v")
+
+            u2d = myg.scratch_array() + u
+            v2d = myg.scratch_array() + v
+
+            self.particles.update_particles(self.dt, u2d, v2d)
         # increment the time
         myd.t += self.dt
         self.n += 1
diff --git a/advection_weno/simulation.py b/advection_weno/simulation.py
@@ -72,6 +72,16 @@ def evolve(self):
 
         rk.compute_final_update()
 
+        if self.particles is not None:
+            myg = self.cc_data.grid
+            u = self.rp.get_param("advection.u")
+            v = self.rp.get_param("advection.v")
+
+            u2d = myg.scratch_array() + u
+            v2d = myg.scratch_array() + v
+
+            self.particles.update_particles(self.dt, u2d, v2d)
+
         # increment the time
         myd.t += self.dt
         self.n += 1
diff --git a/compressible/_defaults b/compressible/_defaults
@@ -22,5 +22,5 @@ grav = 0.0                ; gravitational acceleration (in y-direction)
 riemann = HLLC            ; HLLC or CGF
 
 [particles]
-do_particles = 1
+do_particles = 0
 particle_generator = grid
diff --git a/compressible/simulation.py b/compressible/simulation.py
@@ -143,7 +143,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)
+            n_particles = self.rp.get_param("particles.n_particles")
+            particle_generator = self.rp.get_param("particles.particle_generator")
+            self.particles = particles.Particles(self.cc_data, bc, n_particles, particle_generator)
 
         # some auxillary data that we'll need to fill GC in, but isn't
         # really part of the main solution
@@ -226,10 +228,7 @@ def evolve(self):
         ener[:, :] += 0.5*self.dt*(ymom[:, :] + old_ymom[:, :])*grav
 
         if self.particles is not None:
-            u, v = self.cc_data.get_var("velocity")
-
-            self.particles.update_particles(u, v, self.dt)
-            self.particles.enforce_particle_boundaries()
+            self.particles.update_particles(self.dt)
 
         # increment the time
         self.cc_data.t += self.dt
diff --git a/compressible_fv4/simulation.py b/compressible_fv4/simulation.py
@@ -89,6 +89,9 @@ def evolve(self):
 
         rk.compute_final_update()
 
+        if self.particles is not None:
+            self.particles.update_particles(self.dt)
+
         # increment the time
         myd.t += self.dt
         self.n += 1
diff --git a/compressible_react/simulation.py b/compressible_react/simulation.py
@@ -49,9 +49,15 @@ def evolve(self):
 
         self.diffuse(self.dt/2)
 
+        if self.particles is not None:
+            self.particles.update_particles(self.dt/2)
+
         # note: this will do the time increment and n increment
         super().evolve()
 
+        if self.particles is not None:
+            self.particles.update_particles(self.dt/2)
+
         self.diffuse(self.dt/2)
 
         self.burn(self.dt/2)
diff --git a/compressible_rk/simulation.py b/compressible_rk/simulation.py
@@ -90,6 +90,9 @@ def evolve(self):
 
         rk.compute_final_update()
 
+        if self.particles is not None:
+            self.particles.update_particles(self.dt)
+
         # increment the time
         myd.t += self.dt
         self.n += 1
diff --git a/compressible_sdc/simulation.py b/compressible_sdc/simulation.py
@@ -95,6 +95,9 @@ def evolve(self):
         # store the new solution
         self.cc_data.data[:, :, :] = U_knew[-1].data[:, :, :]
 
+        if self.particles is not None:
+            self.particles.update_particles(self.dt)
+
         # increment the time
         myd.t += self.dt
         self.n += 1
diff --git a/diffusion/_defaults b/diffusion/_defaults
@@ -5,7 +5,3 @@ cfl       = 0.8           ; diffusion CFL number
 
 [diffusion]
 k = 1.0      ; conductivity
-
-
-
-
diff --git a/docs/source/particles_basics.rst b/docs/source/particles_basics.rst
@@ -34,7 +34,9 @@ The particles can be initialised in a number of ways:
   each direction, so the spacing will be different in each direction if the
   domain is not square.) The number of particles will be increased/decreased
   in order to fill the whole domain;
-* the user can define their own ``particle_generator_func`` and pass this into the
+* :func:`array_generate_particles <particles.particles.Particle.array_generate_particles>`,
+  generates particles based on array of particle positions passed to the constructor;
+* the user can define their own ``particle_generator`` and pass this into the
   Particles constructor. This function takes the number of particles to be
   generated and returns a dictionary of ``Particle`` objects.
 
@@ -44,19 +46,21 @@ solver's ``Simulation.initialize`` function:
 .. code-block:: python
 
    if self.rp.get_param("particles.do_particles") == 1:
-       self.particles = particles.Particles(self.cc_data, bc, self.rp)
+         n_particles = self.rp.get_param("particles.n_particles")
+         particle_generator = self.rp.get_param("particles.particle_generator")
+         self.particles = particles.Particles(self.cc_data, bc, n_particles, particle_generator)
 
 The particles can then be advanced by inserting the following code after the
 update of the other variables in the solver's ``Simulation.evolve`` function:
 
 .. code-block:: python
 
    if self.particles is not None:
-        self.particles.update_particles(u, v, self.dt)
-        self.particles.enforce_particle_boundaries()
+        self.particles.update_particles(self.dt)
 
-where ``u`` and ``v`` are the ``ArrayIndexer`` objects holding the x-velocity and
-y-velocity on the grid.
+For some problems (e.g. advection), the x- and y- velocities must also be passed
+in as arguments to this function as they cannot be accessed using the standard
+``get_var("velocity")`` command.
 
 We can turn on/off the particles solver using the following runtime paramters:
 
@@ -73,3 +77,37 @@ We can turn on/off the particles solver using the following runtime paramters:
 |                       | option can be overridden by passing a custom generator |
 |                       | function to the ``Particles`` constructor.             |
 +-----------------------+--------------------------------------------------------+
+
+Plotting particles
+------------------
+
+Particles can be plotted by getting their positions using
+
+.. code-block:: python
+
+   particle_positions = particles.get_positions()
+
+In order to track the movement of particles over time, it's useful
+to 'dye' the particles based on their initial positions. This can
+be done by calling
+
+.. code-block:: python
+
+   colors = particles.get_init_positions()
+
+For example, if we color the particles from white to black based on their initial
+x-position, we can plot them on the figure axis ``ax`` using the following code:
+
+.. code-block:: python
+
+      particle_positions = particles.get_positions()
+
+      # dye particles based on initial x-position
+      colors = 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])
diff --git a/incompressible/_defaults b/incompressible/_defaults
@@ -6,5 +6,5 @@ proj_type = 2             ; what are we projecting? 1 includes -Gp term in U*
 cfl = 0.8
 
 [particles]
-do_particles = 1
+do_particles = 0
 particle_generator = grid
diff --git a/incompressible/simulation.py b/incompressible/simulation.py
@@ -45,7 +45,9 @@ def initialize(self):
         self.cc_data = my_data
 
         if self.rp.get_param("particles.do_particles") == 1:
-            self.particles = particles.Particles(self.cc_data, bc, self.rp)
+            n_particles = self.rp.get_param("particles.n_particles")
+            particle_generator = self.rp.get_param("particles.particle_generator")
+            self.particles = particles.Particles(self.cc_data, bc, n_particles, particle_generator)
 
         # now set the initial conditions for the problem
         problem = importlib.import_module("incompressible.problems.{}".format(self.problem_name))
@@ -393,8 +395,7 @@ def evolve(self):
         self.cc_data.fill_BC("y-velocity")
 
         if self.particles is not None:
-            self.particles.update_particles(u, v, self.dt)
-            self.particles.enforce_particle_boundaries()
+            self.particles.update_particles(self.dt)
 
         # increment the time
         if not self.in_preevolve:
diff --git a/particles/particles.py b/particles/particles.py
diff --git a/particles/tests/test_particles.py b/particles/tests/test_particles.py
diff --git a/simulation_null.py b/simulation_null.py
diff --git a/swe/_defaults b/swe/_defaults
diff --git a/swe/simulation.py b/swe/simulation.py
diff --git a/util/io.py b/util/io.py
