refactor: jit significant portions of the water and sediment routing

This makes them available outside the object, which is important in some
cases, and the class just gets them from here.

feat: make shared_tools into a module, not a class.

This is in preparation for jitting, since we want them to be standalone
functions.

feat: break out step that updates water parcel positions

These functions now become helper utility functions in shared_tools, and
are wrapped by update_Q, which handles them all, and passes them what
they need.

refactor: break out sed functions into smaller functions

This modularizes the code and makes it easier to compile.

feat: jit the functions

apply jit to functions, and fix some type problems that came along with
refactoring.

refactor: move next step calculation to jitted functions

The next slowest operation was getting the indexes of the next cells, so
I moved that to a jitted function too.

refactor: jit the calculation of weight array for the water

fix: move random number generation into numba

Copy of what @amoodie did in #27

fix: some broken tests to test the shared tools portion.

fix: add numba to requirements

fix: move all inlet picking into the random_pick function

There's essentially no difference between the random_pick function and
the random_pick_inlet function other than there is no preference given
for the inlet cells, so you can pick whichever one, so I just adjust the
mechanism to use the existing random pick machinery.

fix: water and sediment routing bugs that caused inconsistent behavior

There were two issues with the refactoring. One was that water routing
developed a flat flow field because I wasn't adding each passing water
parcel to the Q field. The other was related to sediment parcel routing,
where I had transcribed the i and j of stepping incorrectly in the
jitted functions.

pyDeltaRCM/init_tools.py

@@ -22,6 +22,9 @@
 import logging
 import time
 import yaml
+
+from . import shared_tools
+
 # tools for initiating deltaRCM model domain
 
 
@@ -97,7 +100,7 @@ def import_files(self):
         if self.seed is not None:
             if self.verbose >= 2:
                 print("setting random seed to %s " % str(self.seed))
-            np.random.seed(self.seed)
+            shared_tools.set_random_seed(self.seed)
 
     def set_constants(self):
 
@@ -113,17 +116,13 @@ def set_constants(self):
                               [-1, 0, 1],
                               [-sqrt05, 0, sqrt05]])
 
-        self.iwalk = np.array([[-1, 0, 1],
-                               [-1, 0, 1],
-                               [-1, 0, 1]])
+        self.iwalk = shared_tools.get_iwalk()
 
         self.jvec = np.array([[-sqrt05, -1, -sqrt05],
                               [0, 0, 0],
                               [sqrt05, 1, sqrt05]])
 
-        self.jwalk = np.array([[-1, -1, -1],
-                               [0, 0, 0],
-                               [1, 1, 1]])
+        self.jwalk = shared_tools.get_jwalk()
         self.dxn_iwalk = [1, 1, 0, -1, -1, -1, 0, 1]
         self.dxn_jwalk = [0, 1, 1, 1, 0, -1, -1, -1]
         self.dxn_dist = \

pyDeltaRCM/sed_tools.py

@@ -18,12 +18,12 @@
 
 from netCDF4 import Dataset
 
-from .shared_tools import shared_tools
+from . import shared_tools
 
 # tools for sediment routing algorithms and deposition/erosion
 
 
-class sed_tools(shared_tools):
+class sed_tools(object):
 
     def sed_route(self):
         """route all sediment"""
@@ -183,51 +183,38 @@ def sed_parcel(self, theta_sed, sed, px, py):
             # choose next with weights
 
             it += 1
-            ind = (px, py)
-
             depth_ind = self.pad_depth[
-                ind[0] - 1 + 1:ind[0] + 2 + 1, ind[1] - 1 + 1:ind[1] + 2 + 1]
+                px - 1 + 1:px + 2 + 1, py - 1 + 1:py + 2 + 1]
             cell_type_ind = self.pad_cell_type[
-                ind[0] - 1 + 1:ind[0] + 2 + 1, ind[1] - 1 + 1:ind[1] + 2 + 1]
+                px - 1 + 1:px + 2 + 1, py - 1 + 1:py + 2 + 1]
 
-            w1 = np.maximum(0, (self.qx[ind] * self.jvec +
-                                self.qy[ind] * self.ivec))
-            w2 = depth_ind ** theta_sed
-            weight = (w1 * w2 / self.distances)
+            w1 = shared_tools.get_flux_wt(self.qx[(px, py)], self.qy[(px, py)], self.ivec, self.jvec)
+            w2 = shared_tools.get_depth_wt(depth_ind, theta_sed)
 
-            weight[depth_ind <= self.dry_depth] = 0.0001
-            weight[cell_type_ind == -2] = 0
+            w3 = shared_tools.get_combined_weight(w1, w2, self.distances)
 
-            if ind[0] == 0:
-                weight[0, :] = 0
+            weights = shared_tools.get_filtered_weight(
+                w3.flatten(), px, depth_ind.flatten(),
+                cell_type_ind.flatten(), self.dry_depth)
 
