diff --git a/pyDeltaRCM/default.yml b/pyDeltaRCM/default.yml
index 65809d9c..e2f70762 100644
--- a/pyDeltaRCM/default.yml
+++ b/pyDeltaRCM/default.yml
@@ -84,7 +84,7 @@ save_discharge_figs:
   default: True
 save_velocity_figs:
   type: 'bool'
-  default: True
+  default: False
 save_eta_grids:
   type: 'bool'
   default: False
@@ -96,10 +96,10 @@ save_depth_grids:
   default: False
 save_discharge_grids:
   type: 'bool'
-  default: True
+  default: False
 save_velocity_grids:
   type: 'bool'
-  default: True
+  default: False
 save_dt:
   type: 'int'
   default: 50
diff --git a/pyDeltaRCM/deltaRCM_tools.py b/pyDeltaRCM/deltaRCM_tools.py
index d8f79998..e8358bd0 100644
--- a/pyDeltaRCM/deltaRCM_tools.py
+++ b/pyDeltaRCM/deltaRCM_tools.py
@@ -39,8 +39,6 @@ def run_one_timestep(self):
 
         for iteration in range(self.itermax):
 
-            self.count = 0
-
             self.init_water_iteration()
             self.run_water_iteration()
 
diff --git a/pyDeltaRCM/sed_tools.py b/pyDeltaRCM/sed_tools.py
index ee5615f2..c7c47084 100644
--- a/pyDeltaRCM/sed_tools.py
+++ b/pyDeltaRCM/sed_tools.py
@@ -35,6 +35,7 @@ def sed_route(self):
         self.Vp_dep_mud[:] = 0
 
         self.sand_route()
+
         self.mud_route()
 
         self.topo_diffusion()
@@ -195,16 +196,16 @@ def sed_parcel(self, theta_sed, sed, px, py):
             weight = (w1 * w2 / self.distances)
 
             weight[depth_ind <= self.dry_depth] = 0.0001
-            weight[cell_type_ind == -2] = np.nan
+            weight[cell_type_ind == -2] = 0
 
             if ind[0] == 0:
-                weight[0, :] = np.nan
+                weight[0, :] = 0
 
-            new_cell = self.random_pick(weight)
+            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**2 + jstep**2)
+            dist = np.sqrt(istep * istep + jstep * jstep)
 
             # deposition and erosion
 
@@ -254,7 +255,6 @@ def sand_route(self):
 
             self.sed_parcel(theta_sed, 'sand', px, py)
 
-#         self.topo_diffusion()
     def topo_diffusion(self):
         """
         Diffuse topography after routing all coarse sediment parcels
diff --git a/pyDeltaRCM/shared_tools.py b/pyDeltaRCM/shared_tools.py
index 0bcd9b86..34ae9af6 100644
--- a/pyDeltaRCM/shared_tools.py
+++ b/pyDeltaRCM/shared_tools.py
@@ -11,19 +11,13 @@ class shared_tools(object):
 
     def random_pick(self, probs):
         """
-        Randomly pick a number weighted by array probs (len 8)
+        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.
         """
 
-        num_nans = sum(np.isnan(probs))
-
-        if np.nansum(probs) == 0:
-            probs[~np.isnan(probs)] = 1
-            probs[1, 1] = 0
-
-        probs[np.isnan(probs)] = 0
-        cutoffs = np.cumsum(probs)
-        idx = cutoffs.searchsorted(np.random.uniform(0, cutoffs[-1]))
+        idx = probs.searchsorted(np.random.uniform(0, probs[-1]))
 
         return idx
 
@@ -36,7 +30,7 @@ def random_pick_inlet(self, choices, probs=None):
         """
 
         if not probs:
-            probs = np.array([1 for i in range(len(choices))])
+            probs = np.ones(len(choices))
 
         cutoffs = np.cumsum(probs)
         idx = cutoffs.searchsorted(np.random.uniform(0, cutoffs[-1]))
