Merge pull request #27 from wrightky/master

Optimize weights computation

dorado/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.4.2"
+__version__ = "2.5.0"
 
 
 from . import lagrangian_walker

dorado/lagrangian_walker.py

@@ -39,56 +39,185 @@ def random_pick_seed(choices, probs=None):
     return choices[idx]
 
 
-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]
+def big_sliding_window(raster):
+    """Creates 3D array organizing local neighbors at every index
+
+    Returns a raster of shape (L,W,9) which organizes (along the third
+    dimension) all of the neighbors in the original raster at a given
+    index, in the order [NW, N, NE, W, 0, E, SW, S, SE]. Outputs are
+    ordered to match np.ravel(), so it functions similarly to a loop
+    applying ravel to the elements around each index.
+    For example, the neighboring values in raster indexed at (i,j) are 
+    raster(i-1:i+2, j-1:j+2).ravel(). These 9 values have been mapped to
+    big_ravel(i,j,:) for ease of computations. Helper function for make_weight.
 
-    # calculate surface slope weights
-    weight_sfc = maximum(0,
-                         (Particles.stage[ind] - stage_ind) /
-                         Particles.distances)
+    **Inputs** :
+
+        raster : `ndarray`
+            2D array of values (e.g. stage, qx)
+
+    **Outputs** :
+
+        big_ravel : `ndarray`
+            3D array which sorts the D8 neighbors at index (i,j) in 
+            raster into the 3rd dimension at (i,j,:)
+
+    """
+    big_ravel = np.zeros((raster.shape[0],raster.shape[1],9))
+    big_ravel[1:-1,1:-1,0] = raster[0:-2,0:-2]
+    big_ravel[1:-1,1:-1,1] = raster[0:-2,1:-1]
+    big_ravel[1:-1,1:-1,2] = raster[0:-2,2:]
+    big_ravel[1:-1,1:-1,3] = raster[1:-1,0:-2]
+    big_ravel[1:-1,1:-1,4] = raster[1:-1,1:-1]
+    big_ravel[1:-1,1:-1,5] = raster[1:-1,2:]
+    big_ravel[1:-1,1:-1,6] = raster[2:,0:-2]
+    big_ravel[1:-1,1:-1,7] = raster[2:,1:-1]
+    big_ravel[1:-1,1:-1,8] = raster[2:,2:]
+
+    return big_ravel
+
+
+def tile_local_array(local_array, L, W):
+    """Take a local array [[NW, N, NE], [W, 0, E], [SW, S, SE]]
+    and repeat it into an array of shape (L,W,9), where L, W are
+    the shape of the domain, and the original elements are ordered
+    along the third axis. Helper function for make_weight.
+
+    **Inputs** :
+
+        local_array : `ndarray`
+            2D array of represnting the D8 neighbors around
+            some index (e.g. ivec, jvec)
+
+        L : `int`
+            Length of the domain
+
+        W : `int`
+            Width of the domain
+
+    **Outputs** :
+
+        tiled_array : `ndarray`
+            3D array repeating local_array.ravel() LxW times
+
+    """
+    return np.tile(local_array.ravel(), (L, W, 1))
+
+
+def tile_domain_array(raster):
+    """Repeat a large 2D array 9 times along the third axis.
+    Helper function for make_weight.
+
+    **Inputs** :
+
+        raster : `ndarray`
+            2D array of values (e.g. stage, qx)
+
+    **Outputs** :
+
+        tiled_array : `ndarray`
+            3D array repeating raster 9 times
+
+    """
+    return np.repeat(raster[:, :, np.newaxis], 9, axis=2)
 
-    # calculate inertial component weights
-    weight_int = maximum(0, ((Particles.qx[ind] * Particles.jvec +
-                              Particles.qy[ind] * Particles.ivec) /
-                         Particles.distances))
 
+def clear_borders(tiled_array):
+    """Set to zero all the edge elements of a vertical stack 
+    of 2D arrays. Helper function for make_weight.
+
+    **Inputs** :
+
+        tiled_array : `ndarray`
+            3D array repeating raster (e.g. stage, qx) 9 times
+            along the third axis
+
+    **Outputs** :
+
+        tiled_array : `ndarray`
+            The same 3D array as the input, but with the borders
+            in the first and second dimension set to 0.
+
+    """
+    tiled_array[0,:,:] = 0.
+    tiled_array[:,0,:] = 0.
+    tiled_array[-1,:,:] = 0.
+    tiled_array[:,-1,:] = 0.
+    return
+
+
+def make_weight(Particles):
+    """Create the weighting array for particle routing
+
+    Function to compute the routing weights at each index, which gets
+    stored inside the :obj:`dorado.particle_track.Particles` object
+    for use when routing. Called when the object gets instantiated.
+
+    **Inputs** :
+
+        Particles : :obj:`dorado.particle_track.Particles`
+            A :obj:`dorado.particle_track.Particles` object
+
+    **Outputs** :
+
+        Updates the weights array inside the
+        :obj:`dorado.particle_track.Particles` object
+
+    """
+    L, W = Particles.stage.shape
+
+    # calculate surface slope weights
+    weight_sfc = (tile_domain_array(Particles.stage) \
+                  - big_sliding_window(Particles.stage))
+    weight_sfc /= tile_local_array(Particles.distances, L, W)
+    weight_sfc[weight_sfc <= 0] = 0
+    clear_borders(weight_sfc)
+
+    # calculate inertial component weights
+    weight_int = (tile_domain_array(Particles.qx)*tile_local_array(Particles.jvec, L, W)) \
+                 + (tile_domain_array(Particles.qy)*tile_local_array(Particles.ivec, L, W))
+    weight_int /= tile_local_array(Particles.distances, L, W)
+    weight_int[weight_int <= 0] = 0
+    clear_borders(weight_int)
+
     # 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]
