tests/testthat/test-vae-covariate-estimate.R

In nlmixr2est: Nonlinear Mixed Effects Models in Population PK/PD, Estimation Routines

## With covariateSelection=FALSE the covariate coefficients WRITTEN in the model
## (linear beta*WT and transformed beta*log(WT/70)) must be estimated by the
## regress M-step regardless of nonMuTheta -- previously they were held fixed
## under "none" and errored under "fix"/"eta".  Covariate identity is read from
## the shared muRefCovariateDataFrame/allCovs (never mutated).  Fast: no fit; the
## end-to-end "coefficient moves" checks live in the slow
## test-vae-covariate-selection.R.

nmTest({
  ## linear (rxode2 records it in muRefCovariateDataFrame) covariate on cl
  .lin <- function() {
    ini({ tka <- 0.45; tcl <- 1; tv <- 3.45; cl.wt <- 0.01; add.err <- 0.7; eta.cl ~ 0.1 })
    model({ ka <- exp(tka); cl <- exp(tcl + cl.wt * WT + eta.cl); v <- exp(tv)
      d/dt(depot) <- -ka * depot; d/dt(center) <- ka * depot - cl / v * center
      cp <- center / v; cp ~ add(add.err) })
  }
  ## transformed (NOT recognized as a mu-ref covariate) effect on cl
  .tr <- function() {
    ini({ tka <- 0.45; tcl <- 1; tv <- 3.45; cl.wt <- 0.1; add.err <- 0.7; eta.cl ~ 0.1 })
    model({ ka <- exp(tka); cl <- exp(tcl + cl.wt * log(WT / 70) + eta.cl); v <- exp(tv)
      d/dt(depot) <- -ka * depot; d/dt(center) <- ka * depot - cl / v * center
      cp <- center / v; cp ~ add(add.err) })
  }
  d <- nlmixr2data::theo_sd

  test_that(".vaeCovariateCoefThetas detects linear and transformed coefficients", {
    expect_equal(.vaeCovariateCoefThetas(rxode2::assertRxUi(.lin())), "cl.wt")
    expect_equal(.vaeCovariateCoefThetas(rxode2::assertRxUi(.tr())), "cl.wt")
  })

  test_that(".vaeCovariateCoefThetas drops a user-fixed coefficient", {
    fx <- function() {
      ini({ tka <- 0.45; tcl <- 1; tv <- 3.45; cl.wt <- fix(0.01); add.err <- 0.7; eta.cl ~ 0.1 })
      model({ ka <- exp(tka); cl <- exp(tcl + cl.wt * WT + eta.cl); v <- exp(tv)
        d/dt(depot) <- -ka * depot; d/dt(center) <- ka * depot - cl / v * center
        cp <- center / v; cp ~ add(add.err) })
    }
    expect_equal(.vaeCovariateCoefThetas(rxode2::assertRxUi(fx())), character(0))
  })

  test_that(".vaeNonMuThetas excludes covariate coefficients (they are not structural)", {
    ## cl.wt must NOT be offered to nonMuTheta eta/fix injection
    expect_false("cl.wt" %in% .vaeNonMuThetas(rxode2::assertRxUi(.lin())))
    expect_false("cl.wt" %in% .vaeNonMuThetas(rxode2::assertRxUi(.tr())))
  })

  test_that("covariateSelection=FALSE regresses the coefficient in every nonMuTheta mode", {
    for (m in c("regress", "none", "fix", "eta")) {
      p <- suppressWarnings(.vaeDataPrep(rxode2::assertRxUi(.tr()), d,
                                         vaeControl(covariateSelection = FALSE, nonMuTheta = m)))
      expect_true("cl.wt" %in% p$regressNames, info = m)
      i <- match("cl.wt", p$regressNames)
      expect_true(is.finite(p$regressLower[i]) && is.finite(p$regressUpper[i]), info = m)
    }
  })

  test_that("covariateSelection=TRUE does NOT force the coefficient into the regress set", {
    ## default nonMuTheta='regress' still picks up genuine non-mu structural thetas
    ## (tka, tv) but the coefficient is left to the selection machinery
    p <- suppressWarnings(.vaeDataPrep(rxode2::assertRxUi(.tr()), d,
                                       vaeControl(covariateSelection = TRUE, nonMuTheta = "none")))
    expect_false("cl.wt" %in% p$regressNames)
  })

  test_that("a fixed coefficient is not regressed", {
    fx <- function() {
      ini({ tka <- 0.45; tcl <- 1; tv <- 3.45; cl.wt <- fix(0.01); add.err <- 0.7; eta.cl ~ 0.1 })
      model({ ka <- exp(tka); cl <- exp(tcl + cl.wt * WT + eta.cl); v <- exp(tv)
        d/dt(depot) <- -ka * depot; d/dt(center) <- ka * depot - cl / v * center
        cp <- center / v; cp ~ add(add.err) })
    }
    p <- suppressWarnings(.vaeDataPrep(rxode2::assertRxUi(fx()), d,
                                       vaeControl(covariateSelection = FALSE, nonMuTheta = "none")))
    expect_false("cl.wt" %in% p$regressNames)
  })

  test_that("the unbounded fallback bound is scale-aware (tighter for a raw covariate)", {
    ## raw WT (~O(70)) => bound ~ 1/max|WT|, much tighter than the log-scale default
    pLin <- suppressWarnings(.vaeDataPrep(rxode2::assertRxUi(.lin()), d,
                                          vaeControl(covariateSelection = FALSE, nonMuTheta = "none")))
    pTr <- suppressWarnings(.vaeDataPrep(rxode2::assertRxUi(.tr()), d,
                                         vaeControl(covariateSelection = FALSE, nonMuTheta = "none")))
    bLin <- pLin$regressUpper[match("cl.wt", pLin$regressNames)]
    bTr <- pTr$regressUpper[match("cl.wt", pTr$regressNames)]
    expect_lt(bLin, bTr)
    expect_equal(bLin, .vaeCovCoefEffect / max(abs(d$WT)))
  })

  test_that("the fit surfaces a 'estimating covariate coef(s)' note and never errors", {
    for (m in c("none", "fix", "eta", "regress")) {
      w <- character(0)
      withCallingHandlers(
        .preProcessVaeNonMuTheta(rxode2::assertRxUi(.tr()), "vae", d,
                                 vaeControl(covariateSelection = FALSE, nonMuTheta = m)),
        warning = function(cnd) { w <<- c(w, conditionMessage(cnd)); invokeRestart("muffleWarning") })
      expect_true(any(grepl("estimating covariate coef", w)), info = m)
    }
  })

  ## ---- pinCovariates: restrict the search to model-declared covariate pairs ----

  test_that(".vaeModelCovariatePairs pairs a transformed continuous covariate", {
    ui <- rxode2::assertRxUi(.tr())               # cl.wt * log(WT/70) on cl
    dd <- as.data.frame(d); names(dd) <- toupper(names(dd))
    cov <- .vaeCovariateSearch(dd, unique(dd$ID))
    pr <- .vaeModelCovariatePairs(ui, cov)
    expect_equal(nrow(pr), 1L)
    expect_equal(pr$coefName, "cl.wt")
    expect_equal(pr$covName, "WT")
    expect_equal(pr$thetaName, "tcl")
    expect_equal(pr$covType, "continuous")
    expect_equal(pr$userCenter, 70)               # log(WT/70) center read from the model
    expect_true(pr$inPool)
  })

  test_that("pinCovariates=TRUE builds a 1-cell mask and zeros the training coef", {
    ui <- rxode2::assertRxUi(.tr())               # single eta (cl)
    p <- suppressWarnings(.vaeDataPrep(ui, d, vaeControl(pinCovariates = TRUE)))
    expect_true(p$pinActive)
    ## WT allowed only on the cl dim; exactly one 1 in the mask
    expect_equal(sum(p$covAllow), 1L)
    k <- match("tcl", .foceiEtaThetaMap(ui)$thetaForEta)
    expect_equal(p$covAllow[k, match("WT", p$covRaw)], 1L)
    ## in-pool coefficient held at 0 for the covariate-free training decoder
    thNames <- ui$iniDf$name[!is.na(ui$iniDf$ntheta)]
    expect_equal(unname(p$th[match("cl.wt", thNames)]), 0)
    ## in-pool coefficient is NOT regressed (the search estimates it)
    expect_false("cl.wt" %in% p$regressNames)
  })

  test_that("a raw-linear effect on a continuous covariate pins as the identity shape", {
    ## Previously untransferable (the search only understood log(cov/center)),
    ## so this was routed to the regress M-step.  `beta*WT` is now the identity
    ## shape, so it pins to the linear-family column and IS searched.
    ui <- rxode2::assertRxUi(.lin())              # cl.wt * WT (linear on continuous WT)
    p <- suppressWarnings(.vaeDataPrep(ui, d, vaeControl(pinCovariates = TRUE)))
    expect_true(p$pinActive)
    expect_equal(p$pinPairs$shape, "identity")
    expect_true(all(p$pinPairs$inPool))
    expect_false("cl.wt" %in% p$regressNames)      # searched, not regressed
    ## exactly one cell allowed, on the linear-family column
    expect_equal(sum(p$covAllow), 1L)
    .j <- which(colSums(p$covAllow) > 0L)
    expect_equal(p$covFamily[.j], "lin")
    ## identity means the column is the raw covariate, so the slope transfers
    .wt <- vapply(unique(d$ID), function(i) d$WT[d$ID == i][1], numeric(1))
    expect_equal(unname(p$covMat[, .j]), unname(.wt))
    expect_equal(p$covPop[.j], 0)
  })

  test_that("a genuinely untransferable form still regresses", {
    ## sqrt(WT) is not one of the shapes the search can write back, so it must
    ## NOT be mistaken for the identity shape just because WT appears in it
    sq <- function() {
      ini({ tka <- 0.45; tcl <- 1; tv <- 3.45; cl.wt <- 0.1; add.err <- 0.7
        eta.cl ~ 0.1 })
      model({ ka <- exp(tka); cl <- exp(tcl + cl.wt * sqrt(WT) + eta.cl)
        v <- exp(tv)
        d/dt(depot) <- -ka * depot
        d/dt(center) <- ka * depot - cl / v * center
        cp <- center / v; cp ~ add(add.err) })
    }
    p <- suppressWarnings(.vaeDataPrep(rxode2::assertRxUi(sq()), d,
                                       vaeControl(pinCovariates = TRUE)))
    expect_false(any(p$pinPairs$inPool))
    expect_true("cl.wt" %in% p$regressNames)
    expect_equal(sum(p$covAllow), 0L)
  })

  test_that("pinCovariates=FALSE turns the search off and regresses model covariates", {
    ui <- rxode2::assertRxUi(.tr())
    w <- character(0)
    p <- withCallingHandlers(
      .vaeDataPrep(ui, d, vaeControl(pinCovariates = FALSE)),
      warning = function(cnd) { w <<- c(w, conditionMessage(cnd)); invokeRestart("muffleWarning") })
    expect_false(p$pinActive)
    expect_equal(ncol(p$covMat), 0L)               # search pool emptied
    expect_true("cl.wt" %in% p$regressNames)
    expect_true(any(grepl("pinCovariates=FALSE", w)))
  })

  test_that("a multi-covariate transformed line still pins (never falls back to full search)", {
    ## two covariate effects in one parameter line; with muRefCovAlg off they stay
    ## as transformed effects.  Each coefficient must be paired with ITS covariate
    ## (not skipped), so pinning stays restrictive.
    multi <- function() {
      ini({ tka <- 0.45; tcl <- 1; tv <- 3.45; wt.cl <- 0.1; age.cl <- 0.1
            add.err <- 0.7; eta.cl ~ 0.1 })
      model({ ka <- exp(tka)
        cl <- exp(tcl + wt.cl * log(WT / 70) + age.cl * log(AGE / 40) + eta.cl); v <- exp(tv)
        d/dt(depot) <- -ka * depot; d/dt(center) <- ka * depot - cl / v * center
        cp <- center / v; cp ~ add(add.err) })
    }
    dd <- as.data.frame(d); names(dd) <- toupper(names(dd))
    .ageById <- stats::setNames(rep(c(30, 40, 55), length.out = length(unique(dd$ID))),
                                as.character(unique(dd$ID)))
    dd$AGE <- .ageById[as.character(dd$ID)]          # subject-constant
    ui <- rxode2::assertRxUi(multi())
    cov <- .vaeCovariateSearch(dd, unique(dd$ID))
    pr <- .vaeModelCovariatePairs(ui, cov)
    ## both coefficients detected and paired with their own covariate
    expect_setequal(pr$coefName, c("wt.cl", "age.cl"))
    expect_equal(pr$covName[pr$coefName == "wt.cl"], "WT")
    expect_equal(pr$covName[pr$coefName == "age.cl"], "AGE")
    expect_true(all(pr$inPool))                     # both are clean log forms
  })

  test_that("one covariate pinned on two parameters adjusts its column only once", {
    ## same covariate declared on two etas: the covariate column must be
    ## un-centered exactly once (a second pass would take log(0) -> -Inf)
    two <- function() {
      ini({ tka <- 0.45; tcl <- 1; tv <- 3.45; wt.ka <- 0.1; wt.cl <- 0.1
            add.err <- 0.7; eta.ka ~ 0.1; eta.cl ~ 0.1 })
      model({ ka <- exp(tka + wt.ka * log(WT / 70) + eta.ka)
        cl <- exp(tcl + wt.cl * log(WT / 70) + eta.cl); v <- exp(tv)
        d/dt(depot) <- -ka * depot; d/dt(center) <- ka * depot - cl / v * center
        cp <- center / v; cp ~ add(add.err) })
    }
    p <- suppressWarnings(.vaeDataPrep(rxode2::assertRxUi(two()), d,
                                       vaeControl(pinCovariates = TRUE,
                                                  muRefCovAlg = FALSE)))
    jWT <- match("WT", p$covRaw)
    expect_true(all(is.finite(p$covMat[, jWT])))     # never -Inf
    expect_equal(p$covPop[jWT], 0)                   # uncentered exactly once
    ## column is log(WT/70) -- the model's centering, applied once
    expect_equal(unname(p$covMat[1, jWT]), log(d$WT[1] / 70))
    ## both coefficients pinned (same center, so no conflict/demotion)
    expect_true(all(p$pinPairs$inPool))
    expect_equal(sum(p$covAllow), 2L)
  })

  test_that("a mu2 centered covariate is pinned via its nlmixrMuDerCov# column", {
    ## wt.cl*(WT/70) is a mu2 reference; the hook rewrites it to a linear
    ## nlmixrMuDerCov# column (centering carried by the mu2 data).  After the
    ## hook the covariate must be encoded LINEARLY (not log) and pinned in place.
    mu2 <- function() {
      ini({ tka <- 0.45; tcl <- 1; tv <- 3.45; wt.cl <- 0.1; add.err <- 0.7; eta.cl ~ 0.1 })
      model({ ka <- exp(tka); cl <- exp(tcl + wt.cl * (WT / 70) + eta.cl); v <- exp(tv)
        d/dt(depot) <- -ka * depot; d/dt(center) <- ka * depot - cl / v * center
        cp <- center / v; cp ~ add(add.err) })
    }
    lst <- suppressWarnings(.uiModifyForCovs(rxode2::assertRxUi(mu2()), as.data.frame(d)))
    p <- suppressWarnings(.vaeDataPrep(lst$ui, lst$data, vaeControl(pinCovariates = TRUE)))
    ## the derived column is linear (categorical), never log-encoded
    jMu <- grep("NLMIXRMUDERCOV", p$covRaw)
    expect_length(jMu, 1L)
    expect_equal(p$covType[jMu], "categorical")
    ## pinned to the cl dim, in place, with no VAE-applied center (mu2 data carries it)
    expect_true(p$pinActive)
    pr <- p$pinPairs
    expect_equal(pr$coefName, "wt.cl")
    expect_true(pr$inPool)
    expect_equal(pr$userCenter, 0)
    ## searched (not regressed) and zeroed in the training theta
    expect_false("wt.cl" %in% p$regressNames)
  })

  ## ---- issue #801: a covariate reaching the coefficient line only through an
  ## intermediate model variable must still be treated as a covariate coefficient
  ## (estimated, never frozen), not mis-classified as a non-mu structural theta. ----
  .ind <- function() {
    ini({ tka <- 0.45; tcl <- 1; tv <- 3.45; cl.wt <- 0.1; add.err <- 0.7; eta.cl ~ 0.1 })
    model({ wt70 <- WT / 70
      ka <- exp(tka); cl <- exp(tcl + cl.wt * log(wt70) + eta.cl); v <- exp(tv)
      d/dt(depot) <- -ka * depot; d/dt(center) <- ka * depot - cl / v * center
      cp <- center / v; cp ~ add(add.err) })
  }

  test_that(".vaeCovariateCoefThetas detects a coefficient behind an intermediate var", {
    ## rxode2 records the coefficient in mu2RefCovariateReplaceDataFrame even when
    ## the covariate reaches the line via an intermediate (`wt70 <- WT/70`), so it
    ## is classified as a covariate coefficient, not a plain non-mu structural theta.
    ui <- rxode2::assertRxUi(.ind())
    expect_true("cl.wt" %in% ui$mu2RefCovariateReplaceDataFrame$covariateParameter)
    expect_equal(.vaeCovariateCoefThetas(ui), "cl.wt")
    ## and it is NOT offered to the nonMuTheta eta/fix injection (it errored before)
    expect_false("cl.wt" %in% .vaeNonMuThetas(ui))
  })

  test_that("a structural theta on/after a covariate line is not a covariate coef", {
    ## Two structural (non-random-effect) thetas that must NOT be swept into the
    ## covariate-coefficient set: `tka`, an additive intercept sharing beta.ka's
    ## covariate line; and `tlag`, a downstream multiplier of the (covariate-
    ## bearing) `ka`.  rxode2 classifies only `beta.ka` as the covariate
    ## coefficient (mu2RefCovariateReplaceDataFrame), so neither structural theta
    ## is mis-detected.
    struct <- function() {
      ini({ tka <- 0.45; beta.ka <- 0.1; tcl <- 1; tv <- 3.45; tlag <- 0.5
            add.err <- 0.7; eta.cl ~ 0.1 })
      model({ wt70 <- WT / 70
        ka <- exp(tka + beta.ka * log(wt70)); cl <- exp(tcl + eta.cl); v <- exp(tv)
        klag <- tlag * ka
        d/dt(depot) <- -klag * depot; d/dt(center) <- klag * depot - cl / v * center
        cp <- center / v; cp ~ add(add.err) })
    }
    ui <- rxode2::assertRxUi(struct())
    expect_equal(.vaeCovariateCoefThetas(ui), "beta.ka")   # only the real coefficient
    expect_true(all(c("tka", "tlag") %in% .vaeNonMuThetas(ui)))  # structural thetas kept
  })

  test_that("exotic covariate transforms are classified via the mu2 derivative check", {
    ## rxode2's mu2/mu3 classifier records a coefficient whenever the derivative
    ## of its additive linear-predictor term is free of the coefficient (the
    ## derivative becomes the nlmixrMuDerCov# column).  This covers arbitrary
    ## transforms, so reading mu2RefCovariateReplaceDataFrame needs no special
    ## cases -- confirm exotic forms are detected, and a non-mu form (a multiplier
    ## outside the transform, whose slope cannot transfer) is correctly excluded.
    mk <- function(clexpr) {
      eval(parse(text = sprintf(
        "function(){ ini({ tcl<-1; tv<-3.45; b<-0.1; eta.cl~0.1; add.err<-0.7 })\n model({ %s; v<-exp(tv)\n d/dt(depot) <- -cl*depot; cp<-depot/v; cp~add(add.err) }) }",
        clexpr)))
    }
    for (e in c("cl <- exp(tcl + b*sqrt(WT) + eta.cl)",
                "cl <- exp(tcl + b*(WT/70)^2 + eta.cl)",
                "cl <- exp(tcl + b*exp(WT/100) + eta.cl)",
                "cl <- exp(tcl + b*(WT - 70) + eta.cl)")) {
      expect_equal(.vaeCovariateCoefThetas(rxode2::assertRxUi(mk(e))), "b", info = e)
    }
    ## multiplier outside the transform: not a mu reference, correctly not a coef
    expect_equal(
      .vaeCovariateCoefThetas(rxode2::assertRxUi(mk("cl <- exp(tcl + eta.cl)*(1 + b*WT)"))),
      character(0))
  })

  test_that("an indirect covariate coefficient is estimated in every nonMuTheta mode", {
    ## the #801 guarantee: declared covariate effects are never frozen, regardless
    ## of nonMuTheta.  Before the fix it was regressed only under "regress"/"grad",
    ## frozen under "none", and errored under "eta"/"fix".
    for (m in c("regress", "none", "eta", "fix")) {
      p <- suppressWarnings(.vaeDataPrep(rxode2::assertRxUi(.ind()), d,
                                         vaeControl(nonMuTheta = m)))
      expect_true("cl.wt" %in% p$regressNames, info = m)
    }
  })

  test_that("pinCovariates=TRUE with no model covariates is a null path (full search)", {
    noCov <- function() {
      ini({ tka <- 0.45; tcl <- 1; tv <- 3.45; add.err <- 0.7; eta.cl ~ 0.1 })
      model({ ka <- exp(tka); cl <- exp(tcl + eta.cl); v <- exp(tv)
        d/dt(depot) <- -ka * depot; d/dt(center) <- ka * depot - cl / v * center
        cp <- center / v; cp ~ add(add.err) })
    }
    p <- suppressWarnings(.vaeDataPrep(rxode2::assertRxUi(noCov), d, vaeControl(pinCovariates = TRUE)))
    expect_false(p$pinActive)
    expect_null(p$covAllow)                         # no mask -> unrestricted search
    ## WT still discovered from data, now as one search column per shape family
    ## (and one per ARM for the hockey family)
    expect_equal(unique(p$covRaw), "WT")
    expect_equal(p$covNames,
                 c("WT_power", "WT_lin", "WT_hockeyLow", "WT_hockeyHi"))
    ## every shape shares an exclusion group, so only one can ever be selected
    expect_equal(length(unique(p$covGroup)), 1L)
    ## and the two hockey arms are one block inside that group
    expect_equal(p$covBlock, c(1L, 2L, 3L, 3L))
  })
})