A Python implementation of the MOD16 terrestrial evapotranspiration algorithm
r'''
The MODIS MOD16 terrestrial evapotranspiration (ET) algorithm. See the README
for full references.

**There are two types of interfaces in the MOD16 Python code.**

User-friendly methods of a `MOD16` instance, parameterized for a single
land-cover type:

- `MOD16.evapotranspiration()`
- `MOD16.transpiration()`
- `MOD16.evaporation_soil()`
- `MOD16.evaporation_wet_canopy()`
- And so on.

There is also a single, vectorized interface implemented as a static method
of the `MOD16` class, which can handle multiple land-cover types:

- `MOD16._evapotranspiration()`
- `MOD16._et()` (alias for the function above)

The user-friendly, instance methods have code blocks that are easy to read and
understand, but those methods might run slow for large spatial domains because
they incur a lot of Python overhead. The values returned by those functions
are in mass-flux units (kg [H2O] m-2 sec-1).

The vectorized interface is faster because it incurs less Python overhead; it
is useful, for example, when calibrating MOD16, because this may require
hundreds or thousands of function evaluations every second. It can also handle
multiple land-cover types. The vectorized interface returns values in energy
units (W m-2), for comparison to eddy covariance tower measurements.

NOTES:

- Air pressure is an input field, which could be taken as given from a
    reanalysis dataset or calculated using the lapse-rate function given in
    the MOD16 C6.1 User's Guide (pp. 7-8), for estimating air pressure from
    elevation.
- MOD16 C6.1 User's Guide has an outdated formula for net radiation
    (Equation 7, page 4); it is no longer based on surface emissivity and is
    instead based on the sum of short-wave and long-wave radiation; see
    `MOD16.radiation_soil()`.
- MOD16 C6.1 User's Guide has an error in the calculation of boundary-layer
    resistance to bare soil evaporation (Equation 19, page 9); instead of what
    is written there: When VPD <= `VPD_open`, then `rbl_max` should be used
    and when VPD >= `VPD_close`, then `rbl_min` should be used.
- MOD16 C6.1 User's Guide suggested that, in calculating the radiation
    received by the soil (Equation 8), only net radiation is modulated by
    bare soil area, (1 - fPAR). In fact, the difference between net
    radiation and the ground heat flux is what is modulated; i.e., the
    correct equation is \(A_{soil} = (1 - fPAR)\times (A - G)\)
- The MERRA2 longwave radiation field `LWGNT` is defined as the "surface net
    downward longwave flux", hence, it is usually negative because of the net
    outward flux of longwave radiation from the Earth's surface.
- For numerical stability, the quantity `r_corr` used in this implementation
    is different from the term defined in the MOD16 C6.1 User's Guide (ca.
    Equation 13). Here, `r_corr` is a large number (greater than 1), equal to
    1/r, where r is the value provided in the User Guide; this avoids
    numerical instability associated with keeping track of a very small
    number.

Also, in MOD16 C6.1, canopy conductance was calculated as:
$$
C_c = \frac{(gl_{sh}(G_S = G_C))}{gl_{sh} + G_S + G_C} \text{LAI}(1 - F_{wet})
$$

However, in the new VIIRS VNP16 and updated MOD16 it is calculated as:
$$
C_c = \frac{(G_0(G_S = G_C))}{G_0 + G_S + G_C}
$$

Where:
$$
G_0 = gl_{sh} \times \text{LAI}(1 - F_{wet})
$$

TODO:

- `sw_rad_night` may not be be needed anywhere
'''

import warnings
import numpy as np
from collections import Counter
from typing import Iterable, Sequence, Tuple
from numbers import Number
from mod17 import linear_constraint

PFT_VALID = (1,2,3,4,5,6,7,8,9,10,12)
STEFAN_BOLTZMANN = 5.67e-8 # Stefan-Boltzmann constant, W m-2 K-4
SPECIFIC_HEAT_CAPACITY_AIR = 1013 # J kg-1 K-1, Monteith & Unsworth (2001)
# Ratio of molecular weight of water vapor to that of dry air (ibid.)
MOL_WEIGHT_WET_DRY_RATIO_AIR = 0.622
TEMP_LAPSE_RATE = 0.0065 # Standard temperature lapse rate [-(deg K) m-1]
GRAV_ACCEL = 9.80665 # Gravitational acceleration [m s-2]
GAS_LAW_CONST = 8.3143 # Ideal gas law constant [m3 Pa (mol)-1 K-1]
AIR_MOL_WEIGHT = 28.9644e-3 # Molecular weight of air [kg (mol)-1]
STD_TEMP_K = 288.15 # Standard temperature at sea level [deg K]
STD_PRESSURE_PASCALS = 101325.0 # Standard pressure at sea level [Pa]
# A pre-determined quantity, not physically meaningful, used in air_pressure()
AIR_PRESSURE_RATE = GRAV_ACCEL / (TEMP_LAPSE_RATE * (GAS_LAW_CONST / AIR_MOL_WEIGHT))

# Calculate the latent heat of vaporization (J kg-1)
latent_heat_vaporization = lambda temp_k: (2.501 - 0.002361 * (temp_k - 273.15)) * 1e6


