diff --git a/pyDeltaRCM/init_tools.py b/pyDeltaRCM/init_tools.py index 8cb14a7f..93867e27 100644 --- a/pyDeltaRCM/init_tools.py +++ b/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 = \ diff --git a/pyDeltaRCM/sed_tools.py b/pyDeltaRCM/sed_tools.py index c7c47084..ece4ebc4 100644 --- a/pyDeltaRCM/sed_tools.py +++ b/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): diff --git a/pyDeltaRCM/shared_tools.py b/pyDeltaRCM/shared_tools.py index 34ae9af6..e6d525e2 100644 --- a/pyDeltaRCM/shared_tools.py +++ b/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(): diff --git a/pyDeltaRCM/water_tools.py b/pyDeltaRCM/water_tools.py index bb26642b..9ea5456a 100644 --- a/pyDeltaRCM/water_tools.py +++ b/pyDeltaRCM/water_tools.py @@ -1,4 +1,3 @@ - import sys import os import re @@ -17,13 +16,12 @@ from scipy import ndimage from netCDF4 import Dataset - -from .shared_tools import shared_tools +from . import shared_tools # tools for water routing algorithms -class water_tools(shared_tools): +class water_tools(object): def init_water_iteration(self): @@ -45,13 +43,16 @@ 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)] + inlet_weights = np.ones_like(self.inlet) + start_indices = [ + self.inlet[shared_tools.random_pick(inlet_weights / sum(inlet_weights))] + 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.indices[:, 0] = start_indices - current_inds = list(start_indices) + current_inds = np.array(start_indices) self.looped[:] = 0 @@ -70,16 +71,31 @@ def run_water_iteration(self): if x != (0, 0) else 4 for x in inds_tuple] - new_inds = self.calculate_new_indices(inds_tuple, new_cells) + new_cells = np.array(new_cells) - dist = self.update_steps(current_inds, new_inds, new_cells) + next_index = shared_tools.calculate_new_ind( + current_inds, + new_cells, + self.iwalk.flatten(), + self.jwalk.flatten(), + self.eta.shape) - new_inds = np.array(new_inds, dtype=np.int) - new_inds[np.array(dist) == 0] = 0 + dist, istep, jstep, astep = shared_tools.get_steps( + new_cells, + self.iwalk.flat[:], + self.jwalk.flat[:]) - self.indices[:, iter] = new_inds + self.update_Q(dist, current_inds, next_index, astep, jstep, istep) - current_inds = self.check_for_loops(new_inds, iter) + current_inds, self.looped, self.free_surf_flag = shared_tools.check_for_loops( + self.indices, + next_index, + iter, + self.L0, + self.looped, + self.eta.shape, + self.CTR, + self.free_surf_flag) current_inds = self.check_for_boundary(current_inds) @@ -193,144 +209,41 @@ def get_water_weight_array(self): 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()) + stage_nbrs = self.pad_stage[i - 1 + 1:i + 2 + 1, j - 1 + 1:j + 2 + 1] + depth_nbrs = self.pad_depth[i - 1 + 1:i + 2 + 1, j - 1 + 1:j + 2 + 1] + ct_nbrs = self.pad_cell_type[i - 1 + 1:i + 2 + 1, j - 1 + 1:j + 2 + 1] + self.water_weights[i, j] = shared_tools.get_weight_at_cell( + (i, j), + stage_nbrs.flatten(), depth_nbrs.flatten(), ct_nbrs.flatten(), + self.stage[i, j], self.qx[i, j], self.qy[i, j], + self.ivec.flatten(), self.jvec.flatten(), self.distances.flatten(), + self.dry_depth, self.gamma, self.theta_water) def get_new_cell(self, ind): weight = self.water_weights[ind[0], ind[1]] - new_cell = self.random_pick(weight) + new_cell = shared_tools.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 update_Q(self, dist, current_inds, next_index, astep, jstep, istep): + + self.qxn = shared_tools.update_dirQfield(self.qxn.flat[:], dist, current_inds, astep, jstep + ).reshape(self.qxn.shape) + self.qyn = shared_tools.update_dirQfield(self.qyn.flat[:], dist, current_inds, astep, istep + ).reshape(self.qyn.shape) + self.qwn = shared_tools.update_absQfield(self.qwn.flat[:], dist, current_inds, astep, self.Qp_water, self.dx + ).reshape(self.qwn.shape) + self.qxn = shared_tools.update_dirQfield(self.qxn.flat[:], dist, next_index, astep, jstep + ).reshape(self.qxn.shape) + self.qyn = shared_tools.update_dirQfield(self.qyn.flat[:], dist, next_index, astep, istep + ).reshape(self.qyn.shape) + self.qwn = shared_tools.update_absQfield(self.qwn.flat[:], dist, next_index, astep, self.Qp_water, self.dx + ).reshape(self.qwn.shape) 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 == 0)] = 1 - self.free_surf_flag[(self.cell_type.flat[inds] == -1) - & (self.free_surf_flag == -1)] = 2 + self.free_surf_flag[(self.cell_type.flat[inds] == -1) & (self.free_surf_flag == -1)] = 2 inds[self.free_surf_flag == 2] = 0 diff --git a/requirements.txt b/requirements.txt index 90684434..d60e3ff1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,3 +4,4 @@ scipy netCDF4 basic_modeling_interface pyyaml +numba diff --git a/setup.py b/setup.py index 3bb8c793..97f73513 100644 --- a/setup.py +++ b/setup.py @@ -13,7 +13,8 @@ include_package_data=True, url='https://github.com/DeltaRCM/pyDeltaRCM', install_requires=['matplotlib', 'netCDF4', - 'basic-modeling-interface', 'scipy', 'numpy', 'pyyaml'], + 'basic-modeling-interface', 'scipy', 'numpy', 'pyyaml', + 'numba'], entry_points={ 'console_scripts': ['run_pyDeltaRCM=pyDeltaRCM.command_line:run_model'], } diff --git a/tests/test_shared_tools.py b/tests/test_shared_tools.py index 6580cd57..4f8939bc 100644 --- a/tests/test_shared_tools.py +++ b/tests/test_shared_tools.py @@ -24,13 +24,4 @@ def test_random_pick(): probs = np.zeros((8,)) probs[0] = 1 # should return first index - assert delta.random_pick(probs) == 0 - - -def test_random_pick_inlet(): - """ - Test for function shared_tools.random_pick_inlet - """ - choices = [0] - probs = np.ones((1,)) - assert delta.random_pick_inlet(choices, probs) == 0 + assert shared_tools.random_pick(probs) == 0 diff --git a/tests/test_water_tools.py b/tests/test_water_tools.py index c6643847..e95bbd0d 100644 --- a/tests/test_water_tools.py +++ b/tests/test_water_tools.py @@ -8,6 +8,7 @@ from pyDeltaRCM.deltaRCM_driver import pyDeltaRCM from pyDeltaRCM import water_tools +from pyDeltaRCM import shared_tools # need to create a simple case of pydeltarcm object to test these functions delta = pyDeltaRCM(input_file=os.path.join(os.getcwd(), 'tests', 'test.yaml')) @@ -76,8 +77,9 @@ def test_calculate_new_ind(): Test for function water_tools.calculate_new_ind """ # assign old index - old_ind = [0, 4] + old_inds = np.array([4, 5]) # assign new cell location - new_cell = 7 + new_cells = np.array([7, 7]) # expect new cell to be in location (1,4) -> 14 - assert delta.calculate_new_ind(old_ind, new_cell) == 14 + new_inds = shared_tools.calculate_new_ind(old_inds, new_cells, delta.iwalk.flatten(), delta.jwalk.flatten(), delta.eta.shape) + assert np.all(new_inds == np.array([14, 15])) diff --git a/tests/test_yaml_parsing.py b/tests/test_yaml_parsing.py index 0703ec29..90da48ce 100644 --- a/tests/test_yaml_parsing.py +++ b/tests/test_yaml_parsing.py @@ -6,6 +6,7 @@ from pyDeltaRCM.deltaRCM_driver import pyDeltaRCM +from pyDeltaRCM.shared_tools import set_random_seed, get_random_uniform # utilities for file writing def create_temporary_file(tmp_path, file_name): @@ -105,13 +106,13 @@ def test_random_seed_settings_value(tmp_path): p, f = create_temporary_file(tmp_path, file_name) write_parameter_to_file(f, 'seed', 9999) f.close() - np.random.seed(9999) - _preval_same = np.random.uniform() - np.random.seed(5) - _preval_diff = np.random.uniform(1000) + set_random_seed(9999) + _preval_same = get_random_uniform(1) + set_random_seed(5) + _preval_diff = get_random_uniform(1000) delta = pyDeltaRCM(input_file=p) assert delta.seed == 9999 - _postval_same = np.random.uniform() + _postval_same = get_random_uniform(1) assert _preval_same == _postval_same assert delta.seed == 9999