-            new_cell = self.random_pick(np.cumsum(weight.flatten()))
+            new_cell = shared_tools.random_pick(weights)
 
-            jstep = self.iwalk.flat[new_cell]
-            istep = self.jwalk.flat[new_cell]
-            dist = np.sqrt(istep * istep + jstep * jstep)
+            dist, istep, jstep, _ = shared_tools.get_steps(new_cell, self.iwalk.flat[:], self.jwalk.flat[:])
 
             # deposition and erosion
 
             if sed == 'sand':  # sand
-                if dist > 0:
-                    # deposition in current cell
-                    self.qs[px, py] = (self.qs[px, py] +
-                                       self.Vp_res / 2 / self.dt / self.dx)
-                px = px + istep
-                py = py + jstep
-
-                if dist > 0:
-                    # deposition in downstream cell
-                    self.qs[px, py] = (self.qs[px, py] +
-                                       self.Vp_res / 2 / self.dt / self.dx)
+
+                depoPart = self.Vp_res / 2 / self.dt / self.dx
+
+                px, py, self.qs = shared_tools.partition_sand(self.qs, depoPart, py, px, dist, istep, jstep)
 
                 self.sand_dep_ero(px, py)
 
             if sed == 'mud':  # mud
 
-                px = px + istep
-                py = py + jstep
+                px = px + jstep
+                py = py + istep
 
                 self.mud_dep_ero(px, py)
 
@@ -240,8 +227,10 @@ def sand_route(self):
         theta_sed = self.theta_sand
 
         num_starts = int(self.Np_sed * self.f_bedload)
-        start_indices = [self.random_pick_inlet(
-            self.inlet) for x in range(num_starts)]
+        inlet_weights = np.ones_like(self.inlet)
+        start_indices = [
+            self.inlet[shared_tools.random_pick(inlet_weights / sum(inlet_weights))]
+            for x in range(num_starts)]
 
         for np_sed in range(num_starts):
 
@@ -281,8 +270,10 @@ def mud_route(self):
         theta_sed = self.theta_mud
 
         num_starts = int(self.Np_sed * (1 - self.f_bedload))
-        start_indices = [self.random_pick_inlet(
-            self.inlet) for x in range(num_starts)]
+        inlet_weights = np.ones_like(self.inlet)
+        start_indices = [
+            self.inlet[shared_tools.random_pick(inlet_weights / sum(inlet_weights))]
+            for x in range(num_starts)]
 
         for np_sed in range(num_starts):
 

pyDeltaRCM/shared_tools.py

@@ -3,39 +3,220 @@
 import numpy as np
 
 import time
+from numba import njit
 
 # tools shared between deltaRCM water and sediment routing
 
 
-class shared_tools(object):
+def get_iwalk():
+    return np.array([[-1, 0, 1],
+                     [-1, 0, 1],
+                     [-1, 0, 1]])
 
-    def random_pick(self, probs):
-        """
-        Randomly pick a number weighted by array probabilities (len 9)
-        Return the index of the selected weight in array probs
-        Takes a numpy array that is the precalculated cumulative probability
-        around the cell flattened to 1D.
-        """
 
-        idx = probs.searchsorted(np.random.uniform(0, probs[-1]))
+def get_jwalk():
+    return np.array([[-1, -1, -1],
+                     [0, 0, 0],
+                     [1, 1, 1]])
 
-        return idx
 
-    def random_pick_inlet(self, choices, probs=None):
-        """
-        Randomly pick a number from array choices weighted by array probs
-        Values in choices are column indices
+@njit
+def set_random_seed(_seed):
+    np.random.seed(_seed)
 
-        Return a tuple of the randomly picked index for row 0
-        """
 
-        if not probs:
-            probs = np.ones(len(choices))
+@njit
+def get_random_uniform(N):
+    return np.random.uniform(0, 1)
 
-        cutoffs = np.cumsum(probs)
-        idx = cutoffs.searchsorted(np.random.uniform(0, cutoffs[-1]))
 
