---
title: "Neural networks in nlmixr2 models"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{Neural networks in nlmixr2 models}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

```{r, include = FALSE}
knitr::opts_chunk$set(collapse = TRUE, comment = "#>")
## NB: the package must be loaded before probing the backend -- the symbol lives
## in its DLL, so probing first silently reports "no torch" and skips every
## fitting chunk, leaving a vignette that builds in seconds and demonstrates
## nothing.
.hasRx    <- requireNamespace("rxode2", quietly = TRUE)
.hasNn    <- .hasRx && requireNamespace("nlmixr2nn", quietly = TRUE)
.hasEst   <- .hasNn && requireNamespace("nlmixr2est", quietly = TRUE)
## The fits below take minutes, and `R CMD check` runs them TWICE: once in
## `R CMD build`, then again when it re-builds the vignette outputs.  CI sets
## NLMIXR2NN_VIGNETTE_FITS=false on the platforms where that is a wall-clock
## risk.  Unset means run them, so a local build still shows real output.
.wantFits <- !identical(tolower(Sys.getenv("NLMIXR2NN_VIGNETTE_FITS", "true")),
                        "false")
## Ask whether this build can TRAIN, not whether it linked libtorch: a build
## using the builtin Adam/SGD trains perfectly well and must still run the fits
## below.  Testing the libtorch probe alone would ship a vignette that skips
## every fit on exactly the platforms the fallback exists for.
.nnBk     <- tryCatch(.Call("_nlmixr2nn_nnBackend", PACKAGE = "nlmixr2nn"),
                      error = function(e) NA_character_)
.canTrain <- identical(.nnBk, "builtin") ||
  tryCatch(isTRUE(.Call("_nlmixr2nn_nnTorchAvailable", PACKAGE = "nlmixr2nn")),
           error = function(e) FALSE)
.hasTorch <- .hasEst && .wantFits && .canTrain
```

## What `nn()` is

`nn()` puts a neural network inside a model. You write it where you would write
any other term, and everything else is an ordinary `rxode2`/`nlmixr2` workflow:
the model solves, `nlmixr2()` fits it, `predict()` works, `saveRDS()` works.

```{r, eval = .hasRx}
library(nlmixr2nn)

mod <- function() {
  ini({
    lka <- 0.5
    lVc <- 1
  })
  model({
    ka <- exp(lka)
    Vc <- exp(lVc)
    d/dt(depot)   <- -ka * depot
    ## a 6-unit network of the central amount enters the elimination
    d/dt(central) <-  ka * depot - nn(central, nHidden = 6)
    cp <- central / Vc
  })
}

rxode2::rxSetSeed(42)
ui <- rxode2::rxode2(mod)
```

The network is **opaque**: it becomes a single compiled call, so the model text
stays the same size no matter how big the network is. Its value and its exact
input derivatives are computed in fast thread-safe C, and those derivatives are
what make a random effect passed *into* a network identifiable.

```{r, eval = .hasRx}
cat(paste(ui$lstChr, collapse = "\n"))
```

The weights were drawn when the model was parsed, using rxode2's threefry
generator — so `rxSetSeed()` reproduces a network exactly, in this session or any
other, and the model is ready to solve immediately. Threefry is also the
generator whose stream is defined across threads, so a network is reproducible
in the same places a parallel solve is. Drawing weights disturbs neither the R
stream nor the rxode2 seed.

```{r, eval = .hasRx}
head(nnWeights(ui), 4)
```

Everything below is a single `nlmixr2()` call. There is no setup step.

## Simulating a model you have not fitted

Because the model carries its own weights, it solves straight away. This is the
quickest way to check that a model does what you think before fitting anything.

```{r, eval = .hasRx}
ev <- rxode2::et(amt = 100, cmt = "depot") |> rxode2::et(0, 24, by = 4)
head(rxode2::rxSolve(ui, ev), 4)
```

## Learning a term you do not know (UDE)

Keep the mechanism you trust and let a small network absorb the part you do not.
Here the truth is a saturating (Michaelis-Menten) elimination that varies between
subjects; the model knows only that elimination is *some* bounded function of the
amount, and a latent random effect `eta.nn` enters the network as an input.

