```#!/usr/bin/env python3
"""Solve advection numerically using Total Variation Diminishing method."""

import numpy as np

from landlab.field.errors import FieldError

[docs]

For all links, return ID of upwind link, defined based on the sign of `u`.
If `u` is zero, the upwind link is found as though `u` were positive.

For instance (see examples below), consider a 3x4 raster grid with link
numbering::

.-14-.-15-.-16-.
|    |    |    |
10  11   12   13
|    |    |    |
.--7-.--8-.--9-.
|    |    |    |
3    4    5    6
|    |    |    |
.--0-.--1-.--2-.

There are at most 7 active links (4, 5, 7, 8, 9, 11, 12).
If `u` is positive everywhere, then the upwind links are::

.----.-14-.-15-.
|    |    |    |
3    4    5    6
|    |    |    |
.----.--7-.--8-.
|    |    |    |
|    |    |    |
|    |    |    |
.----.--0-.--1-.

If `u` is negative everywhere, then the upwind links are::

.-15-.-16-.----.
|    |    |    |
|    |    |    |
|    |    |    |
.--8-.--9-.----.
|    |    |    |
10  11   12   13
|    |    |    |
.--1-.--2-.----.

Parameters
----------
grid : RasterModelGrid or HexModelGrid
A landlab grid.
u : float or (n_links,) ndarray
Array of *at-link* values used to determine which node is
upwind.

Returns
-------

Examples
--------
>>> import numpy as np
>>> from landlab import RasterModelGrid
>>> grid = RasterModelGrid((3, 4))

array([[-1, -1, -1, -1],
[ 3,  4,  5,  6]])
array([[-1,  0,  1],
[-1,  7,  8],
[-1, 14, 15]])

array([[10, 11, 12, 13],
[-1, -1, -1, -1]])
array([[ 1,  2, -1],
[ 8,  9, -1],
[15, 16, -1]])

>>> u[4:6] = -1
>>> u = -1
>>> u[8:10] = 1
>>> u[11:13] = 1
array([[ 0., -1., -1.,  0.],
[ 0.,  1.,  1.,  0.]])
array([[ 0.,  0.,  0.],
[-1.,  1.,  1.],
[ 0.,  0.,  0.]])
array([[-1, 11, 12, -1],
[ 3,  4, 5,   6]])
array([[-1,  0,  1],
[ 8,  7,  8],
[-1, 14, 15]])
"""

cols = np.choose(np.broadcast_to(u, len(pll)) >= 0, [1, 0])
uwl = pll[np.arange(len(pll)), cols]

return uwl

[docs]
"""Calculate and return ratio of upwind to local gradient in v.

In Total Variation Diminishing (TVD) numerical schemes, this
ratio is input to a flux limiter to calculate the weighting factor
for higher-order vs. lower-order terms.

Parameters
----------
grid : RasterModelGrid or HexModelGrid
A landlab grid.
Array of upwind links for every link (as returned, for example, by
If provided, place output into this array. Otherwise, create a new array.

Returns
-------
is set to one. For links that do not have an upwind link, the ratio is also
set to one.
"""
if out is None:
else:
out[:] = 1.0

np.divide(
local_diff[uwll], local_diff, where=(uwll != -1) & (local_diff != 0.0), out=out
)

return out

[docs]
"""Numerical solution for advection using a Total Variation Diminishing method.

The component is restricted to regular grids (e.g., Raster or Hex).
the last one listed.

Parameters
----------
grid : RasterModelGrid or HexModelGrid
A Landlab grid object.
fields_to_advect : field name or list or (n_nodes,) array (default None)
A node field of scalar values that will be advected, or list of fields.
If not given, the component creates a generic field, initialized to zeros,
applied to each field in it.
Indicates whether the directions of advection are expected to remain
steady throughout a run. If True, some computation time is saved
by calculating upwind links only once.

Examples
--------
>>> import numpy as np
>>> from landlab import RasterModelGrid
>>> grid = RasterModelGrid((3, 7))
>>> s[9:12] = np.array([1.0, 2.0, 1.0])
>>> for _ in range(5):
...
>>> np.argmax(s[7:14])
4
"""

_unit_agnostic = True

_info = {
"dtype": float,
"intent": "out",
"optional": True,
"units": "-",
"mapping": "node",
},
"dtype": float,
"intent": "out",
"optional": True,
"units": "m2/y",
},
"dtype": float,
"intent": "in",
"optional": False,
"units": "m/y",
},
}

[docs]
def __init__(
self,
grid,
):

# Call base class methods to check existence of input fields,
# create output fields, etc.
super().__init__(grid)
self.initialize_output_fields()

self._scalars = []  # list of fields to advect
self._fluxes = []  # list of flux fields
try:
except KeyError:
self._scalars.append(
)
try:
except KeyError:
flux_counter = 0
self._scalars.append(return_array_at_node(self.grid, field))
if isinstance(field, str):
flux_name = "flux_of_" + field
else:
flux_counter += 1
try:
except FieldError:
self._fluxes.append(flux)
else:
else:
try:
except FieldError:
self._fluxes.append(flux)

] = -1

[docs]
def calc_rate_of_change_at_nodes(self, scalar, flux, dt, update_upwind_links=False):
"""Calculate time rate of change in the advected quantity at nodes.

Parameters
----------
scalar : (n_nodes, ) array
Scalar at-node field of values to be advected.
dt : float
Time-step duration. Needed to calculate the Courant number.
update_upwind_links : bool (optional; default False)
If True, upwind links will be updated (set to True if the
direction of advection is changing through time; if there are
multiple advected quantities, it only needs to be True for the
first one updated, and the update will be used for the others)
"""
] = -1
scalar, dt * self._vel / self.grid.length_of_link
)
psi = flux_lim_vanleer(r)
)
return -self.grid.calc_flux_div_at_node(flux)

[docs]
def update(self, dt):
"""Update the solution by one time step dt.

Same as :meth:`~.run_one_step`.

Parameters
----------
dt : float
Time-step duration. Needed to calculate the Courant number.
"""
for i in range(len(self._scalars)):  # update each of the advected scalars
scalar = self._scalars[i]
flux = self._fluxes[i]
roc = self.calc_rate_of_change_at_nodes(
)
scalar[self.grid.core_nodes] += roc[self.grid.core_nodes] * dt
update_upwinds = False  # should always be False after 1st scalar done

[docs]
def run_one_step(self, dt):
"""Update the solution by one time step dt.

Same as :meth:`~.update`.

Parameters
----------
dt : float
Time-step duration. Needed to calculate the Courant number.
"""
self.update(dt)

```