-        return choices[idx]
+@njit
+def get_flux_wt(qx, qy, ivec, jvec):
+    return np.maximum(0, (qx * jvec + qy * ivec))
+
+
+@njit
+def get_depth_wt(depth_ind, theta):
+    return depth_ind ** theta
+
+
+@njit
+def get_combined_weight(w1, w2, distances):
+    return w1 * w2 / distances
+
+
+@njit
+def get_filtered_weight(weight, px, depth_ind, cell_type_ind, dry_depth):
+    dry = depth_ind <= dry_depth
+    weight[dry] = 0.0001
+    walls = cell_type_ind == -2
+    weight[walls] = 0
+    if px == 0:
+        weight[:3] = 0
+    # weight = np.cumsum(weight)
+    # weight.flatten()
+    return weight / np.sum(weight)
+
+
+@njit
+def partition_sand(qs, depoPart, py, px, dist, istep, jstep):
+    if dist > 0:
+        # deposition in current cell
+        qs[px, py] += depoPart
+
+    px = px + jstep
+    py = py + istep
+
+    if dist > 0:
+        # deposition in downstream cell
+        qs[px, py] += depoPart
+    return px, py, qs
+
+
+@njit
+def get_steps(new_cells, iwalk, jwalk):
+
+    istep = iwalk[new_cells]
+    jstep = jwalk[new_cells]
+    dist = np.sqrt(istep * istep + jstep * jstep)
+
+    astep = dist != 0
+
+    return dist, istep, jstep, astep
+
+
+@njit
+def update_dirQfield(qfield, dist, inds, astep, dirstep):
+    for i, ii in enumerate(inds):
+        if astep[i]:
+            qfield[ii] += dirstep[i] / dist[i]
+    return qfield
+
+
+@njit
+def update_absQfield(qfield, dist, inds, astep, Qp_water, dx):
+    for i, ii in enumerate(inds):
+        if astep[i]:
+            qfield[ii] += Qp_water / dx / 2
+    return qfield
+
+
+@njit
+def random_pick(prob):
+    """
+    Randomly pick a number weighted by array probabilities (len 9)
+    Return the index of the selected weight in array probs
+    Takes a numpy array that is the precalculated cumulative probability
+    around the cell flattened to 1D.
+    """
+
+    arr = np.arange(len(prob))
+    return arr[np.searchsorted(np.cumsum(prob), get_random_uniform(1))]
+
+
+@njit
+def custom_unravel(i, shape):
+    x = i // shape[1]
+    y = i % shape[1]
+    return x, y
+
+
+@njit
+def custom_ravel(tup, shape):
+    x = tup[0] * shape[1]
+    y = tup[1]
+    return x + y
+
+
+@njit
+def check_for_loops(indices, inds, it, L0, loopedout, domain_shape, CTR, free_surf_flag):
+
+    looped = [len(i[i > 0]) != len(set(i[i > 0])) for i in indices]
+
+    travel = (0, it)
+
+    for n in range(indices.shape[0]):
+
+        ind = inds[n]
+
+        if looped[n] and (ind > 0) and (max(travel) > L0):
+
+            loopedout[n] += 1
+
+            it = custom_unravel(ind, domain_shape)
+
+            px, py = it
+
+            Fx = px - 1
+            Fy = py - CTR
+
+            Fw = np.sqrt(Fx**2 + Fy**2)
+
+            if Fw != 0:
+                px = px + np.round(Fx / Fw * 5.)
+                py = py + np.round(Fy / Fw * 5.)
+
+            px = max(px, L0)
+            px = int(min(domain_shape[0] - 2, px))
+
+            py = max(1, py)
+            py = int(min(domain_shape[1] - 2, py))
+
+            nind = custom_ravel((px, py), domain_shape)
+
+            inds[n] = nind
+
+            free_surf_flag[n] = -1
+
+    return inds, loopedout, free_surf_flag
+
+
+@njit
+def calculate_new_ind(indices, new_cells, iwalk, jwalk, domain_shape):
+    newbies = []
+    for p, q in zip(indices, new_cells):
+        if q != 4:
+            ind_tuple = custom_unravel(p, domain_shape)
+            new_ind = (ind_tuple[0] + jwalk[q],
+                       ind_tuple[1] + iwalk[q])
+            newbies.append(custom_ravel(new_ind, domain_shape))
+        else:
+            newbies.append(0)
+
+    return np.array(newbies)
+
+
+@njit
+def get_weight_at_cell(ind, stage_nbrs, depth_nbrs, ct_nbrs, stage, qx, qy,
+                       ivec, jvec, distances, dry_depth, gamma, theta_water):
+
+    weight_sfc = np.maximum(0, (stage - stage_nbrs) / distances)
+
+    weight_int = np.maximum(0, (qx * jvec + qy * ivec) / distances)
+
+    if ind[0] == 0:
+        weight_sfc[:3] = 0
+        weight_int[:3] = 0
+
+    drywall = (depth_nbrs <= dry_depth) | (ct_nbrs == -2)
+    weight_sfc[drywall] = 0
+    weight_int[drywall] = 0
+
+    if np.nansum(weight_sfc) > 0:
+        weight_sfc = weight_sfc / np.nansum(weight_sfc)
+
+    if np.nansum(weight_int) > 0:
+        weight_int = weight_int / np.nansum(weight_int)
+
+    weight = gamma * weight_sfc + (1 - gamma) * weight_int
+    weight = depth_nbrs ** theta_water * weight
+    weight[depth_nbrs <= dry_depth] = 0
+
+    if np.sum(weight) != 0:
+        weight = weight / np.sum(weight)
+    else:
+        weight = weight
+    return weight / np.sum(weight)
 
 
 def _get_version():