class MOD16(object):
    '''
    The MODIS MxD16 Evapotranspiration model. The required model parameters are:

    - `tmin_close`: Temperature at which stomata are almost completely
        closed due to (minimum) temperature stress (deg C)
    - `tmin_open`: Temperature at which stomata are completely open, i.e.,
        there is no effect of temperature on transpiration (deg C)
    - `vpd_open`: The VPD at which stomata are completely open, i.e.,
        there is no effect of water stress on transpiration (Pa)
    - `vpd_close`: The VPD at which stomata are almost completely closed
        due to water stress (Pa)
    - `gl_sh`: Leaf conductance to sensible heat per unit LAI
        (m s-1 LAI-1);
    - `gl_wv`: Leaf conductance to evaporated water per unit LAI
        (m s-1 LAI-1);
    - `g_cuticular`: Leaf cuticular conductance (m s-1);
    - `csl`: Mean potential stomatal conductance per unit leaf area (m s-1);
    - `rbl_min`: Minimum leaf boundary layer resistance (s m-1);
    - `rbl_max`: Maximum leaf boundary layer resistance (s m-1);
    - `beta`: Factor in soil moisture constraint on potential soil
        evaporation, i.e., (VPD / beta); from Bouchet (1963)

    Parameters
    ----------
    params : dict
        Dictionary of model parameters
    '''
    required_parameters = [
        'tmin_close', 'tmin_open', 'vpd_open', 'vpd_close', 'gl_sh', 'gl_wv',
        'g_cuticular', 'csl', 'rbl_min', 'rbl_max', 'beta'
    ]

    def __init__(self, params: dict):
        self.params = params
        for key in self.required_parameters:
            setattr(self, key, params[key])

    @staticmethod
    def _et(
            params, lw_net_day, lw_net_night, sw_rad_day, sw_rad_night,
            sw_albedo, temp_day, temp_night, temp_annual, tmin, vpd_day,
            vpd_night, pressure, fpar, lai, f_wet = None, tiny = 1e-9,
            r_corr_list = None
        ) -> Number:
        '''
        Optimized ET code, intended for use in model calibration ONLY.
        Returns combined day and night ET.

        Parameters
        ----------
        params : list
            A list of arrays, each array representing a different parameter,
            in the order specified by `MOD16.required_parameters`. Each array
            should be a (1 x N) array, where N is the number of sites/pixels.
        *drivers
            Every subsequent argument is a separate (T x N) where T is the
            number of time steps and N is the number of sites/pixels.

        Returns
        -------
        numpy.ndarray
            The total latent heat flux [W m-2] for each site/pixel
        '''
        day, night = MOD16._evapotranspiration(
            params, lw_net_day, lw_net_night, sw_rad_day, sw_rad_night,
            sw_albedo, temp_day, temp_night, temp_annual, tmin, vpd_day,
            vpd_night, pressure, fpar, lai, f_wet = None, tiny = 1e-9,
            r_corr_list = r_corr_list)
        return np.add(day, night)

    @staticmethod
    def _evapotranspiration(
            params, lw_net_day, lw_net_night, sw_rad_day, sw_rad_night,
            sw_albedo, temp_day, temp_night, temp_annual, tmin, vpd_day,
            vpd_night, pressure, fpar, lai, f_wet = None, tiny = 1e-9,
            r_corr_list = None
        ) -> Iterable[Tuple[Sequence, Sequence]]:
        '''
        Optimized ET code, intended for use in model calibration ONLY. The
        `params` are expected to be given in the order specified by
        `MOD16.required_parameters`. NOTE: total ET values returned are in
        [W m-2], for comparison to tower ET values. Divide by the latent heat
        of vaporization (J kg-1) to obtain a mass flux (kg m-2 s-1).

        Parameters
        ----------
        params : list
            A list of arrays, each array representing a different parameter,
            in the order specified by `MOD16.required_parameters`. Each array
            should be a (1 x N) array, where N is the number of sites/pixels.
        *drivers
            Every subsequent argument is a separate (T x N) where T is the
            number of time steps and N is the number of sites/pixels.

        Returns
        -------
        list
            A 2-element list of (day, night) latent heat flux [W m-2] totals
        '''
        # Radiation intercepted by the soil surface --
        rad_net_day = sw_rad_day * (1 - sw_albedo) + lw_net_day
        rad_net_night = lw_net_night # At night, SW radiation should be zero

        # Soil heat flux
        g_soil = []
        condition = np.logical_and( # MOD16 UG Equation 9
            np.logical_and(
                temp_annual < (273.15 + 25),
                temp_annual > (273.15 + params[1])
            ), (temp_day - temp_night) >= 5)
        for rad_i, temp_i in zip(
                (rad_net_day, rad_net_night), (temp_day, temp_night)):
            g = np.where(condition, (4.73 * (temp_i - 273.15)) - 20.87, 0)
            # Modify soil heat flux under these conditions...
            g = np.where(np.abs(g) > (0.39 * np.abs(rad_i)), 0.39 * rad_i, g)
            g_soil.append(g)
        g_soil_day, g_soil_night = g_soil
        # -- End soil heat flux calculation

        # Radiation received by the soil, see Section 2.2 of User Guide
        g_soil_day = np.where(
            np.logical_and(rad_net_day - g_soil_day < 0, rad_net_day > 0),
            rad_net_day, g_soil_day)
        g_soil_night = np.where(
            np.logical_and(
                rad_net_day > 0,
                (rad_net_night - g_soil_night) < (-0.5 * rad_net_day)),
            rad_net_night + (0.5 * rad_net_day), g_soil_night)
        rad_soil_day = (1 - fpar) * (rad_net_day - g_soil_day)
        rad_soil_night = (1 - fpar) * (rad_net_night - g_soil_night)

        # Compute day and night components of each ET component
        grouped_drivers = zip(
            (temp_day, temp_night),
            (vpd_day, vpd_night),
            (sw_rad_day, sw_rad_night),
            (lw_net_day, lw_net_night),
            (rad_soil_day, rad_soil_night))
        e_canopy = list()
        e_soil = list()
        transpiration = list()
        et_total = list()
        for i, group in enumerate(grouped_drivers):
            daytime = (i == 0)
            temp_k, vpd, sw_rad, lw_net, rad_soil = group
            # Net radiation to surface, based on down-welling short-wave
            #   radiation and net long-wave radiation
            rad_net = sw_rad * (1 - sw_albedo) + lw_net
            # Radiation intercepted by the canopy
            rad_canopy = fpar * rad_net
            # Compute wet surface fraction and other quantities
            # -- Requiring relative humidity to be calculated from VPD:
            #   VPD = VPsat - VPactual; RH = VPactual / VPsat
            #     --> RH = (VPsat - VPD) / VPsat
            _svp = svp(temp_k)
            rhumidity = (_svp - vpd) / _svp
            if f_wet is None:
                f_wet = np.where(rhumidity < 0.7, 0, rhumidity**4)
            # Slope of saturation vapor pressure curve
            s = svp_slope(temp_k, _svp)
            # Latent heat of vaporization (J kg-1)
            lhv = latent_heat_vaporization(temp_k)
            # Psychrometric constant
            gamma = psychrometric_constant(pressure, temp_k)
            # Correction for atmospheric temperature and pressure
            #   (Equation 13, MOD16 C6.1 User's Guide)
            if r_corr_list is None:
                r_corr = (101300 / pressure) * (temp_k / 293.15)**1.75
            else:
                r_corr = r_corr_list[i]
            # Compute evaporation from wet canopy
            rho = MOD16.air_density(temp_k, pressure, rhumidity) # Air density
            # -- Resistance to radiative heat transfer through air ("rrc")
            r_r = (rho * SPECIFIC_HEAT_CAPACITY_AIR) / (
                4 * STEFAN_BOLTZMANN * temp_k**3)
            with warnings.catch_warnings():
                warnings.simplefilter('ignore')
                # -- Wet canopy resistance to sensible heat ("rhc")
                r_h = 1 / (params[4] * lai * f_wet)
                # -- Wet canopy resistance to evaporated water on the surface
                #   ("rvc")
                r_e = 1 / (params[5] * lai * f_wet)
                # -- Aerodynamic resistance to evaporated water on the wet
                #   canopy surface ("rhrc")
                r_a_wet = np.divide(r_h * r_r, r_h + r_r) # (s m-1)

            # EVAPORATION FROM WET CANOPY
            e = np.divide(
                f_wet * ((s * rad_canopy) + (
                    rho * SPECIFIC_HEAT_CAPACITY_AIR * fpar * vpd * 1/r_a_wet
                )),
                s + ((pressure * SPECIFIC_HEAT_CAPACITY_AIR * r_e) *\
                    1/(lhv * MOL_WEIGHT_WET_DRY_RATIO_AIR * r_a_wet)))
            e_canopy.append(np.where(lai * f_wet <= tiny, 0, e))

            # -- Surface conductance (zero at nighttime)
            g_surf = 0
            if daytime:
                m_tmin = linear_constraint(params[0], params[1])
                m_vpd = linear_constraint(params[2], params[3], 'reversed')
                g_surf = (params[7] * m_tmin(tmin - 273.15) * m_vpd(vpd))
            g_surf /= r_corr
            g_cuticular = params[6] / r_corr
            # -- Canopy conductance, should be zero when LAI or f_wet are zero;
            #   updated calculation for conductance to sensible heat, see User
            #   Guide's Equation 15
            gl_sh = params[4] * lai * (1 - f_wet)
            g = ((gl_sh * (g_surf + g_cuticular)) / (
                gl_sh + g_surf + g_cuticular))
            g_canopy = np.where(
                np.logical_and(lai > 0, (1 - f_wet) > 0), g, tiny)
            # -- Aerodynamic resistance to heat, water vapor from dry canopy
            #   surface into the air (Equation 16, MOD16 C6.1 User's Guide)
            r_a_dry = (1/params[4] * r_r) / (1/params[4] + r_r) # (s m-1)

            # PLANT TRANSPIRATION
            if np.any(g_surf > 0):
                t = (1 - f_wet) * ((s * rad_canopy) +\
                    (rho * SPECIFIC_HEAT_CAPACITY_AIR * fpar * (vpd / r_a_dry)))
                if daytime:
                    t /= (s + gamma * (1 + (1 / g_canopy) / r_a_dry))
                else:
                    # Simplify calculation because g_canopy := 0 at night
                    t /= (s + gamma)
            else:
                t = 0
            transpiration.append(t)

            # BARE SOIL EVAPORATION
            # -- Total aerodynamic resistance as a function of VPD and the
            #   atmospheric boundary layer resistance...
            r_tot = np.where(vpd <= params[2], params[8],
                np.where(vpd >= params[3], params[9],
                params[9] - (
                    (params[9] - params[8]) * (params[3] - vpd))\
                        / (params[3] - params[2])))
            # ...CORRECTED for atmospheric temperature, pressure
            r_tot = r_tot / r_corr
            # -- Aerodynamic resistance at the soil surface
            r_as = (r_tot * r_r) / (r_tot + r_r)
            # -- Terms common to evaporation and potential evaporation
            numer = (s * rad_soil) +\
                (rho * SPECIFIC_HEAT_CAPACITY_AIR * (1 - fpar) * (vpd / r_as))
            denom = (s + gamma * (r_tot / r_as))
            # -- Evaporation from "wet" soil (saturated fraction)
            evap_sat = (numer * f_wet) / denom
            # -- (Potential) Evaporation from unsaturated fraction
            evap_unsat = (numer * (1 - f_wet)) / denom
            # -- Finally, apply the soil moisture constraint from Fisher et al.
            #   (2008); see MOD16 C6.1 User's Guide, pp. 9-10
            e = evap_sat + evap_unsat * rhumidity**(vpd / params[10])
            e_soil.append(e)

            # Result is the sum of the three components
            et_total.append((transpiration[i] + e_canopy[i] + e_soil[i]))
        return et_total

    @staticmethod
    def air_density(
            temp_k: Number, pressure: Number, rhumidity: Number
        ) -> Number:
        '''
        NIST simplified air density formula with buoyancy correction from:

            National Physical Laboratory (2021),
              "Buoyancy Correction and Air Density Measurement."
              http://resource.npl.co.uk/docs/science_technology/
                mass_force_pressure/clubs_groups/instmc_weighing_panel/
                buoycornote.pdf

        Parameters
        ----------
        temp_k : int or float or numpy.ndarray
            Air temperature (degrees K)
        pressure : int or float or numpy.ndarray
            Air pressure (Pa)
        rhumidity : int or float or numpy.ndarray
            Relative humidity, on the interval [0, 1]

        Returns
        -------
        int or float or numpy.ndarray
            Air density (kg m-3)
        '''
        return np.divide( # Convert Pa to mbar, RH to RH% (percentage)
            0.348444 * (pressure / 100) - (rhumidity * 100) *\
                (0.00252 * (temp_k - 273.15) - 0.020582),
            temp_k) # kg m-3

    @staticmethod
    def air_pressure(elevation_m: Number) -> Number:
        r'''
        Atmospheric pressure as a function of elevation. From the discussion on
        atmospheric statics (p. 168) in:

            Iribane, J.V., and W.L. Godson, 1981. Atmospheric Thermodynamics
                2nd Edition. D. Reidel Publishing Company, Dordrecht,
                The Netherlands.

        It is calculated:

        $$
        P_a = P_{\text{std}}\times \left(
            1 - \ell z T_{\text{std}}^{-1}
          \right)^{5.256}
        $$

        Where \(\ell\) is the standard temperature lapse rate (-0.0065 deg K
        per meter), \(z\) is the elevation in meters, and \(P_{\text{std}}\),
        \(T_{\text{std}}\) are the standard pressure (101325 Pa) and
        temperature (288.15 deg K) at sea level.

        Parameters
        ----------
        elevation_m : Number

        Returns
        -------
        Number
            Air pressure in Pascals
        '''
        temp_ratio = 1 - ((TEMP_LAPSE_RATE * elevation_m) / STD_TEMP_K)
        return STD_PRESSURE_PASCALS * np.pow(temp_ratio, AIR_PRESSURE_RATE)

    @staticmethod
    def vpd(qv10m: Number, pressure: Number, tmean: Number) -> Number:
        r'''
        Computes vapor pressure deficit (VPD) from surface meteorology:

        $$
        \text{VPD} = \left(610.7\times \text{exp}\left(
          \frac{17.38\times T}{239 + T}
        \right) - \text{AVP}\right)
        $$

        Where \(T\) is the temperature in deg K and the actual vapor pressure
        (AVP) is given:

        $$
        \text{AVP} = \frac{\text{QV10M}\times
            \text{P}}{0.622 + 0.379\times \text{QV10M}}
        $$

        Where P is the air pressure in Pascals and QV10M is the water vapor
        mixing ratio at 10-meter height.

        Parameters
        ----------
        qv10m : int or float or numpy.ndarray
            Water vapor mixing ratio at 10-meter height (Pa)
        pressure : int or float or numpy.ndarray
            Atmospheric pressure (Pa)
        tmean : int or float or numpy.ndarray
            Mean daytime temperature (degrees K)

        Returns
        -------
        int or float or numpy.ndarray
        '''
        temp_c = tmean - 273.15
        # Actual vapor pressure (Gates 1980, Biophysical Ecology, p.311)
        avp = (qv10m * pressure) / (0.622 + (0.379 * qv10m))
        # Saturation vapor pressure (similar to FAO formula)
        svp = 610.7 * np.exp((17.38 * temp_c) / (239 + temp_c))
        return svp - avp

    @staticmethod
    def rhumidity(temp_k: Number, vpd: Number) -> Number:
        r'''
        Calculates relative humidity as the difference between the saturation
        vapor pressure (SVP) and VPD, normalized by the SVP:

        $$
        \text{RH} = \frac{\text{SVP} - \text{VPD}}{\text{SVP}}
        $$

        Parameters
        ----------
        temp_k : int or float or numpy.ndarray
        vpd : int or float or numpy.ndarray

        Returns
        -------
        int or float or numpy.ndarray
        '''
        # Requiring relative humidity to be calculated from VPD:
        #   VPD = VPsat - VPactual; RH = VPactual / VPsat
        #     --> RH = (VPsat - VPD) / VPsat
        esat = svp(temp_k)
        avp = esat - vpd
        rh = avp / esat
        return np.where(avp < 0, 0, np.where(rh > 1, 1, rh))

    def evapotranspiration(
            self, lw_net_day: Number, lw_net_night: Number,
            sw_rad_day: Number, sw_rad_night: Number, sw_albedo: Number,
            temp_day: Number, temp_night: Number, temp_annual: Number,
            tmin: Number, vpd_day: Number, vpd_night: Number,
            pressure: Number, fpar: Number, lai: Number,
            f_wet: Number = None, separate: bool = False
        ) -> Iterable[Tuple[Sequence, Sequence]]:
        '''
        Instaneous evapotranspiration (ET) [kg m-2 s-1], the sum of bare soil
        evaporation, wet canopy evaporation, and canopy transpiration.

        Parameters
        ----------
        lw_net_day : Number
            Net downward long-wave radiation integrated during daylight hours
        lw_net_night : Number
            Net downward long-wave radiation integrated during nighttime hours
        sw_rad_day : Number
            Down-welling short-wave radiation integrated during daylight hours
        sw_rad_night : Number
            Down-welling short-wave radiation integrated during night-time hours
        sw_albedo : Number
            Down-welling short-wave albedo (under "black-sky" conditions)
        temp_day : Number
            Daytime prevailing air temperature at 10-meter height (deg K)
        temp_night : Number
            Nighttime prevailing air temperature at 10-meter height (deg K)
        temp_annual : Number
            Mean annual air temperature at 10-meter height (deg K)
        tmin : Number
            Minimum daily air temperature at 10-meter height (deg K)
        vpd_day : Number
            Daytime mean vapor pressure deficit (VPD) (Pa)
        vpd_night : Number
            Nighttime mean VPD (Pa)
        pressure : Number
            Air pressure (Pa)
        fpar : Number
            Fraction of photosynthetically active radiation (PAR) absorbed by
            the vegetation canopy
        lai : Number
            Leaf area index (LAI)
        f_wet : Number
            (Optional) Fraction of surface that is saturated with water
        separate : bool
            True to return the separate ET components; False to return their
            sum (Default: False)

        Returns
        -------
        tuple
            Sequence of (daytime, nighttime) ET flux(es)
        '''
        # Get radiation intercepted by the soil surface
        rad_soil_all = self.radiation_soil(
            lw_net_day, lw_net_night, sw_rad_day, sw_rad_night, sw_albedo,
            temp_day, temp_night, temp_annual, fpar)
        # Compute day and night components
        grouped_drivers = zip(
            (temp_day, temp_night),
            (vpd_day, vpd_night),
            (sw_rad_day, sw_rad_night),
            (lw_net_day, lw_net_night),
            rad_soil_all)
        e_canopy = list()
        e_soil = list()
        transpiration = list()
        et_total = list()
        for i, group in enumerate(grouped_drivers):
            daytime = (i == 0)
            temp_k, vpd, sw_rad, lw_net, rad_soil = group
            # Net radiation to surface, based on down-welling short-wave
            #   radiation and net long-wave radiation
            rad_net = sw_rad * (1 - sw_albedo) + lw_net
            # Radiation intercepted by the canopy
            rad_canopy = fpar * rad_net
            # Compute wet surface fraction and other quantities
            # -- Requiring relative humidity to be calculated from VPD:
            #   VPD = VPsat - VPactual; RH = VPactual / VPsat
            #     --> RH = (VPsat - VPD) / VPsat
            rhumidity = MOD16.rhumidity(temp_k, vpd)
            if f_wet is None:
                f_wet = np.where(rhumidity < 0.7, 0, np.power(rhumidity, 4))
            # Slope of saturation vapor pressure curve
            s = svp_slope(temp_k)
            # Latent heat of vaporization (J kg-1)
            lhv = latent_heat_vaporization(temp_k)
            # Correction for atmospheric temperature and pressure
            #   (Equation 13, MOD16 C6.1 User's Guide)
            r_corr = (101300 / pressure) * (temp_k / 293.15)**1.75
            # EVAPORATION FROM WET CANOPY
            e_canopy.append(self.evaporation_wet_canopy(
                pressure, temp_k, vpd, lai, fpar, rad_canopy, lhv, rhumidity,
                f_wet))
            # EVAPORATION FROM BARE SOIL
            e_soil.append(self.evaporation_soil(
                pressure, temp_k, vpd, fpar, rad_soil, r_corr, lhv, rhumidity,
                f_wet))
            # PLANT TRANSPIRATION
            transpiration.append(
                self.transpiration(
                    pressure, temp_k, vpd, lai, fpar, rad_canopy, tmin,
                    r_corr, lhv, rhumidity, f_wet, daytime = daytime))
            if separate:
                et_total.append((e_canopy[i], e_soil[i], transpiration[i]))
            else:
                et_total.append((e_canopy[i] + e_soil[i] + transpiration[i]))
        return et_total

    def evaporation_soil(
            self, pressure: Number, temp_k: Number, vpd: Number, fpar: Number,
            rad_soil: Number, r_corr: Number, lhv: Number = None,
            rhumidity: Number = None, f_wet: Number = None
        ) -> Number:
        r'''
        Evaporation from bare soil, calculated as the sum of evaporation from
        the saturated and unsaturated soil surfaces, as:

        $$
        \lambda E_{\mathrm{sat}} = F_{\text{wet}}
          \frac{s A_{\mathrm{soil}} + \rho C_p (1 - F_c) D r_a^{-1}}{s + \gamma(1 + r_t r_a^{-1})}
        $$
        $$
        \lambda E_{\mathrm{unsat}} = (1 - F_{\text{wet}})
          \frac{s A_{\mathrm{soil}} + \rho C_p (1 - F_c) D r_a^{-1}}{s + \gamma(1 + r_t r_a^{-1})}
        \phi^{D\beta^{-1}}
        $$

        NOTE: The `r_corr` argument to this function is different from that
        defined in Equation 13 in the MOD16 C61 User's Guide. For numerical
        stability, `r_corr` equals 1/r where r is the quantity defined in
        Equation 13. See `MOD16.evapotranspiration()` for how `r_corr` is
        calculated.

        Parameters
        ----------
        pressure : float or numpy.ndarray
            The air pressure in Pascals
        temp_k : float or numpy.ndarray
            The air temperature in degrees K
        vpd : float or numpy.ndarray
            The vapor pressure deficit in Pascals
        fpar : float or numpy.ndarray
            Fraction of photosynthetically active radiation (PAR) absorbed by
            the vegetation canopy
        rad_soil : float or numpy.ndarray
            Net radiation to the soil surface (J m-2 s-1)
        r_corr : float or numpy.ndarray or None
            The temperature and pressure correction factor for surface
            conductance
        lhv : float or numpy.ndarray or None
            (Optional) The latent heat of vaporization
        rhumidity : float or numpy.ndarray or None
            (Optional) The relative humidity
        f_wet : float or numpy.ndarray or None
            (Optional) The fraction of the surface that has standing water

        Returns
        -------
        float or numpy.ndarray
            Evaporation from bare soil (kg m-2 s-1)
        '''
        if lhv is None:
            lhv = latent_heat_vaporization(temp_k)
        if rhumidity is None:
            rhumidity = MOD16.rhumidity(temp_k, vpd)
        if f_wet is None:
            f_wet = np.where(rhumidity < 0.7, 0, np.power(rhumidity, 4))
        # Slope of saturation vapor pressure curve
        s = svp_slope(temp_k)
        # Air density
        rho = MOD16.air_density(temp_k, pressure, rhumidity)
        # Psychrometric constant
        gamma = psychrometric_constant(pressure, temp_k)
        # Resistance to radiative heat transfer through air ("rrc" or "rrs")
        r_r = (rho * SPECIFIC_HEAT_CAPACITY_AIR) / (
            4 * STEFAN_BOLTZMANN * temp_k**3)
        # Total aerodynamic resistance as a function of VPD and the
        #   atmospheric boundary layer resistance...
        r_tot = np.where(vpd <= self.vpd_open, self.rbl_min,
            np.where(vpd >= self.vpd_close, self.rbl_max,
            self.rbl_max - (
                (self.rbl_max - self.rbl_min) * (self.vpd_close - vpd))\
                    / (self.vpd_close - self.vpd_open)))
        # ...CORRECTED for atmospheric temperature, pressure
        r_tot = r_tot / r_corr
        # -- Aerodynamic resistance at the soil surface
        r_as = (r_tot * r_r) / (r_tot + r_r)
        # -- Terms common to evaporation and potential evaporation
        numer = (s * rad_soil) +\
            (rho * SPECIFIC_HEAT_CAPACITY_AIR * (1 - fpar) * (vpd / r_as))
        denom = (s + gamma * (r_tot / r_as))
        # -- Evaporation from "wet" soil (saturated fraction)
        evap_sat = (numer * f_wet) / denom
        # -- (Potential) Evaporation from unsaturated fraction
        evap_unsat = (numer * (1 - f_wet)) / denom
        # BARE SOIL EVAPORATION
        # -- Finally, apply the soil moisture constraint from Fisher et al.
        #   (2008); see MOD16 C6.1 User's Guide, pp. 9-10
        e = np.where(evap_sat < 0, 0, evap_sat)
        e += np.where(
            evap_unsat < 0, 0,
            evap_unsat * np.power(rhumidity, vpd / self.beta))
        # Divide by the latent heat of vaporization (J kg-1) to obtain mass
        #   flux (kg m-2 s-1)
        return e / lhv

    def evaporation_wet_canopy(
            self, pressure: Number, temp_k: Number, vpd: Number, lai: Number,
            fpar: Number, rad_canopy: Number, lhv: Number = None,
            rhumidity: Number = None, f_wet: Number = None, tiny: float = 1e-9
        ) -> Number:
        r'''
        Wet canopy evaporation calculated according to the MODIS MOD16
        framework:

        $$
        \lambda E = F_{wet} \frac{
            s\, A_c\, F_c + \rho\, C_p\, D\, (F_c / r_w)
          }{s + (P_a\, C_p\, r_e)(\lambda\, \varepsilon\, r_w)^{-1}}
        $$

        Where:

        - \(s\) is the slope of the saturation vapor pressure curve (see
            `mod16.svp_slope()`)
        - \(A_c\) is the radiation intercepted by the canopy [J m-2 s-1]
        - \(F_c\) is the canopy fraction of ground cover (i.e., fPAR)
        - \(\rho\) is the density of air [kg m-3]
        - \(C_p\) is the specific heat capacity of air
        - \(D\) is the vapor pressure deficit
        - \(F_{wet}\) is the fraction of the land surface that is saturated
        - \(P_a\) is the air pressure (Pa)
        - \(\lambda\) is the latent heat of vaporization
        - \(\varepsilon\) is the ratio of molecular weight of water vapor to that
            of dry air

        Parameters
        ----------
        pressure : float or numpy.ndarray
            The air pressure in Pascals
        temp_k : float or numpy.ndarray
            The air temperature in degrees K
        vpd : float or numpy.ndarray
            The vapor pressure deficit in Pascals
        lai : float or numpy.ndarray
            The leaf area index (LAI)
        fpar : float or numpy.ndarray
            Fraction of photosynthetically active radiation (PAR) absorbed by
            the vegetation canopy
        rad_canopy : float or numpy.ndarray
            Net radiation to the canopy (J m-2 s-1)
        lhv : float or numpy.ndarray or None
            (Optional) The latent heat of vaporization
        rhumidity : float or numpy.ndarray or None
            (Optional) The relative humidity
        f_wet : float or numpy.ndarray or None
            (Optional) The fraction of the surface that has standing water
        tiny : float
            (Optional) A very small number, the numerical tolerance for zero
            (Default: 1e-9)

        Returns
        -------
        float or numpy.ndarray
            Evaporation from the wet canopy surface (kg m-2 s-1)
        '''
        if lhv is None:
            lhv = latent_heat_vaporization(temp_k)
        if rhumidity is None:
            rhumidity = MOD16.rhumidity(temp_k, vpd)
        if f_wet is None:
            f_wet = np.where(rhumidity < 0.7, 0, np.power(rhumidity, 4))
        # MODPR16_main.c L3893; zero LAI/ f_wet -> zero evaporation
        f_wet = np.where(f_wet == 0, f_wet + tiny, f_wet)
        lai = np.where(lai == 0, lai + tiny, lai)
        # Slope of saturation vapor pressure curve
        s = svp_slope(temp_k)
        # Air density
        rho = MOD16.air_density(temp_k, pressure, rhumidity)
        with warnings.catch_warnings():
            warnings.simplefilter('ignore')
            # Wet canopy resistance to sensible heat ("rhc")
            r_h = 1 / (self.gl_sh * lai * f_wet)
            # Wet canopy resistance to evaporated water on the surface ("rvc")
            r_e = 1 / (self.gl_wv * lai * f_wet)
            # Resistance to radiative heat transfer through air ("rrc")
            r_r = (rho * SPECIFIC_HEAT_CAPACITY_AIR) / (
                4 * STEFAN_BOLTZMANN * temp_k**3)
        # Aerodynamic resistance to evaporated water on the wet canopy
        #   surface ("rhrc")
        r_a_wet = np.divide(r_h * r_r, r_h + r_r) # (s m-1)
        numer = f_wet * ((s * rad_canopy) +\
            (rho * SPECIFIC_HEAT_CAPACITY_AIR * fpar * (vpd * 1/r_a_wet)))
        denom = s + ((pressure * SPECIFIC_HEAT_CAPACITY_AIR * r_e) *\
            1/(lhv * MOL_WEIGHT_WET_DRY_RATIO_AIR * r_a_wet))
        # Mu et al. (2011), Equation 17; PET (J m-2 s-1) is divided by the
        #   latent heat of vaporization (J kg-1) to obtain mass flux
        #   (kg m-2 s-1)
        evap = np.where(numer < 0, 0, (numer / denom) / lhv)
        # If f_wet or LAI are ~zero, then wet canopy evaporation is zero
        return np.where(np.logical_or(f_wet <= tiny, lai <= tiny), 0, evap)

    def radiation_soil(
            self, lw_net_day: Number, lw_net_night: Number,
            sw_rad_day: Number, sw_rad_night: Number, sw_albedo: Number,
            temp_day: Number, temp_night: Number, temp_annual: Number,
            fpar: Number
        ) -> Iterable[Tuple[Sequence, Sequence]]:
        r'''
        Net radiation received by the soil surface, calculated as the
        difference between the net radiation intercepted by the soil and the
        soil heat flux. Net incoming radiation to the land surface (soil and
        non-soil components) during the daytime is calculated:

        $$
        A = (1-\alpha)R_{S\downarrow} + R_L
        $$

        Where \(\alpha\) is the short-wave albedo (under "black-sky"
        conditions); \(R_{S\downarrow}\) is the (daytime) down-welling short-
        wave radidation; \(R_L\) is the (daytime) net downward long-wave
        radiation.

        At nighttime, net incoming radiation is simply the net downward long-
        wave radiation during nighttime hours:

        $$
        A = R_L
        $$

        And the net radiation received by the soil is modulated by the
        fractional area covered by vegetation, \(F_C\), such that areas with
        100% vegetation cover have no net radiation to the soil:

        $$
        A_{\text{soil}} = (1 - F_C)\times (A - G)
        $$

        Parameters
        ----------
        lw_net_day : int or float or numpy.ndarray
            Net downward long-wave radiation integrated during daylight hours
        lw_net_night : int or float or numpy.ndarray
            Net downward long-wave radiation integrated during nighttime hours
        sw_rad_day : int or float or numpy.ndarray
            Down-welling short-wave radiation integrated during daylight hours
        sw_rad_night : int or float or numpy.ndarray
            Down-welling short-wave radiation integrated during night-time hours
        sw_albedo : int or float or numpy.ndarray
            Down-welling short-wave albedo ("black-sky" albedo)
        temp_day : int or float or numpy.ndarray
            Average temperature (degrees K) during daylight hours
        temp_night : int or float or numpy.ndarray
            Average temperature (degrees K) during nighttime hours
        temp_annual : int or float or numpy.ndarray
            Annual average daily temperature (degrees K)
        fpar : float or numpy.ndarray
            Fraction of photosynthetically active radiation (PAR) absorbed by
            the vegetation canopy

        Returns
        -------
        tuple
            A 2-element tuple of (daytime, night-time) radiation received by the
            soil surface
        '''
        # At night, SW radiation should be zero (i.e., sw_rad_night = 0) which
        #   means that only LW radiation contributes (L3852 of MODPR16_main.c)
        rad_net_day = sw_rad_day * (1 - sw_albedo) + lw_net_day
        rad_net_night = lw_net_night
        g_soil_day, g_soil_night = self.soil_heat_flux(
            rad_net_day, rad_net_night, temp_day, temp_night, temp_annual, fpar)
        # Section 2.8 of MOD16 Collection 6.1 User's Guide (pp. 10-11);
        #   if soil heat flux is greater than net incoming radiation...
        g_soil_day = np.where(
            np.logical_and(rad_net_day - g_soil_day < 0, rad_net_day > 0),
            rad_net_day, g_soil_day)
        g_soil_night = np.where(
            np.logical_and(
                rad_net_day > 0,
                (rad_net_night - g_soil_night) < (-0.5 * rad_net_day)),
            rad_net_night + (0.5 * rad_net_day), g_soil_night)
        # "In the improved algorithm, there will be no soil heat flux
        #   interaction between the soil and atmosphere if the ground is 100%
        #   covered with vegetation" (Mu et al. 2011, RSE)
        # Equation 7, Mu et al. (2011)
        rad_soil_day = (1 - fpar) * (rad_net_day - g_soil_day)
        rad_soil_night = (1 - fpar) * (rad_net_night - g_soil_night)
        return (rad_soil_day, rad_soil_night)

    def soil_heat_flux(
            self, rad_net_day: Number, rad_net_night: Number,
            temp_day: Number, temp_night: Number, temp_annual: Number,
            fpar: Number
        ) -> Iterable[Tuple[Sequence, Sequence]]:
        r'''
        Soil heat flux [MJ m-2], based on surface energy balance. If the mean
        annual temperature is between `Tmin_close` and 25 deg C and the
        contrast between daytime and nighttime air temperatures is greater
        than 5 deg C:

        $$
        G_{\text{soil}} = 4.73 T - 20.87
            \quad\text{iff}\quad T_{\text{min,close}}
            \le T_{\text{annual}} < 25 ,\,
                T_{\text{day}} - T_{\text{night}} \ge 5
        $$

        Otherwise, soil heat flux is zero.

        Finally, soil heat flux should be no greater than 39% of the net
        radiation to the land surface.

        $$
        G_{\text{soil}} = 0.39 A \quad\text{iff}\quad G_{soil} > 0.39 |A|
        $$

        Parameters
        ----------
        rad_net_day : int or float or numpy.ndarray
            Net radiation to the land surface during daylight hours; i.e.,
            integrated while the sun is up [MJ m-2]
        rad_net_night : int or float or numpy.ndarray
            Net radiation to the land surface during nighttime hours; i.e.,
            integrated while the sun is down [MJ m-2]
        temp_day : int or float or numpy.ndarray
            Average temperature (degrees K) during daylight hours
        temp_night : int or float or numpy.ndarray
            Average temperature (degrees K) during nighttime hours
        temp_annual : int or float or numpy.ndarray
            Annual average daily temperature (degrees K)
        fpar : float or numpy.ndarray
            Fraction of photosynthetically active radiation (PAR) absorbed by
            the vegetation canopy

        Returns
        -------
        tuple
            A 2-element tuple of (daytime, nighttime) soil heat flux
        '''
        # See ca. Line 4355 in MODPR16_main.c
        g_soil = []
        condition = np.logical_and(
            np.logical_and(
                temp_annual < (273.15 + 25),
                temp_annual > (273.15 + self.tmin_close)
            ), (temp_day - temp_night) >= 5)
        for rad_i, temp_i in zip(
                (rad_net_day, rad_net_night), (temp_day, temp_night)):
            g = np.where(condition, (4.73 * (temp_i - 273.15)) - 20.87, 0)
            # Modify soil heat flux under these conditions...
            g = np.where(np.abs(g) > (0.39 * np.abs(rad_i)), 0.39 * rad_i, g)
            # g_soil is the soil heat flux when fPAR == 0
            # NOTE: We do not scale the result by (1 - fPAR) here, as in
            #   Mu et al. (2011), because this is accounted for when
            #   calculating net radidation received by the soil surface
            g_soil.append(g)
        return g_soil

    def surface_conductance(self, tmin: Number, vpd_day: Number) -> Number:
        r'''
        Surface conductance to transpiration (m s-1), during daylight hours,
        NOT corrected for atmospheric temperature and pressure. Surface
        conductance at night is assumed to be zero (as non-CAM plants close
        stomata at night to prevent water loss when there is no photosynthetic
        carbon gain).

        $$
        G_{\text{surf}} = C_L \times f(T_{\text{min}}) \times f(\text{VPD})
        $$

        Where \(C_L\) is the mean potential stomatal conductance per unit leaf
        area.

        Parameters
        ----------
        tmin : int or float or numpy.ndarray
            Daily minimum temperature (degrees K)
        vpd_day : int or float or numpy.ndarray
            Daytime VPD (Pa)

        Returns
        -------
        int or float or numpy.ndarray
            The daytime surface conductance to transpiration (m s-1)
        '''
        m_tmin = linear_constraint(self.tmin_close, self.tmin_open)
        m_vpd = linear_constraint(self.vpd_open, self.vpd_close, 'reversed')
        return (self.csl * m_tmin(tmin - 273.15) * m_vpd(vpd_day))

    def transpiration(
            self, pressure: Number, temp_k: Number, vpd: Number, lai: Number,
            fpar: Number, rad_canopy: Number, tmin: Number, r_corr: Number,
            lhv: Number = None, rhumidity: Number = None,
            f_wet: Number = None, daytime: bool = True, tiny = 1e-9
        ) -> Number:
        r'''
        Plant transpiration over daytime or night-time hours, in mass flux
        units (kg m-2 s-1), according to:

        $$
        \lambda E_T = \frac{s A_c + \rho C_p F_c D r_d^{-1}}
          {s + \gamma(1 + r_s r_d^{-1})}
        $$

        NOTE: The `r_corr` argument to this function is different from that
        defined in Equation 13 in the MOD16 C61 User's Guide. For numerical
        stability, `r_corr` equals 1/r where r is the quantity defined in
        Equaiton 13. See `MOD16.evapotranspiration()` for how `r_corr` is
        calculated.

        Parameters
        ----------
        pressure : float or numpy.ndarray
            The air pressure in Pascals
        temp_k : float or numpy.ndarray
            The air temperature in degrees K
        vpd : float or numpy.ndarray
            The vapor pressure deficit in Pascals
        lai : float or numpy.ndarray
            The leaf area index (LAI)
        fpar : float or numpy.ndarray
            Fraction of photosynthetically active radiation (PAR) absorbed by
            the vegetation canopy
        rad_canopy : float or numpy.ndarray
            Net radiation to the canopy (J m-2 s-1)
        tmin : float or numpy.ndarray
            Minimum daily temperature (degrees K)
        r_corr : float or numpy.ndarray or None
            The temperature and pressure correction factor for surface
            conductance
        lhv : float or numpy.ndarray or None
            (Optional) The latent heat of vaporization
        rhumidity : float or numpy.ndarray or None
            (Optional) The relative humidity
        f_wet : float or numpy.ndarray or None
            (Optional) The fraction of the surface that has standing water
        daytime : bool
            (Optional) True to calculate daytime (sun-up) transpiration, False
            to calculate night-time (sun-down) transpiration (Default: True)
        tiny : float
            (Optional) A very small number, the numerical tolerance for zero
            (Default: 1e-9)

        Returns
        -------
        float or numpy.ndarray
            Transpiration (kg m-2 s-1)
        '''
        if lhv is None:
            lhv = latent_heat_vaporization(temp_k)
        if rhumidity is None:
            rhumidity = MOD16.rhumidity(temp_k, vpd)
        if f_wet is None:
            f_wet = np.where(rhumidity < 0.7, 0, np.power(rhumidity, 4))
        # Slope of saturation vapor pressure curve
        s = svp_slope(temp_k)
        # Air density
        rho = MOD16.air_density(temp_k, pressure, rhumidity)
        # Psychrometric constant
        gamma = psychrometric_constant(pressure, temp_k)
        # Resistance to radiative heat transfer through air ("rrc")
        r_r = (rho * SPECIFIC_HEAT_CAPACITY_AIR) / (
            4 * STEFAN_BOLTZMANN * temp_k**3)
        # Surface conductance (zero at nighttime)
        g_surf = 0 # (Equation 13, MOD16 C6.1 User's Guide)
        if daytime:
            # NOTE: C_L (self.csl) is included in self.surface_conductance()
            g_surf = self.surface_conductance(tmin, vpd) / r_corr
        g_cuticular = self.g_cuticular / r_corr
        # -- Canopy conductance, should be zero when LAI or f_wet are zero;
        #   updated calculation for conductance to sensible heat, see User
        #   Guide's Equation 15
        gl_sh = self.gl_sh * lai * (1 - f_wet)
        g = ((gl_sh * (g_surf + g_cuticular)) / (
            gl_sh + g_surf + g_cuticular))
        g_canopy = np.where(np.logical_and(lai > 0, (1 - f_wet) > 0), g, tiny)
        # -- Aerodynamic resistance to heat, water vapor from dry canopy
        #   surface into the air (Equation 16, MOD16 C6.1 User's Guide)
        r_a_dry = (1/self.gl_sh * r_r) / (1/self.gl_sh + r_r) # (s m-1)
        # If canopy radiation is negative, drop (s * A_c) term in the
        #   transpriration calculation (L4046 in MODPR16_main.c)
        rad_canopy = np.where(rad_canopy < 0, 0, rad_canopy)
        # PLANT TRANSPIRATION
        t = (1 - f_wet) * ((s * rad_canopy) +\
            (rho * SPECIFIC_HEAT_CAPACITY_AIR * fpar * (vpd / r_a_dry)))
        if daytime:
            t /= (s + gamma * (1 + (1 / g_canopy) / r_a_dry))
        else:
            # At nighttime, g_canopy is zero, so simplify the calculation
            t /= (s + gamma) # (Equation 17, MOD16 C6.1 User's Guide)
        # Divide by the latent heat of vaporization (J kg-1) to obtain mass
        #   flux (kg m-2 s-1)
        return np.where(g_canopy <= tiny, 0, t / lhv)


