tests/testthat/test-adjoint-sens.R

In rxode2: Facilities for Simulating from ODE-Based Models

rxTest({
  # Adjoint (backward) sensitivity analysis.
  #
  # .rxAdjoint() symbolically generates the continuous-adjoint ODE system
  #   costate:    d/dt(lambda_k)   = -J^T lambda_k
  #   quadrature: d/dt(sens_k_p)   = -lambda_k^T df/dp
  # reusing the rx__df_*_dy_*__ elemental derivatives that .rxJacobian
  # materialises (no automatic differentiation).  These tests verify (a) the
  # generated equations have the correct transpose-Jacobian / quadrature
  # structure and (b) that integrating them backward reproduces dy_k(T)/dp to
  # the same tolerance as the forward-sensitivity path and a finite difference.

  .nm <- function(state, p) paste0("rx__sens_", state, "_BY_", p, "__")

  mText <- "d/dt(depot)  = -ka*depot\nd/dt(center) =  ka*depot - (cl/v)*center"
  vars  <- c("ka", "cl", "v")
  p0    <- c(ka = 1.1, cl = 3.0, v = 20.0)
  Tfin  <- 12
  outState <- "center"

  model <- rxode2::rxS(rxode2::rxGetModel(mText), TRUE, promoteLinSens = FALSE)
  st    <- rxode2::rxStateOde(model)
  invisible(rxode2::.rxJacobian(model, c(st, vars)))
  s1    <- rxode2::.rxSens(model, vars)
  adj   <- rxode2::.rxAdjoint(model, vars, outState)

  test_that(".rxAdjoint emits one costate per (out-state,state) and one quadrature per param", {
    lam <- grep("^d/dt\\(rx__adjLambda_", adj, value = TRUE)
    quad <- grep("^d/dt\\(rx__sens_", adj, value = TRUE)
    expect_equal(length(lam), length(st))          # one lambda per state
    expect_equal(length(quad), length(vars))        # one dy/dp per param
    # quadrature reuses the forward-sensitivity output names exactly
    expect_true(all(vapply(vars, function(p)
      any(grepl(.nm(outState, p), quad, fixed = TRUE)), logical(1))))
    # costate for a decaying central compartment is +cl/v*lambda (transpose sign)
    expect_true(any(grepl("rx__adjLambda_center_center__)=cl*rx__adjLambda_center_center__/v",
                          adj, fixed = TRUE)))
    expect_null(NULL)
  })

  # ---- forward-sensitivity + finite-difference references --------------------
  fwdMod <- rxode2::rxode2(paste(c(mText, s1), collapse = "\n"))
  ev     <- rxode2::et(amt = 100, cmt = "depot") |> rxode2::et(Tfin)
  fwdRow <- as.data.frame(rxode2::rxSolve(fwdMod, params = p0, ev,
                                          returnType = "data.frame",
                                          atol = 1e-12, rtol = 1e-12))
  fwdRow <- fwdRow[fwdRow$time == Tfin, ]
  yT     <- c(depot = fwdRow$depot, center = fwdRow$center)
  refFwd <- vapply(vars, function(p) fwdRow[[.nm(outState, p)]], numeric(1))

  solveY <- function(p) {
    d <- as.data.frame(rxode2::rxSolve(rxode2::rxode2(mText), params = p, ev,
                                       returnType = "data.frame",
                                       atol = 1e-12, rtol = 1e-12))
    d[d$time == Tfin, outState]
  }
  refFD <- vapply(vars, function(p) {
    h <- p0[[p]] * 1e-5
    pp <- p0; pm <- p0; pp[p] <- pp[p] + h; pm[p] <- pm[p] - h
    (solveY(pp) - solveY(pm)) / (2 * h)
  }, numeric(1))

  # ---- adjoint via reverse-time (s = Tfin - t) reconstruction ----------------
  # dy/ds = -f(y); costate/quadrature flip sign from the backward-in-t form.
  negLine <- function(line) {
    m <- regmatches(line, regexec("^\\s*(d/dt\\([^)]*\\))\\s*=\\s*(.*)$", line))[[1]]
    paste0(m[2], "=-(", m[3], ")")
  }
  revPrimal <- vapply(strsplit(mText, "\n")[[1]], negLine, character(1),
                      USE.NAMES = FALSE)
  revAdj    <- vapply(adj, negLine, character(1), USE.NAMES = FALSE)
  revMod    <- rxode2::rxode2(paste(c(revPrimal, revAdj), collapse = "\n"))

  inits <- c(depot = yT[["depot"]], center = yT[["center"]])
  for (i in st) inits[paste0("rx__adjLambda_", outState, "_", i, "__")] <-
                  as.numeric(i == outState)
  for (p in vars) inits[.nm(outState, p)] <- 0
  revEv <- rxode2::et(seq(0, Tfin, length.out = 2001))
  revEnd <- as.data.frame(rxode2::rxSolve(revMod, params = p0, revEv,
                                          inits = inits, returnType = "data.frame",
                                          atol = 1e-12, rtol = 1e-12))
  revEnd  <- revEnd[nrow(revEnd), ]
  adjSens <- vapply(vars, function(p) revEnd[[.nm(outState, p)]], numeric(1))

  test_that("adjoint sensitivities match the forward-sensitivity path", {
    expect_equal(unname(adjSens), unname(refFwd), tolerance = 1e-6)
  })

  test_that("adjoint sensitivities match a central finite difference", {
    expect_equal(unname(adjSens), unname(refFD), tolerance = 1e-5)
  })

  # ---- full-trajectory adjoint (.rxAdjointSolve) over a long window ----------
  # Robustness: a 24 h window with fast absorption would make reverse-primal
  # reconstruction blow up like exp(+||J||*T); the checkpoint-interpolation
  # path stays bounded and matches forward sensitivities at every output time.
  fullOutTimes <- seq(1, 24, by = 1)
  fullP  <- c(ka = 1.5, cl = 4.0, v = 30.0)
  fullEv <- rxode2::et(amt = 100, cmt = "depot")

  fmod <- rxode2::rxode2(paste(c(mText, s1), collapse = "\n"))
  fev  <- rxode2::et(amt = 100, cmt = "depot") |> rxode2::et(fullOutTimes)
  fref <- as.data.frame(rxode2::rxSolve(fmod, params = fullP, fev,
                                        returnType = "data.frame",
                                        atol = 1e-11, rtol = 1e-11,
                                        addDosing = FALSE))
  fref <- fref[fref$time %in% fullOutTimes, ]
  allCols <- unlist(lapply(st, function(k)
    vapply(vars, function(p) .nm(k, p), character(1))))

  relTrajErr <- function(denseBy) {
    adjT <- rxode2::.rxAdjointSolve(mText, fullP, fullEv, vars, fullOutTimes,
                                    denseBy = denseBy)
    max(vapply(allCols, function(c)
      max(abs(adjT[[c]] - fref[[c]]) / (abs(fref[[c]]) + 1e-6)), numeric(1)))
  }

  test_that("full-trajectory adjoint matches forward sens over a long window", {
    expect_lt(relTrajErr(0.01), 1e-4)   # bounded (no reverse-primal blow-up)
  })

  test_that("full-trajectory adjoint error converges as the grid refines", {
    e_coarse <- relTrajErr(0.04)
    e_fine   <- relTrajErr(0.01)
    # O(denseBy^2) covariate-interpolation error: refining 4x cuts it markedly
    expect_lt(e_fine, e_coarse)
    expect_lt(e_fine, 5e-5)
  })

  # ---- functional-gradient adjoint (.rxAdjointGrad): the genuine win ---------
  # ONE backward sweep yields dG/dtheta for ALL theta, where
  #   G = sum_i 1/2 * (h(y(t_i), theta) - obs_i)^2,  h = center/v.
  test_that("adjoint objective gradient matches a central finite difference", {
    gP   <- c(ka = 1.2, cl = 3.5, v = 25.0)
    gCS  <- c("ka", "cl", "v")
    gPred <- "center/v"
    gObsT <- c(0.5, 1, 2, 4, 6, 8, 12, 18, 24)
    gEv  <- rxode2::et(amt = 100, cmt = "depot")

    truthMod <- rxode2::rxode2(paste0(mText, "\ncp=", gPred))
    gFev <- gEv |> rxode2::et(gObsT)
    truth <- as.data.frame(rxode2::rxSolve(truthMod, params = gP, gFev,
                                           returnType = "data.frame", addDosing = FALSE))
    gObs <- truth$cp[truth$time %in% gObsT] * 1.1 + 0.05   # nonzero residuals

    objG <- function(p) {
      d <- as.data.frame(rxode2::rxSolve(truthMod, params = p, gFev,
                                         returnType = "data.frame", addDosing = FALSE))
      sum(0.5 * (d$cp[d$time %in% gObsT] - gObs)^2)
    }
    gFD <- vapply(gCS, function(p) {
      h <- gP[[p]] * 1e-6; pp <- gP; pm <- gP
      pp[p] <- pp[p] + h; pm[p] <- pm[p] - h
      (objG(pp) - objG(pm)) / (2 * h)
    }, numeric(1))

    gAdj <- rxode2::.rxAdjointGrad(mText, gP, gEv, gCS, gPred, gObsT, gObs)
    expect_equal(unname(gAdj), unname(gFD), tolerance = 1e-3)
  })

  # ---- F (bioavailability) dose-jump dual ------------------------------------
  # A parameter entering ONLY through f(depot)=Fbio gets its whole gradient from
  # the dose-jump term  lambda_depot(t0+)*amt*dF/dtheta.
  test_that("adjoint gradient handles bioavailability (F) dose-jump dual", {
    fText <- paste0(mText, "\nf(depot)=Fbio")
    fP    <- c(ka = 1.2, cl = 3.5, v = 25.0, Fbio = 0.7)
    fCS   <- c("ka", "cl", "v", "Fbio")
    fPred <- "center/v"
    fObsT <- c(0.5, 1, 2, 4, 6, 8, 12, 18, 24)
    fEv   <- rxode2::et(amt = 100, cmt = "depot")

    truthMod <- rxode2::rxode2(paste0(fText, "\ncp=", fPred))
    fFev <- fEv |> rxode2::et(fObsT)
    truth <- as.data.frame(rxode2::rxSolve(truthMod, params = fP, fFev,
                                           returnType = "data.frame", addDosing = FALSE))
    fObs <- truth$cp[truth$time %in% fObsT] * 1.1 + 0.05

    objG <- function(p) {
      d <- as.data.frame(rxode2::rxSolve(truthMod, params = p, fFev,
                                         returnType = "data.frame", addDosing = FALSE))
      sum(0.5 * (d$cp[d$time %in% fObsT] - fObs)^2)
    }
    gFD <- vapply(fCS, function(p) {
      h <- fP[[p]] * 1e-6; pp <- fP; pm <- fP
      pp[p] <- pp[p] + h; pm[p] <- pm[p] - h
      (objG(pp) - objG(pm)) / (2 * h)
    }, numeric(1))

    gAdj <- rxode2::.rxAdjointGrad(fText, fP, fEv, fCS, fPred, fObsT, fObs)
    # Fbio gradient is entirely from the dose term and clearly nonzero
    expect_gt(abs(gAdj[["Fbio"]]), 1)
    expect_equal(unname(gAdj), unname(gFD), tolerance = 1e-3)
    # F dose dual also runs through the C++ sweep
    Bf <- rxode2::.rxAdjointGradBuild(fText, fCS, fPred, fEv)
    gCf <- rxode2::.rxAdjointGradEvalC(Bf, fP, fObsT, fObs)
    gRf <- rxode2::.rxAdjointGradEval(Bf, fP, fObsT, fObs)
    expect_equal(unname(gCf), unname(gRf), tolerance = 1e-4)
  })

  # ---- modeled lag (alag) transversality dual --------------------------------
  # A parameter entering through alag(depot)=tlag gets its gradient from the
  # time-triggered transversality term at the (lagged) dose time.
  test_that("adjoint gradient handles modeled lag (alag) transversality dual", {
    lText <- paste0(mText, "\nalag(depot)=tlag")
    lP    <- c(ka = 1.2, cl = 3.5, v = 25.0, tlag = 0.8)
    lCS   <- c("ka", "cl", "v", "tlag")
    lPred <- "center/v"
    lObsT <- c(1, 2, 4, 6, 8, 12, 18, 24)
    lEv   <- rxode2::et(amt = 100, cmt = "depot")

    truthMod <- rxode2::rxode2(paste0(lText, "\ncp=", lPred))
    lFev <- lEv |> rxode2::et(lObsT)
    truth <- as.data.frame(rxode2::rxSolve(truthMod, params = lP, lFev,
                                           returnType = "data.frame", addDosing = FALSE))
    lObs <- truth$cp[truth$time %in% lObsT] * 1.1 + 0.05

    objG <- function(p) {
      d <- as.data.frame(rxode2::rxSolve(truthMod, params = p, lFev,
                                         returnType = "data.frame", addDosing = FALSE))
      sum(0.5 * (d$cp[d$time %in% lObsT] - lObs)^2)
    }
    gFD <- vapply(lCS, function(p) {
      h <- lP[[p]] * 1e-6; pp <- lP; pm <- lP
      pp[p] <- pp[p] + h; pm[p] <- pm[p] - h
      (objG(pp) - objG(pm)) / (2 * h)
    }, numeric(1))

    gAdj <- rxode2::.rxAdjointGrad(lText, lP, lEv, lCS, lPred, lObsT, lObs, denseBy = 0.005)
    # the tlag transversality term is exact (point values); check it tightly
    expect_equal(gAdj[["tlag"]], gFD[["tlag"]], tolerance = 1e-3)
    # structural params limited by covariate interpolation of the sharp peak
    expect_equal(unname(gAdj), unname(gFD), tolerance = 1e-2)
    # lag transversality dual also runs through the C++ sweep
    Bl <- rxode2::.rxAdjointGradBuild(lText, lCS, lPred, lEv)
    gCl <- rxode2::.rxAdjointGradEvalC(Bl, lP, lObsT, lObs, denseBy = 0.005)
    gRl <- rxode2::.rxAdjointGradEval(Bl, lP, lObsT, lObs, denseBy = 0.005)
    expect_equal(unname(gCl), unname(gRl), tolerance = 1e-4)
  })

  # ---- replace(evid5) / multiply(evid6) costate jumps ------------------------
  # The costate jump is essential for correct structural-param gradients even
  # with a constant replace value (resetting/scaling lambda_c at the event).
  test_that("adjoint gradient handles replace/multiply costate jumps", {
    eP    <- c(ka = 1.2, cl = 3.5, v = 25.0)
    eCS   <- c("ka", "cl", "v")
    ePred <- "center/v"
    eObsT <- c(1, 2, 4, 6, 8, 12)
    tmod  <- rxode2::rxode2(paste0(mText, "\ncp=", ePred))

    chk <- function(eEv) {
      eFev <- eEv |> rxode2::et(eObsT)
      truth <- as.data.frame(rxode2::rxSolve(tmod, params = eP, eFev,
                                             returnType = "data.frame", addDosing = FALSE))
      eObs <- truth$cp[truth$time %in% eObsT] * 1.1 + 0.05
      objG <- function(p) {
        d <- as.data.frame(rxode2::rxSolve(tmod, params = p, eFev,
                                           returnType = "data.frame", addDosing = FALSE))
        sum(0.5 * (d$cp[d$time %in% eObsT] - eObs)^2)
      }
      gFD <- vapply(eCS, function(p) {
        h <- eP[[p]] * 1e-6; pp <- eP; pm <- eP
        pp[p] <- pp[p] + h; pm[p] <- pm[p] - h
        (objG(pp) - objG(pm)) / (2 * h)
      }, numeric(1))
      gAdj <- rxode2::.rxAdjointGrad(mText, eP, eEv, eCS, ePred, eObsT, eObs,
                                     denseBy = 0.005)
      expect_equal(unname(gAdj), unname(gFD), tolerance = 5e-3)
      # replace/multiply costate jumps also run through the C++ sweep
      B <- rxode2::.rxAdjointGradBuild(mText, eCS, ePred, eEv)
      gC <- rxode2::.rxAdjointGradEvalC(B, eP, eObsT, eObs, denseBy = 0.005)
      expect_equal(unname(gC), unname(gAdj), tolerance = 1e-4)
    }
    chk(rxode2::et(amt = 100, cmt = "depot") |>
          rxode2::et(time = 3, amt = 40, cmt = "center", evid = 5))   # replace
    chk(rxode2::et(amt = 100, cmt = "depot") |>
          rxode2::et(time = 3, amt = 0.5, cmt = "center", evid = 6))  # multiply
  })

  # ---- modeled-rate infusion continuous-forcing + moving-boundary dual -------
  # rate(depot)=Rin adds +R over [0, amt/R]; dG/dRin combines a continuous
  # forcing integral of lambda_depot and a moving-boundary term (amt/R moves).
  test_that("adjoint gradient handles modeled-rate infusion (forcing + boundary)", {
    iText <- paste0(mText, "\nrate(depot)=Rin")
    iP    <- c(ka = 1.2, cl = 3.5, v = 25.0, Rin = 50)
    iCS   <- c("ka", "cl", "v", "Rin")
    iPred <- "depot"                              # observe during infusion (strong signal)
    iObsT <- c(0.5, 1, 1.5, 2.5, 3, 4)
    iEv   <- rxode2::et(amt = 100, rate = -1, cmt = "depot")

    tmod <- rxode2::rxode2(paste0(iText, "\ncp=", iPred))
    iFev <- iEv |> rxode2::et(iObsT)
    truth <- as.data.frame(rxode2::rxSolve(tmod, params = iP, iFev,
                                           returnType = "data.frame", addDosing = FALSE))
    iObs <- truth$cp[truth$time %in% iObsT] * 1.1 + 0.5

    objG <- function(p) {
      d <- as.data.frame(rxode2::rxSolve(tmod, params = p, iFev,
                                         returnType = "data.frame", addDosing = FALSE))
      sum(0.5 * (d$cp[d$time %in% iObsT] - iObs)^2)
    }
    gFD <- vapply(iCS, function(p) {
      h <- iP[[p]] * 1e-5; pp <- iP; pm <- iP
      pp[p] <- pp[p] + h; pm[p] <- pm[p] - h
      (objG(pp) - objG(pm)) / (2 * h)
    }, numeric(1))

    gAdj <- rxode2::.rxAdjointGrad(iText, iP, iEv, iCS, iPred, iObsT, iObs, denseBy = 0.002)
    # Rin gradient (~ -3.2) is dominated by the infusion forcing+boundary dual
    expect_gt(abs(gAdj[["Rin"]]), 1)
    expect_equal(unname(gAdj), unname(gFD), tolerance = 5e-3)
    # infusion forcing + moving-boundary duals also run through the C++ sweep
    Bi <- rxode2::.rxAdjointGradBuild(iText, iCS, iPred, iEv)
    gCi <- rxode2::.rxAdjointGradEvalC(Bi, iP, iObsT, iObs, denseBy = 0.002)
    expect_equal(unname(gCi), unname(gAdj), tolerance = 1e-4)
  })

  # ---- modeled-duration infusion (rate = amt/D) ------------------------------
  test_that("adjoint gradient handles modeled-duration infusion", {
    dText <- paste0(mText, "\ndur(depot)=Dd")
    dP    <- c(ka = 1.2, cl = 3.5, v = 25.0, Dd = 2)
    dCS   <- c("ka", "cl", "v", "Dd")
    dPred <- "depot"
    dObsT <- c(0.5, 1, 1.5, 2.5, 3, 4)
    dEv   <- rxode2::et(amt = 100, rate = -2, cmt = "depot")

    tmod <- rxode2::rxode2(paste0(dText, "\ncp=", dPred))
    dFev <- dEv |> rxode2::et(dObsT)
    truth <- as.data.frame(rxode2::rxSolve(tmod, params = dP, dFev,
                                           returnType = "data.frame", addDosing = FALSE))
    dObs <- truth$cp[truth$time %in% dObsT] * 1.1 + 0.5

    objG <- function(p) {
      d <- as.data.frame(rxode2::rxSolve(tmod, params = p, dFev,
                                         returnType = "data.frame", addDosing = FALSE))
      sum(0.5 * (d$cp[d$time %in% dObsT] - dObs)^2)
    }
    gFD <- vapply(dCS, function(p) {
      h <- dP[[p]] * 1e-5; pp <- dP; pm <- dP
      pp[p] <- pp[p] + h; pm[p] <- pm[p] - h
      (objG(pp) - objG(pm)) / (2 * h)
    }, numeric(1))

    gAdj <- rxode2::.rxAdjointGrad(dText, dP, dEv, dCS, dPred, dObsT, dObs, denseBy = 0.002)
    expect_gt(abs(gAdj[["Dd"]]), 1)
    expect_equal(unname(gAdj), unname(gFD), tolerance = 5e-3)
    # modeled-dur duals also run through the C++ sweep
    Bd <- rxode2::.rxAdjointGradBuild(dText, dCS, dPred, dEv)
    gCd <- rxode2::.rxAdjointGradEvalC(Bd, dP, dObsT, dObs, denseBy = 0.002)
    expect_equal(unname(gCd), unname(gAdj), tolerance = 1e-4)
  })

  # ---- FOCEi -2LL objective gradient (structural + residual-error params) ----
  # The gradient FOCEi propagates: dG/dtheta of G = sum r^2/v + log(v),
  # v = add^2 + (prop*f)^2, over structural AND residual-error parameters, all
  # from ONE backward sweep.
  test_that("adjoint gradient of the FOCEi -2LL objective matches finite differences", {
    lP    <- c(ka = 1.2, cl = 3.5, v = 25.0, add = 0.3, prop = 0.1)
    lCS   <- c("ka", "cl", "v", "add", "prop")
    lPred <- "center/v"
    lObsT <- c(0.5, 1, 2, 4, 6, 8, 12, 18, 24)
    lEv   <- rxode2::et(amt = 100, cmt = "depot")
    tmod  <- rxode2::rxode2(paste0(mText, "\ncp=", lPred))
    lFev  <- lEv |> rxode2::et(lObsT)
    truth <- as.data.frame(rxode2::rxSolve(tmod, params = lP, lFev,
                                           returnType = "data.frame", addDosing = FALSE))
    set.seed(2)
    lObs <- truth$cp[truth$time %in% lObsT] * (1 + stats::rnorm(length(lObsT), 0, 0.1)) +
      stats::rnorm(length(lObsT), 0, 0.2)

    m2ll <- function(p) {
      d <- as.data.frame(rxode2::rxSolve(tmod, params = p, lFev,
                                         returnType = "data.frame", addDosing = FALSE))
      f <- d$cp[d$time %in% lObsT]; v <- p[["add"]]^2 + (p[["prop"]] * f)^2
      sum((f - lObs)^2 / v + log(v))
    }
    gFD <- vapply(lCS, function(p) {
      h <- abs(lP[[p]]) * 1e-6; pp <- lP; pm <- lP
      pp[p] <- pp[p] + h; pm[p] <- pm[p] - h
      (m2ll(pp) - m2ll(pm)) / (2 * h)
    }, numeric(1))
    gAdj <- rxode2::.rxAdjointGrad(mText, lP, lEv, lCS, lPred, lObsT, lObs,
                                   errModel = list(add = "add", prop = "prop"))
    expect_equal(unname(gAdj), unname(gFD), tolerance = 1e-3)
  })

  # ---- C++ eval (build once, evaluate many): matches R eval + FD -------------
  # The pure-numeric backward sweep runs in C++ (rxAdjointSweepC) for the
  # continuous case, matching the R eval to solver precision and the FOCEi -2LL
  # finite difference, with no symbolic work at evaluation time.
  test_that("C++ adjoint eval matches the R eval and finite differences", {
    cPred <- "center/v"
    cCS   <- c("ka", "cl", "v", "add", "prop")
    cObsT <- c(0.5, 1, 2, 4, 6, 8, 12, 18, 24)
    cEv   <- rxode2::et(amt = 100, cmt = "depot")
    B <- rxode2::.rxAdjointGradBuild(mText, cCS, cPred, cEv,
                                     errModel = list(add = "add", prop = "prop"))
    tmod <- rxode2::rxode2(paste0(mText, "\ncp=", cPred))
    cFev <- cEv |> rxode2::et(cObsT)
    for (cP in list(c(ka = 1.2, cl = 3.5, v = 25, add = 0.3, prop = 0.1),
                    c(ka = 0.8, cl = 5.0, v = 18, add = 0.5, prop = 0.05))) {
      truth <- as.data.frame(rxode2::rxSolve(tmod, params = cP, cFev,
                                             returnType = "data.frame", addDosing = FALSE))
      cObs <- truth$cp[truth$time %in% cObsT] * 1.15 + 0.1
      m2ll <- function(p) {
        d <- as.data.frame(rxode2::rxSolve(tmod, params = p, cFev,
                                           returnType = "data.frame", addDosing = FALSE))
        f <- d$cp[d$time %in% cObsT]; v <- p[["add"]]^2 + (p[["prop"]] * f)^2
        sum((f - cObs)^2 / v + log(v))
      }
      gFD <- vapply(cCS, function(p) {
        h <- abs(cP[[p]]) * 1e-6; pp <- cP; pm <- cP
        pp[p] <- pp[p] + h; pm[p] <- pm[p] - h
        (m2ll(pp) - m2ll(pm)) / (2 * h)
      }, numeric(1))
      gR <- rxode2::.rxAdjointGradEval(B, cP, cObsT, cObs, denseBy = 0.01)
      gC <- rxode2::.rxAdjointGradEvalC(B, cP, cObsT, cObs, denseBy = 0.01)
      expect_equal(unname(gC), unname(gR), tolerance = 1e-5)  # RK4 vs adaptive
      expect_equal(unname(gC), unname(gFD), tolerance = 2e-3)
    }
  })

  # ---- C++ full-trajectory sweep (rxode2AdjointTrajSweep / .rxAdjointSolveEvalC) --
  # dy_k(t_i)/dp for EVERY state of interest, output time, and param -- the
  # in-engine counterpart of forward sensitivity's full output structure.
  # IMPORTANT: this requires ONE independent backward sweep PER output time
  # (see the code comment in src/adjoint.cpp) -- it is NOT the O(1) win that
  # .rxAdjointGrad's scalar objective sweep is; it exists for output-structure
  # parity / validation, matching the plan's honest scoping note.
  # For a structurally-zero sensitivity (e.g. a state with no dependence on a
  # given param), the analytic answer is exactly 0 but a finite difference has
  # roundoff noise; the floor absorbs that noise before computing a relative
  # error so a genuinely tiny/zero true value does not blow up the ratio.
  relOrAbsErr <- function(cc, fd, absFloor = 0.02) {
    pmax(abs(cc - fd) - absFloor, 0) / (abs(fd) + absFloor)
  }

  test_that("C++ full-trajectory adjoint sweep matches the R reference (continuous)", {
    tCS <- c("ka", "cl", "v"); tP <- c(ka = 1.2, cl = 3.5, v = 25)
    tObsT <- c(1, 2, 4, 6, 8, 12, 18, 24)
    tEv <- rxode2::et(amt = 100, cmt = "depot")
    B <- rxode2::.rxAdjointSolveBuild(mText, tCS, tEv)
    resC <- rxode2::.rxAdjointSolveEvalC(B, tP, tObsT, denseBy = 0.01)
    resR <- rxode2::.rxAdjointSolve(mText, tP, tEv, tCS, tObsT, denseBy = 0.01)
    tCols <- setdiff(names(resC), "time")
    for (cc in tCols) {
      expect_lt(max(relOrAbsErr(resC[[cc]], resR[[cc]])), 1e-3)
    }
  })

  test_that("C++ full-trajectory adjoint sweep handles F and modeled-lag duals (vs FD)", {
    tOb <- c(1, 2, 4)
    tEv <- rxode2::et(amt = 100, cmt = "depot")
    chkTraj <- function(txt, cs, pp) {
      B <- rxode2::.rxAdjointSolveBuild(txt, cs, tEv)
      resC <- rxode2::.rxAdjointSolveEvalC(B, pp, tOb, denseBy = 0.005)
      tmod <- rxode2::rxode2(txt); tFev <- tEv |> rxode2::et(tOb)
      solveAt <- function(p) as.data.frame(rxode2::rxSolve(tmod, params = p, tFev,
                                                            returnType = "data.frame",
                                                            addDosing = FALSE))
      for (.stt in c("depot", "center")) for (.pp in cs) {
        h <- abs(pp[[.pp]]) * 1e-6; pp1 <- pp; pp2 <- pp
        pp1[.pp] <- pp1[.pp] + h; pp2[.pp] <- pp2[.pp] - h
        d1 <- solveAt(pp1); d2 <- solveAt(pp2)
        fd <- (d1[[.stt]] - d2[[.stt]]) / (2 * h)
        cc <- resC[[paste0("rx__sens_", .stt, "_BY_", .pp, "__")]]
        expect_lt(max(relOrAbsErr(cc, fd)), 2e-2)
      }
    }
    chkTraj(paste0(mText, "\nf(depot)=Fbio"), c("ka", "cl", "v", "Fbio"),
            c(ka = 1.2, cl = 3.5, v = 25, Fbio = 0.7))
    chkTraj(paste0(mText, "\nalag(depot)=tlag"), c("ka", "cl", "v", "tlag"),
            c(ka = 1.2, cl = 3.5, v = 25, tlag = 0.8))
  })

  # ---- rxSolveAdjoint(): drop-in rxSolve() wrapper with adjoint sens columns --
  # Same column names/structure as forward-sensitivity rxSolve(calcSens=), just
  # computed via the backward (adjoint) path -- output-structure parity.
  test_that("rxSolveAdjoint matches forward-sensitivity rxSolve output", {
    aP <- c(ka = 1.2, cl = 3.5, v = 25); aCS <- c("ka", "cl", "v")
    aEv <- rxode2::et(amt = 100, cmt = "depot") |> rxode2::et(c(1, 2, 4, 6, 8, 12))
    res <- rxode2::rxSolveAdjoint(mText, aP, aEv, aCS)

    model <- rxode2::rxS(rxode2::rxGetModel(mText), TRUE, promoteLinSens = FALSE)
    st <- rxode2::rxStateOde(model)
    invisible(rxode2::.rxJacobian(model, c(st, aCS)))
    s1 <- rxode2::.rxSens(model, aCS)
    fmod <- rxode2::rxode2(paste(c(mText, s1), collapse = "\n"))
    fref <- as.data.frame(rxode2::rxSolve(fmod, params = aP, aEv,
                                          returnType = "data.frame"))

    sensCols <- grep("^rx__sens_", names(fref), value = TRUE)
    expect_true(all(sensCols %in% names(res)))          # same column names
    expect_equal(res$depot, fref$depot, tolerance = 1e-4)     # same primal solve
    expect_equal(res$center, fref$center, tolerance = 1e-4)
    for (cc in sensCols) {
      expect_equal(res[[cc]], fref[[cc]], tolerance = 1e-3)   # same sens values
    }
  })

  # ---- population objective gradient (.rxAdjointGradPop) ---------------------
  # dOFV/dtheta = sum over subjects of each subject's adjoint -2LL gradient --
  # the exact quantity a gradient-based population estimator minimises, with one
  # backward sweep per subject regardless of the number of population parameters.
  test_that("adjoint population gradient matches a finite difference of the population -2LL", {
    pTh <- c(ka = 1.2, cl = 3.5, v = 25, add = 0.4, prop = 0.08)
    pCS <- c("ka", "cl", "v", "add", "prop")
    pErr <- list(add = "add", prop = "prop")
    pEv <- rxode2::et(amt = 100, cmt = "depot") |>
      rxode2::et(c(0.5, 1, 2, 4, 6, 8, 12, 18, 24))
    tmod <- rxode2::rxode2(paste0(mText, "\ncp=center/v"))
    set.seed(5); nsub <- 5
    truth <- as.data.frame(rxode2::rxSolve(tmod, params = pTh, pEv,
                                           returnType = "data.frame", addDosing = FALSE))
    data <- do.call(rbind, lapply(seq_len(nsub), function(i) {
      vv <- pTh[["add"]]^2 + (pTh[["prop"]] * truth$cp)^2
      data.frame(id = i, time = truth$time,
                 dv = truth$cp + stats::rnorm(length(truth$cp), 0, sqrt(vv)))
    }))
    m2ll <- function(th) {
      tot <- 0
      for (i in seq_len(nsub)) {
        di <- data[data$id == i, ]
        d <- as.data.frame(rxode2::rxSolve(tmod, params = th, pEv,
                                           returnType = "data.frame", addDosing = FALSE))
        f <- d$cp; vv <- th[["add"]]^2 + (th[["prop"]] * f)^2
        tot <- tot + sum((f - di$dv)^2 / vv + log(vv))
      }
      tot
    }
    gFD <- vapply(pCS, function(pn) {
      h <- abs(pTh[[pn]]) * 1e-6; p1 <- pTh; p2 <- pTh
      p1[pn] <- p1[pn] + h; p2[pn] <- p2[pn] - h
      (m2ll(p1) - m2ll(p2)) / (2 * h)
    }, numeric(1))
    gPop <- rxode2::.rxAdjointGradPop(mText, pTh, pEv, pCS, "center/v", data,
                                      pErr, denseBy = 0.005)
    expect_equal(unname(gPop), unname(gFD), tolerance = 5e-3)
  })

  test_that("adjoint population gradient handles per-subject covariates", {
    cText <- paste0("cl = exp(tcl + bwt*lwt)\n", "v  = exp(tv + bage*lage)\n", mText)
    cTh <- c(tcl = 1.2, tv = 3.2, ka = 1.1, bwt = 0.3, bage = -0.1, add = 0.4, prop = 0.08)
    cCS <- c("tcl", "tv", "ka", "bwt", "bage", "add", "prop")
    cErr <- list(add = "add", prop = "prop")
    cEv <- rxode2::et(amt = 100, cmt = "depot") |>
      rxode2::et(c(0.5, 1, 2, 4, 6, 8, 12, 18, 24))
    tmod <- rxode2::rxode2(paste0(cText, "\ncp=center/v"))
    set.seed(3); nsub <- 6
    covs <- data.frame(id = seq_len(nsub), lwt = stats::rnorm(nsub), lage = stats::rnorm(nsub))
    data <- do.call(rbind, lapply(seq_len(nsub), function(i) {
      p <- c(cTh, lwt = covs$lwt[i], lage = covs$lage[i])
      d <- as.data.frame(rxode2::rxSolve(tmod, params = p, cEv,
                                         returnType = "data.frame", addDosing = FALSE))
      vv <- p[["add"]]^2 + (p[["prop"]] * d$cp)^2
      data.frame(id = i, time = d$time,
                 dv = d$cp + stats::rnorm(length(d$cp), 0, sqrt(vv)),
                 lwt = covs$lwt[i], lage = covs$lage[i])
    }))
    m2ll <- function(th) {
      tot <- 0
      for (i in seq_len(nsub)) {
        di <- data[data$id == i, ]
        p <- c(th, lwt = di$lwt[1], lage = di$lage[1])
        d <- as.data.frame(rxode2::rxSolve(tmod, params = p, cEv,
                                           returnType = "data.frame", addDosing = FALSE))
        f <- d$cp; vv <- th[["add"]]^2 + (th[["prop"]] * f)^2
        tot <- tot + sum((f - di$dv)^2 / vv + log(vv))
      }
      tot
    }
    gFD <- vapply(cCS, function(pn) {
      h <- abs(cTh[[pn]]) * 1e-6; p1 <- cTh; p2 <- cTh
      p1[pn] <- p1[pn] + h; p2[pn] <- p2[pn] - h
      (m2ll(p1) - m2ll(p2)) / (2 * h)
    }, numeric(1))
    gPop <- rxode2::.rxAdjointGradPop(cText, cTh, cEv, cCS, "center/v", data,
                                      cErr, denseBy = 0.005)
    expect_equal(unname(gPop), unname(gFD), tolerance = 1e-2)
  })

  # ---- capstone: adjoint gradient drives a gradient-based fit (nlm-style) -----
  # Proves the functional-gradient adjoint is usable as the ONLY gradient source
  # for a BFGS optimisation that recovers the data-generating parameters.
  test_that("adjoint objective gradient drives a BFGS fit to recover parameters", {
    truePar <- c(ka = 1.2, cl = 3.5, v = 25.0)
    oPred <- "center/v"
    oObsT <- c(0.5, 1, 2, 4, 6, 8, 12, 18, 24)
    oEv   <- rxode2::et(amt = 100, cmt = "depot")
    tmod  <- rxode2::rxode2(paste0(mText, "\ncp=", oPred))
    oFev  <- oEv |> rxode2::et(oObsT)
    set.seed(1)
    truth <- as.data.frame(rxode2::rxSolve(tmod, params = truePar, oFev,
                                           returnType = "data.frame", addDosing = FALSE))
    oObs  <- truth$cp[truth$time %in% oObsT] * (1 + stats::rnorm(length(oObsT), 0, 0.02))

    obj <- function(lp) {
      p <- stats::setNames(exp(lp), c("ka", "cl", "v"))
      d <- as.data.frame(rxode2::rxSolve(tmod, params = p, oFev,
                                         returnType = "data.frame", addDosing = FALSE))
      sum(0.5 * (d$cp[d$time %in% oObsT] - oObs)^2)
    }
    gr <- function(lp) {
      p <- stats::setNames(exp(lp), c("ka", "cl", "v"))
      rxode2::.rxAdjointGrad(mText, p, oEv, c("ka", "cl", "v"), oPred, oObsT, oObs,
                             denseBy = 0.02) * p   # chain rule for log-params
    }
    fit <- stats::optim(log(c(2, 5, 15)), obj, gr, method = "BFGS",
                        control = list(reltol = 1e-8))
    expect_equal(unname(exp(fit$par)), unname(truePar), tolerance = 0.05)
  })
})