Numba for speed up

.travis.yml

@@ -15,6 +15,8 @@ jobs:
   include:
     - stage: coverage
       python: "3.8"
+      before_install:
+        - NUMBA_DISABLE_JIT=1
       install: 
         - pip install -r requirements.txt 
         - pip install -e .

pyDeltaRCM/default.yml

@@ -81,10 +81,10 @@ save_depth_figs:
   default: False
 save_discharge_figs:
   type: 'bool'
-  default: True
+  default: False
 save_velocity_figs:
   type: 'bool'
-  default: True
+  default: False
 save_eta_grids:
   type: 'bool'
   default: False

pyDeltaRCM/deltaRCM_tools.py

@@ -187,7 +187,7 @@ def output_data(self):
             if self.save_eta_figs:
 
                 plt.pcolor(self.eta)
-                plt.clim(self.clim_eta[0], self.clim_eta[1])
+                # plt.clim(self.clim_eta[0], self.clim_eta[1])
                 plt.colorbar()
                 plt.axis('equal')
                 self.save_figure(self.prefix + "eta_" + str(timestep))

pyDeltaRCM/init_tools.py

@@ -9,6 +9,7 @@
 from math import floor, sqrt, pi
 import numpy as np
 from random import shuffle
+import numba
 
 import matplotlib
 from matplotlib import pyplot as plt
@@ -22,6 +23,10 @@
 import logging
 import time
 import yaml
+
+from . import utils
+
+
 # tools for initiating deltaRCM model domain
 
 
@@ -97,7 +102,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)
+            utils.set_random_seed(self.seed)
 
     def set_constants(self):
 
@@ -310,6 +315,8 @@ def create_domain(self):
         self.cell_type[:self.L0, channel_inds:y_channel_max] = cell_channel
 
         self.inlet = list(np.unique(np.where(self.cell_type == 1)[1]))
+        self.inlet_typed = numba.typed.List()
+        [self.inlet_typed.append(x) for x in self.inlet]
         self.eta[:] = self.stage - self.depth
 
         self.clim_eta = (-self.h0 - 1, 0.05)

pyDeltaRCM/sed_tools.py

@@ -19,6 +19,7 @@
 from netCDF4 import Dataset
 
 from .shared_tools import shared_tools
+from . import utils
 
 # tools for sediment routing algorithms and deposition/erosion
 
@@ -200,7 +201,7 @@ def sed_parcel(self, theta_sed, sed, px, py):
             if ind[0] == 0:
                 weight[0, :] = np.nan
 
-            new_cell = self.random_pick(weight)
+            new_cell = utils.random_pick(weight)
 
             jstep = self.iwalk.flat[new_cell]
             istep = self.jwalk.flat[new_cell]
@@ -239,8 +240,8 @@ 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)]
+        typed_inlet = self.inlet_typed
+        start_indices = [utils.random_pick_inlet(typed_inlet) for x in range(num_starts)]
 
         for np_sed in range(num_starts):
 
@@ -281,8 +282,8 @@ 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)]
+        typed_inlet = self.inlet_typed
+        start_indices = [utils.random_pick_inlet(typed_inlet) for x in range(num_starts)]
 
         for np_sed in range(num_starts):
 

pyDeltaRCM/utils.py

@@ -0,0 +1,70 @@
+
+import numpy as np
+
+import time
+import numba
+
+# utilities used in various places in the model and docs
+
+
+@numba.njit
+def set_random_seed(_seed):
+    np.random.seed(_seed)
+
+
+@numba.njit
+def get_random_uniform(N):
+    return np.random.uniform(0, 1, N)
+
+
+@numba.njit
+def random_pick(probs):
+    """
+    Randomly pick a number weighted by array probs (len 8)
+    Return the index of the selected weight in array probs
+    """
+
+    # catch bad prob arrays
+    if np.nansum(probs) == 0:
+        probs[1, 1] = 0
+        for i, prob in np.ndenumerate(probs):
+            if np.isnan(prob):
+                probs[i] = 1
+
+    # prep array
+    for i, prob in np.ndenumerate(probs):
+        if np.isnan(prob):
+            probs[i] = 0
+
+    cutoffs = np.cumsum(probs)
+    v = np.random.uniform(0, cutoffs[-1])
+    idx = np.searchsorted(cutoffs, v)
+
+    return int(idx)
+
+
+@numba.njit
+def random_pick_inlet(choices, probs=None):
+    """
+    Randomly pick a number from array choices weighted by array probs
+    Values in choices are column indices
+
+    Return a tuple of the randomly picked index for row 0
+    """
+
+    if probs is None:
+        probs = np.array([1. for i in range(len(choices))])
+
+    cutoffs = np.cumsum(probs)
+    v = np.random.uniform(0, cutoffs[-1])
+    idx = np.searchsorted(cutoffs, v)
+
+    return choices[idx]
+
+
+def _get_version():
+    """
+    Extract version number from single file, and make it availabe everywhere.
+    """
+    from . import _version
+    return _version.__version__()