def psychrometric_constant(pressure: Number, temp_k: Number) -> Number:
    r'''
    The psychrometric constant, which relates the vapor pressure to the air
    temperature. Calculation derives from the "Handbook of Hydrology" by D.R.
    Maidment (1993), Section 4.2.28.

    $$
    \gamma = \frac{C_p \times P}{\lambda\times 0.622}
    $$

    Where \(C_p\) is the specific heat capacity of air, \(P\) is air pressure,
    and \(\lambda\) is the latent heat of vaporization. The \(C_p\) of air
    varies with its saturation, so the value 1.013e-3 [MJ kg-1 (deg C)-1]
    reflects average atmospheric conditions.

    Parameters
    ----------
    pressure : float or numpy.ndarray
        The air pressure in Pascals
    temp_k : float or numpy.ndarray
        The air temperature in degrees K

    Returns
    -------
    float or numpy.ndarray
        The psychrometric constant at this pressure, temperature (Pa K-1)
    '''
    lhv = latent_heat_vaporization(temp_k) # Latent heat of vaporization (J kg-1)
    return (SPECIFIC_HEAT_CAPACITY_AIR * pressure) /\
        (lhv * MOL_WEIGHT_WET_DRY_RATIO_AIR)


def radiation_net(
        sw_rad: Number, sw_albedo: Number, temp_k: Number) -> Number:
    r'''
    DEPRECATED. Net incoming radiation to the land surface, calculated
    according to the MOD16 algorithm and Cleugh et al. (2007); see
    Equation 7 in the MODIS MOD16 Collection 6.1 User's Guide.

    - Cleugh, H. A., Leuning, R., Mu, Q., & Running, S. W. (2007).
      Regional evaporation estimates from flux tower and MODIS satellite data.
      *Remote Sensing of Environment*, 106(3), 285–304.

    $$
    R_{net} = (1 - \alpha)\times R_{S\downarrow} +
        (\varepsilon_a - \varepsilon_s) \times \sigma \times T^4
    \quad\mbox{where}\quad \varepsilon_s = 0.97
    $$

    Where \(\alpha\) is the MODIS albedo, `R_S` is down-welling short-wave
    radiation, \(\sigma\) is the Stefan-Boltzmann constant, and:

    $$
    \varepsilon_a = 1 - 0.26\,\mathrm{exp}\left(
      -7.77\times 10^{-4}\times (T - 273.15)^2
    \right)
    $$

    Parameters
    ----------
    sw_rad : float or numpy.ndarray
        Incoming short-wave radiation (W m-2)
    sw_albedo : float or numpy.ndarray
        Black-sky albedo, e.g., from MODIS MCD43A3
    temp_k : float or numpy.ndarray
        Air temperature in degrees K

    Returns
    -------
    float or numpy.ndarray
        Net incoming radiation to the land surface (W m-2)
    '''
    # Mu et al. (2011), Equation 5
    emis_surface = 0.97
    emis_atmos = 1 - 0.26 * np.exp(-7.77e-4 * np.power(temp_k - 273.15, 2))
    return sw_rad * (1 - sw_albedo) +\
        STEFAN_BOLTZMANN * (emis_atmos - emis_surface) * np.power(temp_k, 4)


