Repository: passaH2O/dorado

Issue:
Further Speeding Up Routing Weight Calculations
From #16, @wrightky said:

> So, looking at this code, I see that the main structure of the weight computation hasn't changed. We're still constructing small sub-arrays at each index, doing a few quick operations (max, cleaning, multiplications, summing), and saving the 9 resulting weights in a weight array. Surprised we hadn't tried this sooner given how much better it performs.

> It looks like the key reason the runtime scales better in this model isn't anything about how this computation works, it's how many times we do it. Originally, we performed these operations locally once for every particle at every iteration. So, the runtime scaled as Np_tracer * iterations_per_particle. Now, we perform this once for every cell, so it scales with domain size (L-2) * (W-2). I bet if you checked the example cases you benchmarked, the ratio of these values should give roughly the speedup you observed.

> One thing I wonder, though, is whether we could obtain even faster runtimes by modifying the structure of this computation itself, by switching from local array operations to global. Whatever is causing this function to take so long must be due to the fact that we're repeatedly constructing many sub-arrays in a loop and operating on them, instead of performing a few big global array operations (inside which we'd be taking advantage of all the fast numpy broadcasting stuff). In principal, if we broke up these operations (max, cleaning, multiplication, summing) into a loop over each of the D8 directions, with each being a global matrix operation, instead of a loop over each cell, we could reduce the amount of overhead repeatedly calling these functions. I don't know exactly how that scales, but it'd be the difference between time(np.max(small array)) * (L-2) * (W-2) and time(np.max(big array)) * 9. Does that scale better? Not sure.

So that is a potential route for further speeding up the routing weight calculation.


Discussion:

Patch:
diff --git a/docs/source/userguide/index.rst b/docs/source/userguide/index.rst
index 7427ad5..8efb28d 100644
--- a/docs/source/userguide/index.rst
+++ b/docs/source/userguide/index.rst
@@ -53,14 +53,6 @@ Defining the `Particles`
 
 Defining a :obj:`dorado.particle_track.Particles` class is a key step in using `dorado` to perform particle routing. To define a set of particles, the model parameters must first be defined as described above. The `Particles` class is initialized using an instance of the model parameters. From there, particles can be generated and routed.
 
-.. Note::
-    When :obj:`dorado.particle_track.Particles` is initialized, all of the
-    routing weights are automatically calculated. This may take some time for
-    larger model domains, but allows for faster particle routing when particles
-    are actually moved through the domain. There is a progress bar associated
-    with this process so you don't feel like Python has gotten stuck in the
-    object initialization.
-
 Particle Generation and Routing
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
diff --git a/dorado/__init__.py b/dorado/__init__.py
index 2c2564f..5f2a210 100644
--- a/dorado/__init__.py
+++ b/dorado/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "2.4.0"
+__version__ = "2.4.1"
 
 
 from . import lagrangian_walker
diff --git a/dorado/lagrangian_walker.py b/dorado/lagrangian_walker.py
index 225ca49..63f3901 100644
--- a/dorado/lagrangian_walker.py
+++ b/dorado/lagrangian_walker.py
@@ -9,7 +9,7 @@ from builtins import range, map
 from math import cos
 import numpy as np
 from numpy.random import random
-from tqdm import tqdm
+from numpy import maximum, nansum
 
 
 def random_pick_seed(choices, probs=None):
@@ -39,76 +39,63 @@ def random_pick_seed(choices, probs=None):
     return choices[idx]
 
 
-def make_weight(Particles):
-    """Make an array with the routing weights."""
-    # local namespace function imports
-    from numpy import maximum
-    from numpy import nansum
-    # init the weight array
-    L, W = np.shape(Particles.stage)
-    Particles.weight = np.zeros((L, W, 9))
-    # do weighting calculation for each cell
-    print('Calculating routing weights ...')
-    for i in tqdm(list(range(1, L-1)), ascii=True):
-        for j in list(range(1, W-1)):
-            # weights for each location in domain
-            # get stage values for neighboring cells
-            stage_ind = Particles.stage[i-1:i+2, j-1:j+2]
+def make_weight(Particles, ind):
+    """Update weighting array with weights at this index"""
+    # get stage values for neighboring cells
+    stage_ind = Particles.stage[ind[0]-1:ind[0]+2, ind[1]-1:ind[1]+2]
 