pyDeltaRCM/water_tools.py

@@ -6,6 +6,8 @@
 import logging
 import time
 
+import numba
+
 from math import floor, sqrt, pi
 import numpy as np
 from random import shuffle
@@ -19,10 +21,90 @@
 from netCDF4 import Dataset
 
 from .shared_tools import shared_tools
+from . import utils
 
 # tools for water routing algorithms
 
 
+@numba.njit
+def _index_slicer(kernel, ind):
+    """Get a slice out of a kernel for an ind.
+    """
+
+    return kernel[ind[0] - 1 + 1:ind[0] + 2 + 1, ind[1] - 1 + 1:ind[1] + 2 + 1]
+
+
+@numba.njit
+def get_weights_kernels(ind, pad_stage, stage, distances, qx, qy,
+                        ivec, jvec, pad_depth, pad_cell_type):
+    """Get weights and kernels for an ind.
+    """
+    stage_ind = _index_slicer(pad_stage, ind)
+    weight_sfc = np.maximum(0,
+                            (stage[ind] - stage_ind) / distances)
+
+    weight_int = np.maximum(0, (qx[ind] * jvec +
+                            qy[ind] * ivec) / distances)
+    depth_ind = _index_slicer(pad_depth, ind)
+    ct_ind = _index_slicer(pad_cell_type, ind)
+
+    return weight_sfc, weight_int, depth_ind, ct_ind
+
+
+@numba.njit
+def get_new_cell(ind, weight_sfc, weight_int, depth_ind, ct_ind,
+                 dry_depth, gamma, theta_water):
+    """Get new_cell based on weights and kernels for an ind.
+    """
+
+    if ind[0] == 0:
+        weight_sfc[0, :] = 0
+        weight_int[0, :] = 0
+
+    for (i, u), (j, v) in zip(np.ndenumerate(depth_ind),
+                              np.ndenumerate(ct_ind)):
+        if (u <= dry_depth) or (v == -2):
+            weight_sfc[i] = 0
+            weight_int[i] = 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_ind ** theta_water * _weight
+    for i, u in np.ndenumerate(depth_ind):
+        if u <= dry_depth:
+            _weight[i] = np.nan
+
+    new_cell = utils.random_pick(_weight)
+
+    return new_cell
+
+
+@numba.njit
+def get_new_cells(inds_tuple_typed, pad_stage, stage, distances, qx, qy,
+                  ivec, jvec, pad_depth, pad_cell_type, dry_depth, gamma,
+                  theta_water):
+    """Get new_cells based on current model state.
+    """
+    new_cells = [np.int(x) for x in range(0)]
+    for ind in inds_tuple_typed:
+        if ind != (0, 0):
+            weight_sfc, weight_int, depth_ind, ct_ind = get_weights_kernels(
+                ind, pad_stage, stage, distances, qx, qy, ivec, jvec,
+                pad_depth, pad_cell_type)
+            new_cells.append(get_new_cell(ind, weight_sfc, weight_int,
+                                          depth_ind, ct_ind, dry_depth,
+                                          gamma, theta_water))
+        else:
+            new_cells.append(int(4))
+
+    return new_cells
+
+
 class water_tools(shared_tools):
 
     def update_flow_field(self, iteration):
@@ -142,9 +224,10 @@ def init_water_iteration(self):
 
     def run_water_iteration(self):
 
