pars_condEffect.qmd

# Average Marginal Effect {#sec-averageMarginalEffect}

In this section, we assess the impact of individual covariates on the growth, survival, and recruitment models.
We employ the Average Marginal Effect (AME) to summarize the overall effect size of each covariate $x$ on the response $y$ across the dataset.
For instance, we can use the AME to measure the relative effect of competition and climate on the survival function.
In practice, the AME is estimated using partial derivatives.
This involves evaluating the model's output when one covariate is perturbed while keeping the others constant.
To perturb each covariate $x$, we begin by normalizing it within the range of $[0-1]$ using all observations across the 31 species as follows:

$$
  x' = \frac{x - \text{min}(x)}{\text{max}(x) - \text{min}(x)}
$$

where $x'$ is the normalized vector of values for the covariate $x$.
We then increase the normalized values by a perturbation size $p$ before transforming it back to the original scale:

$$
\begin{align*}
&x_{\epsilon} = (x' + p) (\text{max}(x) -\text{min}(x)) + \text{min(x)} \\
&h = x_{\epsilon} - x
\end{align*}
$$

Where $x_{\epsilon}$ is the perturbed covariate in its original scale and $p$ is the perturbation size set to 0.01.
For each covariate $x$, the average marginal effect (AME) across all observations $i$ is estimated as follows:

$$
  AME = \frac{1}{n} \sum_{i=1}^n \frac{\mid f(x_i + h, z_i, \theta) - f(x_i, z_i, \theta) \mid}{h}
$$

$f$ denotes one of the growth, survival, or recruitment demographic models, $z$ the remaining covariates retained at their original values, and $\theta$ is the collection of parameters.
This approach is valuable because it captures the locally linearized effect of each observation. 
Furthermore, computing the average effect across the observed dataset considers the diversity in the distribution of each covariate, giving more weight to regions where observations are more prevalent.

```{r,include=FALSE,echo=FALSE}
Echo=FALSE
Eval=TRUE
Cache=TRUE
Warng=FALSE
Msg=FALSE
library(tidyverse)
library(cmdstanr)
library(posterior)
library(ggdist)
library(ggpubr)
library(ggrepel)
library(data.table)
```

```{r re_loadSpData,echo=Echo,eval=Eval,cache=Cache,warning=Warng,message=Msg}
data_path <- readLines('_data.path')
pars_path <- file.path(data_path, 'output_sim_processed')
models <- c(
  'growth' = 'intcpt_plot_comp_clim',
  'mort' = 'intcpt_plot_comp_clim',
  'recruit' = 'intcpt_plot_comp_clim'
)

## Species if info
spIds <- read_csv(
  file.path(
    data_path, 'species_id.csv'
  )
) |>
mutate(
  shade = factor(shade, levels = c('tolerant', 'intermediate', 'intolerant')),
  shade_sylvics = factor(shade_sylvics, levels = c('very-tolerant',
                                            'tolerant',
                                            'intermediate',
                                            'intolerant',
                                            'very-intolerant'
                                            )),
  succession = factor(succession, levels = c('pioneer',
                                              'intermediate',
                                              'climax'))
) |>
filter(sp_to_analyze)

# climate range used to scale and unscale climate variables
climate_range = readRDS(file.path(data_path, 'climate_scaleRange.RDS'))

# dataset
treeData <- readRDS(file.path(data_path, 'treeData.RDS')) |>
  filter(species_id %in% spIds$species_id_old)

fec_dt <- readRDS(file.path(data_path, 'fec_dt.RDS')) |>
  filter(species_id %in% spIds$species_id_old)
```

```{r ImportPlotRandomEffects,echo=Echo,eval=FALSE,cache=Cache,warning=Warng,message=Msg}
# I set eval to FALSE so this is ran a single time only - keeping here just for documentation
dir.create(file.path('data', 'parameters_randomEffects'), recursive = TRUE)

sample_draws = 4000

for(sp in spIds$species_id_old)
{
  # growth
  #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  readRDS(
    paste0(
      pars_path, '/growth/', models['growth'], '/posterior_plot_', sp, '.RDS'
    )
  ) |>
  select(iter, value, plot_id) |>
  pivot_wider(names_from = plot_id, values_from = value) |>
  select(!iter) |>
  write_csv(
    file = paste0('data/parameters_randomEffects/growth_', sp, '.csv')
  )

  # survival
  #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  readRDS(
    paste0(
      pars_path, '/mort/', models['mort'], '/posterior_plot_', sp, '.RDS'
    )
  ) |>
  select(iter, value, plot_id) |>
  pivot_wider(names_from = plot_id, values_from = value) |>
  select(!iter) |>
  write_csv(
    file = paste0('data/parameters_randomEffects/mort_', sp, '.csv')
  )


  # recruitment
  #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  readRDS(
    paste0(
      pars_path, '/recruit/', models['recruit'], '/posterior_plot_', sp, '.RDS'
    )
  ) |>
  select(iter, value, plot_id) |>
  pivot_wider(names_from = plot_id, values_from = value) |>
  select(!iter) |>
  write_csv(
    file = paste0('data/parameters_randomEffects/recruit_', sp, '.csv')
  )

  cat('species', which(sp == spIds$species_id_old), 'of', nrow(spIds))
}
```

```{r Funcs,echo=Echo,eval=Eval,cache=Cache,warning=Warng,message=Msg}
# Function to import parameters given a specific draw from posterior
getPars_sp <- function(
  sp,
  vitalRate = 'growth', # or 'mort/rec'
  rowID = 1, # posterior draw (1:8000)
  randomEffects = FALSE, # if TRUE, return random effects parameters only
  addressPars = data_path,
  model_dir = models
){
  fixed_or_random = ifelse(randomEffects, 'parameters_randomEffects/', 'output_sim_processed/')

  addressPars_rf = 'data/'
  
  if(randomEffects) {
    colNames = system(
      paste0(
        'perl -nle "print && exit if $. == 1" ',
        addressPars_rf, fixed_or_random, vitalRate, '_', sp, '.csv'
      ),
      intern = TRUE
    )

    vals = system(
      paste0(
        'perl -nle "print && exit if $. == ', rowID + 1, '" ',
        addressPars_rf, fixed_or_random, vitalRate, '_', sp, '.csv'
      ),
      intern = TRUE
    )

    setNames(
        as.numeric(str_split_1(vals, pattern = ',')),
        str_split_1(colNames, pattern = ',')
      ) |>
      rbind() |>
      as_tibble() ->
    pars_fixed
  
  }else{
    readRDS(
        paste0(addressPars, fixed_or_random, vitalRate, '/', model_dir[vitalRate], '/posterior_pop_', sp, '.RDS')
      ) |>
      pivot_wider(names_from = par, values_from = value) |>
      filter(iter == rowID) |>
      select(!c(iter, lp__)) ->
    pars_fixed
  }

  return( pars_fixed )
}

# Function to generate small variation in covariate to compute derivative
cov_perturbation <- function(cov, per = 0.0001) {
  h = (abs(mean(cov, na.rm = TRUE)) + per) * per
  return( cov + h )
}

# Growth function
growth_f <- function(dbh, deltaTime, comp_cons, comp_het, temp, prec, r, plot_re, Lmax, beta, theta, opt_temp, tau_temp, opt_prec, tau_prec) {
  r_plot = exp(
    r +
    plot_re +
    beta * (comp_cons + theta * comp_het) +
    -tau_temp * (temp - opt_temp)^2 +
    -tau_prec * (prec - opt_prec)^2
  )
  y = dbh * exp(-r_plot * deltaTime) + Lmax * (1 -exp(-r_plot * deltaTime))

  return(y)
}

# Mortality function
mort_f <- function(dbh, deltaTime, comp_cons, comp_het, temp, prec, psi, plot_re, beta, theta, opt_temp, tau_temp, opt_prec, tau_prec) {
  longev_log = 1/(1 + exp(
        -(
          psi + 
          plot_re +
          beta * (comp_cons + theta * comp_het) +
          -tau_temp * (temp - opt_temp)^2 +
          -tau_prec * (prec - opt_prec)^2
        )
      )
    )
  y = 1 - (longev_log^deltaTime)
  
  return( y )
}

# Recruitment function
rec_f <- function(deltaYear_plot, plot_size, BA_adult_sp, BA_adult, temp, prec, mPop_log, plot_re, optimal_BA, sigma_BA, p_log, beta_p, opt_temp, tau_temp, opt_prec, tau_prec) {
  m = exp(
      mPop_log + plot_re +
      (-1/sigma_BA^2) * (BA_adult_sp - optimal_BA)^2 +
      -tau_temp * (temp - opt_temp)^2 +
      -tau_prec * (prec - opt_prec)^2
    )
  p = exp(
      -exp(
        p_log
      ) +
      BA_adult * -beta_p
    )
  y = m * plot_size * (1 - p^deltaYear_plot)/(1 - p)
  
  return( y )
}
```

```{r FuncMarginal,echo=Echo,eval=Eval,cache=Cache,warning=Warng,message=Msg}
# function to get the growth marginal effect for a specific species posterior draw
# this is so ungly, sorry if anyone wants to read this
marginEffect <- function(
  Sp, # species id
  rowID, # posterior draw (1:4000)
  vitalRate,
  dataset,
  randomEffects = TRUE, # if false, random effects is set to zero
  at_range = 0.1, # quantile probability to classify plot at latitude ranges (if NULL all plots)
  pertub = 0.01
){
  # focal species data
  if(vitalRate == 'recruit') {
    dataset |>
      filter(species_id == Sp) |>
      select(species_id, plot_id, latitude, deltaYear_plot, plot_size, BA_adult_sp, BA_adult, contains('_scl'), bio_01_mean, bio_12_mean) ->
    td_sp
  }else{
    dataset |>
      filter(species_id == Sp & dbh >= 127) |>
      select(species_id, plot_id, latitude, deltaTime, dbh0, dbh, BA_comp_sp, BA_comp_inter, BA_sp, BA_plot, contains('_scl'), bio_01_mean, bio_12_mean) ->
    td_sp
  }

  # deal if calculation is for all plots together or plots from latitudinal range
  if(!is.null(at_range)) {
    td_sp |>
      mutate(
        plot_position = 'center',
        plot_position = ifelse(
          latitude < quantile(latitude, probs = at_range, na.rm = TRUE),
          'lower',
          plot_position
        ),
        plot_position = ifelse(
          latitude > quantile(latitude, probs = 1 - at_range, na.rm = TRUE),
          'upper',
          plot_position
        )
      ) ->
    td_sp
  }else{
    td_sp$plot_position = 'na' # not using a real NA because I `drop_na()` latter
  }

  # get a draw from parameters posterior for focal species
  parsFixed_sp <- getPars_sp(
    Sp,
    vitalRate = vitalRate,
    rowID = rowID,
    randomEffects = FALSE
  )
  parsRandom_sp <- getPars_sp(
      Sp,
      vitalRate = vitalRate,
      rowID = rowID,
      randomEffects = TRUE
    ) |>
    pivot_longer(
      cols = everything(),
      names_to = 'plot_id',
      values_to = 'plot_re'
    )

  if(!randomEffects)
    parsRandom_sp$plot_re = 0

  # Covariate perturbation
  if(vitalRate == 'recruit') {
    td_sp |>
      # add parameters
      bind_cols(parsFixed_sp) |>
      left_join(parsRandom_sp, by = 'plot_id') |>
      drop_na() |>
        mutate(
        # perturbation in covariate
        ba_sp = scale_vars(BAsp_scl + pertub, cov = 'BAsp', 'unscale'),
        ba_tot = scale_vars(BAtot_scl + pertub, cov = 'BAtot', 'unscale'),
        temp = bio_01_mean_scl + pertub,
        prec = bio_12_mean_scl + pertub,
        # compute natural scale for perturbed climate variables
        temp_nonScl = scale_vars(temp, cov = 'temp', 'unscale'),
        prec_nonScl = scale_vars(prec, cov = 'prec', 'unscale')
      ) ->
    td_sp
  }else{
    td_sp |>
      # add parameters
      bind_cols(parsFixed_sp) |>
      left_join(parsRandom_sp, by = 'plot_id') |>
      filter(deltaTime > 0) |>
      drop_na() |>
      mutate(
        # perturbation in  covariate
        ba_con = scale_vars(BAcon_scl + pertub, cov = 'BAcon', 'unscale'),
        ba_het = scale_vars(BAhet_scl + pertub, cov = 'BAhet', 'unscale'),
        temp = bio_01_mean_scl + pertub,
        prec = bio_12_mean_scl + pertub,
        # compute natural scale for perturbed climate variables
        temp_nonScl = scale_vars(temp, cov = 'temp', 'unscale'),
        prec_nonScl = scale_vars(prec, cov = 'prec', 'unscale')
      ) ->
    td_sp
  }

  # Compute demographic model for natural and perturbed covariate and get their difference
  if(vitalRate == 'growth') {
    td_sp |>
      mutate(
        # base line growth
        y = growth_f(dbh0, deltaTime, BA_comp_sp, BA_comp_inter, bio_01_mean_scl, bio_12_mean_scl, r, plot_re, Lmax, Beta, theta, optimal_temp, tau_temp, optimal_prec, tau_prec),
        # growth for each perturbed covariate
        y_bacon = growth_f(dbh0, deltaTime, ba_con, BA_comp_inter, bio_01_mean_scl, bio_12_mean_scl, r, plot_re, Lmax, Beta, theta, optimal_temp, tau_temp, optimal_prec, tau_prec),
        y_bahet = growth_f(dbh0, deltaTime, BA_comp_sp, ba_het, bio_01_mean_scl, bio_12_mean_scl, r, plot_re, Lmax, Beta, theta, optimal_temp, tau_temp, optimal_prec, tau_prec),
        y_temp = growth_f(dbh0, deltaTime, BA_comp_sp, BA_comp_inter, temp, bio_12_mean_scl, r, plot_re, Lmax, Beta, theta, optimal_temp, tau_temp, optimal_prec, tau_prec),
        y_prec = growth_f(dbh0, deltaTime, BA_comp_sp, BA_comp_inter, bio_01_mean_scl, prec, r, plot_re, Lmax, Beta, theta, optimal_temp, tau_temp, optimal_prec, tau_prec),
        # compute absolute deviation from perturbation
        ydiff_bacon = abs(y_bacon - y)/(ba_con - BA_comp_sp),
        ydiff_bahet = abs(y_bahet - y)/(ba_het - BA_comp_inter),
        ydiff_temp = abs(y_temp - y)/(temp_nonScl - bio_01_mean),
        ydiff_prec = abs(y_prec - y)/(prec_nonScl - bio_12_mean)
      ) |>
      group_by(plot_position) |>
      reframe(across(contains('ydiff'), mean)) ->
    out
  }else if(vitalRate == 'mort') {
    td_sp |>
      mutate(
        # base line mortality
        y = mort_f(dbh0, deltaTime, BA_comp_sp, BA_comp_inter, bio_01_mean_scl, bio_12_mean_scl, psi, plot_re, Beta, theta, optimal_temp, tau_temp, optimal_prec, tau_prec),
        # growth for each perturbed covariate
        y_bacon = mort_f(dbh0, deltaTime, ba_con, BA_comp_inter, bio_01_mean_scl, bio_12_mean_scl, psi, plot_re, Beta, theta, optimal_temp, tau_temp, optimal_prec, tau_prec),
        y_bahet = mort_f(dbh0, deltaTime, BA_comp_sp, ba_het, bio_01_mean_scl, bio_12_mean_scl, psi, plot_re, Beta, theta, optimal_temp, tau_temp, optimal_prec, tau_prec),
        y_temp = mort_f(dbh0, deltaTime, BA_comp_sp, BA_comp_inter, temp, bio_12_mean_scl, psi, plot_re, Beta, theta, optimal_temp, tau_temp, optimal_prec, tau_prec),
        y_prec = mort_f(dbh0, deltaTime, BA_comp_sp, BA_comp_inter, bio_01_mean_scl, prec, psi, plot_re, Beta, theta, optimal_temp, tau_temp, optimal_prec, tau_prec),
        # compute absolute deviation from perturbation
        ydiff_bacon = abs(y_bacon - y)/(ba_con - BA_comp_sp),
        ydiff_bahet = abs(y_bahet - y)/(ba_het - BA_comp_inter),
        ydiff_temp = abs(y_temp - y)/(temp_nonScl - bio_01_mean),
        ydiff_prec = abs(y_prec - y)/(prec_nonScl - bio_12_mean)
      ) |>
      group_by(plot_position) |>
      reframe(across(contains('ydiff'), mean)) ->
    out
  }else if(vitalRate == 'recruit') {
    td_sp |>
      mutate(
        # base line recruitment
        y = rec_f(deltaYear_plot, plot_size, BA_adult_sp, BA_adult, bio_01_mean_scl, bio_12_mean_scl, mPop_log, plot_re, optimal_BA, sigma_BA, p_log, beta_p, optimal_temp, tau_temp, optimal_prec, tau_prec),
        # recruitment for each perturbed covariate
        y_basp = rec_f(deltaYear_plot, plot_size, ba_sp, BA_adult, bio_01_mean_scl, bio_12_mean_scl, mPop_log, plot_re, optimal_BA, sigma_BA, p_log, beta_p, optimal_temp, tau_temp, optimal_prec, tau_prec),
        y_batot = rec_f(deltaYear_plot, plot_size, BA_adult_sp, ba_tot, bio_01_mean_scl, bio_12_mean_scl, mPop_log, plot_re, optimal_BA, sigma_BA, p_log, beta_p, optimal_temp, tau_temp, optimal_prec, tau_prec),
        y_temp = rec_f(deltaYear_plot, plot_size, BA_adult_sp, BA_adult, temp, bio_12_mean_scl, mPop_log, plot_re, optimal_BA, sigma_BA, p_log, beta_p, optimal_temp, tau_temp, optimal_prec, tau_prec),
        y_prec = rec_f(deltaYear_plot, plot_size, BA_adult_sp, BA_adult, bio_01_mean_scl, prec, mPop_log, plot_re, optimal_BA, sigma_BA, p_log, beta_p, optimal_temp, tau_temp, optimal_prec, tau_prec),
        # compute absolute deviation from perturbation
        ydiff_basp = abs(y_basp - y)/(ba_sp - BA_adult_sp),
        ydiff_batot = abs(y_batot - y)/(ba_tot - BA_adult),
        ydiff_temp = abs(y_temp - y)/(temp_nonScl - bio_01_mean),
        ydiff_prec = abs(y_prec - y)/(prec_nonScl - bio_12_mean)
      ) |>
      group_by(plot_position) |>
      reframe(across(contains('ydiff'), mean)) ->
    out
  }
}
```

```{r computeAME,echo=Echo,eval=FALSE,cache=Cache,warning=Warng,message=Msg}

# prepare dataset with transition info
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

## compute de deltaTime between measures of dbh
treeData[,
  deltaTime := year_measured - lag(year_measured, 1),
  by = tree_id
]

## define previous measure
treeData[,
  dbh0 := lag(dbh, 1),
  by = tree_id
]

## Fill the NA first measures with their non lag information
treeData[is.na(deltaTime), deltaTime := 0]
treeData[deltaTime == 0, dbh0 := dbh]


# scale variables so perturbation is at the same scale for all covariates
cov_ranges <- append(
  list(
    'BAcon' = treeData[, range(BA_comp_sp, na.rm = TRUE)],
    'BAhet' = treeData[, range(BA_comp_inter, na.rm = TRUE)],
    'BAsp' = fec_dt[, range(BA_adult_sp, na.rm = TRUE)],
    'BAtot' = fec_dt[, range(BA_adult, na.rm = TRUE)]
  ),
  setNames(
    climate_range,
    c('temp', 'prec')
  )
)

# functions to scale OR unscale climate variables
scale_vars <- function(
  value, # covariate value
  cov, # either 'BAcon', 'BAhet', 'BAsp', 'BAtot', 'temp', or 'prec'
  direction, # either 'scale' or 'unscale'
  range_dt = cov_ranges
){
  if(cov %in% c('BAcon', 'BAhet', 'BAsp', 'BAtot', 'temp', 'prec')){
    cov_rg = range_dt[[cov]]
  }else{
    stop("`cov` must be one of 'BAcon', 'BAhet', 'BAsp', 'BAtot', 'temp', or 'prec' character.")
  }

  min_v = cov_rg[1]
  max_v = cov_rg[2]

  if(direction == 'scale') {
    return( (value - min_v)/(max_v - min_v) )
  }else if(direction == 'unscale') {
    return( value * (max_v - min_v) + min_v )
  }else {
    stop('`direction` argument must be either `scale` or `unscale`.')
  }
}

treeData |>
  mutate(
    BAcon_scl = scale_vars(BA_comp_sp, 'BAcon', direction = 'scale'),
    BAhet_scl = scale_vars(BA_comp_inter, 'BAhet', direction = 'scale')
  ) ->
treeData

fec_dt |>
  mutate(
    BAsp_scl = scale_vars(BA_adult_sp, 'BAsp', direction = 'scale'),
    BAtot_scl = scale_vars(BA_adult, 'BAtot', direction = 'scale')
  ) ->
fec_dt


# Compute marginal effect growth
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

sampled_rows <- sample(1:4000, 500)

# Growth
AME_growth <- list()
for(sp in spIds$species_id_old) {
    AME_growth[[sp]] <- map_dfr(
      sampled_rows,
      ~marginEffect(
        Sp = sp,
        rowID = .x,
        dataset = treeData,
        vitalRate = 'growth',
        randomEffects = TRUE,
        at_range = NULL,
        pertub = 0.01
      )
    )
  cat(' Species', which(sp == spIds$species_id_old), 'of 31\r')
}

AME_growth |>
  bind_rows(.id = 'species_id') |>
  select(!plot_position) ->
AME_growth

# map_dfr(
#   spIds$species_id_old,
#   ~map2_dfr(
#     sampled_rows,
#     .x,
#     ~marginEffect(Sp = .y, rowID = .x, treeData, vitalRate = 'growth', randomEffects = TRUE, at_range = 0.1) |>
#       bind_cols(species_id = .y)
#   ),
#   .progress = TRUE
#  ) ->
# AME_growth_rg

# Survival
AME_mort <- list()
for(sp in spIds$species_id_old) {
    AME_mort[[sp]] <- map_dfr(
      sampled_rows,
      ~marginEffect(
        Sp = sp,
        rowID = .x,
        dataset = treeData,
        vitalRate = 'mort',
        randomEffects = TRUE,
        at_range = NULL,
        pertub = 0.01
      )
    )
  cat(' Species', which(sp == spIds$species_id_old), 'of 31\r')
}

AME_mort |>
  bind_rows(.id = 'species_id') |>
  select(!plot_position) ->
AME_mort

# map_dfr(
#   spIds$species_id_old,
#   ~map2_dfr(
#     sampled_rows,
#     .x,
#     ~marginEffect(Sp = .y, rowID = .x, treeData, vitalRate = 'mort', randomEffects = TRUE, at_range = 0.1) |>
#       bind_cols(species_id = .y)
#   ),
#   .progress = TRUE
#  ) ->
# AME_mort_rg

# Recruitment
AME_rec <- list()
for(sp in spIds$species_id_old) {
    AME_rec[[sp]] <- map_dfr(
      sampled_rows,
      ~marginEffect(
        Sp = sp,
        rowID = .x,
        dataset = fec_dt,
        vitalRate = 'recruit',
        randomEffects = TRUE,
        at_range = NULL,
        pertub = 0.01
      )
    )
  cat(' Species', which(sp == spIds$species_id_old), 'of 31\r')
}

AME_rec |>
  bind_rows(.id = 'species_id') |>
  select(!plot_position) ->
AME_rec

# map_dfr(
#   spIds$species_id_old,
#   ~map2_dfr(
#     sampled_rows,
#     .x,
#     ~marginEffect(Sp = .y, rowID = .x, fec_dt, vitalRate = 'recruit', randomEffects = TRUE, at_range = 0.1) |>
#       bind_cols(species_id = .y)
#   ),
#   .progress = TRUE
#  ) ->
# AME_rec_rg

# save computed AME
dir.create('data')
saveRDS(AME_growth, file.path('data', 'AME_growth.RDS'))
# saveRDS(AME_growth_rg, file.path('data', 'AME_growth_rg.RDS'))
saveRDS(AME_mort, file.path('data', 'AME_mort.RDS'))
# saveRDS(AME_mort_rg, file.path('data', 'AME_mort_rg.RDS'))
saveRDS(AME_rec, file.path('data', 'AME_rec.RDS'))
# saveRDS(AME_rec_rg, file.path('data', 'AME_rec_rg.RDS'))
```


### AME per species

```{r plot_AME_gmr,echo=Echo,eval=Eval,cache=Cache,warning=Warng,message=Msg}
#| label: fig-AMEall
#| fig-width: 9
#| fig-height: 8
#| fig-cap: "Average Marginal effect of each covariate on the growth and survival vital rates. The larger the values, the stronger the marginal effect of the covariate across all observation for the species. Note that size (d.b.h.) is intrinsic to the growth equation model where annual growth rate reduces exponentialy with size reaching a growth rate of zero when size reaches $\\xi_{\\infty}$, so the size AME is exactly ln(1) for all species."

AME_growth <- readRDS(file.path('data', 'AME_growth.RDS'))
AME_mort <- readRDS(file.path('data', 'AME_mort.RDS'))
AME_rec <- readRDS(file.path('data', 'AME_rec.RDS'))

# AME per species per covariate (growth and survival)
AME_growth |>
  pivot_longer(cols = !species_id) |>
  bind_cols(vr = 'Growth') |>
  bind_rows(
    AME_mort |>
      pivot_longer(cols = !species_id) |>
      bind_cols(vr = 'Survival') 
  ) |>
  bind_rows(
    AME_rec |>
      pivot_longer(cols = !species_id) |>
      bind_cols(vr = 'Recruitment') 
  ) |>
  left_join(
    spIds,
    by = c('species_id'= 'species_id_old')
  ) |>
  mutate(
    name = gsub('ydiff_', '', name),
    name = case_match(
      name,
      'bacon' ~ 'BA cons',
      'bahet' ~ 'BA het',
      'basp' ~ 'BA cons',
      'batot' ~ 'BA total',
      'temp' ~ 'Temperature',
      'prec' ~ 'Precipitation'
    )
  ) |>
  ggplot() +
  aes(log(value), species_name) +
  aes(color = name) +
  facet_grid(~vr, scales = 'free') +
  stat_pointinterval() +
  scale_color_manual(
    values = c('#7fc97f', '#beaed4', '#ffff99', '#386cb0', '#fdc086')
  ) +
  theme_classic() +
  theme(
    legend.position = 'top',
    axis.text.y = element_text(face = "italic")
  ) +
  labs(
    color = 'Covariate',
    x = 'ln(Average Marginal effect)',
    y = ''
  )

# AME_rec |>
#   pivot_longer(cols = !species_id) |>
#   left_join(
#     spIds,
#     by = c('species_id'= 'species_id_old')
#   ) |>
#   mutate(
#     name = gsub('ydiff_', '', name),
#     name = case_match(
#       name,
#       'basp' ~ 'BA cons',
#       'batot' ~ 'BA total'
#     )
#   ) |>
#   filter(name == 'BA total') |>
#   ggplot() +
#   aes(shade_sylvics, value) +
#   aes(fill = shade_sylvics) +
#   geom_boxplot() +
#   labs(
#     x = 'Shade tolerance',
#     y = 'Marginal effect of total basal area on recruitment survival'
#   ) +
#   scale_fill_manual(
#     values = c("#20bc45", "#87bc45", "#edbf33", "#ea5545", "#ba0000")
#   ) +
#   theme_classic() +
#   theme(legend.position = 'none')
```


### Summary of AME across species

```{r plot_AME_summ,echo=Echo,eval=Eval,cache=Cache,warning=Warng,message=Msg}
#| label: fig-AMEsumm
#| fig-width: 9
#| fig-height: 5
#| fig-cap: "Average Marginal effect of each covariate grouped by species and their respective uncertainty. The larger the values, the stronger the marginal effect of the covariate across all observations for the species."

ggarrange(
  AME_growth |>
    group_by(species_id) |>
    reframe(across(contains('ydiff'), mean)) |>
    pivot_longer(cols = !species_id) |>
    left_join(
      spIds,
      by = c('species_id'= 'species_id_old')
    ) |>
    mutate(
      name = gsub('ydiff_', '', name),
      name = case_match(
        name,
        'bacon' ~ 'BA cons',
        'bahet' ~ 'BA het',
        'size' ~ 'd.b.h',
        'temp' ~ 'Temperature',
        'prec' ~ 'Precipitation'
      )
    ) |>
    ggplot() +
    aes(x = fct_reorder(name, log(value), .desc = TRUE), y = log(value)) +
    aes(fill = name) +
    geom_boxplot() +
    scale_fill_manual(
      values = c('#7fc97f', '#ffff99', '#386cb0', '#fdc086')
    ) +
    theme_classic() +
    labs(
      x = '',
      y = 'ln(Average Marginal Effect)',
      subtitle = 'Growth'
    ) +
    theme(legend.position = 'none'),
  AME_mort |>
    group_by(species_id) |>
    reframe(across(contains('ydiff'), mean)) |>
    pivot_longer(cols = !species_id) |>
    left_join(
      spIds,
      by = c('species_id'= 'species_id_old')
    ) |>
    mutate(
      name = gsub('ydiff_', '', name),
      name = case_match(
        name,
        'bacon' ~ 'BA cons',
        'bahet' ~ 'BA het',
        'temp' ~ 'Temperature',
        'prec' ~ 'Precipitation'
      )
    ) |>
    ggplot() +
    aes(x = fct_reorder(name, log(value), .desc = TRUE), y = log(value)) +
    aes(fill = name) +
    geom_boxplot() +
    scale_fill_manual(
      values = c('#7fc97f', '#ffff99', '#386cb0', '#fdc086')
    ) +
    theme_classic() +
    labs(
      y = '',
      x = '',
      subtitle = 'Survival'
    ) +
    theme(legend.position = 'none'),
  AME_rec |>
    group_by(species_id) |>
    reframe(across(contains('ydiff'), mean)) |>
    pivot_longer(cols = !species_id) |>
    left_join(
      spIds,
      by = c('species_id'= 'species_id_old')
    ) |>
    mutate(
      name = gsub('ydiff_', '', name),
      name = case_match(
        name,
        'basp' ~ 'BA cons',
        'batot' ~ 'BA total',
        'temp' ~ 'Temperature',
        'prec' ~ 'Precipitation'
      )
    ) |>
    ggplot() +
    aes(x = fct_reorder(name, log(abs(value)), .desc = TRUE), y = log(abs(value))) +
    aes(fill = name) +
    geom_boxplot() +
    scale_fill_manual(
      values = c('#7fc97f', '#ffff99', '#386cb0', '#fdc086')
    ) +
    theme_classic() +
    labs(
      x = '',
      y = '',
      subtitle = 'Recruitment'
    ) +
    theme(legend.position = 'none'),
  ncol = 3
)
```


### AME of competition covariates grouped by shade tolerance

```{r plot_AME_compShade,echo=Echo,eval=Eval,cache=Cache,warning=Warng,message=Msg}
#| label: fig-AMEsummComp
#| fig-width: 8.5
#| fig-height: 5
#| fig-cap: "Average Marginal effect of the competition covariates grouped by species shade tolerance and their respective uncertainty. The larger the values, the stronger the marginal effect of the covariate across all observations for the species."

AME_growth |>
  group_by(species_id) |>
  reframe(across(contains('ydiff'), mean)) |>
  pivot_longer(cols = !species_id) |>
  bind_cols(vr = 'Growth') |>
  bind_rows(
    AME_mort |>
      group_by(species_id) |>
      reframe(across(contains('ydiff'), mean)) |>
      pivot_longer(cols = !species_id) |>
      bind_cols(vr = 'Survival') 
  ) |>
  bind_rows(
    AME_rec |>
      group_by(species_id) |>
      reframe(across(contains('ydiff'), mean)) |>
      pivot_longer(cols = !species_id) |>
      bind_cols(vr = 'Recruitment') 
  ) |>
  left_join(
    spIds,
    by = c('species_id'= 'species_id_old')
  ) |>
  mutate(
    name = gsub('ydiff_', '', name)
  ) |>
  filter(name %in% c('bacon', 'bahet', 'basp', 'batot')) |>
  mutate(
    name = case_match(
      name,
      'bacon' ~ 'BA conspecific',
      'bahet' ~ 'BA heterospecific',
      'basp' ~ 'BA conspecific',
      'batot' ~ 'BA total'
    )
  ) |>
  mutate(
    vr = factor(vr, levels = c('Growth', 'Survival', 'Recruitment')),
    value = log(abs(value))
  ) |>
  ggplot() +
  aes(fct_reorder(name, value, .desc = TRUE), value) +
  facet_wrap(~vr, scales = 'free') +
  aes(fill = shade_sylvics) +
  geom_boxplot() +
  scale_fill_manual(
    values = c("#20bc45", "#87bc45", "#edbf33", "#ea5545", "#ba0000")
  ) +
  theme_classic() +
  labs(
    x = '',
    y = 'ln(Average Marginal effect)',
    fill = ''
  ) +
  theme(
    strip.background = element_blank(),
    strip.text = element_text(hjust = 0, vjust = 1.1, margin = margin(l=0), size = rel(1.1)),
    legend.position = 'bottom'
  )
```


### AME of temperature given the median latitude distribution of the species

```{r plot_AME_growthVsLatitude,echo=Echo,eval=Eval,cache=Cache,warning=Warng,message=Msg}
#| label: fig-AMEsummTemp
#| fig-width: 9
#| fig-height: 5
#| fig-cap: "Average Marginal effect of Temperature in the function of the average latitudinal position of the species. The blue line is 50% quantile and the gray lines are the 25% and 75% quantiles."

AME_growth |>
  select(species_id, ydiff_temp) |>
  bind_cols(vr = 'Growth') |>
  bind_rows(
    AME_mort |>
      select(species_id, ydiff_temp) |>
      bind_cols(vr = 'Survival')
  ) |>
  bind_rows(
    AME_rec |>
      select(species_id, ydiff_temp) |>
      bind_cols(vr = 'Recruitment')
  ) |>
  left_join(
    spIds,
    by = c('species_id'= 'species_id_old')
  ) |>
  left_join(
    treeData |>
      group_by(species_id) |>
      reframe(
        range_pos = median(latitude, na.rm = TRUE)
      )
  ) |>
  group_by(species_name, vr) |>
  reframe(
    range_pos = mean(range_pos),
    value = mean(ydiff_temp)
  ) ->
mean_rangeValue

AME_growth |>
  pivot_longer(cols = !species_id) |>
  bind_cols(vr = 'Growth') |>
  bind_rows(
    AME_mort |>
    pivot_longer(cols = !species_id) |>
    bind_cols(vr = 'Survival') 
  ) |>
  bind_rows(
    AME_rec |>
    pivot_longer(cols = !species_id) |>
    bind_cols(vr = 'Recruitment') 
  ) |>
  left_join(
    spIds,
    by = c('species_id'= 'species_id_old')
  ) |>
  left_join(
    treeData |>
      group_by(species_id) |>
      reframe(
        range_pos = median(latitude, na.rm = TRUE)
      )
  ) |>
  mutate(
    name = gsub('ydiff_', '', name)
  ) |>
  filter(name %in% c('temp')) |>
  ggplot() +
  aes(range_pos, log(value)) +
  stat_pointinterval() +
  facet_wrap(~vr) +
  geom_text_repel(
    data = mean_rangeValue,
    aes(x = range_pos, y = log(value), label = species_name),
    alpha = 0.8,
    size = 2,
    fontface = 'italic'
  ) +
  geom_quantile(quantiles = 0.5, color = 'blue') +
  geom_quantile(quantiles = c(0.25, 0.75), color = 'black', alpha = 0.5) +
  theme_classic() +
  labs(
    x = 'Median latitude distribution',
    y = 'ln(Average Marginal effect)'
  )
```

```{r sim_int_effect,eval=FALSE,echo=FALSE}
# code to test the AME when slope is the same but the intercept is different
md <- function(x, beta, int)
  return( exp(log(int) + beta * x) )

int_a = 3
int_b = 2.1
beta = -1
x_i = seq(0, 6, 0.1)

plot(
  x_i,
  md(x_i,beta, int_a),
  type = 'l'
);abline(h = 0, col = 'grey', lty = 2);abline(h = int_a, col = 1, lty = 3)
points(
  x_i,
  md(x_i,beta, int_b),
  col = 2,
  type = 'l'
);abline(h = int_b, col = 2, lty = 3)

# sample individuals
ind_x <- rlnorm(100, log(0.8), log(2.5))
segments( x0 = ind_x, y0 = 0, x1 = ind_x, y1 = -1, col = rgb(0, 0, 0, 0.5))

# size of derivative
h = (mean(ind_x) + 0.0001) * 0.0001

# predict for both models (A and B) for the observed x predictor
y_A = md(ind_x, beta, int_a)
y_B = md(ind_x, beta, int_b)

plot(sort(ind_x), y_A[order(ind_x)], type = 'o');points(sort(ind_x), y_B[order(ind_x)], type = 'o', col = 2)

# Now predict after small perturbation
y_Ad = md(ind_x + h, beta, int_a)
y_Bd = md(ind_x + h, beta, int_b)

# mean[dy_dx] or Absolute Average Marginal Effect
mean(abs(y_A - y_Ad)/h)
mean(abs(y_B - y_Bd)/h)

# samething but controled for the intercept
mean(abs(y_A - y_Ad)/h/int_a)
mean(abs(y_B - y_Bd)/h/int_b)

```