---
title: "Case-crossover models with hpolcc"
author: "Patrick Brown"
date: "`r Sys.Date()`"
output:
  rmarkdown::html_vignette:
    toc: true
vignette: >
  %\VignetteIndexEntry{Case-crossover models with hpolcc}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

## Introduction

This vignette fits a time-stratified case-crossover model to daily mortality and
PM10 in London (2002–2006), using the example dataset from
[Tobias et al. (2024)](https://github.com/aureliotobias/casecrossover).

Strata are defined by day-of-week within month and year via
`dirichlet_multinom(..., by = "year_month_dow")`. PM10 enters through an
integrated Wiener process smooth (`iwp(pm10, ...)`), with relative humidity as
a linear fixed effect.

```{r setup, include=FALSE}
# adlaplace must load before hpolcc/data.table so CppAD OpenMP setup runs first
library(adlaplace)
library(hpolcc)
library(data.table)
```

## Data

```{r data}
london <- fread(system.file("extdata", "london.csv", package = "hpolcc"))
london$date <- as.Date(ISOdate(london$year, london$month, london$day))
london$year_month_dow <- format(london$date, "%Y-%m-%a")
head(london)
quantile(london$pm10, na.rm = TRUE)
quantile(london$tmean, na.rm = TRUE)
```

## Model specification

```{r model-spec}
london$temp_div10 <- london$tmean / 10
cc_formula <- dirichlet_multinom(
    all,
    by = "year_month_dow",
    init = 0.2,
    parscale = 5
) ~ rh + iwp(
    pm10,
    p = 2,
    knots = seq(0, 100, by = 10),
    ref_value = 40,
    init = 0.001, parscale = 100
) + iwp(
    tmean,
    p = 3, knots = seq(-2, 30, by = 2),
    init = 0.01, parscale = 100,
    ref_value = 20
)
```


## Fit

```{r fit}
fit <- hnlm(
    formula = cc_formula,
    data = london,
    config = list(
        transform_theta = TRUE,
        num_threads = 2L,
        num_shards = 50L,
        verbose = FALSE
    ),
    control = list(
        maxit = 200, trace = 0, REPORT = 1,
        factr = 1e3,
        pgtol = 1e-8
    )
)
```

```{r summary}
fit
summary(fit)
```


```{r plot_results}
par(mfrow = c(1, 2), mar = c(2.25, 2.25, 0, 0), bty = "l", mgp = c(1, 0.5, 0))
for (D in names(fit$sample$sim)) {
    the_xlim <- range(london[[D]], na.rm = TRUE)
    sample_in_xlim <- fit$sample$sim[[D]][
        fit$sample$x[[D]] > min(the_xlim) &
            fit$sample$x[[D]] < max(the_xlim),
    ]
    matplot(fit$sample$x[[D]], fit$sample$sim[[D]],
        type = "l", lty = 1, col = "#00000010",
        xlab = D, ylab = "RR", xaxs = "i", yaxs = "i",
        xlim = the_xlim,
        ylim = quantile(sample_in_xlim, prob = c(0.005, 0.99))
    )
    matlines(fit$sample$x[[D]], fit$sample$envelope$common[[D]], lty = c(2, 1, 2), col = "red")
}
```


## AD development check

`for_dev = TRUE` builds `model_data`, the combined `ad_fun`, and `config` without
running outer optimization.

```{r for_dev}
config_dev <- list(
    transform_theta = TRUE,
    num_threads = 2L,
    num_shards = 50L,
    verbose = FALSE
)
forres <- hnlm(
    cc_formula,
    data = london,
    config = config_dev,
    for_dev = TRUE
)
slotNames(forres$ad_fun)
dim(forres$config$shards)
forres$model_data$data$info$parameters
bob <- stats::optim(
    par = forres$config$opt$init,
    fn = adlaplace::outer_fn,
    gr = adlaplace::outer_gr,
    method = "L-BFGS-B",
    control = forres$control,
    lower = forres$config$opt$lower,
    upper = forres$config$opt$upper,
    config = forres$config,
    ad_fun = forres$ad_fun,
    cache = forres$cache,
    control_inner = forres$control_inner,
    hessian = TRUE
)
```

## Joint log density

At starting values, evaluate the full log-density and its gradient over all
parameters (`beta`, inner `gamma`, and `theta`).

```{r joint-log-dens}
sz <- forres$ad_fun@sizes
x_full <- c(
    forres$config$opt$init[seq_len(sz["beta"])],
    forres$cache$gamma,
    forres$config$opt$init[seq(sz["beta"] + 1L, length.out = sz["theta"])]
)

shards <- seq.int(from = 0L, length.out = adlaplace:::n_groups(forres$ad_fun@ptr))
by_shard <- vapply(
    shards,
    function(s) joint_log_dens(forres$ad_fun, x_full, shards = s),
    numeric(1)
)
by_shard[1:5]
by_shard[seq(to = length(by_shard), length.out = 5)]

sum(by_shard)
joint_log_dens(forres$ad_fun, x_full)
grad(forres$ad_fun, x_full)
bob <- hessian(forres$ad_fun, x_full)
```

## Likelihood

With outer parameters fixed at their starting values, `log_lik_laplace()`
optimizes random effects and returns the  log-likelihood.

```{r likelihood}
x_outer <- forres$model_data$data$info$parameters$init

lik <- log_lik_laplace(
    x = x_outer,
    ad_fun = forres$ad_fun,
    config = modifyList(forres$config, list(verbose = TRUE)),
    gamma = forres$cache$gamma,
    control = list(report.level = 2, report.freq = 1),
    deriv = TRUE
)

lik$neg_log_lik
lik$grad
head(lik$opt$solution)
```

`neg_log_lik` is the negative profile log-likelihood after inner optimization
over random effects.

Conditional simulations of IWP smooths use `adlaplaceHgp::cond_sim_iwp()` with
the Laplace output and model-data bundle:

```{r cond-sim}
sim <- adlaplaceHgp::cond_sim_iwp(
    fit = lik,
    model_data = forres$model_data,
    n = 25
)
names(sim)
# inspect assembled inputs:
# str(adlaplaceHgp::cond_sim_iwp_inputs(lik, forres$model_data))
```