---
title: "Loaloa prevalence with FEM Matérn"
author: "Patrick Brown"
date: "`r Sys.Date()`"
output:
  rmarkdown::html_vignette:
    toc: true
vignette: >
  %\VignetteIndexEntry{Loaloa prevalence with FEM Matérn}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

```{r setup, include=FALSE}
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>",
  fig.width = 7,
  fig.height = 5
)
has_geostatsp <- requireNamespace("geostatsp", quietly = TRUE) &&
  requireNamespace("terra", quietly = TRUE)
```

```{r pkgs}
library(adlaplace)
library(adlaplaceGrf)
```

This vignette fits a binomial geostatistical model for loaloa prevalence
using a village-level IID effect and a FEM Matérn spatial field
(`shape = 2`, ν = 2), matching the Matérn shape used in geostatsp
examples. When **INLA** is available, the same model is also fit with
`geostatsp::glgm` for comparison.

```{r skip-check, echo=FALSE}
if (!has_geostatsp) {
  knitr::knit_exit()
}
```

## Data

```{r data}
data("loaloa", package = "geostatsp")

loaloa <- terra::unwrap(loaloa)
elev_r <- terra::unwrap(elevationLoa)
evi_r <- terra::unwrap(eviLoa)

elev_s_r <- (elev_r - 500) / 100
evi_s_r <- (
  evi_r - mean(terra::values(evi_r), na.rm = TRUE)
) / sd(terra::values(evi_r), na.rm = TRUE)

loaloa_for_fit <- geostatsp:::gm.dataSpatial(
  formula = y ~ elev + evi,
  data = loaloa,
  covariates = list(elev = elev_s_r, evi = evi_s_r),
  grid = geostatsp::squareRaster(loaloa, cells = 200, buffer = 1e4)
)
```

## Knot lines

Use a coarse knot raster over a buffered region around the villages; pass it
directly as `knots` to `matern()` (extent endpoints plus interior cell centers; see
`knots_from_spatraster()`).
(`geostatsp` / `terra` may mask `matern`, so call it with `::`.)



## Fit

geostatsp Matérn **shape = 2** is ν = 2 (SPDE α = 3 in 2D), so we pass
`shape = 2` to `matern()` (cubic B-splines / `random_fem_ssq_3`). Locations come
from the `geometry` column (WKT) in the model data frame.

```{r model_setup, cache=FALSE}
knots_grid <- geostatsp::squareRaster(loaloa, cells = 6, buffer = 300 * 1000)
loaloa_df <- as.data.frame(loaloa_for_fit$data, geom = "WKT")

loaloa_model <- adlaplace::binomial(y, size = N) ~
  elev + evi +
  adlaplace::iid(villageID, init = 0.2) +
  adlaplaceGrf::matern(
    geometry,
    knots = knots_grid,
    shape = 2L,
    # range = practical rho = sqrt(8*nu)/kappa (m); sd = field SD
    init = c(100 * 1000, 0.2),
    lower = c(1 * 1000, 0.001)
  )
```



```{r fit}
fit <- adlaplace::adlaplace(
  loaloa_model,
  data = loaloa_df,
  num_shards = 20L,
  control = list(maxit = 200, trace = 3, REPORT = 1),
  config = list(num_threads = 2L),
  verbose = FALSE
)
fit$optim$convergence
coef(fit)
```

```{r fit-glgm}
loaFit <- geostatsp::glgm(
  y ~ elev_s + evi_s +
    f(
      villageID,
      prior = "pc.prec",
      param = c(1, 0.5),
      model = "iid"
    ),
  loaloa,
  grid = 100,
  covariates = list(elev_s = elev_s_r, evi_s = evi_s_r),
  family = "binomial",
  Ntrials = loaloa$N,
  shape = 2,
  buffer = 25000,
  prior = list(sd = 1, range = 500 * 1000),
  control.inla = list(strategy = "gaussian")
)
```

## Summary

```{r summary}
logLik(fit)

cf <- coef(fit)
ad_pars <- c(
  intercept = unname(cf[["intercept"]]),
  elev = unname(cf[["elev_linear"]]),
  evi = unname(cf[["evi_linear"]]),
  village_sd = unname(cf[["villageID_iid"]]),
  # matern: practical range rho = sqrt(8*nu)/kappa (m); sd = field SD
  range = unname(cf[["geometry_matern_log_range"]]),
  sd = unname(cf[["geometry_matern_log_sd"]])
)
ad_pars
```