@@ -49,4 +43,4 @@ def _get_version():
     Extract version number from single file, and make it availabe everywhere.
     """
     from . import _version
-    return _version.__version__()
\ No newline at end of file
+    return _version.__version__()
diff --git a/pyDeltaRCM/water_tools.py b/pyDeltaRCM/water_tools.py
index 6dd86ea3..bb26642b 100644
--- a/pyDeltaRCM/water_tools.py
+++ b/pyDeltaRCM/water_tools.py
@@ -25,104 +25,6 @@
 
 class water_tools(shared_tools):
 
-    def update_flow_field(self, iteration):
-        """
-        Update water discharge after one water iteration
-        """
-
-        timestep = self._time
-
-        dloc = (self.qxn**2 + self.qyn**2)**(0.5)
-
-        qwn_div = np.ones((self.L, self.W))
-        qwn_div[dloc > 0] = self.qwn[dloc > 0] / dloc[dloc > 0]
-
-        self.qxn *= qwn_div
-        self.qyn *= qwn_div
-
-        if timestep > 0:
-
-            omega = self.omega_flow_iter
-
-            if iteration == 0:
-                omega = self.omega_flow
-
-            self.qx = self.qxn * omega + self.qx * (1 - omega)
-            self.qy = self.qyn * omega + self.qy * (1 - omega)
-
-        else:
-            self.qx = self.qxn.copy()
-            self.qy = self.qyn.copy()
-
-        self.qw = (self.qx**2 + self.qy**2)**(0.5)
-
-        self.qx[0, self.inlet] = self.qw0
-        self.qy[0, self.inlet] = 0
-        self.qw[0, self.inlet] = self.qw0
-
-    def update_velocity_field(self):
-        """
-        Update the flow velocity field after one water iteration
-        """
-
-        mask = (self.depth > self.dry_depth) * (self.qw > 0)
-
-        self.uw[mask] = np.minimum(
-            self.u_max, self.qw[mask] / self.depth[mask])
-        self.uw[~mask] = 0
-        self.ux[mask] = self.uw[mask] * self.qx[mask] / self.qw[mask]
-        self.ux[~mask] = 0
-        self.uy[mask] = self.uw[mask] * self.qy[mask] / self.qw[mask]
-        self.uy[~mask] = 0
-
-    def flooding_correction(self):
-        """
-        Flood dry cells along the shore if necessary
-
-        Check the neighbors of all dry cells. If any dry cells have wet
-        neighbors, check that their stage is not higher than the bed elevation
-        of the center cell.
-        If it is, flood the dry cell.
-        """
-
-        wet_mask = self.depth > self.dry_depth
-        wet_mask_nh = self.get_wet_mask_nh()
-        wet_mask_nh_sum = np.sum(wet_mask_nh, axis=0)
-
-        # makes wet cells look like they have only dry neighbors
-        wet_mask_nh_sum[wet_mask] = 0
-
-        # indices of dry cells with wet neighbors
-        shore_ind = np.where(wet_mask_nh_sum > 0)
-
-        stage_nhs = self.build_weight_array(self.stage)
-        eta_shore = self.eta[shore_ind]
-
-        for i in range(len(shore_ind[0])):
-
-            # pretends dry neighbor cells have stage zero
-            # so they cannot be > eta_shore[i]
-
-            stage_nh = wet_mask_nh[:, shore_ind[0][i], shore_ind[1][i]] * \
-                stage_nhs[:, shore_ind[0][i], shore_ind[1][i]]
-
-            if (stage_nh > eta_shore[i]).any():
-                self.stage[shore_ind[0][i], shore_ind[1][i]] = max(stage_nh)
-
-    def finalize_water_iteration(self, timestep, iteration):
-        """
-        Finish updating flow fields
-        Clean up at end of water iteration
-        """
-
-        self.update_water(timestep, iteration)
-
-        self.stage[:] = np.maximum(self.stage, self.H_SL)
-        self.depth[:] = np.maximum(self.stage - self.eta, 0)
-
-        self.update_flow_field(iteration)
-        self.update_velocity_field()
-
     def init_water_iteration(self):
 
         self.qxn[:] = 0
@@ -143,17 +45,18 @@ 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)]
+        start_indices = [self.random_pick_inlet(self.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.
+        self.qwn.flat[start_indices] += self.Qp_water / self.dx / 2
 
         self.indices[:, 0] = start_indices
         current_inds = list(start_indices)
 
         self.looped[:] = 0
 
+        self.get_water_weight_array()
+
         while (sum(current_inds) > 0) & (iter < self.itmax):
 
             iter += 1
@@ -161,17 +64,15 @@ def run_water_iteration(self):
             self.check_size_of_indices_matrix(iter)
 
             inds = np.unravel_index(current_inds, self.depth.shape)
-            inds_tuple = [(inds[0][i], inds[1][i])
-                          for i in range(len(inds[0]))]
+            inds_tuple = list(zip(*inds))
 
-            new_cells = [self.get_weight(x)
-                         if x != (0, 0) else 4 for x in inds_tuple]
+            new_cells = [self.get_new_cell(x)
+                         if x != (0, 0) else 4
+                         for x in inds_tuple]
 
-            new_inds = list(map(lambda x, y: self.calculate_new_ind(x, y)
-                                if y != 4 else 0, inds_tuple, new_cells))
+            new_inds = self.calculate_new_indices(inds_tuple, new_cells)
 
-            dist = list(map(lambda x, y, z: self.step_update(x, y, z) if x > 0
-                            else 0, current_inds, new_inds, new_cells))
+            dist = self.update_steps(current_inds, new_inds, new_cells)
 
             new_inds = np.array(new_inds, dtype=np.int)
             new_inds[np.array(dist) == 0] = 0
@@ -186,57 +87,6 @@ def run_water_iteration(self):
 
             current_inds[self.free_surf_flag > 0] = 0
 
-    def check_for_boundary(self, inds):
-
-        self.free_surf_flag[(self.cell_type.flat[inds] == -1)
-                            & (self.free_surf_flag == 0)] = 1
-
-        self.free_surf_flag[(self.cell_type.flat[inds] == -1)
-                            & (self.free_surf_flag == -1)] = 2
-
-        inds[self.free_surf_flag == 2] = 0
-
-        return inds
-
-    def check_for_loops(self, inds, it):
-
-        looped = [len(i[i > 0]) != len(set(i[i > 0])) for i in self.indices]
-
-        for n in range(self.Np_water):
-
-            ind = inds[n]
-            # revised 'it' to 'np.max(it)' to match python 2 > assessment
-            if (looped[n]) and (ind > 0) and (np.max(it) > self.L0):
-
-                self.looped[n] += 1
-
-                it = np.unravel_index(ind, self.depth.shape)
-
-                px, py = it
-
-                Fx = px - 1
-                Fy = py - self.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, self.L0)
-                px = int(min(self.L - 2, px))
-
-                py = max(1, py)
-                py = int(min(self.W - 2, py))
-
-                nind = np.ravel_multi_index((px, py), self.depth.shape)
-
-                inds[n] = nind
-
-                self.free_surf_flag[n] = -1
-
-        return inds
-
     def free_surf(self, it):
         """calculate free surface after routing one water parcel"""
 
@@ -252,7 +102,6 @@ def free_surf(self, it):
             if ((self.cell_type[xs[-1], ys[-1]] == -1) and
                     (self.looped[n] == 0)):
 
-                self.count += 1
                 Hnew[xs[-1], ys[-1]] = self.H_SL
                 # if cell is in ocean, H = H_SL (downstream boundary condition)
 
@@ -302,8 +151,190 @@ def free_surf(self, it):
                     self.sfc_visit[i, j] = self.sfc_visit[i, j] + 1
                     # add up # of cell visits
 
-                    self.sfc_sum[i, j] = self.sfc_sum[i, j] + Hnew[i, j]
                     # sum of all water surface elevations
+                    self.sfc_sum[i, j] = self.sfc_sum[i, j] + Hnew[i, j]
+
+    def finalize_water_iteration(self, timestep, iteration):
+        """
+        Finish updating flow fields
+        Clean up at end of water iteration
+        """
+
+        self.update_water(timestep, iteration)
+
+        self.stage[:] = np.maximum(self.stage, self.H_SL)
+        self.depth[:] = np.maximum(self.stage - self.eta, 0)
+
+        self.update_flow_field(iteration)
+        self.update_velocity_field()
+
+    def check_size_of_indices_matrix(self, it):
+        if it >= self.indices.shape[1]:
+            """
+            Initial size of self.indices is half of self.itmax
+            because the number of iterations doesn't go beyond
+            that for many timesteps.
+
+            Once it reaches it > self.itmax/2 once, make the size
+            self.iter for all further timesteps
+            """
+
+            if self.verbose >= 2:
+                self.logger.info('Increasing size of self.indices')
+
+            indices_blank = np.zeros(
+                (np.int(self.Np_water), np.int(self.itmax / 4)), dtype=np.int)
+
+            self.indices = np.hstack((self.indices, indices_blank))
+
+    def get_water_weight_array(self):
+
+        self.water_weights = np.zeros(shape = (self.L, self.W, 9))
+
+        for i in range(self.L):
+            for j in range(self.W):
+                self.water_weights[i, j] = self.get_weight_at_cell((i, j))
+
+    def get_weight_at_cell(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)
+
+        weight = self.gamma * weight_sfc + (1 - self.gamma) * weight_int
+        weight = depth_ind ** self.theta_water * weight
+        weight[depth_ind <= self.dry_depth] = 0
+        if np.sum(weight) != 0:
+            weight = weight / np.sum(weight)
+        else:
+            weight = weight
+
+        return np.cumsum(weight.flatten())
+
+    def get_new_cell(self, ind):
+        weight = self.water_weights[ind[0], ind[1]]
+        new_cell = self.random_pick(weight)
+        return new_cell
+
+    def calculate_new_ind(self, ind, new_cell):
+
+        new_ind = (ind[0] + self.jwalk.flat[new_cell],
+                   ind[1] + self.iwalk.flat[new_cell])
+        # added wrap mode to fct to resolve ValueError due to negative numbers
+        new_ind_flat = np.ravel_multi_index(
+            new_ind, self.depth.shape, mode='wrap')
+
+        return new_ind_flat
+
+    def calculate_new_indices(self, inds_tuple, new_cells):
+        newbies = []
+        for i, j in zip(inds_tuple, new_cells):
+            if j != 4:
+                newbies.append(self.calculate_new_ind(i, j))
+            else:
+                newbies.append(0)
+        return newbies
+
+    def step_update(self, ind, new_ind, new_cell):
+
+        istep = self.iwalk.flat[new_cell]
+        jstep = self.jwalk.flat[new_cell]
+        dist = np.sqrt(istep * istep + jstep * jstep)
+
+        if dist > 0:
+
+            self.qxn.flat[ind] += jstep / dist
+            self.qyn.flat[ind] += istep / dist
+            self.qwn.flat[ind] += self.Qp_water / self.dx / 2
+
+            self.qxn.flat[new_ind] += jstep / dist
+            self.qyn.flat[new_ind] += istep / dist
+            self.qwn.flat[new_ind] += self.Qp_water / self.dx / 2
+
+        return dist
+
+    def update_steps(self, current_inds, new_inds, new_cells):
+        newbies = []
+        for i, j, k in zip(current_inds, new_inds, new_cells):
+            if i > 0:
+                newbies.append(self.step_update(i, j, k))
+            else:
+                newbies.append(0)
+        return newbies
+
+    def check_for_loops(self, inds, it):
+
+        looped = [len(i[i > 0]) != len(set(i[i > 0])) for i in self.indices]
+
+        for n in range(self.Np_water):
+
+            ind = inds[n]
+            # revised 'it' to 'np.max(it)' to match python 2 > assessment
+            if (looped[n]) and (ind > 0) and (np.max(it) > self.L0):
+
+                self.looped[n] += 1
+
+                it = np.unravel_index(ind, self.depth.shape)
+
+                px, py = it
+
+                Fx = px - 1
+                Fy = py - self.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, self.L0)
+                px = int(min(self.L - 2, px))
+
+                py = max(1, py)
+                py = int(min(self.W - 2, py))
+
+                nind = np.ravel_multi_index((px, py), self.depth.shape)
+
+                inds[n] = nind
+
+                self.free_surf_flag[n] = -1
+
+        return inds
+
+    def check_for_boundary(self, inds):
+
+        self.free_surf_flag[(self.cell_type.flat[inds] == -1)
+                            & (self.free_surf_flag == 0)] = 1
+
+        self.free_surf_flag[(self.cell_type.flat[inds] == -1)
+                            & (self.free_surf_flag == -1)] = 2
+
+        inds[self.free_surf_flag == 2] = 0
+
+        return inds
 
     def update_water(self, timestep, itr):
         """Update surface after routing all parcels.
