Merge pull request #24 from elbeejay/faster-generation

Faster "exact" generation

Author: elbeejay

docs/source/userguide/index.rst

@@ -64,7 +64,7 @@ Defining a :obj:`dorado.particle_track.Particles` class is a key step in using `
 Particle Generation and Routing
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-* Particles can be generated using the `particle_generator` function, this allows for random and exact placement of particles
+* Particles can be generated using the :obj:`dorado.particle_track.Particles.generate_particles` function, this allows for random and exact placement of particles
 
 * Particles can be routed using either the High-Level or Low-Level API functionalities described below
 

dorado/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.2.0"
+__version__ = "2.3.0"
 
 from . import lagrangian_walker
 from . import parallel_routing

dorado/particle_track.py

@@ -459,21 +459,16 @@ def generate_particles(self,
                                   in list(range(Np_tracer))]
 
         elif method == 'exact':
-            # step through seed lists to create exact list of start locations
-            # if number of particles exceeds length of seed lists, loop back
-            # to beginning of the seed lists and continue from there
-            Np_to_seed = Np_tracer
-            new_start_xindices = []
-            new_start_yindices = []
-            while Np_to_seed > 0:
-                for i in range(0, len(seed_xloc)):
-                    # condition if we run out of particles within this loop
-                    if Np_to_seed <= 0:
-                        break
-                    # otherwise keep on seeding
-                    new_start_xindices = new_start_xindices + [seed_xloc[i]]
-                    new_start_yindices = new_start_yindices + [seed_yloc[i]]
-                    Np_to_seed -= 1  # reduce number of particles left to seed
+            # create exact start indices based on x, y locations and np_tracer
+            # get number of particles per location and remainder left out
+            Np_divy = divmod(Np_tracer, len(seed_xloc))
+            # assign the start locations for the evenly divided particles
+            new_start_xindices = seed_xloc * Np_divy[0]
+            new_start_yindices = seed_yloc * Np_divy[0]
+            # add the remainder ones to the list one by one via looping
+            for i in range(0, Np_divy[1]):
+                new_start_xindices = new_start_xindices + [seed_xloc[i]]
+                new_start_yindices = new_start_yindices + [seed_yloc[i]]
 
         # Now initialize vectors that will create the structured list
         new_xinds = [[new_start_xindices[i]] for i in

tests/test_particle_track.py

@@ -133,6 +133,17 @@ def test_exact_locations_overflow():
     assert len(all_walk_data['xinds']) == num_ps
     assert len(all_walk_data['yinds']) == num_ps
 
+def test_exact_locations_multidims():
+    """Test case where seed locs given as multi-dimensional arrays."""
+    num_ps = 5
+    particle = particle_track.Particles(params)
+    x_seed = np.array(((0, 0), (1, 1)))
+    y_seed = np.array(((0, 0), (1, 1)))
+    particle.generate_particles(num_ps, x_seed, y_seed, method='exact')
+    all_walk_data = particle.run_iteration()
+    assert len(all_walk_data['xinds']) == num_ps
+    assert len(all_walk_data['yinds']) == num_ps
+
 def test_no_explicit_generation():
     """Test reading of Np_tracer from self if some walk_data exists."""
     # create some walk data (would exist from a previous run)