---
title: "Parameter fitting with automatic differentiation"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{Parameter fitting with automatic differentiation}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

```{r, include = FALSE}
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>"
)
```

A distinctive feature of `odelia` is that ODE systems are templated on their
scalar type, so the solver can run with an *automatic differentiation* (AD)
number type as well as ordinary `double`. This lets the solver return the exact
gradient of a loss function — the mismatch between the solution and a set of
target observations — with respect to the system's parameters. Those gradients
can then be handed to a gradient-based optimiser to calibrate the model.

This vignette demonstrates the workflow by recovering known Lorenz parameters
from a reference trajectory.

```{r setup}
library(odelia)
```

## Generate a reference trajectory

First we solve the system with known "true" parameters to produce the data we
will later try to recover.

```{r reference}
true_pars <- c(sigma = 10.0, R = 28.0, b = 8.0 / 3.0)

lz <- LorenzSystem$new(sigma = true_pars[1], R = true_pars[2], b = true_pars[3])
lz$set_initial_state(c(1, 1, 1), t0 = 0)

ctrl <- OdeControl$new()

runner <- Lorenz_Solver$new(lz$ptr, ctrl$ptr)
runner$advance_adaptive(c(0, 1))

times <- runner$times()
hist <- runner$history()

# The internal solver times that coincide with stored output rows
obs_index <- which(times %in% hist$time)
target_vals <- as.matrix(hist[, c("x", "y", "z")])
```

## Register the calibration target

We move the system parameters away from the truth (this is our starting guess),
build a solver, and register the calibration target. A solver takes its own copy
of the system when it is built, so the parameters it carries are the guess as
set here; changing `lz` afterwards would not reach `ad_runner`, whose own
`$set_params()` does. The reference run's `times` are the schedule every fit
replays, so each evaluation takes the same steps and the loss is a smooth
function of the parameters.

```{r ad-setup}
initial_guess <- c(sigma = 12.0, R = 30.0, b = 3.0)
lz$set_params(initial_guess)

ad_runner <- Lorenz_Solver$new(lz$ptr, ctrl$ptr)
ad_runner$set_target(times, target_vals, obs_index)
```

The `$loss_and_gradient()` method solves the system for a given parameter vector and returns
both the loss and its exact gradient. Underneath, it replays the schedule in
`times` step for step while keeping one record per step, sets the adjoint at
each observed index with the residual against the target, and sweeps the
recording backwards once, which yields the derivative with respect to every
parameter (and, when asked, every initial condition) from that single pass. The
`times` must therefore be a step schedule rather than an output grid, which is
why `obs_index` picks the observations out of it. The article
[Reverse mode](articles/reverse-mode.html) describes the machinery.

```{r fit-once}
res0 <- ad_runner$loss_and_gradient(params = initial_guess)
res0$loss
res0$gradient
```

## Optimise

Because each `$loss_and_gradient()` call returns both loss and gradient, we cache the most
recent result so the optimiser does not have to re-run the solver to obtain the
gradient separately.

```{r memoise}
memoise_fit <- function(runner) {
  cache_p <- NULL
  cache_res <- NULL

  fn <- function(p) {
    cache_res <<- runner$loss_and_gradient(params = p)
    cache_p <<- p
    cache_res$loss
  }
  gr <- function(p) {
    if (!identical(p, cache_p)) fn(p) # populate cache if stale
    cache_res$gradient
  }
  list(obj = fn, grad = gr)
}

helper <- memoise_fit(ad_runner)

res <- optim(
  par = initial_guess,
  fn = helper$obj,
  gr = helper$grad,
  method = "L-BFGS-B",
  control = list(maxit = 100)
)

res$convergence # 0 indicates successful convergence
res$value       # final loss, should be ~0
```

## Check the recovered parameters

```{r recovered}
rbind(true = true_pars, recovered = res$par)
```

The optimiser recovers the true parameters to high precision, confirming that
the gradients produced by AD are correct and usable for calibration.

`$loss_and_gradient()` evaluates at the parameters it is given and stores
nothing, so the solver still carries the initial guess. To integrate at the
fitted values, set them on the solver and start it again from its initial
state:

```{r refit}
ad_runner$set_params(res$par)
ad_runner$reset()
ad_runner$advance_adaptive(c(0, 1))
max(abs(as.matrix(ad_runner$history()[, c("x", "y", "z")]) - target_vals))
```