@@ -361,70 +392,50 @@ def update_water(self, timestep, itr):
 
         self.flooding_correction()
 
-    def step_update(self, ind, new_ind, new_cell):
-
-        istep = self.iwalk.flat[new_cell]
-        jstep = self.jwalk.flat[new_cell]
-        dist = np.sqrt(istep**2 + jstep**2)
-
-        if dist > 0:
-
-            self.qxn.flat[ind] += jstep / dist
-            self.qyn.flat[ind] += istep / dist
-            self.qwn.flat[ind] += self.Qp_water / self.dx / 2.
-
-            self.qxn.flat[new_ind] += jstep / dist
-            self.qyn.flat[new_ind] += istep / dist
-            self.qwn.flat[new_ind] += self.Qp_water / self.dx / 2.
-
-        return dist
-
-    def calculate_new_ind(self, ind, new_cell):
-
-        new_ind = (ind[0] + self.jwalk.flat[new_cell], ind[1] +
-                   self.iwalk.flat[new_cell])
-        # added wrap mode to fct to resolve ValueError due to negative numbers
-        new_ind_flat = np.ravel_multi_index(
-            new_ind, self.depth.shape, mode='wrap')
-
-        return new_ind_flat
+    def flooding_correction(self):
+        """
+        Flood dry cells along the shore if necessary
 
-    def get_weight(self, ind):
+        Check the neighbors of all dry cells. If any dry cells have wet
+        neighbors, check that their stage is not higher than the bed elevation
+        of the center cell.
+        If it is, flood the dry cell.
+        """
 