```{r summary-compare}
if (length(loaFit$parameters)) {
  glgm_sum <- loaFit$parameters$summary
  knitr::kable(glgm_sum[, c("mean", "0.025quant", "0.975quant")], digits = 3)

  # Map ad_pars names to rows of loaFit$parameters$summary
  row_for <- function(...) {
    pats <- c(...)
    hits <- unique(unlist(lapply(pats, function(p) {
      grep(p, rownames(glgm_sum), ignore.case = TRUE)
    })))
    if (!length(hits)) {
      return(NA_character_)
    }
    rownames(glgm_sum)[hits[1L]]
  }
  map <- c(
    intercept = row_for("^\\(Intercept\\)$", "^Intercept$"),
    elev = row_for("^elev"),
    evi = row_for("^evi"),
    village_sd = row_for("^sd villageID$", "sd.?villageID", "villageID.*[Ss]d"),
    range = row_for("^range/1000$", "^range$"),
    sd = row_for("^sd$", "SD for space")
  )

  glgm_pars <- setNames(
    vapply(map, function(nm) {
      if (is.na(nm)) {
        return(NA_real_)
      }
      as.numeric(glgm_sum[nm, "mean"])
    }, numeric(1)),
    names(map)
  )
  # summary may store range in km as "range/1000"
  if (!is.na(map[["range"]]) && identical(map[["range"]], "range/1000")) {
    glgm_pars[["range"]] <- glgm_pars[["range"]] * 1000
  }

  compare <- rbind(
    adlaplace = ad_pars[names(map)],
    glgm = glgm_pars
  )
  knitr::kable(compare, digits = 3)
}
```

## Estimated field

```{r est}
my_sim <- matern_est(fit, loaloa_for_fit$grid, n = 10)
```

Posterior mean and SD of the spatial field from both fits (glgm when available):

```{r plot-surfaces-setup, include=FALSE}
my_col <- mapmisc::colourScale(
  my_sim[[grep("sd", names(my_sim), invert = TRUE)]],
  style = "equal", digits = 2, breaks = 9, col = "Spectral"
)
sd_col <- mapmisc::colourScale(
  my_sim[[grep("sd", names(my_sim))]],
  style = "equal", digits = 2, breaks = 9
)
plot_compare <- exists("loaFit") && length(loaFit$parameters)
plot_alone <- !plot_compare
```

```{r plot-surfaces-compare, fig.width=8, fig.height=8, fig.cap="Posterior mean and SD of the spatial field (glgm vs adlaplace)."}
data("worldMap", package = "mapmisc")
worldMap <- terra::unwrap(worldMap)
worldMap <- terra::crop(worldMap, terra::ext(c(-2, 10, -10, 5) * 1e6))
worldMapT <- terra::project(worldMap, loaloa)
mapmisc::map.new(loaloa)
if (plot_compare) {
  terra::plot(
    loaFit$raster[["random.mean"]],
    add = TRUE, col = my_col$col, breaks = my_col$breaks, legend = FALSE
  )
}
terra::plot(worldMapT, add = TRUE)
mapmisc::legendBreaks("right", my_col)

mapmisc::map.new(loaloa)
terra::plot(
  my_sim[["mean"]],
  add = TRUE, col = my_col$col, breaks = my_col$breaks, legend = FALSE
)
terra::plot(worldMapT, add = TRUE)
mapmisc::legendBreaks("right", my_col)

mapmisc::map.new(loaloa)
if (plot_compare) {
  terra::plot(
    loaFit$raster[["random.sd"]],
    add = TRUE, col = sd_col$col, breaks = sd_col$breaks, legend = FALSE
  )
}
terra::plot(worldMapT, add = TRUE)
mapmisc::legendBreaks("right", sd_col)

mapmisc::map.new(loaloa)
terra::plot(
  my_sim[["sd"]],
  add = TRUE, col = sd_col$col, breaks = sd_col$breaks, legend = FALSE
)
terra::plot(worldMapT, add = TRUE)
mapmisc::legendBreaks("right", sd_col)
```




Four posterior draws of the adlaplace Matérn field:

```{r plot-sims, fig.width=8, fig.height=8}
sim_layers <- paste0("sim", 1:6)
terra::plot(my_sim[[sim_layers]], nc = 2, col = my_col$col, breaks = my_col$breaks, legend = FALSE)
mapmisc::legendBreaks("right", my_col, bty = "n")
```