def svp(temp_k: Number) -> Number:
    r'''
    The saturation vapor pressure, based on [
    the Food and Agriculture Organization's (FAO) formula, Equation 13
    ](http://www.fao.org/3/X0490E/x0490e07.htm).

    $$
    \mathrm{SVP} = 1\times 10^3\left(
    0.6108\,\mathrm{exp}\left(
      \frac{17.27 (T - 273.15)}{T - 273.15 + 237.3}
      \right)
    \right)
    $$

    This is also referred to as the August-Roche-Magnus equation.

    Parameters
    ----------
    temp_k : float or numpy.ndarray
        The air temperature in degrees K

    Returns
    -------
    float or numpy.ndarray
    '''
    temp_c = temp_k - 273.15
    # And convert from kPa to Pa
    return 1e3 * 0.6108 * np.exp((17.27 * temp_c) / (temp_c + 237.3))


def svp_slope(temp_k: Number, s: Number = None) -> Number:
    r'''
    The slope of the saturation vapour pressure curve, which describes the
    relationship between saturation vapor pressure and temperature. This
    approximation is taken from the MOD16 C source code. An alternative is
    based on [the Food and Agriculture Organization's (FAO) formula,
    Equation 13 ](http://www.fao.org/3/X0490E/x0490e07.htm).

    $$
    \Delta = 4098\times [\mathrm{SVP}]\times (T - 273.15 + 237.3)^{-2}
    $$

    Parameters
    ----------
    temp_k : float or numpy.ndarray
        The air temperature in degrees K
    s : float or numpy.ndarray or None
        Saturation vapor pressure, if already known (Optional)

    Returns
    -------
    float or numpy.ndarray
        The slope of the saturation vapor pressure curve in Pascals per
        degree K (Pa degK-1)
    '''
    if s is None:
        s = svp(temp_k)
    return 17.38 * 239.0 * s / (239.0 + temp_k - 273.15)**2