-        stage_ind = self.pad_stage[
-            ind[0] - 1 + 1:ind[0] + 2 + 1, ind[1] - 1 + 1:ind[1] + 2 + 1]
+        wet_mask = self.depth > self.dry_depth
+        wet_mask_nh = self.get_wet_mask_nh()
+        wet_mask_nh_sum = np.sum(wet_mask_nh, axis=0)
 
-        weight_sfc = np.maximum(0,
-                                (self.stage[ind] - stage_ind) / self.distances)
+        # makes wet cells look like they have only dry neighbors
+        wet_mask_nh_sum[wet_mask] = 0
 
-        weight_int = np.maximum(0, (self.qx[ind] * self.jvec +
-                                    self.qy[ind] * self.ivec) / self.distances)
+        # indices of dry cells with wet neighbors
+        shore_ind = np.where(wet_mask_nh_sum > 0)
 
-        if ind[0] == 0:
-            weight_sfc[0, :] = 0
-            weight_int[0, :] = 0
+        stage_nhs = self.build_weight_array(self.stage)
+        eta_shore = self.eta[shore_ind]
 
-        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]
+        for i in range(len(shore_ind[0])):
 
-        weight_sfc[(depth_ind <= self.dry_depth) | (ct_ind == -2)] = 0
-        weight_int[(depth_ind <= self.dry_depth) | (ct_ind == -2)] = 0
+            # pretends dry neighbor cells have stage zero
+            # so they cannot be > eta_shore[i]
 
