From ef18f96fe4310af48e9471617f3eef7a7f9f9bed Mon Sep 17 00:00:00 2001 From: Eric Barefoot Date: Mon, 25 May 2020 20:51:47 -0500 Subject: [PATCH 1/8] 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 | 13 +-- pyDeltaRCM/sed_tools.py | 61 +++++----- pyDeltaRCM/shared_tools.py | 223 +++++++++++++++++++++++++++++++++---- pyDeltaRCM/water_tools.py | 195 +++++++++----------------------- requirements.txt | 1 + setup.py | 3 +- tests/test_shared_tools.py | 11 +- tests/test_water_tools.py | 8 +- tests/test_yaml_parsing.py | 11 +- 9 files changed, 303 insertions(+), 223 deletions(-) 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 From a91aea32cc5a50f55279006c247326c34aae0d12 Mon Sep 17 00:00:00 2001 From: Andrew Moodie Date: Tue, 26 May 2020 08:16:25 -0500 Subject: [PATCH 2/8] Revert "fix tests from other branch" This reverts commit b5e3bcf35e7fed373e516802afe5186979152deb. --- tests/test_deltaRCM_driver.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/test_deltaRCM_driver.py b/tests/test_deltaRCM_driver.py index ea7c2408..c9697f23 100644 --- a/tests/test_deltaRCM_driver.py +++ b/tests/test_deltaRCM_driver.py @@ -20,11 +20,8 @@ def test_init(): assert delta._is_finalized == False -delta = pyDeltaRCM(input_file=os.path.join(os.getcwd(), 'tests', 'test.yaml')) - - def test_update(): - + delta = pyDeltaRCM(input_file=os.path.join(os.getcwd(), 'tests', 'test.yaml')) delta.update() assert delta._time == 1.0 assert delta._is_finalized == False From b0a4d9ba50fb84e0221fd5eff56ca5a6ccefae9a Mon Sep 17 00:00:00 2001 From: Andrew Moodie Date: Tue, 26 May 2020 08:20:50 -0500 Subject: [PATCH 3/8] update loop check to typed list --- pyDeltaRCM/shared_tools.py | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/pyDeltaRCM/shared_tools.py b/pyDeltaRCM/shared_tools.py index e6d525e2..8d6758f7 100644 --- a/pyDeltaRCM/shared_tools.py +++ b/pyDeltaRCM/shared_tools.py @@ -3,7 +3,7 @@ import numpy as np import time -from numba import njit +from numba import njit, jit, typed # tools shared between deltaRCM water and sediment routing @@ -131,22 +131,18 @@ def custom_ravel(tup, shape): @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] - + 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 - - it = custom_unravel(ind, domain_shape) - - px, py = it - + px, py = custom_unravel(ind, domain_shape) Fx = px - 1 Fy = py - CTR @@ -171,6 +167,7 @@ def check_for_loops(indices, inds, it, L0, loopedout, domain_shape, CTR, free_su return inds, loopedout, free_surf_flag + @njit def calculate_new_ind(indices, new_cells, iwalk, jwalk, domain_shape): newbies = [] From 404f29b6d75c0ea9f0eab02580cae645db9d7615 Mon Sep 17 00:00:00 2001 From: Andrew Moodie Date: Tue, 26 May 2020 08:22:30 -0500 Subject: [PATCH 4/8] fix missing delta in old test --- tests/test_deltaRCM_driver.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_deltaRCM_driver.py b/tests/test_deltaRCM_driver.py index c9697f23..2dd83eaa 100644 --- a/tests/test_deltaRCM_driver.py +++ b/tests/test_deltaRCM_driver.py @@ -28,6 +28,7 @@ def test_update(): def test_finalize(): + delta = pyDeltaRCM(input_file=os.path.join(os.getcwd(), 'tests', 'test.yaml')) delta.finalize() assert delta._is_finalized == True From 092d98ca418b93c6d72694042fdeaa28eaeb26af Mon Sep 17 00:00:00 2001 From: Eric Barefoot Date: Tue, 26 May 2020 13:54:22 -0500 Subject: [PATCH 5/8] fix: repair bugs for water and sediment routing Updated sediment routing to use the same weight choosing function as water, and fixed the water so that any non-disallowed cell is given equal weight if all weight is zero --- pyDeltaRCM/sed_tools.py | 20 +++++++------- pyDeltaRCM/shared_tools.py | 53 +++++++++++--------------------------- pyDeltaRCM/water_tools.py | 3 +-- 3 files changed, 25 insertions(+), 51 deletions(-) diff --git a/pyDeltaRCM/sed_tools.py b/pyDeltaRCM/sed_tools.py index ece4ebc4..ad384903 100644 --- a/pyDeltaRCM/sed_tools.py +++ b/pyDeltaRCM/sed_tools.py @@ -183,19 +183,17 @@ def sed_parcel(self, theta_sed, sed, px, py): # choose next with weights it += 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 = 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) + 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] - w3 = shared_tools.get_combined_weight(w1, w2, self.distances) - - weights = shared_tools.get_filtered_weight( - w3.flatten(), px, depth_ind.flatten(), - cell_type_ind.flatten(), self.dry_depth) + 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) new_cell = shared_tools.random_pick(weights) diff --git a/pyDeltaRCM/shared_tools.py b/pyDeltaRCM/shared_tools.py index 8d6758f7..06ea0311 100644 --- a/pyDeltaRCM/shared_tools.py +++ b/pyDeltaRCM/shared_tools.py @@ -30,34 +30,6 @@ def get_random_uniform(N): return np.random.uniform(0, 1) -@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: @@ -185,19 +157,19 @@ def calculate_new_ind(indices, new_cells, iwalk, jwalk, domain_shape): @njit def get_weight_at_cell(ind, stage_nbrs, depth_nbrs, ct_nbrs, stage, qx, qy, - ivec, jvec, distances, dry_depth, gamma, theta_water): + 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] = 0 - weight_int[:3] = 0 + weight_sfc[:3] = np.nan + weight_int[:3] = np.nan drywall = (depth_nbrs <= dry_depth) | (ct_nbrs == -2) - weight_sfc[drywall] = 0 - weight_int[drywall] = 0 + weight_sfc[drywall] = np.nan + weight_int[drywall] = np.nan if np.nansum(weight_sfc) > 0: weight_sfc = weight_sfc / np.nansum(weight_sfc) @@ -206,14 +178,19 @@ def get_weight_at_cell(ind, stage_nbrs, depth_nbrs, ct_nbrs, stage, qx, qy, 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 ** theta * weight weight[depth_nbrs <= dry_depth] = 0 - if np.sum(weight) != 0: - weight = weight / np.sum(weight) + nanWeight = np.isnan(weight) + + if np.any(weight[~nanWeight] != 0): + weight = weight / np.nansum(weight) + weight[nanWeight] = 0 else: - weight = weight - return weight / np.sum(weight) + weight[~nanWeight] = 1 / len(weight[~nanWeight]) + weight[nanWeight] = 0 + return weight + def _get_version(): diff --git a/pyDeltaRCM/water_tools.py b/pyDeltaRCM/water_tools.py index 9ea5456a..0add31cf 100644 --- a/pyDeltaRCM/water_tools.py +++ b/pyDeltaRCM/water_tools.py @@ -221,8 +221,7 @@ def get_water_weight_array(self): def get_new_cell(self, ind): weight = self.water_weights[ind[0], ind[1]] - new_cell = shared_tools.random_pick(weight) - return new_cell + return shared_tools.random_pick(weight) def update_Q(self, dist, current_inds, next_index, astep, jstep, istep): From 9fb7c91649fdb50421e6608d8991103d40ae3cd4 Mon Sep 17 00:00:00 2001 From: Eric Barefoot Date: Tue, 26 May 2020 14:23:32 -0500 Subject: [PATCH 6/8] docs: very basic effort to get docs compliant also repair broken test --- docs/source/reference/shared_tools/index.rst | 11 +++++------ pyDeltaRCM/shared_tools.py | 11 ++++++++++- tests/test_sed_tools.py | 2 ++ 3 files changed, 17 insertions(+), 7 deletions(-) diff --git a/docs/source/reference/shared_tools/index.rst b/docs/source/reference/shared_tools/index.rst index 053cfe34..6d232a68 100644 --- a/docs/source/reference/shared_tools/index.rst +++ b/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 diff --git a/pyDeltaRCM/shared_tools.py b/pyDeltaRCM/shared_tools.py index 06ea0311..a2885dda 100644 --- a/pyDeltaRCM/shared_tools.py +++ b/pyDeltaRCM/shared_tools.py @@ -32,6 +32,8 @@ def get_random_uniform(N): @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 @@ -47,7 +49,8 @@ def partition_sand(qs, depoPart, py, px, dist, istep, jstep): @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) @@ -59,6 +62,8 @@ def get_steps(new_cells, iwalk, jwalk): @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] @@ -67,6 +72,8 @@ def update_dirQfield(qfield, dist, inds, astep, dirstep): @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 @@ -88,6 +95,8 @@ def random_pick(prob): @njit def custom_unravel(i, shape): + '''Spread sand between two cells + ''' x = i // shape[1] y = i % shape[1] return x, y diff --git a/tests/test_sed_tools.py b/tests/test_sed_tools.py index fb577328..29c21e06 100644 --- a/tests/test_sed_tools.py +++ b/tests/test_sed_tools.py @@ -21,6 +21,8 @@ def test_sed_route(): test the function sed_tools.sed_route """ delta.pad_cell_type = np.pad(delta.cell_type, 1, 'edge') + delta.pad_stage = np.pad(delta.stage, 1, 'edge') + delta.pad_depth = np.pad(delta.depth, 1, 'edge') delta.sed_route() [a, b] = np.shape(delta.pad_depth) From b9076d992fba7a8aabfd3e7f602a85a8c7f70e41 Mon Sep 17 00:00:00 2001 From: Eric Barefoot Date: Tue, 2 Jun 2020 17:50:43 -0500 Subject: [PATCH 7/8] tests: add tests and IndexErrors to unravelling functions --- pyDeltaRCM/shared_tools.py | 11 ++++++++--- tests/test_shared_tools.py | 35 +++++++++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/pyDeltaRCM/shared_tools.py b/pyDeltaRCM/shared_tools.py index a2885dda..758d95fa 100644 --- a/pyDeltaRCM/shared_tools.py +++ b/pyDeltaRCM/shared_tools.py @@ -95,8 +95,10 @@ def random_pick(prob): @njit def custom_unravel(i, shape): - '''Spread sand between two cells - ''' + """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 @@ -104,6 +106,10 @@ def custom_unravel(i, shape): @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 @@ -148,7 +154,6 @@ def check_for_loops(indices, inds, it, L0, loopedout, domain_shape, CTR, free_su return inds, loopedout, free_surf_flag - @njit def calculate_new_ind(indices, new_cells, iwalk, jwalk, domain_shape): newbies = [] diff --git a/tests/test_shared_tools.py b/tests/test_shared_tools.py index 4f8939bc..f68ae266 100644 --- a/tests/test_shared_tools.py +++ b/tests/test_shared_tools.py @@ -25,3 +25,38 @@ def test_random_pick(): probs[0] = 1 # should return first index assert shared_tools.random_pick(probs) == 0 +def test_custom_unravel_square(): + arr = np.arange(9).reshape((3, 3)) + # test upper left corner + x, y = shared_tools.custom_unravel(0, arr.shape) + assert x == 0 + assert y == 0 + # test center + x, y = shared_tools.custom_unravel(4, arr.shape) + assert x == 1 + assert y == 1 + # test off-center + x, y = shared_tools.custom_unravel(5, arr.shape) + assert x == 1 + assert y == 2 + # test lower right corner + x, y = shared_tools.custom_unravel(8, arr.shape) + assert x == 2 + assert y == 2 + + +def test_custom_unravel_rectangle(): + arr = np.arange(50).reshape((5, 10)) + # test a few spots + x, y = shared_tools.custom_unravel(19, arr.shape) + assert x == 1 + assert y == 9 + x, y = shared_tools.custom_unravel(34, arr.shape) + assert x == 3 + assert y == 4 + + +@pytest.mark.xfail(raises=IndexError, strict=True) +def test_custom_unravel_exceed_error(): + arr = np.arange(9).reshape((3, 3)) + x, y = shared_tools.custom_unravel(99, arr.shape) From 56a36ca83a695238e621ad329b4ae135d16bb05e Mon Sep 17 00:00:00 2001 From: Eric Barefoot Date: Tue, 2 Jun 2020 17:51:14 -0500 Subject: [PATCH 8/8] tests: add tests for all of shared tools and some for model init tools --- tests/test_deltaRCM_driver.py | 53 ++++++++++- tests/test_shared_tools.py | 169 ++++++++++++++++++++++++++++++++++ 2 files changed, 221 insertions(+), 1 deletion(-) diff --git a/tests/test_deltaRCM_driver.py b/tests/test_deltaRCM_driver.py index 2dd83eaa..92440731 100644 --- a/tests/test_deltaRCM_driver.py +++ b/tests/test_deltaRCM_driver.py @@ -38,9 +38,60 @@ def test_multifinalization_error(): err_delta = pyDeltaRCM(input_file=os.path.join(os.getcwd(), 'tests', 'test.yaml')) err_delta.update() # test will Fail if any assertion is wrong - assert err_delta._time == 1.0 + assert err_delta._time == 1.0 assert err_delta._is_finalized == False err_delta.finalize() assert err_delta._is_finalized == True # next line should throw RuntimeError and test will xFail err_delta.update() + + +def test_getters_setters(): + + delta = pyDeltaRCM(input_file=os.path.join(os.getcwd(), 'tests', 'test.yaml')) + + assert np.all(delta.sea_surface_elevation == 0) + assert delta.water_depth[0, 2] == 0 + assert delta.water_depth[0, 3] == 1 + assert delta.water_depth[4, 4] == 1 + assert delta.bed_elevation[0, 2] == 0 + assert delta.bed_elevation[0, 3] == -1 + assert delta.bed_elevation[4, 4] == -1 + + assert delta.sea_surface_mean_elevation == 0 + delta.sea_surface_mean_elevation = 0.5 + assert delta.sea_surface_mean_elevation == 0.5 + + assert delta.sea_surface_elevation_change == 0.001 + delta.sea_surface_elevation_change = 0.002 + assert delta.sea_surface_elevation_change == 0.002 + + assert delta.bedload_fraction == 0.5 + delta.bedload_fraction = 0.25 + assert delta.bedload_fraction == 0.25 + + delta = pyDeltaRCM(input_file=os.path.join(os.getcwd(), 'tests', 'test.yaml')) + + assert delta.channel_flow_velocity == 1 + delta.channel_flow_velocity = 3 + assert delta.channel_flow_velocity == 3 + assert delta.u_max == 6 + + delta = pyDeltaRCM(input_file=os.path.join(os.getcwd(), 'tests', 'test.yaml')) + + assert delta.channel_width == 2 + delta.channel_width = 10 + assert delta.channel_width == 10 + assert delta.N0 == 10 + + delta = pyDeltaRCM(input_file=os.path.join(os.getcwd(), 'tests', 'test.yaml')) + + assert delta.channel_flow_depth == 1 + delta.channel_flow_depth = 2 + assert delta.channel_flow_depth == 2 + + delta = pyDeltaRCM(input_file=os.path.join(os.getcwd(), 'tests', 'test.yaml')) + + assert delta.influx_sediment_concentration == 0.1 + delta.influx_sediment_concentration = 2 + assert delta.C0 == 0.02 diff --git a/tests/test_shared_tools.py b/tests/test_shared_tools.py index f68ae266..d29d496c 100644 --- a/tests/test_shared_tools.py +++ b/tests/test_shared_tools.py @@ -16,6 +16,85 @@ # delta._delta.**shared_tools_function** +def test_set_random_assignments(): + """ + Test for function shared_tools.get_random_uniform and + test for function shared_tools.set_random_seed + """ + shared_tools.set_random_seed(delta.seed) + got = shared_tools.get_random_uniform(1) + _exp = 0.5488135039273248 + assert got == pytest.approx(_exp) + + +def test_sand_partition(): + """ + Test for function shared_tools.partition_sand + """ + nx, ny, qsn = shared_tools.partition_sand( + delta.qs, 1, 4, 4, 1, 0, 1 + ) + assert nx == 5 + assert ny == 4 + assert qsn[4, 4] == 1 + assert qsn[5, 4] == 1 + assert np.all(qsn[qsn != 1] == 0) + + +def test_get_steps(): + """ + Test for function shared_tools.get_steps + """ + new_cells = np.arange(9) + iwalk = shared_tools.get_iwalk() + jwalk = shared_tools.get_jwalk() + + d, i, j, a = shared_tools.get_steps(new_cells, iwalk.flatten(), jwalk.flatten()) + + d_exp = np.array([1.41421356, 1., 1.41421356, 1., 0., + 1., 1.41421356, 1., 1.41421356]) + i_exp = np.array([-1, 0, 1, -1, 0, 1, -1, 0, 1]) + j_exp = np.array([-1, -1, -1, 0, 0, 0, 1, 1, 1]) + + assert np.all(np.delete(a, 4)) + assert ~a[4] + assert i == pytest.approx(i_exp) + assert j == pytest.approx(j_exp) + assert d == pytest.approx(d_exp) + + +def test_update_dirQfield(): + """ + Test for function shared_tools.update_dirQfield + """ + np.random.seed(delta.seed) + qx = np.random.uniform(0, 10, 9) + d = np.array([1, np.sqrt(2), 0]) + astep = np.array([True, True, False]) + inds = np.array([3, 4, 5]) + stepdir = np.array([1, 1, 0]) + qxn = shared_tools.update_dirQfield(np.copy(qx), d, inds, astep, stepdir) + qxdiff = qxn - qx + qxdiff_exp = np.array([1, np.sqrt(2) / 2, 0]) + assert np.all(qxdiff[3:6] == pytest.approx(qxdiff_exp)) + + +def test_update_absQfield(): + """ + Test for function shared_tools.update_absQfield + """ + np.random.seed(delta.seed) + qw = np.random.uniform(0, 10, 9) + d = np.array([1, np.sqrt(2), 0]) + astep = np.array([True, True, False]) + inds = np.array([3, 4, 5]) + qwn = shared_tools.update_absQfield(np.copy(qw), d, inds, astep, delta.Qp_water, delta.dx) + qwdiff = qwn - qw + diffelem = delta.Qp_water / delta.dx / 2 + qwdiff_exp = np.array([diffelem, diffelem, 0]) + assert np.all(qwdiff[3:6] == pytest.approx(qwdiff_exp)) + + def test_random_pick(): """ Test for function shared_tools.random_pick @@ -25,6 +104,31 @@ def test_random_pick(): probs[0] = 1 # should return first index assert shared_tools.random_pick(probs) == 0 + + +def test_random_pick_anybut_first(): + """ + Test for function shared_tools.random_pick + """ + shared_tools.set_random_seed(delta.seed) + probs = (1 / 7) * np.ones((3, 3), dtype=np.float64) + probs[0, 0] = 0 + probs[1, 1] = 0 + # should never return first index + _rets = np.zeros((100,)) + for i in range(100): + _rets[i] = shared_tools.random_pick(probs.flatten()) + assert np.all(_rets != 0) + assert np.all(_rets != 4) # THIS LINE NEEDS TO PASS!! + assert np.sum(_rets == 1) > 0 + assert np.sum(_rets == 2) > 0 + assert np.sum(_rets == 3) > 0 + assert np.sum(_rets == 5) > 0 + assert np.sum(_rets == 6) > 0 + assert np.sum(_rets == 7) > 0 + assert np.sum(_rets == 8) > 0 + + def test_custom_unravel_square(): arr = np.arange(9).reshape((3, 3)) # test upper left corner @@ -60,3 +164,68 @@ def test_custom_unravel_rectangle(): def test_custom_unravel_exceed_error(): arr = np.arange(9).reshape((3, 3)) x, y = shared_tools.custom_unravel(99, arr.shape) + + +def test_check_for_loops(): + + idxs = np.array( + [[0, 11, 12, 13, 23, 22, 12], + [0, 1, 2, 3, 4, 5, 16]]) + nidx = np.array([21, 6]) + itt = 6 + free = np.array([1, 1]) + CTR = 4 + L0 = 1 + looped = np.array([0, 0]) + + nidx, looped, free = shared_tools.check_for_loops( + idxs, nidx, itt, L0, looped, (10, 10), CTR, free) + + assert np.all(nidx == [41, 6]) + assert np.all(looped == [1, 0]) + assert np.all(free == [-1, 1]) + + +def test_calculate_new_ind(): + + cidx = np.array([12, 16, 16]) + ncel = np.array([6, 1, 4]) + iwalk = shared_tools.get_iwalk() + jwalk = shared_tools.get_jwalk() + + nidx = shared_tools.calculate_new_ind(cidx, ncel, iwalk.flatten(), jwalk.flatten(), (10, 10)) + + nidx_exp = np.array([21, 6, 0]) + assert np.all(nidx == nidx_exp) + + +def test_get_weight_at_cell(): + + ind = (0, 4) + np.random.seed(delta.seed) + stage = np.random.uniform(0.5, 1, 9) + eta = np.random.uniform(0, 0.85, 9) + depth = stage - eta + depth[depth < 0] = 0 + celltype = np.array([-2, -2, -2, 1, 1, -2, 0, 0, 0]) + qx = 1 + qy = 1 + ivec = delta.ivec.flatten() + jvec = delta.jvec.flatten() + dists = delta.distances.flatten() + dry_thresh = 0.1 + gamma = 0.001962 + theta = 1 + + wts = shared_tools.get_weight_at_cell(ind, stage, depth, celltype, stage, qx, qy, + ivec, jvec, dists, dry_thresh, gamma, theta) + assert np.all(wts[[0, 1, 2, 5]] == 0) + assert wts[4] == 0 + assert np.any(wts[[3, 6, 7, 8]] != 0) + + +def test_version_is_valid(): + v = shared_tools._get_version() + assert type(v) is str + dots = [i for i, c in enumerate(v) if c == '.'] + assert len(dots) == 2