+    depth_ind = big_sliding_window(Particles.depth)
+    ct_ind = big_sliding_window(Particles.cell_type)
 
     # 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 weight is above 0 normalize by sum of weights
-    if nansum(weight_int) > 0:
-        weight_int = weight_int / nansum(weight_int)
+    norm_sfc = np.nansum(weight_sfc, axis=2)
+    idx_sfc = tile_domain_array((norm_sfc > 0))
+    weight_sfc[idx_sfc] /= tile_domain_array(norm_sfc)[idx_sfc]
+
+    norm_int = np.nansum(weight_int, axis=2)
+    idx_int = tile_domain_array((norm_int > 0))
+    weight_int[idx_int] /= tile_domain_array(norm_int)[idx_int]
 
     # define actual weight by using gamma, and weight components
     weight = Particles.gamma * weight_sfc + \
-        (1 - Particles.gamma) * weight_int
-
+             (1 - Particles.gamma) * weight_int
+    
     # 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 set weight to 0
+    weight[(depth_ind <= Particles.dry_depth) | (ct_ind == 2)] = 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
+    zero_weights = tile_domain_array((np.nansum(weight, axis=2) <= 0))
+    deepest_cells = (depth_ind == tile_domain_array(np.max(depth_ind,axis=2)))
+    choose_deep_cells = (zero_weights & deepest_cells)
+    weight[choose_deep_cells] = 1.0
+
+    # Final checks, eliminate invalid choices
+    clear_borders(weight)
 
     # set weight in the true weight array
-    Particles.weight[ind[0], ind[1], :] = weight.ravel()
+    Particles.weight = weight
 
 
 def get_weight(Particles, ind):
@@ -111,9 +240,6 @@ 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

@@ -402,7 +402,7 @@ def __init__(self, params):
         self.walk_data = None
 
         # initialize routing weights array
-        self.weight = np.zeros((self.stage.shape[0], self.stage.shape[1], 9))
+        lw.make_weight(self)
 
 
     # function to clear walk data if you've made a mistake while generating it

dorado/routines.py

