text: # -*- coding: utf-8 -*-
  """Landlab component to calculate drainage density."""
  from warnings import warn
  
  import numpy as np
  
  from landlab import Component
  
  
  class DrainageDensity(Component):
  
      r"""Calculate drainage density over a DEM.
  
      Landlab component that implements the distance to channel algorithm of
      Tucker et al., 2001.
  
      This component requires EITHER a channel__mask array with 1's
      where channels exist and 0's elsewhere, OR a set of coefficients
      and exponents for a slope-area relationship and a
      channelization threshold to compare against that relationship.
  
      If an array is provided it MUST be of type np.uint8. See the example below
      for how to make such an array.
  
      The channel__mask array will be assigned to an at-node field with the name
      `channel__mask`. If the channel__mask was originaly created from a passed
      array, a user can update this array to change the mask.
  
      If the channel__mask is created using an area coefficent,
      slope coefficient, area exponent, slope exponent, and channelization
      threshold, the location of the mask will be re-update when
      calculate_drainage_density is called.
  
      If an area coefficient, :math:`C_A`, a slope coefficent, :math:`C_S`, an
      area exponent, :math:`m_r`, a slope exponent, :math:`n_r`, and
      channelization threshold :math:`T_C` are provided, nodes that meet the
      criteria
  
      .. math::
  
          C_A A^{m_r} C_s S^{n_r} > T_c
  
      where :math:`A` is the drainage density and :math:`S` is the local slope,
      will be marked as channel nodes.
  
     ``calculate_drainage_density`` function returns drainage density for the model
      domain.
  
      ``calculate_drainage_density`` calculates the distance from every node to the
      nearest channel node :math:`L` along the flow line of steepest descent
      (assuming D8 routing).
  
  
      This component stores this distance a field, called:
      ``surface_to_channel__minimum_distance``.
  
      The drainage density is then calculated (after Tucker et al., 2001):
  
      .. math::
  
          D_d = \frac{1}{2\overline{L}}
  
      where :math:`\overline{L}` is the mean L for the model domain.
  
      """
  
      _name = "DrainageDensity"
  
      _info = {
          "area_coefficient": {
              "dtype": float,
              "intent": "in",
              "optional": True,
              "units": "-",
              "mapping": "node",
              "doc": "Area coefficient to define channels.",
          },
          "area_exponent": {
              "dtype": float,
              "intent": "in",
              "optional": True,
              "units": "-",
              "mapping": "node",
              "doc": "Area exponent to define channels.",
          },
          "channel__mask": {
              "dtype": np.uint8,
              "intent": "in",
              "optional": True,
              "units": "-",
              "mapping": "node",
              "doc": "Logical map of at which grid nodes channels are present",
          },
          "channelization_threshold": {
              "dtype": float,
              "intent": "in",
              "optional": True,
              "units": "-",
              "mapping": "node",
              "doc": "Channelization threshold for use with area and slope coefficients and exponents.",
          },
          "flow__link_to_receiver_node": {
              "dtype": int,
              "intent": "in",
              "optional": False,
              "units": "-",
              "mapping": "node",
              "doc": "ID of link downstream of each node, which carries the discharge",
          },
          "flow__receiver_node": {
              "dtype": int,
              "intent": "in",
              "optional": False,
              "units": "-",
              "mapping": "node",
              "doc": "Node array of receivers (node that receives flow from current node)",
          },
          "flow__upstream_node_order": {
              "dtype": int,
              "intent": "in",
              "optional": False,
              "units": "-",
              "mapping": "node",
              "doc": "Node array containing downstream-to-upstream ordered list of node IDs",
          },
          "slope_coefficient": {
              "dtype": float,
              "intent": "in",
              "optional": True,
              "units": "-",
              "mapping": "node",
              "doc": "Slope coefficient to define channels.",
          },
          "slope_exponent": {
              "dtype": float,
              "intent": "in",
              "optional": True,
              "units": "-",
              "mapping": "node",
              "doc": "Slope exponent to define channels.",
          },
          "surface_to_channel__minimum_distance": {
              "dtype": float,
              "intent": "out",
              "optional": False,
              "units": "m",
              "mapping": "node",
              "doc": "Distance from each node to the nearest channel",
          },
          "topographic__steepest_slope": {
              "dtype": float,
              "intent": "in",
              "optional": False,
              "units": "-",
              "mapping": "node",
              "doc": "The steepest *downhill* slope",
          },
      }
  
      def __init__(
          self,
          grid,
          channel__mask=None,
          area_coefficient=None,
          slope_coefficient=None,
          area_exponent=None,
          slope_exponent=None,
          channelization_threshold=None,
      ):
          """Initialize the DrainageDensity component.
  
          Parameters
          ----------
          grid : ModelGrid
          channel__mask : Array that holds 1's where
              channels exist and 0's elsewhere
          area_coefficient : coefficient to multiply drainage area by,
              for calculating channelization threshold
          slope_coefficient : coefficient to multiply slope by,
              for calculating channelization threshold
          area_exponent : exponent to raise drainage area to,
              for calculating channelization threshold
          slope_exponent : exponent to raise slope to,
              for calculating channelization threshold
          channelization_threshold : threshold value above
              which channels exist
  
          Examples
          --------
          >>> import numpy as np
          >>> from landlab import RasterModelGrid
          >>> from landlab.components import FlowAccumulator, FastscapeEroder
          >>> mg = RasterModelGrid((10, 10))
          >>> _ = mg.add_zeros('node', 'topographic__elevation')
          >>> np.random.seed(50)
          >>> noise = np.random.rand(100)
          >>> mg.at_node['topographic__elevation'] += noise
          >>> mg.at_node['topographic__elevation'] # doctest: +NORMALIZE_WHITESPACE
          array([ 0.49460165,  0.2280831 ,  0.25547392,  0.39632991,  0.3773151 ,
              0.99657423,  0.4081972 ,  0.77189399,  0.76053669,  0.31000935,
              0.3465412 ,  0.35176482,  0.14546686,  0.97266468,  0.90917844,
              0.5599571 ,  0.31359075,  0.88820004,  0.67457307,  0.39108745,
              0.50718412,  0.5241035 ,  0.92800093,  0.57137307,  0.66833757,
              0.05225869,  0.3270573 ,  0.05640164,  0.17982769,  0.92593317,
              0.93801522,  0.71409271,  0.73268761,  0.46174768,  0.93132927,
              0.40642024,  0.68320577,  0.64991587,  0.59876518,  0.22203939,
              0.68235717,  0.8780563 ,  0.79671726,  0.43200225,  0.91787822,
              0.78183368,  0.72575028,  0.12485469,  0.91630845,  0.38771099,
              0.29492955,  0.61673141,  0.46784623,  0.25533891,  0.83899589,
              0.1786192 ,  0.22711417,  0.65987645,  0.47911625,  0.07344734,
              0.13896007,  0.11230718,  0.47778497,  0.54029623,  0.95807105,
              0.58379231,  0.52666409,  0.92226269,  0.91925702,  0.25200886,
              0.68263261,  0.96427612,  0.22696165,  0.7160172 ,  0.79776011,
              0.9367512 ,  0.8537225 ,  0.42154581,  0.00543987,  0.03486533,
              0.01390537,  0.58890993,  0.3829931 ,  0.11481895,  0.86445401,
              0.82165703,  0.73749168,  0.84034417,  0.4015291 ,  0.74862   ,
              0.55962945,  0.61323757,  0.29810165,  0.60237917,  0.42567684,
              0.53854438,  0.48672986,  0.49989164,  0.91745948,  0.26287702])
          >>> fr = FlowAccumulator(mg, flow_director='D8')
          >>> fsc = FastscapeEroder(mg, K_sp=.01, m_sp=.5, n_sp=1)
          >>> for x in range(100):
          ...     fr.run_one_step()
          ...     fsc.run_one_step(dt = 10.0)
          ...     mg.at_node['topographic__elevation'][mg.core_nodes] += .01
          >>> channels = np.array(mg.at_node['drainage_area'] > 5, dtype=np.uint8)
          >>> dd = DrainageDensity(mg, channel__mask=channels)
          >>> mean_drainage_density = dd.calculate_drainage_density()
          >>> np.isclose(mean_drainage_density, 0.3831100571)
          True
  
          Alternatively you can pass a set of coefficients to identify the channel
          mask. Next shows the same example as above, but with these coefficients
          provided.
  
          >>> mg = RasterModelGrid((10, 10))
          >>> _ = mg.add_zeros('node', 'topographic__elevation')
          >>> np.random.seed(50)
          >>> noise = np.random.rand(100)
          >>> mg.at_node['topographic__elevation'] += noise
          >>> fr = FlowAccumulator(mg, flow_director='D8')
          >>> fsc = FastscapeEroder(mg, K_sp=.01, m_sp=.5, n_sp=1)
          >>> for x in range(100):
          ...     fr.run_one_step()
          ...     fsc.run_one_step(dt = 10.0)
          ...     mg.at_node['topographic__elevation'][mg.core_nodes] += .01
          >>> channels = np.array(mg.at_node['drainage_area'] > 5, dtype=np.uint8)
          >>> dd = DrainageDensity(mg,
          ...                      area_coefficient=1.0,
          ...                      slope_coefficient=1.0,
          ...                      area_exponent=1.0,
          ...                      slope_exponent=0.0,
          ...                      channelization_threshold=5)
          >>> mean_drainage_density = dd.calculate_drainage_density()
          >>> np.isclose(mean_drainage_density, 0.3831100571)
          True
          """
          super(DrainageDensity, self).__init__(grid)
  
          if grid.at_node["flow__receiver_node"].size != grid.size("node"):
              msg = (
                  "A route-to-multiple flow director has been "
                  "run on this grid. The landlab development team has not "
                  "verified that DrainageDensity is compatible with "
                  "route-to-multiple methods. Please open a GitHub Issue "
                  "to start this process."
              )
              raise NotImplementedError(msg)
  
          if channel__mask is not None:
              if area_coefficient is not None:
                  warn(
                      "Channel mask and area "
                      "coefficient supplied. Defaulting "
                      "to channel mask, ignoring area "
                      "coefficient."
                  )
              if slope_coefficient is not None:
                  warn(
                      "Channel mask and slope "
                      "coefficient supplied. Defaulting "
                      "to channel mask, ignoring slope "
                      "coefficient."
                  )
              if area_exponent is not None:
                  warn(
                      "Channel mask and area "
                      "exponent supplied. Defaulting "
                      "to channel mask, ignoring area "
                      "exponent."
                  )
              if slope_exponent is not None:
                  warn(
                      "Channel mask and slope "
                      "exponent supplied. Defaulting "
                      "to channel mask, ignoring slope "
                      "exponent."
                  )
              if channelization_threshold is not None:
                  warn(
                      "Channel mask and channelization "
                      "threshold supplied. Defaulting "
                      "to channel mask, ignoring "
                      "threshold."
                  )
              if grid.number_of_nodes != len(channel__mask):
                  raise ValueError(
                      "Length of channel mask is not equal to " "number of grid nodes"
                  )
  
              if "channel__mask" in grid.at_node:
                  warn("Existing channel__mask grid field was overwritten.")
  
              if channel__mask.dtype.type is not np.uint8:
                  raise ValueError("mask must by np.uint8")
  
              self._mask_as_array = True
              self._update_channel_mask = self._update_channel_mask_array
              grid.at_node["channel__mask"] = channel__mask
  
          if channel__mask is None:
              if area_coefficient is None:
                  raise ValueError(
                      "No channel mask and no area "
                      "coefficient supplied. Either "
                      "a channel mask or all 5 threshold "
                      "parameters are needed."
                  )
              if slope_coefficient is None:
                  raise ValueError(
                      "No channel mask and no slope "
                      "coefficient supplied. Either "
                      "a channel mask or all 5 threshold "
                      "parameters are needed."
                  )
              if area_exponent is None:
                  raise ValueError(
                      "No channel mask and no area "
                      "exponent supplied. Either "
                      "a channel mask or all 5 threshold "
                      "parameters are needed."
                  )
              if slope_exponent is None:
                  raise ValueError(
                      "No channel mask and no slope "
                      "exponent supplied. Either "
                      "a channel mask or all 5 threshold "
                      "parameters are needed."
                  )
              if channelization_threshold is None:
                  raise ValueError(
                      "No channel mask and no channelization "
                      "threshold supplied. Either "
                      "a channel mask or all 5 threshold "
                      "parameters are needed."
                  )
  
              self._mask_as_array = False
              self._update_channel_mask = self._update_channel_mask_values
              self._area_coefficient = area_coefficient
              self._slope_coefficient = slope_coefficient
              self._area_exponent = area_exponent
              self._slope_exponent = slope_exponent
              self._channelization_threshold = channelization_threshold
  
              self._update_channel_mask()
  
          # for this component to work with Cython acceleration,
          # the channel_network must be uint8, not bool...
          self._channel_network = grid.at_node["channel__mask"]
  
          # Flow receivers
          self._flow_receivers = grid.at_node["flow__receiver_node"]
  
          # Links to receiver nodes
          self._stack_links = grid.at_node["flow__link_to_receiver_node"]
  
          # Upstream node order
          self._upstream_order = grid.at_node["flow__upstream_node_order"]
  
          # Distance to channel
          if "surface_to_channel__minimum_distance" in grid.at_node:
              self._distance_to_channel = grid.at_node[
                  "surface_to_channel__minimum_distance"
              ]
          else:
              self._distance_to_channel = grid.add_zeros(
                  "surface_to_channel__minimum_distance", at="node", dtype=float
              )
  
      def _update_channel_mask_array(self):
          raise NotImplementedError(
              (
                  "If you provided a channel mask to "
                  "DrainageDensity, update it by updating the "
                  "model grid field channel__mask"
              )
          )
  
      def _update_channel_mask_values(self):
          channel__mask = (
              self._area_coefficient
              * np.power(self._grid.at_node["drainage_area"], self._area_exponent)
              * self._slope_coefficient
              * np.power(
                  self._grid.at_node["topographic__steepest_slope"], self._slope_exponent
              )
          ) > self._channelization_threshold
          self._grid.at_node["channel__mask"] = channel__mask.astype(np.uint8)
  
      def calculate_drainage_density(self):
          """Calculate drainage density. \
  
          If the channel mask is defined based on slope and area coefficients,
          it will be update based on the current drainage area and slope fields.
  
          Returns
          -------
          landscape_drainage_density : float (1/m)
              Drainage density over the model domain.
          """
          from .cfuncs import _calc_dists_to_channel
  
          if self._mask_as_array is False:
              self._update_channel_mask()
  
          _calc_dists_to_channel(
              self._channel_network,
              self._flow_receivers,
              self._upstream_order,
              self._grid.length_of_d8,
              self._stack_links,
              self._distance_to_channel,
              self._grid.number_of_nodes,
          )
          landscape_drainage_density = 1.0 / (
              2.0
              * np.mean(
                  self._grid.at_node["surface_to_channel__minimum_distance"][
                      self._grid.core_nodes
                  ]
              )
          )
          # this is THE drainage density
          return landscape_drainage_density
  