Repository: passaH2O/dorado

Issue:
Docs Build Hard-Set to Python 3.8
Once the matplotlib dependencies support Python 3.9, we should reset our docs build to use the latest version of Python (in the workflow yaml this would be `'3.x'` as opposed to `'3.8'`). 

See [here](https://github.com/nucleic/kiwi/issues/92) for the upstream issue that once closed, should allow us to move back to `'3.x'`.

Discussion:

Patch:
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 1279613..2f84ed9 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -116,10 +116,10 @@ jobs:
       uses: actions/checkout@v2
       with:
         persist-credentials: false
-    - name: Set up Python 3.8
+    - name: Set up Python 3.x
       uses: actions/setup-python@v2
       with:
-        python-version: '3.8'
+        python-version: '3.x'
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
diff --git a/docs/source/paper/index.rst b/docs/source/paper/index.rst
index 1d62d56..8d29903 100644
--- a/docs/source/paper/index.rst
+++ b/docs/source/paper/index.rst
@@ -4,4 +4,4 @@
 Paper
 =====
 
-A paper accompanying this software package is currently being submitted to the Journal of Open Source Software. The paper Markdown file has been rendered as a ReStructured Text file (although not all formatting translates well) for inclusion in the documentation prior to publication. It can be viewed :ref:`here<paperrst>`.
+A paper accompanying this software package has been published in the Journal of Open Source Software. It can be viewed `here <https://joss.theoj.org/papers/10.21105/joss.02585>`_
diff --git a/docs/source/paper/paper.rst b/docs/source/paper/paper.rst
deleted file mode 100644
index a6ef419..0000000
--- a/docs/source/paper/paper.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-.. _paperrst:
-
-==================
-RST-Rendered Paper
-==================
-
-.. mdinclude:: ../../../joss/paper.md
diff --git a/docs/source/userguide/index.rst b/docs/source/userguide/index.rst
index da25023..b515c46 100644
--- a/docs/source/userguide/index.rst
+++ b/docs/source/userguide/index.rst
@@ -53,6 +53,14 @@ Defining the `Particles`
 
 Defining a :obj:`dorado.particle_track.Particles` class is a key step in using `dorado` to perform particle routing. To define a set of particles, the model parameters must first be defined as described above. The `Particles` class is initialized using an instance of the model parameters. From there, particles can be generated and routed.
 
+.. Note::
+    When :obj:`dorado.particle_track.Particles` is initialized, all of the
+    routing weights are automatically calculated. This may take some time for
+    larger model domains, but allows for faster particle routing when particles
+    are actually moved through the domain. There is a progress bar associated
+    with this process so you don't feel like Python has gotten stuck in the
+    object initialization.
+
 Particle Generation and Routing
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
diff --git a/dorado/__init__.py b/dorado/__init__.py
index 883e7ea..49a4acf 100644
--- a/dorado/__init__.py
+++ b/dorado/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "2.1.1"
+__version__ = "2.2.0"
 
 from . import lagrangian_walker
 from . import parallel_routing
diff --git a/dorado/lagrangian_walker.py b/dorado/lagrangian_walker.py
index 26b6869..225ca49 100644
--- a/dorado/lagrangian_walker.py
+++ b/dorado/lagrangian_walker.py
@@ -9,6 +9,7 @@ from builtins import range, map
 from math import cos
 import numpy as np
 from numpy.random import random
+from tqdm import tqdm
 
 
 def random_pick_seed(choices, probs=None):
@@ -38,11 +39,76 @@ def random_pick_seed(choices, probs=None):
     return choices[idx]
 
 
+def make_weight(Particles):
+    """Make an array with the routing weights."""
+    # local namespace function imports
+    from numpy import maximum
+    from numpy import nansum
+    # init the weight array
+    L, W = np.shape(Particles.stage)
+    Particles.weight = np.zeros((L, W, 9))
+    # do weighting calculation for each cell
+    print('Calculating routing weights ...')
+    for i in tqdm(list(range(1, L-1)), ascii=True):
+        for j in list(range(1, W-1)):
+            # weights for each location in domain
+            # get stage values for neighboring cells
+            stage_ind = Particles.stage[i-1:i+2, j-1:j+2]
+
+            # calculate surface slope weights
+            weight_sfc = maximum(0,
+                                 (Particles.stage[i, j]-stage_ind) /
+                                 Particles.distances)
+
+            # calculate inertial component weights
+            weight_int = maximum(0, ((Particles.qx[i, j] * Particles.jvec +
+                                      Particles.qy[i, j] * Particles.ivec) /
+                                 Particles.distances))
+
+            # get depth and cell types for neighboring cells
+            depth_ind = Particles.depth[i-1:i+2, j-1:j+2]
+            ct_ind = Particles.cell_type[i-1:i+2, j-1:j+2]
+
+            # set weights for cells that are too shallow, or invalid 0
+            weight_sfc[(depth_ind <= Particles.dry_depth) | (ct_ind == 2)] = 0
+            weight_int[(depth_ind <= Particles.dry_depth) | (ct_ind == 2)] = 0
+
+            # if sum of weights is above 0 normalize by sum of weights
+            if nansum(weight_sfc) > 0:
+                weight_sfc = weight_sfc / nansum(weight_sfc)
+
+            # if sum of weight is above 0 normalize by sum of weights
+            if nansum(weight_int) > 0:
+                weight_int = weight_int / nansum(weight_int)
+
+            # define actual weight by using gamma, and weight components
+            weight = Particles.gamma * weight_sfc + \
+                (1 - Particles.gamma) * weight_int
+
+            # modify the weight by the depth and theta weighting parameter
+            weight = depth_ind ** Particles.theta * weight
+
+            # if the depth is below the minimum depth then location is not
+            # considered therefore set the associated weight to nan
+            weight[(depth_ind <= Particles.dry_depth) | (ct_ind == 2)] \
+                = np.nan
+
+            # if it's a dead end with only nans and 0's, choose deepest cell
+            if nansum(weight) <= 0:
+                weight = np.zeros_like(weight)
+                weight[depth_ind == np.max(depth_ind)] = 1.0
+
+            # set weight in the true weight array
+            Particles.weight[i, j, :] = weight.ravel()
+
+    print('Finished routing weight calculation.')
+
+
 def get_weight(Particles, ind):
-    """Assign weights to cells surrounding current index.
+    """Choose new cell location given an initial location.
 
-    Function to assign weights to the surrounding 8 cells around the
-    current index and randomly choose one of those cells.
+    Function to randomly choose 1 of the surrounding 8 cells around the
+    current index using the pre-calculated routing weights.
 
     **Inputs** :
 
@@ -58,54 +124,12 @@ def get_weight(Particles, ind):
             New location given as a value between 1 and 8 (inclusive)
 
     """
-    # pull surrounding cell values from stage array
-    stage_ind = Particles.stage[ind[0]-1:ind[0]+2, ind[1]-1:ind[1]+2]
-    # define water surface gradient weight component (minimum of 0)
-    weight_sfc = np.maximum(0,
-                            (Particles.stage[ind] - stage_ind) /
-                            Particles.distances)
-
-    # define flow inertial weighting component (minimum of 0)
-    weight_int = np.maximum(0, (Particles.qx[ind] * Particles.jvec +
-                                Particles.qy[ind] * Particles.ivec) /
-                            Particles.distances)
-
-    # pull surrounding cell values from depth and cell type arrays
-    depth_ind = Particles.depth[ind[0]-1:ind[0]+2, ind[1]-1:ind[1]+2]
-    ct_ind = Particles.cell_type[ind[0]-1:ind[0]+2, ind[1]-1:ind[1]+2]
-
-    # if the depth is below minimum depth for cell to be weight
-    # or it is a cell type that is not water, then make it impossible for
-    # the parcel to travel there by setting associated weight to 0
-    weight_sfc[(depth_ind <= Particles.dry_depth) | (ct_ind == 2)] = 0
-    weight_int[(depth_ind <= Particles.dry_depth) | (ct_ind == 2)] = 0
-
-    # if sum of weights is above 0 normalize by sum of weights
-    if np.nansum(weight_sfc) > 0:
-        weight_sfc = weight_sfc / np.nansum(weight_sfc)
-    # if sum of weight is above 0 normalize by sum of weights
-    if np.nansum(weight_int) > 0:
-        weight_int = weight_int / np.nansum(weight_int)
-
-    # define actual weight by using gamma, and defined weight components
-    Particles.weight = Particles.gamma * weight_sfc + \
-        (1 - Particles.gamma) * weight_int
-    # modify the weight by the depth and theta weighting parameter
-    Particles.weight = depth_ind ** Particles.theta * Particles.weight
-    # if the depth is below the minimum depth then location is not
-    # considered therefore set the associated weight to nan
-    Particles.weight[(depth_ind <= Particles.dry_depth) | (ct_ind == 2)] \
-        = np.nan
-    # if it's a dead end with only nans and 0's, choose deepest cell
-    if np.nansum(Particles.weight) <= 0:
-        Particles.weight = np.zeros_like(Particles.weight)
-        Particles.weight[depth_ind == np.max(depth_ind)] = 1.0
     # randomly pick the new cell for the particle to move to using the
-    # random_pick function and the set of weights just defined
+    # random_pick function and the set of weights
     if Particles.steepest_descent is not True:
-        new_cell = random_pick(Particles.weight)
+        new_cell = random_pick(Particles.weight[ind[0], ind[1], :])
     elif Particles.steepest_descent is True:
-        new_cell = steep_descent(Particles.weight)
+        new_cell = steep_descent(Particles.weight[ind[0], ind[1], :])
 
     return new_cell
 
@@ -311,8 +335,8 @@ def steep_descent(probs):
 
     """
     max_val = np.nanmax(probs)
-    # remove location 1,1 from consideration
-    probs[1, 1] = 0
+    # remove initial location (index 4) from consideration
+    probs[4] = 0
     # remove any locations from consideration beneath max value
     probs[probs < max_val] = 0
     # any nans become ignored too
diff --git a/dorado/particle_track.py b/dorado/particle_track.py
index 8f6954a..1b96fc6 100644
--- a/dorado/particle_track.py
+++ b/dorado/particle_track.py
@@ -379,6 +379,9 @@ class Particles():
         # initialize the walk_data
         self.walk_data = None
 
+        # create weights - this might take a bit of time for large domains
+        lw.make_weight(self)
+
     # function to clear walk data if you've made a mistake while generating it
     def clear_walk_data(self):
         """Manually reset self.walk_data back to None."""


Test patch:
diff --git a/tests/test_lagrangian_walker.py b/tests/test_lagrangian_walker.py
index ed7a2d4..b43cb9e 100644
--- a/tests/test_lagrangian_walker.py
+++ b/tests/test_lagrangian_walker.py
@@ -222,3 +222,201 @@ def test_get_weight_deep():
     np.random.seed(0)
     # make assertion
     assert lw.get_weight(particles, ind) == 8
+
+
+def test_make_weight_deep():
+    '''
+    Test for function make_weight within lagrangian_walker
+    '''
+    tools = pt.modelParams()
+    # define a bunch of expected values
+    tools.stage = np.ones((5, 5))
+    tools.cell_type = np.zeros_like(tools.stage)
+    tools.qy = tools.stage.copy()
+    tools.qx = np.zeros((5, 5))
+    tools.ivec = ivec
+    tools.jvec = jvec
+    tools.distances = distances*np.nan
+    tools.dry_depth = 0.1
+    tools.gamma = 0.02
+    tools.theta = 1
+    tools.steepest_descent = True
+    tools.depth = tools.stage.copy()
+    tools.depth[2,2] = 10.0  # define index 8 as the deepest
+    tools.seed_xloc = [1]
+    tools.seed_yloc = [1]
+    tools.Np_tracer = 1
+    tools.dx = 1
+    # define particles
+    particles = pt.Particles(tools)
+    # set the current index
+    ind = (1,1)
+    # set seed
+    np.random.seed(0)
+    # make assertion
+    assert lw.get_weight(particles, ind) == 8
+
+
+def test_make_weight_shallow():
+    '''
+    Test calculating weights when all neighboring cells are dry
+    '''
+    tools = pt.modelParams()
+    # define a bunch of expected values
+    tools.stage = np.zeros((5, 5))
+    tools.cell_type = np.zeros_like(tools.stage)
+    tools.qy = tools.stage.copy()
+    tools.qx = np.zeros((5, 5))
+    tools.ivec = ivec
+    tools.jvec = jvec
+    tools.distances = distances*np.nan
+    tools.dry_depth = 0.1
+    tools.gamma = 0.02
+    tools.theta = 1
+    tools.steepest_descent = True
+    tools.depth = tools.stage.copy()
+    tools.depth[1, 1] = 10.0  # particle location will be wet
+    tools.seed_xloc = [1]
+    tools.seed_yloc = [1]
+    tools.Np_tracer = 1
+    tools.dx = 1
+    # define particles
+    particles = pt.Particles(tools)
+    # set the current index
+    ind = (1, 1)
+    # set seed
+    np.random.seed(0)
+    # make assertions about weights
+    # at index, index[4] (self) will be 1 while neighbors will be 0
+    assert particles.weight[1, 1, 4] == 1.0
+    assert np.sum(particles.weight[1, 1, :]) == 1.0
+    # weights at boundary cells should be 0
+    assert np.all(np.sum(particles.weight[0, :, 4]) == 0.0)
+    assert np.all(np.sum(particles.weight[-1, :, 4]) == 0.0)
+    assert np.all(np.sum(particles.weight[:, 0, 4]) == 0.0)
+    assert np.all(np.sum(particles.weight[:, -1, 4]) == 0.0)
+
+
+def test_make_weight_equal_opportunity():
+    '''
+    Test calculating weights
+    '''
+    tools = pt.modelParams()
+    # define a bunch of expected values
+    tools.stage = np.zeros((5, 5))
+    tools.cell_type = np.zeros_like(tools.stage)
+    tools.qy = tools.stage.copy()
+    tools.qx = np.zeros((5, 5))
+    tools.ivec = ivec
+    tools.jvec = jvec
+    tools.distances = distances*np.nan
+    tools.dry_depth = 0.1
+    tools.gamma = 0.02
+    tools.theta = 1
+    tools.steepest_descent = True
+    tools.depth = tools.stage.copy()
+    tools.depth[1, 1] = 10.0  # same depth at some neighboring cells
+    tools.depth[1, 2] = 10.0  # same depth at some neighboring cells
+    tools.depth[2, 2] = 10.0  # same depth at some neighboring cells
+    tools.seed_xloc = [1]
+    tools.seed_yloc = [1]
+    tools.Np_tracer = 1
+    tools.dx = 1
+    # define particles
+    particles = pt.Particles(tools)
+    # set the current index
+    ind = (1, 1)
+    # set seed
+    np.random.seed(0)
+    # make assertions about weights
+    # at index, 3 neighbors will be equiprobable
+    assert np.sum(particles.weight[1, 1, :]) == 3.0
+    # at other wet locations, 3 neighbors will be equiprobable
+    assert np.sum(particles.weight[1, 2, :]) == 3.0
+    assert np.sum(particles.weight[2, 2, :]) == 3.0
+    # weights at boundary cells should be 0
+    assert np.all(np.sum(particles.weight[0, :, 4]) == 0.0)
+    assert np.all(np.sum(particles.weight[-1, :, 4]) == 0.0)
+    assert np.all(np.sum(particles.weight[:, 0, 4]) == 0.0)
+    assert np.all(np.sum(particles.weight[:, -1, 4]) == 0.0)
+
+
+def test_make_weight_unequal_opportunity():
+    '''
+    Test calculating weights w/ different depths in cells
+    '''
+    tools = pt.modelParams()
+    # define a bunch of expected values
+    tools.stage = np.zeros((5, 5))
+    tools.cell_type = np.zeros_like(tools.stage)
+    tools.qy = tools.stage.copy()
+    tools.qx = np.zeros((5, 5))
+    tools.ivec = ivec
+    tools.jvec = jvec
+    tools.distances = distances*np.nan
+    tools.dry_depth = 0.1
+    tools.gamma = 0.02
+    tools.theta = 1
+    tools.steepest_descent = True
+    tools.depth = tools.stage.copy()
+    tools.depth[1, 1] = 10.0  # same depth at some neighboring cells
+    tools.depth[1, 2] = 5.0  # less depth at some neighboring cells
+    tools.depth[2, 2] = 5.0  # less depth at some neighboring cells
+    tools.seed_xloc = [1]
+    tools.seed_yloc = [1]
+    tools.Np_tracer = 1
+    tools.dx = 1
+    # define particles
+    particles = pt.Particles(tools)
+    # set the current index
+    ind = (1, 1)
+    # set seed
+    np.random.seed(0)
+    # make assertions about weights
+    # at index, staying put index[4] higher probability than neighbors
+    assert particles.weight[1, 1, 4] > particles.weight[1, 1, 5]
+    assert particles.weight[1, 1, 4] > particles.weight[1, 1, 8]
+    # but the two neighbors should be equiprobable
+    assert particles.weight[1, 1, 5] == particles.weight[1, 1, 8]
+    # weights at boundary cells should be 0
+    assert np.all(np.sum(particles.weight[0, :, 4]) == 0.0)
+    assert np.all(np.sum(particles.weight[-1, :, 4]) == 0.0)
+    assert np.all(np.sum(particles.weight[:, 0, 4]) == 0.0)
+    assert np.all(np.sum(particles.weight[:, -1, 4]) == 0.0)
+
+
+def test_wet_boundary_no_weight():
+    '''
+    Confirm that even if the boundary cell is deep enough, its weight == 0.
+    '''
+    tools = pt.modelParams()
+    # define a bunch of expected values
+    tools.stage = np.ones((5, 5))
+    tools.cell_type = np.zeros_like(tools.stage)
+    tools.qy = tools.stage.copy()
+    tools.qx = np.zeros((5, 5))
+    tools.ivec = ivec
+    tools.jvec = jvec
+    tools.distances = distances*np.nan
+    tools.dry_depth = 0.1
+    tools.gamma = 0.02
+    tools.theta = 1
+    tools.steepest_descent = True
+    tools.depth = tools.stage.copy() * 10.0
+    tools.seed_xloc = [1]
+    tools.seed_yloc = [1]
+    tools.Np_tracer = 1
+    tools.dx = 1
+    # define particles
+    particles = pt.Particles(tools)
+    # set seed
+    np.random.seed(0)
+    # assert weights at boundary cells should be 0
+    assert np.all(np.sum(particles.weight[0, :, 4]) == 0.0)
+    assert np.all(np.sum(particles.weight[-1, :, 4]) == 0.0)
+    assert np.all(np.sum(particles.weight[:, 0, 4]) == 0.0)
+    assert np.all(np.sum(particles.weight[:, -1, 4]) == 0.0)
+    # assert that weights everywhere else are not 0
+    assert np.all(np.sum(particles.weight[1:-1, 1:-1, 4]) != 0.0)
+    # assert that depths everywhere are 10.0
+    assert np.all(particles.depth == 10.0)
diff --git a/tests/test_particle_track.py b/tests/test_particle_track.py
index bfd3c47..b17e2da 100644
--- a/tests/test_particle_track.py
+++ b/tests/test_particle_track.py
@@ -124,6 +124,31 @@ def test_exact_locations():
     assert all_walk_data['xinds'][0] == seed_xloc
     assert all_walk_data['yinds'][0] == seed_yloc
 
+def test_exact_locations_overflow():
+    """Test where number of particles exceeds number of seed locations."""
+    num_ps = 3
+    particle = particle_track.Particles(params)
+    particle.generate_particles(num_ps, [0, 1], [0, 1], method='exact')
+    all_walk_data = particle.run_iteration()
+    assert len(all_walk_data['xinds']) == num_ps
+    assert len(all_walk_data['yinds']) == num_ps
+
+def test_no_explicit_generation():
+    """Test reading of Np_tracer from self if some walk_data exists."""
+    # create some walk data (would exist from a previous run)
+    walk_data = {'xinds': [[1], [1]], 'yinds': [[1], [1]],
+                 'travel_times': [[0.0], [0.0]]}
+    # init particle class and stick walk data in it
+    # init defines Np_tracer as 0, but it is 2 in walk_data
+    particle = particle_track.Particles(params)
+    particle.walk_data = walk_data
+    # assert that Np_tracer is 0
+    assert particle.Np_tracer == 0
+    # run an iteration and see if Np_tracer is corrected
+    particle.run_iteration()
+    # assert that number of particles has been correctly identified
+    assert particle.Np_tracer == 2
+
 def test_travel_time():
     # Particle doesn't travel in the 3x3 space due to the 'sticky' edge
     # conditions so check that travel time is 0 and particle hasn't moved