-            # calculate surface slope weights
-            weight_sfc = maximum(0,
-                                 (Particles.stage[i, j]-stage_ind) /
-                                 Particles.distances)
+    # calculate surface slope weights
+    weight_sfc = maximum(0,
+                         (Particles.stage[ind] - stage_ind) /
+                         Particles.distances)
 
-            # calculate inertial component weights
-            weight_int = maximum(0, ((Particles.qx[i, j] * Particles.jvec +
-                                      Particles.qy[i, j] * Particles.ivec) /
-                                 Particles.distances))
+    # calculate inertial component weights
+    weight_int = maximum(0, ((Particles.qx[ind] * Particles.jvec +
+                              Particles.qy[ind] * Particles.ivec) /
+                         Particles.distances))
 
-            # get depth and cell types for neighboring cells
-            depth_ind = Particles.depth[i-1:i+2, j-1:j+2]
-            ct_ind = Particles.cell_type[i-1:i+2, j-1:j+2]
+    # get depth and cell types for neighboring cells
+    depth_ind = Particles.depth[ind[0]-1:ind[0]+2, ind[1]-1:ind[1]+2]
+    ct_ind = Particles.cell_type[ind[0]-1:ind[0]+2, ind[1]-1:ind[1]+2]
 
-            # set weights for cells that are too shallow, or invalid 0
-            weight_sfc[(depth_ind <= Particles.dry_depth) | (ct_ind == 2)] = 0
-            weight_int[(depth_ind <= Particles.dry_depth) | (ct_ind == 2)] = 0
+    # set weights for cells that are too shallow, or invalid 0
+    weight_sfc[(depth_ind <= Particles.dry_depth) | (ct_ind == 2)] = 0
+    weight_int[(depth_ind <= Particles.dry_depth) | (ct_ind == 2)] = 0
 
-            # if sum of weights is above 0 normalize by sum of weights
-            if nansum(weight_sfc) > 0:
-                weight_sfc = weight_sfc / nansum(weight_sfc)
+    # if sum of weights is above 0 normalize by sum of weights
+    if nansum(weight_sfc) > 0:
+        weight_sfc = weight_sfc / nansum(weight_sfc)
 
-            # if sum of weight is above 0 normalize by sum of weights
-            if nansum(weight_int) > 0:
-                weight_int = weight_int / nansum(weight_int)
+    # if sum of weight is above 0 normalize by sum of weights
+    if nansum(weight_int) > 0:
+        weight_int = weight_int / nansum(weight_int)
 
-            # define actual weight by using gamma, and weight components
-            weight = Particles.gamma * weight_sfc + \
-                (1 - Particles.gamma) * weight_int
+    # define actual weight by using gamma, and weight components
+    weight = Particles.gamma * weight_sfc + \
+        (1 - Particles.gamma) * weight_int
 
-            # modify the weight by the depth and theta weighting parameter
-            weight = depth_ind ** Particles.theta * weight
+    # modify the weight by the depth and theta weighting parameter
+    weight = depth_ind ** Particles.theta * weight
 
-            # if the depth is below the minimum depth then location is not
-            # considered therefore set the associated weight to nan
-            weight[(depth_ind <= Particles.dry_depth) | (ct_ind == 2)] \
-                = np.nan
+    # if the depth is below the minimum depth then location is not
+    # considered therefore set the associated weight to nan
+    weight[(depth_ind <= Particles.dry_depth) | (ct_ind == 2)] \
+        = np.nan
 
-            # if it's a dead end with only nans and 0's, choose deepest cell
-            if nansum(weight) <= 0:
-                weight = np.zeros_like(weight)
-                weight[depth_ind == np.max(depth_ind)] = 1.0
+    # if it's a dead end with only nans and 0's, choose deepest cell
+    if nansum(weight) <= 0:
+        weight = np.zeros_like(weight)
+        weight[depth_ind == np.max(depth_ind)] = 1.0
 