mod16/__init__.py
arthur-e/MOD16
'''
Calibration of MOD16 against a representative, global eddy covariance (EC)
flux tower network. The model calibration is based on Markov-Chain Monte
Carlo (MCMC). Example use:

    python calibration.py tune --pft=1
    python calibration.py tune --pft=1 --config=MOD16_calibration_config.json

The default configuration file provided in the repository,
`mod16/data/MOD16_calibration_config.json`, should be copied and modified to
suit your needs.

The general calibration protocol used here involves:

1. Check how well the chain(s) are mixing by running short:
`python calibration.py tune 1 --draws=5000`
2. If any chain is "sticky," run a short chain while tuning the jump scale:
`python calibration.py tune 1 --tune=scaling --draws=5000`
3. Using the trace plot from Step (2) as a reference, experiment with
different jump scales to try and achieve the same (optimal) mixing when
tuning on `lambda` (default) instead, e.g.:
`python calibration.py tune 1 --scaling=1e-2 --draws=5000`
4. When the right jump scale is found, run a chain at the desired length.

Once a good mixture is obtained, it is necessary to prune the samples to
eliminate autocorrelation, e.g., in Python:

    sampler = MOD16StochasticSampler(...)
    sampler.plot_autocorr(burn = 1000, thin = 10)
    trace = sampler.get_trace(burn = 1000, thin = 10)

A thinned posterior can be exported from the command line:

    python calibration.py export-bplut output.csv --burn=1000 --thin=10
'''

import datetime
import json
import os
import numpy as np
import h5py
import pymc as pm
import aesara.tensor as at
import arviz as az
import mod16
from multiprocessing import get_context, set_start_method
from pathlib import Path
from typing import Sequence
from scipy import signal
from matplotlib import pyplot
from mod16 import MOD16
from mod16.utils import restore_bplut, pft_dominant
from mod17 import PFT_VALID
from mod17.calibration import BlackBoxLikelihood, StochasticSampler

MOD16_DIR = os.path.dirname(mod16.__file__)