@@ -74,25 +74,29 @@ def steady_plots(particle, num_iter,
         # Do particle iterations
         walk_data = particle.run_iteration()
         if save_output:
-            x0, y0, t0 = get_state(walk_data, 0)
+            # Get current location
             xi, yi, ti = get_state(walk_data)
-
-            fig = plt.figure(dpi=200)
-            ax = fig.add_subplot(111)
-            im = ax.imshow(particle.depth)
-            plt.title('Depth - Particle Iteration ' + str(i))
-            cax = fig.add_axes([ax.get_position().x1+0.01,
-                                ax.get_position().y0,
-                                0.02,
-                                ax.get_position().height])
-            cbar = plt.colorbar(im, cax=cax)
-            cbar.set_label('Water Depth [m]')
-            ax.scatter(y0, x0, c='b', s=0.75)
-            ax.scatter(yi, xi, c='r', s=0.75)
+            if i == 0:
+                # Initialize figure
+                x0, y0, t0 = get_state(walk_data, 0)
+                fig = plt.figure(dpi=200)
+                ax = fig.add_subplot(111)
+                im = ax.imshow(particle.depth)
+                cax = fig.add_axes([ax.get_position().x1+0.01,
+                                    ax.get_position().y0,
+                                    0.02,
+                                    ax.get_position().height])
+                cbar = plt.colorbar(im, cax=cax)
+                cbar.set_label('Water Depth [m]')
+                orig = ax.scatter(y0, x0, c='b', s=0.75)
+                newloc = ax.scatter(yi, xi, c='r', s=0.75)
+            else:
+                # Update figure with new locations
+                newloc.set_offsets(np.array([yi,xi]).T)
+            ax.set_title('Depth - Particle Iteration ' + str(i))
             plt.savefig(folder_name+os.sep +
                         'figs'+os.sep+'output'+str(i)+'.png',
                         bbox_inches='tight')
-            plt.close()
 
     if save_output:
         # save data as json text file - technically human readable
@@ -165,6 +169,7 @@ def unsteady_plots(dx, Np_tracer, seed_xloc, seed_yloc, num_steps, timestep,
     # init params
     params = modelParams()
     params.dx = dx
+    params.verbose = False
     # make directory to save the data
     if folder_name is None:
         folder_name = os.getcwd()
@@ -237,25 +242,30 @@ def unsteady_plots(dx, Np_tracer, seed_xloc, seed_yloc, num_steps, timestep,
 
         walk_data = particle.run_iteration(target_time=target_times[i])
 
-        x0, y0, t0 = get_state(walk_data, 0)
         xi, yi, ti = get_state(walk_data)
-
-        # make and save plots and data
-        fig = plt.figure(dpi=200)
-        ax = fig.add_subplot(111)
-        ax.scatter(y0, x0, c='b', s=0.75)
-        ax.scatter(yi, xi, c='r', s=0.75)
-        im = ax.imshow(params.depth)
-        plt.title('Depth at Time ' + str(target_times[i]))
-        cax = fig.add_axes([ax.get_position().x1+0.01,
-                            ax.get_position().y0,
-                            0.02,
-                            ax.get_position().height])
-        cbar = plt.colorbar(im, cax=cax)
-        cbar.set_label('Water Depth [m]')
+        if i == 0:
+            x0, y0, t0 = get_state(walk_data, 0)
+            # Initialize figure
+            fig = plt.figure(dpi=200)
+            ax = fig.add_subplot(111)
+            im = ax.imshow(params.depth)
+            cax = fig.add_axes([ax.get_position().x1+0.01,
+                                ax.get_position().y0,
+                                0.02,
+                                ax.get_position().height])
+            cbar = plt.colorbar(im, cax=cax)
+            cbar.set_label('Water Depth [m]')
+            orig = ax.scatter(y0, x0, c='b', s=0.75)
+            newloc = ax.scatter(yi, xi, c='r', s=0.75)
+        else:
+            # Update figure with new locations
+            im.set_data(params.depth)
+            im.set_clim(np.min(params.depth), np.max(params.depth))
+            newloc.set_offsets(np.array([yi,xi]).T)
+            plt.draw()
+        ax.set_title('Depth at Time ' + str(target_times[i]))
         plt.savefig(folder_name+os.sep+'figs'+os.sep+'output'+str(i)+'.png',
                     bbox_inches='tight')
-        plt.close()
 
     # save data as a json text file - technically human readable
     fpath = folder_name+os.sep+'data'+os.sep+'data.txt'
@@ -312,30 +322,36 @@ def time_plots(particle, num_iter, folder_name=None):
         walk_data = particle.run_iteration()
 
         # collect latest travel times
-        x0, y0, t0 = get_state(walk_data, 0)
         xi, yi, temptimes = get_state(walk_data)
 
-        # set colorbar using 10th and 90th percentile values
-        cm = matplotlib.cm.colors.Normalize(vmax=np.percentile(temptimes, 90),
-                                            vmin=np.percentile(temptimes, 10))
-
-        fig = plt.figure(dpi=200)
-        ax = plt.gca()
-        plt.title('Depth - Particle Iteration ' + str(i))
-        ax.scatter(y0, x0, c='b', s=0.75)
-        sc = ax.scatter(yi, xi, c=temptimes, s=0.75, cmap='coolwarm', norm=cm)
-        divider = make_axes_locatable(ax)
-        cax = divider.append_axes("right", size="5%", pad=0.05)
-        cbar = plt.colorbar(sc, cax=cax)
-        cbar.set_label('Particle Travel Times [s]')
-        im = ax.imshow(particle.depth)
-        divider = make_axes_locatable(ax)
-        cax = divider.append_axes("bottom", size="5%", pad=0.5)
-        cbar2 = plt.colorbar(im, cax=cax, orientation='horizontal')
-        cbar2.set_label('Water Depth [m]')
+        if i == 0:
+            x0, y0, t0 = get_state(walk_data, 0)
+            # Initialize figure
+            fig = plt.figure(dpi=200)
+            ax = plt.gca()
+            im = ax.imshow(particle.depth)
+            orig = ax.scatter(y0, x0, c='b', s=0.75)
+            sc = ax.scatter(yi, xi, c=temptimes, s=0.75, cmap='coolwarm')
+            sc.set_clim(np.percentile(temptimes,90),
+                        np.percentile(temptimes,10))
+            divider = make_axes_locatable(ax)
+            cax = divider.append_axes("right", size="5%", pad=0.05)
+            cbar = plt.colorbar(sc, cax=cax)
+            cbar.set_label('Particle Travel Times [s]')
+            divider = make_axes_locatable(ax)
+            cax = divider.append_axes("bottom", size="5%", pad=0.5)
+            cbar2 = plt.colorbar(im, cax=cax, orientation='horizontal')
+            cbar2.set_label('Water Depth [m]')
+        else:
+            # Update figure with new locations
+            sc.set_offsets(np.array([yi,xi]).T) # Location
+            sc.set_array(np.array(temptimes)) # Color values
+            sc.set_clim(np.percentile(temptimes,90),
+                        np.percentile(temptimes,10)) # Color limits
+            plt.draw()
+        ax.set_title('Depth - Particle Iteration ' + str(i))
         plt.savefig(folder_name+os.sep+'figs'+os.sep+'output'+str(i)+'.png',
                     bbox_inches='tight')
-        plt.close()
 
     # save data as a json text file - technically human readable
     fpath = folder_name+os.sep+'data'+os.sep+'data.txt'