-            # set weight in the true weight array
-            Particles.weight[i, j, :] = weight.ravel()
-
-    print('Finished routing weight calculation.')
+    # set weight in the true weight array
+    Particles.weight[ind[0], ind[1], :] = weight.ravel()
 
 
 def get_weight(Particles, ind):
     """Choose new cell location given an initial location.
 
     Function to randomly choose 1 of the surrounding 8 cells around the
-    current index using the pre-calculated routing weights.
+    current index using the routing weights from make_weight.
 
     **Inputs** :
 
@@ -124,6 +111,9 @@ def get_weight(Particles, ind):
             New location given as a value between 1 and 8 (inclusive)
 
     """
+    # Check if weights have been computed for this location:
+    if nansum(Particles.weight[ind[0], ind[1], :]) <= 0:
+        make_weight(Particles, ind)
     # randomly pick the new cell for the particle to move to using the
     # random_pick function and the set of weights
     if Particles.steepest_descent is not True:
diff --git a/dorado/particle_track.py b/dorado/particle_track.py
index 3a6886e..75d5de0 100644
--- a/dorado/particle_track.py
+++ b/dorado/particle_track.py
@@ -377,8 +377,8 @@ class Particles():
         # initialize the walk_data
         self.walk_data = None
 
-        # create weights - this might take a bit of time for large domains
-        lw.make_weight(self)
+        # initialize routing weights array
+        self.weight = np.zeros((self.stage.shape[0], self.stage.shape[1], 9))
 
     # function to clear walk data if you've made a mistake while generating it
     def clear_walk_data(self):


Test patch:
diff --git a/tests/test_lagrangian_walker.py b/tests/test_lagrangian_walker.py
index b43cb9e..a1d52d2 100644
--- a/tests/test_lagrangian_walker.py
+++ b/tests/test_lagrangian_walker.py
@@ -286,6 +286,11 @@ def test_make_weight_shallow():
     ind = (1, 1)
     # set seed
     np.random.seed(0)
+    # do weighting calculation for each cell
+    L, W = np.shape(particles.stage)
+    for i in list(range(1, L-1)):
+        for j in list(range(1, W-1)):
+            lw.make_weight(particles, (i, j))
     # make assertions about weights
     # at index, index[4] (self) will be 1 while neighbors will be 0
     assert particles.weight[1, 1, 4] == 1.0
@@ -328,6 +333,11 @@ def test_make_weight_equal_opportunity():
     ind = (1, 1)
     # set seed
     np.random.seed(0)
+    # do weighting calculation for each cell
+    L, W = np.shape(particles.stage)
+    for i in list(range(1, L-1)):
+        for j in list(range(1, W-1)):
+            lw.make_weight(particles, (i, j))
     # make assertions about weights
     # at index, 3 neighbors will be equiprobable
     assert np.sum(particles.weight[1, 1, :]) == 3.0
@@ -372,6 +382,11 @@ def test_make_weight_unequal_opportunity():
     ind = (1, 1)
     # set seed
     np.random.seed(0)
+    # do weighting calculation for each cell
+    L, W = np.shape(particles.stage)
+    for i in list(range(1, L-1)):
+        for j in list(range(1, W-1)):
+            lw.make_weight(particles, (i, j))
     # make assertions about weights
     # at index, staying put index[4] higher probability than neighbors
     assert particles.weight[1, 1, 4] > particles.weight[1, 1, 5]
@@ -411,6 +426,11 @@ def test_wet_boundary_no_weight():
     particles = pt.Particles(tools)
     # set seed
     np.random.seed(0)
+    # do weighting calculation for each cell
+    L, W = np.shape(particles.stage)
+    for i in list(range(1, L-1)):
+        for j in list(range(1, W-1)):
+            lw.make_weight(particles, (i, j))
     # assert weights at boundary cells should be 0
     assert np.all(np.sum(particles.weight[0, :, 4]) == 0.0)
     assert np.all(np.sum(particles.weight[-1, :, 4]) == 0.0)