```{r, eval = .hasTorch}
rxode2::rxSetSeed(1)   # pins the network
set.seed(1)            # pins the simulated data below
truth <- rxode2::rxode2("d/dt(centr) = -(2*exp(eV))*centr/(3+centr)")
etaTrue <- rnorm(8, 0, sqrt(0.15))
d <- do.call(rbind, lapply(1:8, function(id) {
  s <- rxode2::rxSolve(truth,
         data.frame(id = id, time = c(0, .5, 1, 2, 4, 6, 8, 10),
                    evid = c(1, rep(0, 7)), cmt = 1, amt = c(10, rep(0, 7))),
         params = c(eV = etaTrue[id]), returnType = "data.frame")
  s <- s[s$time > 0, ]
  data.frame(id = id, time = c(0, s$time), evid = c(1, rep(0, nrow(s))), cmt = 1,
             amt = c(10, rep(0, nrow(s))), dv = c(NA, s$centr + rnorm(nrow(s), 0, .1)))
}))

ude <- function() {
  ini({ add.sd <- 0.3; eta.nn ~ 0.2 })
  model({
    g <- nn(centr, eta.nn, nHidden = 3, act = "tanh")
    d/dt(centr) <- -(1 / (1 + exp(-g))) * centr     # a bounded learned rate
    centr ~ add(add.sd)
  })
}

fit <- nlmixr2est::nlmixr2(ude, d, "focei")
```

That is the whole call: no `nn=` argument, and the data was not touched. The
result is an ordinary `nlmixr2` fit, so the random effect recovers the
between-subject variation the way any eta would:

```{r, eval = .hasTorch}
cor(fit$eta$eta.nn[order(fit$eta$ID)], etaTrue)
```

To see what the network *learned*, evaluate it directly:

```{r, eval = .hasTorch}
e <- nnEval(fit, centr = seq(0.5, 10, length.out = 25), eta.nn = 0)
head(e, 3)
```

`plot(e)` draws that function, and `plot(e, true = f)` overlays a reference
curve when you have one.

## Letting the network find a covariate relationship

The same idea with covariates as the inputs: the network learns how they drive a
parameter, and the latent eta carries what they do not explain.

```{r, eval = FALSE}
covmod <- function() {
  ini({ tCL <- 1; tV <- 3.5; add.sd <- 0.5; eta.nn ~ 0.1 })
  model({
    CL <- exp(tCL + nn(WT, EGFR, eta.nn))        # learned covariate map
    V  <- exp(tV)
    d/dt(central) <- -(CL / V) * central
    cp <- central / V
    cp ~ add(add.sd)
  })
}
fit <- nlmixr2est::nlmixr2(covmod, data, "focei")
```

## A model with no between-subject variability

Systems-pharmacology models are often fitted with no random effects at all: one
mechanistic system, with a network standing in for an unknown term. Fit those
with a population optimizer — `"bobyqa"`, `"nlminb"`, `"lbfgsb3c"` and friends —
and there is no eta to estimate.

```{r, eval = .hasTorch}
rxode2::rxSetSeed(1)
set.seed(1)
qspTruth <- rxode2::rxode2("d/dt(centr) = -(2)*centr/(3+centr)")
qd <- do.call(rbind, lapply(1:6, function(id) {
  s <- rxode2::rxSolve(qspTruth,
         data.frame(id = id, time = c(0, .5, 1, 2, 4, 6, 8, 10),
                    evid = c(1, rep(0, 7)), cmt = 1, amt = c(10, rep(0, 7))),
         returnType = "data.frame")
  s <- s[s$time > 0, ]
  data.frame(id = id, time = c(0, s$time), evid = c(1, rep(0, nrow(s))), cmt = 1,
             amt = c(10, rep(0, nrow(s))), dv = c(NA, s$centr + rnorm(nrow(s), 0, .1)))
}))

qsp <- function() {
  ini({ add.sd <- 0.3 })                     # residual error only -- no eta
  model({
    g <- nn(centr, nHidden = 3, act = "tanh")
    d/dt(centr) <- -(1 / (1 + exp(-g))) * centr
    centr ~ add(add.sd)
  })
}

qspFit <- nlmixr2est::nlmixr2(qsp, qd, "bobyqa")
```

The fitted model reproduces the dynamics it was never told about:

```{r, eval = .hasTorch}
nd <- data.frame(time = c(0, .5, 1, 2, 4, 6, 8, 10), evid = c(1, rep(0, 7)),
                 cmt = 1, amt = c(10, rep(0, 7)))
a <- rxode2::rxSolve(qspFit$finalUi, nd, returnType = "data.frame")
b <- rxode2::rxSolve(qspTruth, nd, returnType = "data.frame")
data.frame(time = a$time, fitted = round(a$centr, 2), true = round(b$centr, 2))
```

Passing a population-only optimizer for a model that *does* have a random effect
is an error rather than a silent compromise — the weights and the Omega would
otherwise come from two different objectives.

## Saving, reloading and refitting

The trained weights are carried on the fitted model, so a saved fit is
self-contained: it reloads and re-solves in a fresh session with no external
state. Feeding a fit back in resumes from its trained weights.

```{r, eval = .hasTorch}
f <- tempfile(fileext = ".rds")
saveRDS(qspFit, f)
reloaded <- readRDS(f)
identical(round(rxode2::rxSolve(reloaded$finalUi, nd, returnType = "data.frame")$centr, 8),
          round(a$centr, 8))
unlink(f)
```

```{r, eval = FALSE}
refit <- nlmixr2est::nlmixr2(qspFit, moreData, "bobyqa")   # warm start
```

## When you want to steer the training

The training schedule is inferred from the model, the data and the estimator you
chose: whether the estimator can resume from a partial fit, whether there is a
random effect, and what the residual model is. `nnControl()` overrides any part
of it and leaves the rest inferred.

```{r, eval = FALSE}
nlmixr2est::nlmixr2(ude, d, "focei", nn = nnControl(rounds = 400, lr = 0.01))
```

| setting | inferred as | override when |
|---|---|---|
| `mode` | `"joint"` when the estimator resumes from a partial fit, else `"iter"` | you want fully converged inner fits each round |
| `rounds` | 200 joint, 30 iterative, 60 population | training was still improving when it stopped |
| `cotangent` | the Gaussian closed form for an untransformed `add()`/`prop()` endpoint, the distribution's own score for a count endpoint, the captured exact score otherwise | rarely — a source the endpoint cannot support is an error, not a silent fallback |
| `lr`, `optimizer` | 0.03 (0.01 on a refit), Adam | the objective oscillates or moves too slowly |
| `warmStart` | a population pre-fit, skipped when the model already carries trained weights | you want a longer or shorter pre-fit |
| `l2`, `smooth` | `0` -- no penalty on the weights or on curvature | a fitted network looks wigglier than the data can justify (see below) |
| `kinetic` | `0` | the learned term is stiff, or the augmented solve is slow |

### Regularization, when a network learns too much

A network with more capacity than the data supports will invent structure — a
covariate curve with sharp wiggles in it, or per-subject variation that should
have gone to the random effect. It does this silently: the fit converges, the
objective looks fine, and only the learned shape gives it away.

Two penalties are available for that, and they express different preferences.
`l2` is an L2-norm penalty on the weights — ridge regression, or weight decay in
machine-learning terms — which charges a weight for being far from zero and so
prefers the flattest network the data will allow. `smooth` is a roughness penalty
in the spline sense: it charges for the *second* derivative of the network along
each input, so a monotone slope is free and only the wiggles are paid for. Put
plainly, `l2` asks the network to be small and `smooth` asks it to be smooth,
which is usually what you want from a learned covariate relationship — a shape,
not a straight line and not an interpolation of the noise. Biases are never penalized, and neither is a latent eta's input column —
the eta reaches the model only through the network, so shrinking those weights
would shrink the random effect itself.

Both are a **fraction of the objective**, not an absolute amount: `l2 = 0.05`
means the weight penalty starts at 5% of the objective, so one value means
roughly the same thing whatever the endpoint, estimator or data size.

```{r, eval = FALSE}
fit <- nlmixr2est::nlmixr2(ude, d, "focei", nn = nnControl(l2 = 0.05))
plot(nnEval(fit, centr = seq(0, 10, length.out = 101)))
```

**Both default to `0`.** That is a measured choice, not a cautious one. `l2`
shrinks the network toward the zero function, and in these models the network
*is* the model — so any value large enough to suppress structure the data does
not support is also large enough to attenuate structure it does. On the package's
own recovery tests, every value from `1e-3` up cost a real covariate effect, and
every value that preserved them moved the weight norm by less than half a percent.