class MOD16StochasticSampler(StochasticSampler):
    '''
    A Markov Chain-Monte Carlo (MCMC) sampler for MOD16. The specific sampler
    used is the Differential Evolution (DE) MCMC algorithm described by
    Ter Braak (2008), though the implementation is specific to the PyMC3
    library.

    Parameters
    ----------
    config : dict
        Dictionary of configuration parameters
    model : Callable
        The function to call (with driver data and parameters); this function
        should take driver data as positional arguments and the model
        parameters as a `*Sequence`; it should require no external state.
    observed : Sequence
        Sequence of observed values that will be used to calibrate the model;
        i.e., model is scored by how close its predicted values are to the
        observed values
    params_dict : dict or None
        Dictionary of model parameters, to be used as initial values and as
        the basis for constructing a new dictionary of optimized parameters
    backend : str or None
        Path to a NetCDF4 file backend (Default: None)
    weights : Sequence or None
        Optional sequence of weights applied to the model residuals (as in
        weighted least squares)
    '''
    required_parameters = {
        'ET': MOD16.required_parameters
    }
    required_drivers = {
        'ET': [
            'lw_net_day', 'lw_net_night', 'sw_rad_day', 'sw_rad_night',
            'sw_albedo', 'temp_day', 'temp_night', 'temp_annual', 'tmin',
            'vpd_day', 'vpd_night', 'pressure', 'fpar', 'lai'
        ]
    }

    def compile_et_model(
            self, observed: Sequence, drivers: Sequence) -> pm.Model:
        '''
        Creates a new ET model based on the prior distribution. Model can be
        re-compiled multiple times, e.g., for cross validation.

        Parameters
        ----------
        observed : Sequence
            Sequence of observed values that will be used to calibrate the model;
            i.e., model is scored by how close its predicted values are to the
            observed values
        drivers : list or tuple
            Sequence of driver datasets to be supplied, in order, to the
            model's run function

        Returns
        -------
        pm.Model
        '''
        # Define the objective/ likelihood function
        log_likelihood = BlackBoxLikelihood(
            self.model, observed, x = drivers, weights = self.weights)
        # With this context manager, "all PyMC3 objects introduced in the indented
        #   code block...are added to the model behind the scenes."
        with pm.Model() as model:
            # NOTE: Parameters shared with MOD17 are fixed based on MOD17
            #   re-calibration
            tmin_close = self.params['tmin_close']
            tmin_open = self.params['tmin_open']
            vpd_open = self.params['vpd_open']
            vpd_close = self.params['vpd_close']
            gl_sh =       pm.Triangular('gl_sh', **self.prior['gl_sh'])
            gl_wv =       pm.Triangular('gl_wv', **self.prior['gl_wv'])
            g_cuticular = pm.LogNormal(
                'g_cuticular', **self.prior['g_cuticular'])
            csl =         pm.LogNormal('csl', **self.prior['csl'])
            rbl_min =     pm.Triangular('rbl_min', **self.prior['rbl_min'])
            rbl_max =     pm.Triangular('rbl_max', **self.prior['rbl_max'])
            beta =        pm.Triangular('beta', **self.prior['beta'])
            # (Stochstic) Priors for unknown model parameters
            # Convert model parameters to a tensor vector
            params_list = [
                tmin_close, tmin_open, vpd_open, vpd_close, gl_sh, gl_wv,
                g_cuticular, csl, rbl_min, rbl_max, beta
            ]
            params = at.as_tensor_variable(params_list)
            # Key step: Define the log-likelihood as an added potential
            pm.Potential('likelihood', log_likelihood(params))
        return model


class CalibrationAPI(object):
    '''
    Convenience class for calibrating the MOD17 GPP and NPP models. Meant to
    be used with `fire.Fire()`.
    '''

    def __init__(self, config = None):
        config_file = config
        if config_file is None:
            config_file = os.path.join(
                MOD16_DIR, 'data/MOD16_calibration_config.json')
        with open(config_file, 'r') as file:
            self.config = json.load(file)
        self.hdf5 = self.config['data']['file']

    def _filter(self, raw: Sequence, size: int):
        'Apply a smoothing filter with zero phase offset'
        if size > 1:
            window = np.ones(size) / size
            return np.apply_along_axis(
                lambda x: signal.filtfilt(window, np.ones(1), x), 0, raw)
        return raw # Or, revert to the raw data

    def clean_observed(
            self, raw: Sequence, drivers: Sequence, protocol: str = 'ET',
            filter_length: int = 2) -> Sequence:
        '''
        Cleans observed tower flux data according to a prescribed protocol.

        Parameters
        ----------
        raw : Sequence
        drivers : Sequence
        protocol : str
        filter_length : int
            The window size for the smoothing filter, applied to the observed
            data

        Returns
        -------
        Sequence
        '''
        # Read in the observed data and apply smoothing filter; then mask out
        #   negative latent heat observations
        obs = self._filter(raw, filter_length)
        return np.where(obs < 0, np.nan, obs)

    def export_posterior(
            self, model: str, param: str, output_path: str, thin: int = 10,
            burn: int = 1000, k_folds: int = 1):
        '''
        Exports posterior distribution for a parameter, for each PFT to HDF5.

        Parameters
        ----------
        model : str
            The name of the model ("GPP" or "NPP")
        param : str
            The model parameter to export
        output_path : str
            The output HDF5 file path
        thin : int
            Thinning rate
        burn : int
            The burn-in (i.e., first N samples to discard)
        k_folds : int
            The number of k-folds used in cross-calibration/validation;
            if more than one (default), the folds for each PFT will be
            combined into a single HDF5 file
        '''
        params_dict = restore_bplut(self.config['BPLUT'][model])
        bplut = params_dict.copy()
        # Filter the parameters to just those for the PFT of interest
        post = []
        for pft in PFT_VALID:
            params = dict([(k, v[pft]) for k, v in params_dict.items()])
            backend = self.config['optimization']['backend_template'] %\
                (model, pft)
            post_by_fold = []
            for fold in range(1, k_folds + 1):
                if k_folds > 1:
                    backend = self.config['optimization']['backend_template'] %\
                        (f'{model}-k{fold}', pft)
                # NOTE: This value was hard-coded in the extant version of MOD16
                if 'beta' not in params:
                    params['beta'] = 250
                sampler = MOD16StochasticSampler(
                    self.config, getattr(MOD16, '_%s' % model.lower()), params,
                    backend = backend)
                trace = sampler.get_trace()
                fit = trace.sel(draw = slice(burn, None, thin))['posterior']
                if param in fit:
                    post_by_fold.append(
                        az.extract_dataset(fit, combined = True)[param].values)
                else:
                    # In case there is, e.g., a parameter that takes on a
                    #   constant value for a specific PFT
                    if k_folds > 1:
                        post_by_fold.append(
                            np.ones((1, post[-1].shape[-1])) * np.nan)
                    else:
                        a_key = list(fit.keys())[0]
                        post_by_fold.append(
                            np.ones(fit[a_key].values.shape) * np.nan)
            if k_folds > 1:
                post.append(np.vstack(post_by_fold))
            else:
                post.extend(post_by_fold)
        # If not every PFT's posterior has the same number of samples (e.g.,
        #   when one set of chains was run longer than another)...
        if not all([p.shape == post[0].shape for p in post]):
            max_len = max([p.shape for p in post])[0]
            # ...Reshape all posteriors to match the greatest sample size
            import ipdb
            ipdb.set_trace()#FIXME
            post = [
                np.pad(
                    p.astype(np.float32), (0, max_len - p.size),
                    mode = 'constant', constant_values = (np.nan,))
                for p in post
            ]
        with h5py.File(output_path, 'a') as hdf:
            post = np.stack(post)
            ts = datetime.date.today().strftime('%Y-%m-%d') # Today's date
            dataset = hdf.create_dataset(
                f'{param}_posterior', post.shape, np.float32, post)
            dataset.attrs['description'] = 'CalibrationAPI.export_posterior() on {ts}'

    def tune(
            self, pft: int, plot_trace: bool = False, ipdb: bool = False,
            save_fig: bool = False, **kwargs):
        '''
        Run the MOD16 ET calibration.

        Parameters
        ----------
        pft : int
            The Plant Functional Type (PFT) to calibrate
        plot_trace : bool
            True to plot the trace for a previous calibration run; this will
            also NOT start a new calibration (Default: False)
        ipdb : bool
            True to drop the user into an ipdb prompt, prior to and instead of
            running calibration
        save_fig : bool
            True to save figures to files instead of showing them
            (Default: False)
        **kwargs
            Additional keyword arguments passed to
            `MOD16StochasticSampler.run()`

        NOTE that `MOD16StochasticSampler` inherits methods from the `mod17`
        module, including [run()](https://arthur-e.github.io/MOD17/calibration.html#mod17.calibration.StochasticSampler).
        '''
        assert pft in PFT_VALID, f'Invalid PFT: {pft}'
        # Set var_names to tell ArviZ to plot only the free parameters
        kwargs.update({'var_names': MOD16.required_parameters[4:]})
        # Pass configuration parameters to MOD16StochasticSampler.run()
        for key in ('chains', 'draws', 'tune', 'scaling'):
            if key in self.config['optimization'].keys():
                kwargs[key] = self.config['optimization'][key]
        # Filter the parameters to just those for the PFT of interest
        params_dict = restore_bplut(self.config['BPLUT']['ET'])
        params_dict = dict([(k, v[pft]) for k, v in params_dict.items()])
        # NOTE: This value was hard-coded in the extant version of MOD16
        if np.isnan(params_dict['beta']):
            params_dict['beta'] = 250
        model = MOD16(params_dict)
        with h5py.File(self.hdf5, 'r') as hdf:
            sites = hdf['FLUXNET/site_id'][:].tolist()
            if hasattr(sites[0], 'decode'):
                sites = [s.decode('utf-8') for s in sites]
            # Get dominant PFT
            pft_map = pft_dominant(hdf['state/PFT'][:], site_list = sites)
            # Blacklist various sites
            blacklist = self.config['data']['sites_blacklisted']
            pft_mask = np.logical_and(pft_map == pft, ~np.in1d(sites, blacklist))
            weights = hdf['weights'][pft_mask]
            # Read in tower observations
            tower_obs = hdf['FLUXNET/latent_heat'][:][:,pft_mask]
            # Read the validation mask; mask out observations that are
            #   reserved for validation
            print('Masking out validation data...')
            mask = hdf['FLUXNET/validation_mask'][pft]
            tower_obs[mask] = np.nan
            # Read start and end dates and mask data appropriately
            timestamps = [
                 f'{y}-{str(m).zfill(2)}-{str(d).zfill(2)}'
                 for y, m, d in hdf['time'][:].tolist()
            ]
            start = self.config['data']['dates']['start']
            end = self.config['data']['dates']['end']
            t0 = timestamps.index(start)
            t1 = timestamps.index(end) + 1
            tower_obs = tower_obs[t0:t1]
            # Read in driver datasets
            print('Loading driver datasets...')
            lw_net_day = hdf['MERRA2/LWGNT_daytime'][:][t0:t1,pft_mask]
            lw_net_night = hdf['MERRA2/LWGNT_nighttime'][:][t0:t1,pft_mask]
            if self.config['optimization']['platform'] == 'VIIRS':
                sw_albedo = hdf['VIIRS/VNP43MA3_black_sky_sw_albedo'][:][t0:t1,pft_mask]
            else:
                sw_albedo = hdf['MODIS/MCD43GF_black_sky_sw_albedo'][:][t0:t1,pft_mask]
            sw_albedo = np.nanmean(sw_albedo, axis = -1)
            sw_rad_day = hdf['MERRA2/SWGDN_daytime'][:][t0:t1,pft_mask]
            sw_rad_night = hdf['MERRA2/SWGDN_nighttime'][:][t0:t1,pft_mask]
            temp_day = hdf['MERRA2/T10M_daytime'][:][t0:t1,pft_mask]
            temp_night = hdf['MERRA2/T10M_nighttime'][:][t0:t1,pft_mask]
            tmin = hdf['MERRA2/Tmin'][:][t0:t1,pft_mask]
            # As long as the time series is balanced w.r.t. years (i.e., same
            #   number of records per year), the overall mean is the annual mean
            temp_annual = hdf['MERRA2/T10M'][:][t0:t1,pft_mask].mean(axis = 0)
            vpd_day = MOD16.vpd(
                hdf['MERRA2/QV10M_daytime'][:][t0:t1,pft_mask],
                hdf['MERRA2/PS_daytime'][:][t0:t1,pft_mask],
                temp_day)
            vpd_night = MOD16.vpd(
                hdf['MERRA2/QV10M_nighttime'][:][t0:t1,pft_mask],
                hdf['MERRA2/PS_nighttime'][:][t0:t1,pft_mask],
                temp_night)
            pressure = hdf['MERRA2/PS'][:][t0:t1,pft_mask]
            # Read in fPAR, LAI, and convert from (%) to [0,1]
            prefix = 'MODIS/MOD'
            if self.config['optimization']['platform'] == 'VIIRS':
                prefix = 'VIIRS/VNP'
            fpar = np.nanmean(
                hdf[f'{prefix}15A2HGF_fPAR_interp'][:][t0:t1,pft_mask], axis = -1)
            lai = np.nanmean(
                hdf[f'{prefix}15A2HGF_LAI_interp'][:][t0:t1,pft_mask], axis = -1)
            # Convert fPAR from (%) to [0,1] and re-scale LAI; reshape fPAR and LAI
            fpar /= 100
            lai /= 10
        # Compile driver datasets
        drivers = [
            lw_net_day, lw_net_night, sw_rad_day, sw_rad_night, sw_albedo,
            temp_day, temp_night, temp_annual, tmin, vpd_day, vpd_night,
            pressure, fpar, lai
        ]
        print('Initializing sampler...')
        backend = self.config['optimization']['backend_template'] % ('ET', pft)
        sampler = MOD16StochasticSampler(
            self.config, MOD16._et, params_dict, backend = backend,
            weights = weights)
        if plot_trace or ipdb:
            # This matplotlib setting prevents labels from overplotting
            pyplot.rcParams['figure.constrained_layout.use'] = True
            trace = sampler.get_trace()
            if ipdb:
                import ipdb
                ipdb.set_trace()
            az.plot_trace(trace, var_names = MOD16.required_parameters)
            pyplot.show()
            return
        tower_obs = self.clean_observed(tower_obs, drivers)
        # Get (informative) priors for just those parameters that have them
        with open(self.config['optimization']['prior'], 'r') as file:
            prior = json.load(file)
        prior_params = list(filter(
            lambda p: p in prior.keys(), sampler.required_parameters['ET']))
        prior = dict([
            (p, dict([(k, v[pft]) for k, v in prior[p].items()]))
            for p in prior_params
        ])
        sampler.run(
            tower_obs, drivers, prior = prior, save_fig = save_fig, **kwargs)


