text: #!/usr/bin/env python
  
  import numpy as np
  import six
  
  from landlab import Component, CLOSED_BOUNDARY
  from ...utils.decorators import use_file_name_or_kwds
  
  
  class SoilInfiltrationGreenAmpt(Component):
  
      """Infiltrate surface water into a soil following the Green-Ampt method.
  
      This code is based on an overland flow model by Francis Rengers and
      colleagues, after Julien et al., 1995. The infiltration scheme follows the
      Green and Ampt equation.
  
      It was implemented in Landlab by DEJH, March '16. Please cite
      Rengers et al., 2016, Model Predictions of Water Runoff in Steep
      Catchments after Wildfire, WRR.
  
      Examples
      --------
      >>> from landlab import RasterModelGrid
      >>> from landlab.components import SoilInfiltrationGreenAmpt
      >>> mg = RasterModelGrid((4,3), spacing=10.)
      >>> hydraulic_conductivity = mg.ones('node')*1.e-6
      >>> hydraulic_conductivity.reshape((4,3))[0:2,:] *= 10000.
      >>> h = mg.add_ones('node', 'surface_water__depth')
      >>> h *= 0.01
      >>> d = mg.add_ones('node', 'soil_water_infiltration__depth', dtype=float)
      >>> d *= 0.2
      >>> SI = SoilInfiltrationGreenAmpt(
      ...     mg,hydraulic_conductivity=hydraulic_conductivity)
      >>> for i in range(10):  # 100s total
      ...     SI.run_one_step(10.)
      >>> mg.at_node['surface_water__depth']
      array([  1.00000000e-08,   1.00000000e-08,   1.00000000e-08,
               1.00000000e-08,   1.00000000e-08,   1.00000000e-08,
               9.88530416e-03,   9.88530416e-03,   9.88530416e-03,
               9.88530416e-03,   9.88530416e-03,   9.88530416e-03])
      >>> mg.at_node['soil_water_infiltration__depth']
      array([ 0.20999999,  0.20999999,  0.20999999,  0.20999999,  0.20999999,
              0.20999999,  0.2001147 ,  0.2001147 ,  0.2001147 ,  0.2001147 ,
              0.2001147 ,  0.2001147 ])
      """
  
      _name = 'SoilInfiltrationGreenAmpt'
  
      _input_var_names = (
          'surface_water__depth',
          'soil_water_infiltration__depth'
          )
  
      _output_var_names = (
          'surface_water__depth',
          'soil_water_infiltration__depth',
          )
  
      _var_units = {
          'surface_water__depth': 'm',
          'soil_water_infiltration__depth': 'm'
          }
  
      _var_mapping = {
          'surface_water__depth': 'node',
          'soil_water_infiltration__depth': 'node'
          }
  
      _var_doc = {
          'surface_water__depth': 'Depth of water above the surface',
          'soil_water_infiltration__depth': (
              'Water column height above the surface previously absorbed into '
              'the soil. Note that this is NOT the actual depth of the wetted '
              'front, which also depends on the porosity.')
          }
  
      # This follows mean values from Rawls et al., 1992; lambda then h_b
      SOIL_PROPS = {
          'sand': (0.694, 0.0726),
          'loamy sand': (0.553, 0.0869),
          'sandy loam': (0.378, 0.1466),
          'loam': (0.252, 0.1115),
          'silt loam': (0.234, 0.2076),
          'sandy clay loam': (0.319, 0.2808),
          'clay loam': (0.242, 0.2589),
          'silty clay loam': (0.177, 0.3256),
          'sandy clay': (0.223, 0.2917),
          'silty clay': (0.150, 0.3419),
          'clay': (0.165, 0.3730),
      }
  
      @use_file_name_or_kwds
      def __init__(self, grid, hydraulic_conductivity=0.005,
                   soil_bulk_density=1590., rock_density=2650.,
                   initial_soil_moisture_content=0.15, soil_type='sandy loam',
                   volume_fraction_coarse_fragments=0.2,
                   coarse_sed_flag=False,
                   surface_water_minimum_depth=1.e-8,
                   soil_pore_size_distribution_index=None,
                   soil_bubbling_pressure=None,
                   wetting_front_capillary_pressure_head=None, **kwds):
          """Kinematic wave approximation overland flow component.
  
          Parameters
          ----------
          grid : RasterModelGrid
              A grid.
          hydraulic_conductivity : float, array, or field name (m/s)
              The soil effective hydraulic conductivity.
          soil_bulk_density : float (kg/m**3)
              The dry bulk density of the soil.
          rock_density : float (kg/m**3)
              The density of the soil constituent material (i.e., lacking porosity).
          initial_soil_moisture_content : float (m**3/m**3, 0. to 1.)
              The fraction of the initial pore space filled with water.
          soil_type : {'sand', loamy sand', 'sandy loam', 'loam', 'silt loam',
                       'sandy clay loam', 'clay loam', 'silty clay loam',
                       'sandy clay', 'silty clay', 'clay'}, or None
              A soil type to automatically set soil_pore_size_distribution_index
              and soil_bubbling_pressure, using mean values from Rawls et al.,
              1992.
          volume_fraction_coarse_fragments : float (m**3/m**3, 0. to 1.)
              The fraction of the soil made up of rocky fragments with very
              little porosity, with diameter > 2 mm.
          coarse_sed_flag : boolean, optional
              If this flag is set to true, the fraction of coarse material in the
              soil column with be used as a correction for phi, the porosity factor.
          surface_water_minimum_depth : float (m), optional
              A minimum water depth to stabilize the solutions for surface flood
              modelling. Leave as the default in most normal use cases.
          soil_pore_size_distribution_index : float, optional
              An index describing the distribution of pore sizes in the soil,
              and controlling effective hydraulic conductivity at varying water
              contents, following Brooks and Corey (1964). Can be set by
              soil_type. Typically denoted "lambda".
          soil_bubbling_pressure : float (m), optional
              The bubbling capillary pressure of the soil, controlling effective
              hydraulic conductivity at varying water contents, following Brooks
              and Corey (1964). Can be set by soil_type. Typically denoted "h_b".
          wetting_front_capillary_pressure_head : float (m), optional
              The effective head at the wetting front in the soil driven by
              capillary pressure in the soil pores. If not set, will be
              calculated by the component from the pore size distribution and
              bubbling pressure, following Brooks and Corey.
          """
  
          self._grid = grid
          self.min_water = surface_water_minimum_depth
          self.hydraulic_conductivity = hydraulic_conductivity
  
          if not coarse_sed_flag:
              volume_fraction_coarse_fragments = 0.
  
          self.moisture_deficit = self.calc_moisture_deficit(
              soil_bulk_density=soil_bulk_density,
              rock_density=rock_density,
              volume_fraction_coarse_fragments=volume_fraction_coarse_fragments,
              soil_moisture_content=initial_soil_moisture_content)
  
          if wetting_front_capillary_pressure_head is None:
              self.capillary_pressure = self.calc_soil_pressure(
                  soil_type=soil_type,
                  soil_pore_size_distribution_index=soil_pore_size_distribution_index,
                  soil_bubbling_pressure=soil_bubbling_pressure)
          else:
              self.capillary_pressure = wetting_front_capillary_pressure_head
  
      @staticmethod
      def calc_soil_pressure(soil_type=None,
                             soil_pore_size_distribution_index=1.,
                             soil_bubbling_pressure=0.):
          """Calculate capillary pressure in a soil type.
  
          Parameters
          ----------
          soil_type : str, optional
              The name of a soil type.
          soil_pore_size_distribution_index : float
              Pore-size distribution index [-].
          soil_bubbling_pressure : float
              Bubbling pressure [m].
          """
          if soil_type is None:
              soil_props = (soil_pore_size_distribution_index,
                            soil_bubbling_pressure)
          else:
              try:
                  soil_props = SoilInfiltrationGreenAmpt.SOIL_PROPS[soil_type]
              except KeyError:
                  raise ValueError('{0}: unknown soil type'.format(soil_type))
  
          return SoilInfiltrationGreenAmpt.calc_pressure_head(*soil_props)
  
      @staticmethod
      def calc_pressure_head(lam, h_b):
          """Calculate pressure head.
  
          Pressure head is set using *lambda* and *h_b*, using an
          equation after Brooks-Corey (1964), following Rawls et al., 1992.
  
          Parameters
          ----------
          lam : float, optional
              Pore-size distribution index. Exponent that describes the
              distribution of pore sizes in the soil, and controls
              effective hydraulic conductivity at varying water
              contents, following Brooks and Corey (1964) [-].
          h_b : float (m), optional
              Bubbling pressure. Capillary pressure of the soil,
              controlling effective hydraulic conductivity at varying
              water contents, following Brooks and Corey (1964) [m]
          """
          return ((2. + 3. * lam) / (1. + 3. * lam) * h_b * .5)
  
      @staticmethod
      def calc_moisture_deficit(soil_bulk_density=1590., rock_density=2650.,
                                volume_fraction_coarse_fragments=0.,
                                soil_moisture_content=0.):
          """Calculate the moisture deficit in a soil.
          
          Parameters
          ----------
          soil_bulk_density : float or array of float
              Bulk density of the soil [kg / m3].
          rock_density : float or array of float
              Density of rock [kg / m3].
          volume_fraction_coarse_fragments : float or array of float
              Volume fraction of sediment made up of coarse grains [-].
          soil_moisture_content : float or array of float
              Fraction of soil filled with water [-].
  
          Returns
          -------
          float or array of float
              Moisture deficit.
          """
          if np.any(soil_bulk_density <= 0.):
              raise ValueError('non-positive soil bulk density')
          if np.any(rock_density < soil_bulk_density):
              raise ValueError('soil bulk density greater than rock density')
          if np.any(volume_fraction_coarse_fragments < 0.):
              raise ValueError('negative volume fraction of coarse grains')
          if np.any(volume_fraction_coarse_fragments > 1.):
              raise ValueError('volume fraction of coarse grains')
  
          soil_porosity = 1. - np.true_divide(soil_bulk_density, rock_density)
          soil_porosity *= 1. - volume_fraction_coarse_fragments
  
          if np.any(soil_moisture_content > soil_porosity):
              raise ValueError('soil moisture greater than porosity')
  
          return soil_porosity - soil_moisture_content
  
      @property
      def min_water(self):
          """Minimum surface water depth."""
          return self._min_water
  
      @min_water.setter
      def min_water(self, new_value):
          if np.any(new_value <= 0.):
              raise ValueError('minimum water depth must be positive')
          self._min_water = new_value
  
      @property
      def hydraulic_conductivity(self):
          """Hydraulic conductivity of soil."""
          return self._hydraulic_conductivity
      
      @hydraulic_conductivity.setter
      def hydraulic_conductivity(self, new_value):
          if isinstance(new_value, six.string_types):
              new_value = self.grid.at_node[new_value]
          if np.any(new_value < 0.):
              raise ValueError('hydraulic conductivity must be positive')
          self._hydraulic_conductivity = new_value
  
      @property
      def moisture_deficit(self):
          """Moisture deficit of soil."""
          return self._moisture_deficit
  
      @moisture_deficit.setter
      def moisture_deficit(self, new_value):
          if np.any(new_value < 0.):
              raise ValueError('negative moisture deficit')
          self._moisture_deficit = new_value
  
      @property
      def capillary_pressure(self):
          """Capillary pressure of soil."""
          return self._capillary_pressure
  
      @capillary_pressure.setter
      def capillary_pressure(self, new_value):
          if np.any(new_value < 0.):
              raise ValueError('negative capillary pressure')
          self._capillary_pressure = new_value
  
      def run_one_step(self, dt):
          """Update fields with current hydrologic conditions.
  
          Parameters
          ----------
          dt : float (s)
              The imposed timestep for the model.
          """
          water_depth = self.grid.at_node['surface_water__depth']
          infiltration_depth = self.grid.at_node['soil_water_infiltration__depth']
  
          assert(np.all(infiltration_depth >= 0.))
  
          wettingfront_depth = infiltration_depth / self.moisture_deficit
          
          potential_infilt = dt * self.hydraulic_conductivity * (
              (wettingfront_depth + self.capillary_pressure + water_depth) /
              wettingfront_depth)
          np.clip(potential_infilt, 0., None, out=potential_infilt)
  
          available_water = water_depth - self.min_water
          np.clip(available_water, 0., None, out=available_water)
  
          actual_infiltration = np.choose(potential_infilt > available_water,
                                          (potential_infilt, available_water))
  
          water_depth -= actual_infiltration
          infiltration_depth += actual_infiltration
  