Formula interface and tidybayes support

[jasonelaw]

I implemented a front end function for spatialGEV_fit that is more R idiomatic. It accepts a multi-part formula like the ones in the Formula package, a square dataset, and refactors many arguments into a control and prior argument. I've also implemented a basic function to initialize the parameters. This uses the evd package since it was already imported. You can use it on the simulatedData data set:

library(SpatialGEV)
library(tidyverse)
library(tidybayes)

nlocs <- nrow(simulatedData$locs)
data <- as_tibble(simulatedData$locs) |>
  mutate(
    id = factor(1:nlocs),
    x1 = rnorm(nlocs),
    x2 = rlnorm(nlocs),
    x3 = rlnorm(nlocs),
    z = simulatedData$y
  )
data <- tidyr::unchop(data, z)

You can call the spatial_gev function like this:

f_reg <- spatial_gev(
  formula = z ~ (1 | id) | (1 | id), 
  data = data, 
  coordinates = ~ x + y, 
  reparam_s = "unconstrained", 
  kernel = "matern", 
  control = list(silent = TRUE),
  prior = list(nu = 3/2, s_prior = c(-0.25, 0.25))
)

The formula can be a 3 part formula separated by | where the parts correspond to location, scale, and shape models (y ~ location | scale | shape). Each of the subparts must be a lme4 style random effect term: (x1 + log(x2) | location_id) where the grouping factor corresponds to an identifier for the observed locations. Thus, the model I fit above corresponds to a random = "ab" with X_a and X_b design matrices with an intercept only. You could specify covariates like z ~ (x1 | id) | (x2 | id). In addition, I've supported the Formula package dot shorthand which means you could specify a random = "abs" model with the same regression terms like this y ~ (x1 | id) | . | . and the formula will be expanded to accomodate the b and s terms.

I also implemented a posterior::as_draws_df method which allows us to use spatialGEVsam objects with the tidybayes package.

draws <- spatialGEV_sample(f_reg, n_draw = 100)
draws_par <- draws |>
  recover_types(tbl) |>
  spread_draws(a[id], b[id])

draws_par <- draws_par |> 
  median_qi(a, b) |>
  left_join(tbl |> distinct(id, x, y), by = join_by(id))

draws_par |> 
  ggplot(aes(x, y)) +
  geom_raster(aes(fill = b))

I'm sure the code is somewhat brittle and could use some checks on the sanity of the data.frame that is passed. The control and prior implementations are pretty straightforward and are mostly to simplify the signature of the function and align with idiomatic R conventions. I generally tried to stick to the tidymodels recommendations: https://tidymodels.github.io/model-implementation-principles/function-interfaces.html.

[meixichen]

Wow thanks a lot for this. This is really interesting. Will take a look at this!