if __name__ == '__main__':
    import fire
    import warnings
    with warnings.catch_warnings():
        warnings.simplefilter('ignore')
        fire.Fire(CalibrationAPI)

mod16/calibration.py
arthur-e/MOD16
'''
'''

import json
import os
import numpy as np
import h5py
import mod16
from tqdm import tqdm
from mod16 import MOD16
from mod16.utils import restore_bplut, pft_dominant
from mod17.science import nash_sutcliffe
from SALib.sample import saltelli
from SALib.analyze import sobol

OUTPUT_TPL = '/home/arthur/Workspace/NTSG/projects/Y2021_MODIS-VIIRS/data/MOD16_sensitivity_%s_analysis.json'
MOD16_DIR = os.path.dirname(mod16.__file__)
with open(os.path.join(MOD16_DIR, 'data/MOD16_calibration_config.json'), 'r') as file:
    CONFIG = json.load(file)
BOUNDS = {
    "tmin_close": [-35, 0],
    "tmin_open": [0, 25],
    "vpd_open": [0, 1500],
    "vpd_close": [1500, 8000],
    "gl_sh": [0.01, 0.12],
    "gl_wv": [0.01, 0.12],
    "g_cuticular": [1e-6, 1e-3],
    "csl": [0.001, 0.1],
    "rbl_min": [10, 80],
    "rbl_max": [80, 150],
    "beta": [0, 1000]
}


def main(pft = None):
    # Stratify the data using the validation mask so that an equal number of
    #   samples from each PFT are used
    drivers, tower_obs = load_data(pft = pft, validation_mask_only = pft is None)
    params = MOD16.required_parameters
    problem = {
        'num_vars': len(params),
        'names': params,
        'bounds': [
            BOUNDS[p]
            for p in params
        ]
    }
    # NOTE: Number of samples must be a power of 2
    param_sweep = saltelli.sample(problem, 256 if pft is None else 128)
    Y = np.zeros([param_sweep.shape[0]])
    for i, X in enumerate(tqdm(param_sweep)):
        yhat = MOD16._et(X, *drivers)
        Y[i] = nash_sutcliffe(yhat, tower_obs, norm = True)
    metrics = sobol.analyze(problem, Y)
    filename = OUTPUT_TPL % 'ET'
    if pft is not None:
        filename = OUTPUT_TPL % f'ET-PFT{pft}'
    with open(filename, 'w') as file:
        json.dump(dict([(k, v.tolist()) for k, v in metrics.items()]), file)


def load_data(pft, validation_mask_only = False):
    print('Loading driver datasets...')
    with h5py.File(CONFIG['data']['file'], 'r') as hdf:
        if pft is not None:
            site_list = hdf['FLUXNET/site_id'][:].tolist()
            if hasattr(site_list[0], 'decode'):
                site_list = [s.decode('utf-8') for s in site_list]
            sites = pft_dominant(hdf['state/PFT'][:], site_list = site_list)
            sites = sites == pft
        else:
            shp = hdf['MERRA2/Tmin'].shape
            sites = np.ones(shp[1]).astype(bool)
        lw_net_day = hdf['MERRA2/LWGNT_daytime'][:][:,sites]
        lw_net_night = hdf['MERRA2/LWGNT_nighttime'][:][:,sites]
        sw_albedo = hdf['MODIS/MCD43A3_black_sky_sw_albedo'][:][:,sites]
        sw_rad_day = hdf['MERRA2/SWGDN_daytime'][:][:,sites]
        sw_rad_night = hdf['MERRA2/SWGDN_nighttime'][:][:,sites]
        temp_day = hdf['MERRA2/T10M_daytime'][:][:,sites]
        temp_night = hdf['MERRA2/T10M_nighttime'][:][:,sites]
        tmin = hdf['MERRA2/Tmin'][:][:,sites]
        # As long as the time series is balanced w.r.t. years (i.e., same
        #   number of records per year), the overall mean is the annual mean
        temp_annual = hdf['MERRA2/T10M'][:][:,sites].mean(axis = 0)[None,:]\
            .repeat(tmin.shape[0], axis = 0)
        vpd_day = MOD16.vpd(
            hdf['MERRA2/QV10M_daytime'][:][:,sites],
            hdf['MERRA2/PS_daytime'][:][:,sites],
            temp_day)
        vpd_night = MOD16.vpd(
            hdf['MERRA2/QV10M_nighttime'][:][:,sites],
            hdf['MERRA2/PS_nighttime'][:][:,sites],
            temp_night)
        pressure = hdf['MERRA2/PS'][:][:,sites]
        # Read in fPAR, LAI, and convert from (%) to [0,1]
        fpar = np.nanmean(
            hdf['MODIS/MOD15A2HGF_fPAR_interp'][:][:,sites], axis = -1)
        lai = np.nanmean(
            hdf['MODIS/MOD15A2HGF_LAI_interp'][:][:,sites], axis = -1)
        # Convert fPAR from (%) to [0,1] and re-scale LAI; reshape fPAR and LAI
        fpar /= 100
        lai /= 10
        tower_obs = hdf['FLUXNET/latent_heat'][:][:,sites]
        if pft is None:
            is_test = hdf['FLUXNET/validation_mask'][:].sum(axis = 0).astype(bool)
    # Compile driver datasets
    drivers = [
        lw_net_day, lw_net_night, sw_rad_day, sw_rad_night, sw_albedo,
        temp_day, temp_night, temp_annual, tmin, vpd_day, vpd_night,
        pressure, fpar, lai
    ]
    # Speed things up by focusing only on data points where valid data exist
    mask = ~np.isnan(tower_obs)
    if pft is None and validation_mask_only:
        # Stratify the data using the validation mask so that an equal number
        #   of samples from each PFT are used
        mask = np.logical_and(is_test, mask)
    drivers = [d[mask] for d in drivers]
    return (drivers, tower_obs[mask])


if __name__ == '__main__':
    import fire
    fire.Fire(main)

mod16/sensitivity.py
arthur-e/MOD16
'''
Utilities related to the MOD16 algorithm.
'''

from __future__ import annotations
import csv
import os
import numpy as np
import pandas as pd
import mod16
from collections import Counter
from typing import Callable, Sequence
from pandas._typing import FilePath, ReadCsvBuffer

BPLUT_FIELD_LOOKUP = {
    'Tmin_min(C)':      'tmin_close',
    'Tmin_max(C)':      'tmin_open',
    'VPD_min(Pa)':      'vpd_open',
    'VPD_max(Pa)':      'vpd_close',
    'gl_sh((m/s)':      'gl_sh',
    'gl_e_wv(m/s)':     'gl_wv',
    'g_cuticular(m/s)': 'g_cuticular',
    'Cl(m/s)':          'csl',
    'RBL_MIN(s/m)':     'rbl_min',
    'RBL_MAX(s/m)':     'rbl_max',
    'beta':             'beta'
}

def pft_dominant(
        pft_map: np.ndarray, site_list: list = None,
        valid_pft: list = (1,2,3,4,5,6,7,8,9,10,12)):
    '''
    Returns the dominant PFT type, based on the PFT mode among the sub-grid
    for each site. Note that this is specific to the MOD17 calibration/
    validation (Cal/Val) protocol, i.e., three sites are always classified as
    Deciduous Needleleaf (PFT 3):

        CA-SF2
        CA-SF3
        US-NGC

    One other site, US-A10, is black-listed because it's PI-reported PFT is
    not compatible with our PFT classification.

    Parameters
    ----------
    pft_map : numpy.ndarray
        (N x M) array of PFT classes, where N is the number of sites and
        M is the number of sub-grid cells (N PFT classes are returned)
    site_list : list
        (Optional) List of the site names; must be provided to get PFT
        classes that accurately match the Cal/Val protocol
    valid_pft : list
        (Optional) List of valid PFT classes (Default: `range(1, 13)` but
        without 11)

    Returns
    -------
    numpy.ndarray
        An (N,) array of the dominant PFT classes
    '''
    pft_dom = np.zeros(pft_map.shape[0], np.float32)
    for i in range(0, pft_map.shape[0]):
        try:
            pft_dom[i] = Counter(
                list(filter(lambda x: x in valid_pft, pft_map[i])))\
                .most_common()[0][0]
        except:
            # Skip those sites that have no valid PFTs
            continue
    if site_list is not None:
        pft_dom[site_list.index('US-A10')] = np.nan
        # For PFT==3, DNF, use pre-determined locations
        idx = [site_list.index(sid) for sid in ('CA-SF2', 'CA-SF3', 'US-NGC')]
        pft_dom[idx] = 3
    return pft_dom


