Merge pull request #25 from wrightky/master

Weights speedup v2

Author: wrightky

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
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

dorado/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.4.0"
+__version__ = "2.4.1"
 
 
 from . import lagrangian_walker

dorado/lagrangian_walker.py

@@ -9,7 +9,7 @@
 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:

dorado/particle_track.py

@@ -377,8 +377,8 @@ def __init__(self, params):
         # 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):

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)