R/sw_two_stream.R

In ashiklom/edr-da: ED2 RTM testing functions

#' R implementation of ED2 two-stream radiation model
#'
#' @param czen Cosine of angle of incidence (`cosaio`)
#' @param iota_g Ground (soil + snow) albedo (nwl)
#' @param pft PFT identities of each cohort, as integer (ncoh). Shortest to tallest.
#' @param lai Leaf area index of each cohort (ncoh). Shorest to tallest.
#' @param wai Wood area index of each cohort (ncoh). Shorest to tallest.
#' @param cai Crown area of each cohort (ncoh). Shorest to tallest.
#' @param orient_factor Orient factor (npft)
#' @param clumping_factor Clumping factor (npft)
#' @param leaf_reflect Leaf reflectance spectra (nwl * npft)
#' @param leaf_trans Leaf transmittance spectra (nwl * npft)
#' @param wood_reflect Wood reflectance spectra (nwl * npft)
#' @param wood_trans Wood transmittance spectra (nwl * npft)
#' @param down_sky Normalized diffuse solar spectrum (nwl)
#' @param down0_sky Normalized direct solar spectrum (nwl)
#' @param wavelengths Numeric vector of wavelengths to use, in nm
#'   (nwl). Default is 400:2500.
#' @return
#' @author Alexey Shiklomanov
#' @useDynLib redr
#' @export
sw_two_stream <- function(czen,
                          iota_g,
                          pft,
                          lai, wai, cai,
                          orient_factor, clumping_factor,
                          leaf_reflect, leaf_trans,
                          wood_reflect, wood_trans,
                          down_sky, down0_sky,
                          wavelengths = seq(400, 2500),
                          ssa_type = "sellers") {

  ssa_type <- tolower(ssa_type)
  # Sanity checks
  nwl <- length(wavelengths)
  stopifnot(
    ssa_type %in% c("sellers", "sail"),
    NROW(leaf_reflect) == nwl,
    NROW(leaf_trans) == nwl,
    NROW(wood_reflect) == nwl,
    NROW(wood_trans) == nwl,
    NROW(iota_g) == nwl
  )

  ncoh <- length(pft)
  stopifnot(
    length(lai) == ncoh,
    length(wai) == ncoh,
    length(cai) == ncoh
  )

  stopifnot(abs(czen) <= 1)

  # All spectra are between 0 and 1
  stopifnot(
    !any(iota_g > 1 | iota_g < 0),
    !any(leaf_reflect > 1 | leaf_reflect < 0),
    !any(leaf_trans > 1 | leaf_trans < 0),
    !any(wood_reflect > 1 | wood_reflect < 0),
    !any(wood_trans > 1 | wood_trans < 0)
  )

  # Incident radiation has to sum to 1 across all wavelengths
  stopifnot(
    !any(down_sky > 1 | down_sky < 0),
    !any(down0_sky > 1 | down0_sky < 0),
    all(down_sky + down0_sky == 1)
  )

  ##########
  # Unpack PFT-specific parameters across cohorts for easier vectorization
  leaf_reflect <- leaf_reflect[, pft]
  leaf_trans <- leaf_trans[, pft]
  wood_reflect <- wood_reflect[, pft]
  wood_trans <- wood_trans[, pft]
  orient_factor <- orient_factor[pft]
  clumping_factor <- clumping_factor[pft]

  ##########
  # Begin ED2 supplement

  # Eqns S98 and S99, pg. S48
  phi1 <- 0.5 - orient_factor * (0.633 + 0.33 * orient_factor)
  phi2 <- 0.877 * (1 - 2 * phi1)

  # S97; Goudriaan 1977
  # E(Z, \chi_k) in the text
  proj_area <- phi1 + phi2 * czen

  # S100; Sellers 1985
  mu_bar <- (1 - phi1 * log(1 + phi2 / phi1) / phi2) / phi2
  mu_bar[orient_factor == 0] <- 1
  if (!all(is.finite(mu_bar))) {
    ibad <- which(!is.finite(mu_bar))
    stop(
      "mu_bar calculation failed with orient_factor: ",
      paste(orient_factor[ibad], collapse = ", ")
    )
  }

  # Calculations from sfc_rad
  # S101
  leaf_scatter <- leaf_reflect + leaf_trans
  wood_scatter <- wood_reflect + wood_trans

  # S103
  cos_Ak2 <- (1 + orient_factor)^2
  # S102
  leaf_backscatter <- (leaf_scatter + 0.25 * (leaf_reflect - leaf_trans) * cos_Ak2) /   #nolint
    (2 * leaf_scatter)
  wood_backscatter <- (wood_scatter + 0.25 * (wood_reflect - wood_trans) * cos_Ak2) /   #nolint
    (2 * wood_scatter)

  ##########

  # Loop over bands -- skipping because this is already vectorized

  # Loop over cohorts -- vectorizing here
  elai <- clumping_factor * lai
  etai <- elai + wai

  leaf_weight <- elai / etai
  wood_weight <- 1 - leaf_weight

  # Inverse optical depth of direct radiation
  mu0 <- -etai / log((1 - cai) + cai * exp(-proj_area * etai / (cai * czen)))

  # Inverse optical depth of diffuse radiation
  mu <- -etai / log((1 - cai) + cai * exp(-etai / mu_bar))

  # Backscatter coefficients for diffuse radiation
  # S105
  # Single scattering albedo
  iota_ratio <- 1 / (2 * (1 + phi2 * mu0)) *
    (1 - phi1 * mu0 / (1 + phi2 * mu0) *
       log((1 + (phi1 + phi2) * mu0) / (phi1 * mu0)))
  stopifnot(all(is.finite(iota_ratio)))

  # S104
  beta0 <- iota_ratio * (1 + mu0 / mu)
  stopifnot(all(is.finite(beta0)))

  epsil0 <- 1 - 2 * beta0

  #################
  # Define boundary conditions
  #################
  z <- ncoh + 1
  elai[z] <- 0
  etai[z] <- 0
  leaf_weight[z] <- 0.5
  wood_weight[z] <- 0.5
  proj_area[z] <- 0.5
  mu0[z] <- czen / proj_area[z]
  mu[z] <- 1
  iota_ratio[z] <- 0.5 * (1 - 0.5 * mu0[z] *
                            log(1 / (0.5 * mu0[z]) + 1))
  beta0[z] <- iota_ratio[z] * (mu0[z] + mu[z]) / mu[z]
  epsil0[z] <- 1 - 2 * beta0[z]

  # Direct radiation profile via exponential attentuation
  # Transmissivity of direct radiation; S110
  expm0_minus <- exp(-etai / mu0)
  expm0_minus[z] <- 1
  # S111-S112
  down0 <- matrix(0, nwl, z)
  down0[, z] <- down0_sky
  for (j in seq(ncoh, 1)) {
    down0[, j] <- down0[, j + 1] * expm0_minus[j]
  }

  # Convert scalars to matrices, so I can do elementwise multiplication
  vec2mat <- function(x) matrix(rep(x, nwl), nrow = nwl, byrow = TRUE)
  etai <- vec2mat(etai)
  leaf_weight <- vec2mat(leaf_weight)
  wood_weight <- vec2mat(wood_weight)
  mu <- vec2mat(mu)
  mu0 <- vec2mat(mu0)
  epsil0 <- vec2mat(epsil0)
  expm0_minus <- vec2mat(expm0_minus)

  # Diffuse radiation properties
  # All of these are wavelength-dependent quantities (nwl x ncoh)
  iota <- leaf_weight[, -z] * leaf_scatter + wood_weight[, -z] * wood_scatter
  beta <- leaf_weight[, -z] * leaf_backscatter + wood_weight[, -z] * wood_backscatter

  # Part of S113
  epsil <- 1 - 2 * beta

  lambda <- sqrt((1 - epsil * iota) * (1 - iota)) / mu[, -z]

  # Ancillary variables for right-hand side
  iota_mu <- iota / mu[, -z]
  iota_mu0 <- iota / mu0[, -z]
  down0_mu0 <- down0[, -1] / mu0[, -z]
  mu02 <- mu0[, -z] ^ 2
  lambda2 <- lambda ^ 2

  a_aux <- -((1 - epsil * iota) * iota_mu + epsil0[, -z] * iota_mu0) * down0_mu0
  s_aux <- -((1 - iota) * epsil0[, -z] * iota_mu + iota_mu0) * down0_mu0
  delta <- (a_aux + s_aux) * mu02 / (2 * (1 - lambda2 * mu02))
  upsilon <- (a_aux - s_aux) * mu02 / (2 * (1 - lambda2 * mu02))

  # Upwelling and downwelling radiation
  # S124
  iez <- sqrt((1 - iota) / (1 - epsil * iota))
  gamm_plus <- 0.5 * (1 + iez)
  gamm_minus <- 0.5 * (1 - iez)

  # Transmissivity of diffuse light
  # See S122-123
  expl_plus <- exp(lambda * etai[, -z])
  expl_minus <- exp(-lambda * etai[, -z])

  # Define boundary conditions for above
  iota <- cbind(iota, rep(1, nwl))
  beta <- cbind(beta, rep(0, nwl))
  epsil <- cbind(epsil, 1 - 2 * beta[, z])
  lambda <- cbind(lambda, rep(0, nwl))
  a_aux <- cbind(a_aux, -epsil0[, z] * down0_sky / (mu0[, z] ^ 2))
  s_aux <- cbind(s_aux, -iota[, z] * down0_sky / (mu0[, z] ^ 2))
  delta <- cbind(delta, 0.5 * (a_aux[, z] + s_aux[, z]) * mu0[, z] ^ 2)
  upsilon <- cbind(upsilon, 0.5 *(a_aux[, z] - s_aux[, z]) * mu0[, z] ^ 2)
  gamm_plus <- cbind(gamm_plus, rep(1, nwl))
  gamm_minus <- cbind(gamm_minus, rep(0, nwl))
  expl_plus <- cbind(expl_plus, rep(1, nwl))
  expl_minus <- cbind(expl_minus, rep(1, nwl))

  # Size of solution matrix
  nsiz <- 2 * ncoh + 2

  mmat <- array(0, c(nsiz, nsiz, nwl))
  yvec <- matrix(0, nsiz, nwl)

  # Bottom (1) and top boundary conditions
  mmat[1, 1, ] <- (gamm_minus[, 1] - iota_g * gamm_plus[, 1]) * expl_minus[, 1]
  mmat[1, 2, ] <- (gamm_plus[, 1] - iota_g * gamm_minus[, 1]) * expl_plus[, 1]
  mmat[nsiz, nsiz-1, ] <- gamm_plus[, z]
  mmat[nsiz, nsiz, ] <- gamm_minus[, z]
  yvec[1, ] <- iota_g * down0[, 1] - (upsilon[, 1] - iota_g * delta[, 1]) * expm0_minus[1]
  yvec[nsiz, ] <- down_sky - delta[, z]

  for (k in seq_len(ncoh)) {
    kp1 <- k + 1
    k2 <- 2 * k
    k2m1 <- k2 - 1
    k2p1 <- k2 + 1
    k2p2 <- k2 + 2

    yvec[k2, ] <- delta[, kp1] * expm0_minus[, kp1] - delta[, k]
    yvec[k2p1, ] <- upsilon[, kp1] * expm0_minus[, kp1] - upsilon[, k]

    mmat[k2, k2m1, ] <- gamm_plus[, k]
    mmat[k2, k2, ] <- gamm_minus[, k]
    mmat[k2, k2p1, ] <- -gamm_plus[, kp1] * expl_minus[, kp1]
    mmat[k2, k2p2, ] <- -gamm_minus[, kp1] * expl_plus[, kp1]
    mmat[k2p1, k2m1, ] <- gamm_minus[, k]
    mmat[k2p1, k2, ] <- gamm_plus[, k]
    mmat[k2p1, k2p1, ] <- -gamm_minus[, kp1] * expl_minus[, kp1]
    mmat[k2p1, k2p2, ] <- -gamm_plus[, kp1] * expl_plus[, kp1]
  }

  stopifnot(is.finite(sum(mmat)))

  # Solve the radiation balance at each wavelength
  xvec_t <- solvearray(mmat, yvec)
  xvec <- t(xvec_t)

  # Store the solution in matrices (nwl x (ncoh + 1))
  down <- matrix(0, nwl, ncoh + 1)
  up <- matrix(0, nwl, ncoh + 1)
  for (k in seq_len(ncoh + 1)) {
    k2 <- 2 * k
    k2m1 <- k2 - 1
    # S122
    down[, k] <- xvec[, k2m1] * gamm_plus[, k] * expl_minus[, k] +
      xvec[, k2] * gamm_minus[, k] * expl_plus[, k] +
      delta[, k] * expm0_minus[, k]
    # S123
    up[, k] <- xvec[, k2m1] * gamm_minus[, k] * expl_minus[, k] +
      xvec[, k2] * gamm_plus[, k] * expl_plus[, k] +
      upsilon[, k] * expm0_minus[, k]
  }

  # Integrate light levels
  light_level <- matrix(0, nwl, ncoh)
  light_beam_level <- matrix(0, nwl, ncoh)
  light_diff_level <- matrix(0, nwl, ncoh)

  for (k in seq_len(ncoh)) {
    kp1 <- k + 1
    light_level[, k] <- light_level[, k] + 0.5 * (down[, k] + down[, kp1] + down0[, k] + down0[, kp1])
    light_beam_level[, k] <- light_beam_level[, k] + 0.5 * (down0[, k] + down0[, kp1])
    light_diff_level[, k] <- light_diff_level[, k] + 0.5 * (down[, k] + down[, kp1])
  }

  # Albedo is "up" at top of canopy
  list(
    albedo = up[, z],                    # Albedo is upwelling radiation profile at top of canopy (nwl)
    up = up,                             # Upwelling radiation profile, by cohort + top of canopy (nwl x (ncoh + 1))
    down = down,                         # Downwelling radiation, by cohort + top of canopy (nwl x (ncoh + 1))
    light_level = light_level,           # Light level, by cohort (nwl x ncoh)
    light_beam_level = light_beam_level, # Direct light level, by cohort (nwl x ncoh)
    light_diff_level = light_diff_level  # Diffuse light level, by cohort (nwl x ncoh)
  )
}