-        if np.nansum(weight_sfc) > 0:
-            weight_sfc = weight_sfc / np.nansum(weight_sfc)
+            stage_nh = wet_mask_nh[:, shore_ind[0][i], shore_ind[1][i]] * \
+                stage_nhs[:, shore_ind[0][i], shore_ind[1][i]]
 
-        if np.nansum(weight_int) > 0:
-            weight_int = weight_int / np.nansum(weight_int)
+            if (stage_nh > eta_shore[i]).any():
+                self.stage[shore_ind[0][i], shore_ind[1][i]] = max(stage_nh)
 
-        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
+    def get_wet_mask_nh(self):
+        """
+        Returns np.array((8,L,W)), for each neighbor around a cell
+        with 1 if the neighbor is wet and 0 if dry
+        """
 
-        new_cell = self.random_pick(self.weight)
+        wet_mask = (self.depth > self.dry_depth) * 1
+        wet_mask_nh = self.build_weight_array(wet_mask, fix_edges=True)
 
-        return new_cell
+        return wet_mask_nh
 
     def build_weight_array(self, array, fix_edges=False, normalize=False):
         """
@@ -467,32 +478,52 @@ def build_weight_array(self, array, fix_edges=False, normalize=False):
 
         return wgt_array
 
-    def get_wet_mask_nh(self):
+    def update_flow_field(self, iteration):
         """