def restore_bplut(
        path_or_buffer: FilePath | ReadCsvBuffer | str,
        nrows: int = 11) -> dict:
    '''
    NOTE: I manually exported Maosheng's fixed-width version (fixed-width
    files are a crime) to CSV for easier handling.

    Parameters
    ----------
    path_or_buffer : str or buffer
        File path or buffer representing the BPLUT to be read
    nrows : int
        Number of rows to read

    Returns
    -------
    dict
    '''
    # Remaps Maosheng's PFT order to the actual PFT code from MCD12Q1
    #   LC_Type2
    pft_lookup = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12]
    data = pd.read_csv(path_or_buffer, nrows = nrows)
    # Create a dictionary with an array for every key
    output = dict([
        (k, np.full((13,), np.nan))
        for k in BPLUT_FIELD_LOOKUP.values()
    ])
    # Assumes the first column indexes the parameter/ field names
    field_index = data.columns[0]
    pft_index = list(data.columns)
    pft_index.remove(field_index)
    for k, key in enumerate(data[field_index]):
        if key not in data[data.columns[0]].values:
            continue # e.g., "beta" not included in Collection 5.x
        values = data.loc[data[field_index] == key, pft_index].values.ravel()
        output[BPLUT_FIELD_LOOKUP[key]][pft_lookup] = values
    return output


def write_bplut(params_dict: dict, output_path: str):
    '''
    Writes a BPLUT parameters dictionary to an output CSV file.

    Parameters
    ----------
    params_dict : dict
    output_path : str
        The output CSV file path
    '''
    template = os.path.join(
        os.path.dirname(mod16.__file__), 'data/MOD16_BPLUT_C5.1_05deg_MCD43B_Albedo_MERRA_GMAO.csv')
    with open(template, 'r') as file:
        reader = csv.reader(file)
        for line in reader:
            if reader.line_num > 1:
                break
            header = line
    with open(output_path, 'w') as file:
        writer = csv.writer(file)
        writer.writerow(header)
        for name, key in BPLUT_FIELD_LOOKUP.items():
            values = []
            for pft in mod16.PFT_VALID:
                values.append(params_dict[key][pft])
            writer.writerow((name, *values))

mod16/utils.py
arthur-e/MOD16
'''
Unit tests for the `mod16` Python utilities library.
'''

import os
import unittest
import numpy as np
from mod16 import MOD16, psychrometric_constant, radiation_net, svp_slope

# MOD17_BPLUT = os.path.join(
#     os.path.dirname(mod17.__file__), 'data/MOD17_BPLUT_C5.1_MERRA_NASA.csv')

class CanopyEvaporation(unittest.TestCase):
    '''
    Suite of test cases related to wet canopy evaporation.
    '''

    @classmethod
    def setUp(cls):
        cls.params = dict().fromkeys(MOD16.required_parameters, None)
        cls.params.update({
            'gl_sh': 0.01,
            'gl_wv': 0.01,
            'g_cuticular': 1e-5,
            'tmin_close': -8, # deg C
            'tmin_open': 8, # deg C
            'vpd_open': 650,
            'vpd_close': 3000,
            'rbl_min': 60,
            'rbl_max': 90,
            'csl': 2.4e-3,
            'beta': 250 # Original (hard-coded constant) value
        })
        cls.pressure = 100e3
        cls.temp_k = 273.15 + 30
        cls.tmin = 273.15 + 20
        cls.vpd = 1000
        cls.lai = 1.5
        cls.fpar = 0.5
        cls.rad_canopy = 5000
        cls.rad_soil = 5000
        cls.f_wet = 0.5
        cls.r_corr = (101300 / cls.pressure) * (cls.temp_k / 293.15)**1.75
        # 5 levels of each
        cls._pressure = np.arange(98e3, 103e3, 1e3)
        cls._temp_k = 273.15 + np.array([0, 10, 20, 30, 40])
        cls._vpd = np.arange(0, 5000, 1000)
        cls._lai = np.arange(0.5, 3, 0.5)
        cls._fpar = np.array([0.1, 0.3, 0.5, 0.7, 0.9])
        cls._rad_canopy = np.arange(3e3, 8e3, 1e3)
        cls._f_wet = np.array([0.1, 0.3, 0.5, 0.7, 0.9])

    def test_evaporation_soil(self):
        'Should accurately calculate evaporation from bare soil'
        model = MOD16(self.params)
        evap = 60 * 60 * model.evaporation_soil(
            self.pressure, self.temp_k, self.vpd, self.fpar, self.rad_soil,
            self.r_corr)
        self.assertEqual(evap.round(3), 3.102)

    def test_evaporation_soil_by_fpar(self):
        'Should accurately calc. evap. from bare soil across fPAR gradient'
        model = MOD16(self.params)
        evap = 60 * 60 * model.evaporation_soil(
            self.pressure, self.temp_k, self.vpd, self._fpar, self.rad_soil,
            self.r_corr)
        # <KAE> 2021-12-08 Maosheng indicated that the User Guide has an error
        #   w.r.t. to the scaling of soil heat flux by (1 - fPAR); this has
        #   been removed from MOD16.soil_heat_flux()
        self.assertTrue(np.equal(evap.round(3), np.array([
            3.128, 3.115, 3.102, 3.089, 3.076
        ])).all())

    def test_transpiration_daytime(self):
        'Should accurately calculate transpiration during daytime'
        model = MOD16(self.params)
        trans = 60 * 60 * model.transpiration(
            self.pressure, self.temp_k, self.vpd, self.lai, self.fpar,
            self.rad_canopy, self.tmin, self.r_corr, daytime = True)
        # <KAE> 2021-12-08 The Tmin ramp function parameters had been
        #   inadvertently switched when the tests were first written
        self.assertEqual(trans.round(3), 1.248) # kg m-2 hr-1

    def test_transpiration_nighttime(self):
        'Should accurately calculate transpiration during nighttime'
        model = MOD16(self.params)
        trans = 60 * 60 * model.transpiration(
            self.pressure, self.temp_k, self.vpd, self.lai, self.fpar,
            self.rad_canopy, self.tmin, self.r_corr, daytime = False)
        # <KAE> 2021-12-08 The Tmin ramp function parameters had been
        #   inadvertently switched when the tests were first written
        self.assertEqual(trans.round(3), 3.854) # kg m-2 hr-1

    def test_wet_canopy_evaporation(self):
        'Should accurately calculate wet canopy evaporation (kg m-2 s-1)'
        model = MOD16(self.params)
        evap = model.evaporation_wet_canopy(
            self.pressure, self.temp_k, self.vpd, self.lai, self.fpar,
            self.rad_canopy).round(6)
        # <KAE> 2021-12-08 There was a bug in the r_corr definition
        self.assertEqual(evap, 4.49e-4)

    def test_wet_canopy_evaporation_by_pressure(self):
        'Accurately calculate wet canopy evaporation by pressure gradient'
        model = MOD16(self.params)
        evap = 60 * 60 * model.evaporation_wet_canopy(
            self._pressure, self.temp_k, self.vpd, self.lai, self.fpar,
            self.rad_canopy)
        self.assertTrue(np.equal(evap.round(3), np.array([
            1.623, 1.62 , 1.618, 1.615, 1.612 # kg m-2 hr-1
        ])).all())

    def test_wet_canopy_evaporation_by_temperature(self):
        'Accurately calculate wet canopy evaporation by temperature gradient'
        model = MOD16(self.params)
        evap = 60 * 60 * model.evaporation_wet_canopy(
            self.pressure, self._temp_k, self.vpd, self.lai, self.fpar,
            self.rad_canopy)
        self.assertTrue(np.equal(evap.round(3), np.array([
            0., 0., 0., 1.618, 3.222 # kg m-2 hr-1
        ])).all())

    def test_wet_canopy_evaporation_by_vpd(self):
        'Accurately calculate wet canopy evaporation across VPD gradient'
        model = MOD16(self.params)
        evap = 60 * 60 * model.evaporation_wet_canopy(
            self.pressure, self.temp_k, self._vpd, self.lai, self.fpar,
            self.rad_canopy)
        self.assertTrue(np.equal(evap.round(3), np.array([
            5.382, 1.618, 0, 0, 0 # kg m-2 hr-1
        ])).all())

    def test_wet_canopy_evaporation_by_lai(self):
        'Accurately calculate wet canopy evaporation across LAI gradient'
        model = MOD16(self.params)
        evap = 60 * 60 * model.evaporation_wet_canopy(
            self.pressure, self.temp_k, self.vpd, self._lai, self.fpar,
            self.rad_canopy)
        self.assertTrue(np.equal(evap.round(3), np.array([
            1.174, 1.478, 1.618, 1.699, 1.752 # kg m-2 hr-1
        ])).all())

    def test_wet_canopy_evaporation_by_fpar(self):
        'Accurately calculate wet canopy evaporation across fPAR gradient'
        model = MOD16(self.params)
        evap = 60 * 60 * model.evaporation_wet_canopy(
            self.pressure, self.temp_k, self.vpd, self.lai, self._fpar,
            self.rad_canopy)
        self.assertTrue(np.equal(evap.round(3), np.array([
            1.611, 1.615, 1.618, 1.621, 1.624 # kg m-2 hr-1
        ])).all())

    def test_wet_canopy_evaporation_by_canopy_radiation(self):
        'Accurately calculate wet canopy evaporation by canopy radiation'
        model = MOD16(self.params)
        evap = 60 * 60 * model.evaporation_wet_canopy(
            self.pressure, self.temp_k, self.vpd, self.lai, self.fpar,
            self._rad_canopy)
        self.assertTrue(np.equal(evap.round(3), np.array([
            0.974, 1.296, 1.618, 1.94 , 2.262 # kg m-2 hr-1
        ])).all())

    def test_psychrometric_constant(self):
        'Should accurately calculate the psychrometric constant'
        pressure = np.array((100e3, 80e3, 100e3, 80e3))
        temp_k = 273.15 + np.array((10, 10, 25, 25))
        answer = [65.74, 52.59, 66.69, 53.35]
        for i in range(0, 4):
            self.assertEqual(
                answer[i],
                psychrometric_constant(pressure[i], temp_k[i]).round(2))
        # Example from FAO:
        #   http://www.fao.org/3/X0490E/x0490e07.htm#psychrometric%20constant%20(g)
        self.assertEqual(
            54.55, np.round(psychrometric_constant(81.8e3, 25 + 273.15), 2))

    def test_radiation_net(self):
        'Should accurately calculate net radiation to the land surface'
        swrad = np.array((500, 5000, 500, 5000, 500, 5000, 500, 5000))
        albedo = np.array((0.4, 0.4, 0.8, 0.8, 0.4, 0.4, 0.8, 0.8))
        temp_k = 273.15 + np.array((10, 10, 10, 10, 25, 25, 25, 25))
        answer = [223.3, 2923.3, 23.3, 923.3, 241.8, 2941.8, 41.8, 941.8]
        for i in range(0, 8):
            self.assertEqual(
                answer[i],
                radiation_net(swrad[i], albedo[i], temp_k[i]).round(1))

    def test_svp_slope(self):
        'Should accurately calculate slope of SVP curve'
        self.assertEqual( 82.3, svp_slope(273.15 + 10).round(1))
        self.assertEqual(144.8, svp_slope(273.15 + 20).round(1))
        self.assertEqual(188.8, svp_slope(273.15 + 25).round(1))


if __name__ == '__main__':
    unittest.main()

tests/tests.py
arthur-e/MOD16