Merge pull request #36 from ericbarefoot/numbajit

BUMBANUMBA

docs/source/reference/shared_tools/index.rst

@@ -4,12 +4,11 @@
 shared_tools
 *********************************
 
-The tools are defined in ``pyDeltaRCM.shared_tools``. 
+The tools are defined in ``pyDeltaRCM.shared_tools``.
 
+.. currentmodule:: pyDeltaRCM
 
-.. currentmodule:: pyDeltaRCM.shared_tools
+.. autosummary::
+  :toctree: ../../_autosummary
 
-.. autosummary:: 
-	:toctree: ../../_autosummary
-
-	shared_tools
+  shared_tools

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,36 @@ 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]
-            cell_type_ind = self.pad_cell_type[
-                ind[0] - 1 + 1:ind[0] + 2 + 1, ind[1] - 1 + 1:ind[1] + 2 + 1]
+            stage_nbrs = self.pad_stage[px - 1 + 1:px + 2 + 1, py - 1 + 1:py + 2 + 1]
+            depth_ind = self.pad_depth[px - 1 + 1:px + 2 + 1, py - 1 + 1:py + 2 + 1]
+            cell_type_ind = self.pad_cell_type[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)
+            weights = shared_tools.get_weight_at_cell(
+                (px, py),
+                stage_nbrs.flatten(), depth_ind.flatten(), cell_type_ind.flatten(),
+                self.stage[px, py], self.qx[px, py], self.qy[px, py],
+                self.ivec.flatten(), self.jvec.flatten(), self.distances.flatten(),
+                self.dry_depth, self.gamma, theta_sed)
 
-            weight[depth_ind <= self.dry_depth] = 0.0001
-            weight[cell_type_ind == -2] = 0
+            new_cell = shared_tools.random_pick(weights)
 
-            if ind[0] == 0:
-                weight[0, :] = 0
-
-            new_cell = self.random_pick(np.cumsum(weight.flatten()))
-
-            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 +225,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 +268,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,208 @@
 import numpy as np
 
 import time
+from numba import njit, jit, typed
 
 # 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 partition_sand(qs, depoPart, py, px, dist, istep, jstep):
+    '''Spread sand between two cells
+    '''
+    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):
+    '''find the values giving the next step
+    '''
+    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):
+    '''update unit vector of water flux in x or y
+    '''
+    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):
+    '''Update norm of water flux vector
+    '''
+    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):
+    """Function to unravel indexes for 2D array
+    """
+    if i > (shape[1] * shape[0]):
+        raise IndexError("Index is out of matrix bounds")
+    x = i // shape[1]
+    y = i % shape[1]
+    return x, y
+
+
+@njit
+def custom_ravel(tup, shape):
+    """Function to ravel indexes for 2D array
+    """
+    if tup[0] > shape[0] or tup[1] > shape[1]:
+        raise IndexError("Index is out of matrix bounds")
+    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 = typed.List() # numba typed list for iteration
+    for i in np.arange(len(indices)):
+        row = indices[i, :]
+        v = len(row[row > 0]) != len(set(row[row > 0]))
+        looped.append(v)
+    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
+            px, py = custom_unravel(ind, domain_shape)
+            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):
+
+    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] = np.nan
+        weight_int[:3] = np.nan
+
+    drywall = (depth_nbrs <= dry_depth) | (ct_nbrs == -2)
+    weight_sfc[drywall] = np.nan
+    weight_int[drywall] = np.nan
+
+    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 * weight
+    weight[depth_nbrs <= dry_depth] = 0
+
+    nanWeight = np.isnan(weight)
+
+    if np.any(weight[~nanWeight] != 0):
+        weight = weight / np.nansum(weight)
+        weight[nanWeight] = 0
+    else:
+        weight[~nanWeight] = 1 / len(weight[~nanWeight])
+        weight[nanWeight] = 0
+    return weight
+
 
 
 def _get_version():