-        Returns np.array((8,L,W)), for each neighbor around a cell
-        with 1 if the neighbor is wet and 0 if dry
+        Update water discharge after one water iteration
         """
 
-        wet_mask = (self.depth > self.dry_depth) * 1
-        wet_mask_nh = self.build_weight_array(wet_mask, fix_edges=True)
+        timestep = self._time
 
-        return wet_mask_nh
+        dloc = (self.qxn**2 + self.qyn**2)**(0.5)
 
-    def check_size_of_indices_matrix(self, it):
-        if it >= self.indices.shape[1]:
-            """
-            Initial size of self.indices is half of self.itmax
-            because the number of iterations doesn't go beyond
-            that for many timesteps.
+        qwn_div = np.ones((self.L, self.W))
+        qwn_div[dloc > 0] = self.qwn[dloc > 0] / dloc[dloc > 0]
 
-            Once it reaches it > self.itmax/2 once, make the size
-            self.iter for all further timesteps
-            """
+        self.qxn *= qwn_div
+        self.qyn *= qwn_div
 
-            if self.verbose >= 2:
-                self.logger.info('Increasing size of self.indices')
+        if timestep > 0:
 
-            indices_blank = np.zeros(
-                (np.int(self.Np_water), np.int(self.itmax / 4)), dtype=np.int)
+            omega = self.omega_flow_iter
 
-            self.indices = np.hstack((self.indices, indices_blank))
+            if iteration == 0:
+                omega = self.omega_flow
+
+            self.qx = self.qxn * omega + self.qx * (1 - omega)
+            self.qy = self.qyn * omega + self.qy * (1 - omega)
+
+        else:
+            self.qx = self.qxn.copy()
+            self.qy = self.qyn.copy()
+
+        self.qw = (self.qx**2 + self.qy**2)**(0.5)
+
+        self.qx[0, self.inlet] = self.qw0
+        self.qy[0, self.inlet] = 0
+        self.qw[0, self.inlet] = self.qw0
+
+    def update_velocity_field(self):
+        """
+        Update the flow velocity field after one water iteration
+        """
+
+        mask = (self.depth > self.dry_depth) * (self.qw > 0)
+
+        self.uw[mask] = np.minimum(
+            self.u_max, self.qw[mask] / self.depth[mask])
+        self.uw[~mask] = 0
+        self.ux[mask] = self.uw[mask] * self.qx[mask] / self.qw[mask]
+        self.ux[~mask] = 0
+        self.uy[mask] = self.uw[mask] * self.qy[mask] / self.qw[mask]
+        self.uy[~mask] = 0