So reach for it deliberately, when a fitted curve looks wigglier than the data can
justify. `l2 = 0.05` is where shrinkage becomes substantial — on a test fixture
with no true covariate effect it took the weight norm from 11.5 to 3.9 while the
unpenalized objective *improved*. Always look at what the network learned with
`nnEval()` afterwards rather than trusting the objective alone: real effects
shrink alongside invented ones.

A third penalty, `kinetic`, asks a different question. `l2` and `smooth` are about
the *shape* of the learned function; `kinetic` is about the *dynamics* it
produces. It charges for the network's output along the trajectory the model
actually solves — how hard the network is pushing where the solution really goes,
rather than on a grid of input combinations that may never occur — and so asks
the model to reach the same data by pushing less hard.

```{r, eval = FALSE}
fit <- nlmixr2est::nlmixr2(ude, d, "focei", nn = nnControl(kinetic = 0.05))
```

That makes the learned dynamics easier to integrate, so it can make a fit *faster*
as well as smoother — which matters here because the training solve carries a
sensitivity state for every state-weight pair, so every solver step saved is
multiplied. It defaults to `0` for a reason of its own on top of the one above:
it is the only one of the three that changes the dynamics being solved rather
than only steering the optimizer.

`kinetic` is a fraction of the objective too, though in a weaker sense than `l2`
and `smooth`: `kinetic = 0.05` charges 5% of the objective per unit of the term,
rather than starting at exactly 5% of it. The difference is deliberate and
`vignette("nlmixr2nn-node")` explains it — briefly, the default initialization
starts the network near the zero function on purpose, so normalizing this term
against its own starting value would divide by something four orders of magnitude
too small.

All three shape the optimization only. The reported `objf` — and so AIC and BIC —
stays the unpenalized -2 log-likelihood, so a regularized network model is still
directly comparable to one with an analytic covariate relationship.
`vignette("nlmixr2nn-node")` covers where `kinetic` comes from and which other
ideas from the neural ODE literature do and do not carry over to a model like
this one.

## What is supported

- **Inputs**: 1 to 4, given positionally — states, covariates, or a latent eta.
- **Activations**: `"softplus"` (default), `"tanh"`, `"relu"`, `"gelu"`, `"silu"`.
  The default is smooth with a nonzero second derivative, which matters because
  the model Jacobian and the FOCEi sensitivities are chained through it.
- **Hidden layers**: one. `nHidden` is a single width.
- **Endpoints**: one. Any normal-distribution residual model — `add()`, `prop()`,
  both, or a transformation such as `lnorm()` — and the count endpoints `pois()`
  and `binom()`.

  A count endpoint differs in one way worth knowing. For `y ~ pois(lam)` there is
  no "prediction" to differentiate: the model's prediction *is* the log-density,
  so the network's effect is chained through the distribution's own parameter
  (`lam`) instead, and the score is that distribution's derivative. Two
  consequences: `nnControl(cotangent = "exact")` does not apply and is refused,
  and **the parameter must be positive by construction** — write
  `lam <- exp(nn(...))` rather than letting a rate wander negative, which aborts
  the likelihood evaluation. For `binom(n, p)` the network drives `p`; the size
  `n` must reach the solve, so bind it to a data column (`n <- ntrials`) rather
  than writing a bare constant, which the model compiler folds away.

  Other distributions (`nbinom()`, ordinal, and the rest) are refused with a
  message naming the reason, not fitted wrongly.
- **Networks per model**: several; they are trained together, and `nnEval(net=)`
  selects between them.
- **Estimators**: any standard one. `"focei"`, the Laplace/AGQ methods, the
  importance-sampling methods (`"imp"`, `"impmap"`) and the variational methods
  interleave network training with the population fit, because their outer
  iteration genuinely resumes from where it left off. SAEM (`"saem"`,
  `"fsaem"`), `"qrpem"` and the nonparametric methods (`"npag"`, `"npb"`)
  alternate instead: SAEM's stochastic-approximation gain sequence restarts on
  every call, so chopping it into pieces would not continue the chain.

Network initialization is scaled to each input's magnitude, measured from the
data, so a network reading a concentration of 0.02 and one reading an amount of
500 both start in a usable range.