-        iter = 0
-        start_indices = [self.random_pick_inlet(
-            self.inlet) for x in range(self.Np_water)]
+        _iter = 0
+        typed_inlet = self.inlet_typed
+        start_indices = [utils.random_pick_inlet(typed_inlet)
+                         for x in range(self.Np_water)]
 
         self.qxn.flat[start_indices] += 1
         self.qwn.flat[start_indices] += self.Qp_water / self.dx / 2.
@@ -154,18 +237,25 @@ def run_water_iteration(self):
 
         self.looped[:] = 0
 
-        while (sum(current_inds) > 0) & (iter < self.itmax):
+        while (sum(current_inds) > 0) & (_iter < self.itmax):
 
-            iter += 1
+            _iter += 1
 
-            self.check_size_of_indices_matrix(iter)
+            self.check_size_of_indices_matrix(_iter)
 
             inds = np.unravel_index(current_inds, self.depth.shape)
+            inds_tuple_typed = numba.typed.List()
             inds_tuple = [(inds[0][i], inds[1][i])
                           for i in range(len(inds[0]))]
+            [inds_tuple_typed.append(x) for x in inds_tuple]
 
-            new_cells = [self.get_weight(x)
-                         if x != (0, 0) else 4 for x in inds_tuple]
+            new_cells = get_new_cells(inds_tuple_typed, self.pad_stage,
+                                      self.stage, self.distances,
+                                      self.qx, self.qy,
+                                      self.ivec, self.jvec,
+                                      self.pad_depth, self.pad_cell_type,
+                                      self.dry_depth, self.gamma,
+                                      self.theta_water)
 
             new_inds = list(map(lambda x, y: self.calculate_new_ind(x, y)
                                 if y != 4 else 0, inds_tuple, new_cells))
@@ -176,13 +266,13 @@ def run_water_iteration(self):
             new_inds = np.array(new_inds, dtype=np.int)
             new_inds[np.array(dist) == 0] = 0
 
-            self.indices[:, iter] = new_inds
+            self.indices[:, _iter] = new_inds
 
-            current_inds = self.check_for_loops(new_inds, iter)
+            current_inds = self.check_for_loops(new_inds, _iter)
 
             current_inds = self.check_for_boundary(current_inds)
 
-            self.indices[:, iter] = current_inds
+            self.indices[:, _iter] = current_inds
 
             current_inds[self.free_surf_flag > 0] = 0
 
@@ -389,43 +479,6 @@ def calculate_new_ind(self, ind, new_cell):
 
         return new_ind_flat
 
-    def get_weight(self, ind):
-
-        stage_ind = self.pad_stage[
-            ind[0] - 1 + 1:ind[0] + 2 + 1, ind[1] - 1 + 1:ind[1] + 2 + 1]
-
-        weight_sfc = np.maximum(0,
-                                (self.stage[ind] - stage_ind) / self.distances)
-
-        weight_int = np.maximum(0, (self.qx[ind] * self.jvec +
-                                    self.qy[ind] * self.ivec) / self.distances)
-
-        if ind[0] == 0:
-            weight_sfc[0, :] = 0
-            weight_int[0, :] = 0
-
-        depth_ind = self.pad_depth[
-            ind[0] - 1 + 1:ind[0] + 2 + 1, ind[1] - 1 + 1:ind[1] + 2 + 1]
-        ct_ind = self.pad_cell_type[
-            ind[0] - 1 + 1:ind[0] + 2 + 1, ind[1] - 1 + 1:ind[1] + 2 + 1]
-
-        weight_sfc[(depth_ind <= self.dry_depth) | (ct_ind == -2)] = 0
-        weight_int[(depth_ind <= self.dry_depth) | (ct_ind == -2)] = 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)
-
-        self.weight = self.gamma * weight_sfc + (1 - self.gamma) * weight_int
-        self.weight = depth_ind ** self.theta_water * self.weight
-        self.weight[depth_ind <= self.dry_depth] = np.nan
-
-        new_cell = self.random_pick(self.weight)
-
-        return new_cell
-
     def build_weight_array(self, array, fix_edges=False, normalize=False):
         """
         Create np.array((8,L,W)) of quantity a

requirements.txt

@@ -4,3 +4,4 @@ scipy
 netCDF4
 basic_modeling_interface
 pyyaml
+numba