R/methods.R

In RTMB: 'R' Bindings for 'TMB'

################################################################################
## This file contains:
## - Methods for sparse matrices
## - Methods for dense matrices
## - Hand written distributions (Autogenerated are in 'distributions.R')
################################################################################

setMethod("show", "advector", function(object) print.advector(object) )

setAs("sparseMatrix", "adsparse",
      function(from) {
          x <- from
          x <- as(x, "generalMatrix")
          x <- as(x, "CsparseMatrix")
          new("adsparse", x=advector(x@x), i=x@i, p=x@p, Dim=x@Dim)
      })

setAs("advector", "sparseMatrix",
      function(from) Dense2Sparse(as.matrix(from)) )

################################################################################
## Utilities to reuse methods from the Matrix package
################################################################################
IndexMap <- function(x) new("dgCMatrix",
                            i=x@i, p=x@p, Dim=x@Dim, x=as.numeric(seq_along(x@x)))
ApplyMatrixMethod <- function(method, x, ...) {
    if (is.character(method))
        method <- match.fun(method)
    ans <- method(IndexMap(x), ...)
    if (!is.null(dim(ans))) {
        ans <- as(as(ans, "generalMatrix"), "sparseMatrix")
        new("adsparse", i=ans@i, p=ans@p, Dim=ans@Dim, x=x@x[ans@x])
    }
    else
        c(advector(0), x@x)[ans + 1L]
}
ApplyMatrixReplaceMethod <- function(method, x, ..., value) {
    method <- match.fun(method)
    xmap <- IndexMap(x)
    vmap <- seq_along(value) + length(x@x)
    ans <- method(xmap, ..., value=vmap)
    if (is.null(dim(ans))) stop("unexpected")
    ans <- as(as(ans, "generalMatrix"), "sparseMatrix")
    new("adsparse", i=ans@i, p=ans@p, Dim=ans@Dim,
        x=c(advector(0), x@x, value)[ans@x + 1L])
}

##' @describeIn ADmatrix AD sparse matrix transpose. Re-directs to \link[Matrix]{t,CsparseMatrix-method}.
##' @return Object of class \code{advector} with a dimension attribute for dense matrix operations; Object of class \code{adsparse} for sparse matrix operations.
t.adsparse <- function(x) ApplyMatrixMethod(Matrix::t, x)
##' @describeIn ADmatrix AD sparse matrix subsetting. Re-directs to \link[Matrix]{[-methods}.
"[.adsparse" <- function(x, ...) ApplyMatrixMethod("[", x, ...)
##' @describeIn ADmatrix AD sparse matrix subset assignment. Re-directs to \link[Matrix]{[<--methods}.
##' @param value Replacement value
"[<-.adsparse" <- function(x, ..., value) {
    x <- ApplyMatrixReplaceMethod("[<-", x, ..., value=value)
    x
}
##' @describeIn ADmatrix AD sparse matrix diagonal extract. Re-directs to \link[Matrix]{diag,CsparseMatrix-method}.
setMethod("diag", c("adsparse", "missing", "missing"), function(x) ApplyMatrixMethod("diag", x) )

##setClassUnion("advector_castable", c("advector", "numeric"))

##' @describeIn ADmatrix AD matrix exponential
setMethod("expm", "advector", function(x) math_expm(x))
##' @describeIn ADmatrix AD matrix exponential
setMethod("expm", "adsparse", function(x) math_expm(x))

## Methods sparseMatrix -> adsparse

setMethod("Ops",
          signature("sparseMatrix", "advector"),
          function(e1, e2) callGeneric( as(e1, "adsparse") , e2) )

setMethod("Ops",
          signature("advector", "sparseMatrix"),
          function(e1, e2) callGeneric( e1, as(e2, "adsparse") ) )

setMethod("Ops",
          signature("sparseMatrix", "adsparse"),
          function(e1, e2) callGeneric( as(e1, "adsparse") , e2) )

setMethod("Ops",
          signature("adsparse", "sparseMatrix"),
          function(e1, e2) callGeneric( e1, as(e2, "adsparse") ) )

## Methods adsparse
##' @describeIn ADmatrix AD sparse matrix dimension
setMethod("dim", "adsparse", function(x) x@Dim)

setMethod("Ops",
          signature("ad", "adsparse"),
          function(e1, e2) SparseArith2(advector(e1), e2, .Generic) )
setMethod("Ops",
          signature("adsparse", "ad"),
          function(e1, e2) SparseArith2(e1, advector(e2), .Generic) )
setMethod("Ops",
          signature("adsparse", "adsparse"),
          function(e1, e2) {
              if (!identical(e1@Dim, e2@Dim))
                  stop("non-conformable arguments")
              SparseArith2(e1, e2, .Generic)
          })
setMethod("Ops",
          signature("adsparse", "missing"),
          function(e1, e2) {
              e1@x[] <- callGeneric(e1@x)
              e1
          })
##' @describeIn ADmatrix AD matrix multiply
setMethod("%*%",
          signature("anysparse", "ad"),
          function(x, y) {
              x <- as(x, "adsparse")
              y <- as.matrix(advector(y))
              if ( ncol(x) != nrow(y) )
                  stop("non-conformable arguments")
              SparseArith2(x, y, .Generic)
          })
##' @describeIn ADmatrix AD matrix multiply
setMethod("%*%",
          signature("ad", "anysparse"),
          function(x, y) {
              x <- as.matrix(advector(x))
              y <- as(y, "adsparse")
              if ( ncol(x) != nrow(y) )
                  stop("non-conformable arguments")
              SparseArith2(x, y, .Generic)
          })
##' @describeIn ADmatrix AD matrix multiply
setMethod("%*%",
          signature("adsparse", "adsparse"),
          function(x, y) {
              if ( ncol(x) != nrow(y) )
                  stop("non-conformable arguments")
              SparseArith2(x, y, .Generic)
          })
##' @describeIn ADmatrix AD matrix multiply
setMethod("%*%",
          signature("ad", "ad"),
          function(x, y) {
              y <- as.matrix(y)
              ## Promotion of a vector to a 1-row matrix
              if (is.null(dim(x)) && length(x) == nrow(y))
                  x <- t(x)
              x <- as.matrix(x)
              matmul(advector(x), advector(y))
          })
##' @describeIn ADmatrix AD matrix multiply
setMethod("tcrossprod", signature("ad", "ad."),
          function(x, y) {if (is.null(y)) y <- x; x %*% t(y)} )
##' @describeIn ADmatrix AD matrix multiply
setMethod( "crossprod", signature("ad", "ad."),
          function(x, y) {if (is.null(y)) y <- x; t(x) %*% y} )
##' @describeIn ADmatrix AD matrix cov2cor
##' @param V Covariance matrix
setMethod( "cov2cor", signature("advector"),
          function(V) {
              oldval <- TapeConfig()["comparison"]
              on.exit(TapeConfig(comparison=oldval))
              TapeConfig(comparison="allow")
              stats::cov2cor(V)
          })
##' @describeIn ADmatrix AD matrix inversion and solve
##' @param a matrix
##' @param b matrix, vector or missing
setMethod("solve",
          signature("ad", "ad."),
          function(a, b) {
              a <- as.matrix(a)
              ans <- matinv(advector(a))
              if (!missing(b)) {
                  b <- as.matrix(b)
                  ans <- ans %*% b
              }
              ans
          })
##' @describeIn ADmatrix AD matrix inversion and solve
setMethod("solve", signature("num", "num."),
          function(a, b) {
              base::solve(a, b)
          })
##' @describeIn ADmatrix Sparse AD matrix solve (not yet implemented)
##' @param a matrix
##' @param b matrix, vector or missing
setMethod("solve",
          signature("anysparse", "ad."),
          function(a, b) {
              stop("Sparse AD solve is not yet implemented")
          })
##' @describeIn ADmatrix AD matrix (or array) colsums
##' @param na.rm Logical; Remove NAs while taping.
##' @param dims Same as \link[base]{colSums} and \link[base]{rowSums}.
setMethod("colSums", signature(x="advector"),
          function(x, na.rm, dims) {
              if (dims != 1L) stop("AD version requires dims=1")
              apply(x, seq_len(length(dim(x)))[-1L], sum, na.rm=na.rm)
          } )
##' @describeIn ADmatrix AD matrix (or array) rowsums
setMethod("rowSums", signature("advector"),
          function(x, na.rm, dims) {
              if (dims != 1L) stop("AD version requires dims=1")
              apply(x, 1L, sum, na.rm=na.rm)
          } )
##' @describeIn ADmatrix AD matrix column bind
##' @param ... As \link[base]{cbind}
cbind.advector <- function (...) {
    args <- lapply(list(...), advector)
    ans <- do.call("cbind", lapply(args, unclass))
    class(ans) <- "advector"
    asS4(ans)
}
##' @describeIn ADmatrix AD matrix row bind
rbind.advector <- function (...) {
    args <- lapply(list(...), advector)
    ans <- do.call("rbind", lapply(args, unclass))
    class(ans) <- "advector"
    asS4(ans)
}

## Show general idea which is automated in 'distributions.R'
## First we generate the version we want for AD types (dot signifies 'default argument')
##' @describeIn Distributions AD implementation of \link[stats]{dnorm}
setMethod("dnorm", signature("ad", "ad.", "ad.", "logical."),
          function(x, mean, sd, log) {
              r <- (x - mean) / sd
              ans <- - .5 * r * r - log(sqrt(2*pi)) - log(sd)
              if (log) ans else exp(ans)
          })
## This matches 'too much', so we fix by adding a specialization:
##' @describeIn Distributions Default method
setMethod("dnorm", signature("num", "num.", "num.", "logical."),
          function(x, mean, sd, log) {
              stats::dnorm(x, mean, sd, log)
          })
## For S4 generics we add the OSA version like this:
##' @describeIn Distributions OSA implementation
setMethod("dnorm", "osa", function(x, mean, sd, log) {
    dGenericOSA(.Generic, x=x, mean=mean, sd=sd, log=log)
})
## For S4 generics we add the simref version like this:
##' @describeIn Distributions Simulation implementation. Modifies \code{x} and returns zero.
setMethod("dnorm", "simref", function(x, mean, sd, log) {
    ## works when x, mean or sd are simref
    if (inherits(mean, "simref")) {
        x <- x - mean
        mean <- 0
    }
    if (inherits(sd, "simref")) {
        x <- x / sd
        sd <- 1
    }
    dGenericSim(.Generic, x=x, mean=mean, sd=sd, log=log)
})

##' @describeIn Distributions AD implementation of \link[stats]{dlnorm}.
##' @param meanlog Parameter; Mean on log scale.
##' @param sdlog Parameter; SD on log scale.
setMethod("dlnorm", "ANY", function (x, meanlog, sdlog, log) {
    y <- log(x)
    ans <- dnorm(y, meanlog, sdlog, log=TRUE) - y
    ans <- ans[] ## if 'simref' do complete simulation
    if (log) ans else exp(ans)
})
##' @describeIn Distributions OSA implementation.
setMethod("dlnorm", "osa", function (x, meanlog, sdlog, log) {
    dGenericOSA(.Generic, x=x, meanlog=meanlog, sdlog=sdlog, log=log)
})
##' @describeIn Distributions Default method.
setMethod("dlnorm", signature("num", "num.", "num.", "logical."),
          function(x, meanlog, sdlog, log) {
              stats::dlnorm(x, meanlog, sdlog, log)
          })

##' @describeIn Distributions Minimal AD implementation of \link[stats]{plogis}
setMethod("plogis", c("advector", "missing", "missing", "missing", "missing"),
          function(q) 1 / (1 + exp(-q) ) )
##' @describeIn Distributions Minimal AD implementation of \link[stats]{qlogis}
setMethod("qlogis", c("advector", "missing", "missing", "missing", "missing"),
          function(p) log( p / ( 1 - p ) ) )

## 'diag' needs patching.
## - base::diag works fine for AD matrix input (diagonal extraction and replacement)
## - However, matrix construction has issues

##' @describeIn ADconstruct Equivalent of \link[base]{diag}
##' @param x As \link[base]{diag}
##' @return Object of class \code{"advector"} with a dimension attribute.
setMethod("diag", signature(x="advector", nrow="ANY", ncol="ANY"),
          function(x, nrow, ncol) {
              ## Diagonal extraction: base::diag works fine
              if (length(dim(x)) >= 2)
                  return(callNextMethod())
              ## Matrix creation
              ans <- advector(base::diag(seq_along(x), nrow=nrow, ncol=ncol))
              diag(ans) <- x
              ans
          })

##' @describeIn ADconstruct Equivalent of \link[base]{matrix}
##' @param data As \link[base]{matrix}
##' @param nrow As \link[base]{matrix}
##' @param ncol As \link[base]{matrix}
##' @param byrow As \link[base]{matrix}
##' @param dimnames As \link[base]{matrix}
setMethod("matrix", signature(data="advector"),
          function(data, nrow, ncol, byrow, dimnames) {
              ans <- callNextMethod()
              asS4(structure(ans, class="advector"))
          })
##' @describeIn ADconstruct Equivalent of \link[base]{matrix}
setMethod("matrix", signature(data="num."),
          function(data, nrow, ncol, byrow, dimnames) {
              ans <- callNextMethod()
              if (ad_context())
                  ans <- advector(ans)
              ans
          })
##' @describeIn ADapply As \link[base]{apply}
##' @param X As \link[base]{apply}
##' @param MARGIN As \link[base]{apply}
##' @param FUN As \link[base]{apply}
##' @param ... As \link[base]{apply}
##' @return Object of class \code{"advector"} with a dimension attribute.
setMethod("apply", signature(X="advector"),
          function (X, MARGIN, FUN, ...)  {
              ans <- callNextMethod()
              if (is.complex(ans))
                  class(ans) <- "advector"
              ans
          })
##' @describeIn ADapply As \link[base]{sapply}
##' @param simplify As \link[base]{sapply}
##' @param USE.NAMES As \link[base]{sapply}
setMethod("sapply", signature(X="ANY"),
          function (X, FUN, ..., simplify = TRUE, USE.NAMES = TRUE) {
              ans <- base::sapply(X, FUN, ..., simplify = FALSE, USE.NAMES = TRUE)
              ## Adapted from 'base::sapply':
              if (!isFALSE(simplify)) {
                  cl <- class(ans[[1L]])
                  ans <- simplify2array(ans, higher = (simplify == "array"))
                  if (identical(cl, "advector")) ## FIXME: Test all elements
                      class(ans) <- cl
              }
              ans
          })

##' @describeIn ADvector Equivalent of \link[base]{ifelse}
##' @param test \code{logical} vector
##' @param yes \code{advector}
##' @param no \code{advector}
setMethod("ifelse", signature(test="num", yes="ad", no="ad"),
          function(test, yes, no) {
              yes <- advector(yes)
              no <- advector(no)
              ans <- callNextMethod()
              class(ans) <- "advector"
              ans
          })
##' @describeIn ADvector Default method
setMethod("ifelse", signature(test="num", yes="num", no="num"),
          function(test, yes, no) {
              base::ifelse(test, yes, no)
          })

################################################################################

##' @describeIn Distributions Conway-Maxwell-Poisson. Calculate density.
dcompois <- function(x, mode, nu, log = FALSE) {
    if (inherits(x,"osa"))    return (dGenericOSA( "dcompois" , x=x, mode=mode, nu=nu, log=log ))
    if (inherits(x,"simref")) return (dGenericSim( "dcompois" , x=x, mode=mode, nu=nu, log=log ))
    loglambda <- nu * log(mode)
    ans <- x * loglambda - nu * lfactorial(x)
    ans <- ans - compois_calc_logZ(loglambda, nu)
    if (log) ans else exp(ans)
}
## internal (simref)
rcompois <- function(n, mode, nu) {
    loglambda <- nu * log(mode)
    loglambda <- rep(loglambda, length.out=n)
    mapply(distr_rcompois, loglambda, nu) ## mapply re-cycles
}

##' @describeIn Distributions Conway-Maxwell-Poisson. Calculate density parameterized via the mean.
dcompois2 <- function(x, mean, nu, log = FALSE) {
    if (inherits(x,"osa"))    return (dGenericOSA( "dcompois2" , x=x, mean=mean, nu=nu, log=log ))
    if (inherits(x,"simref")) return (dGenericSim( "dcompois2" , x=x, mean=mean, nu=nu, log=log ))
    logmean <- log(mean)
    loglambda <- compois_calc_loglambda(logmean, nu)
    ans <- x * loglambda - nu * lfactorial(x)
    ans <- ans - compois_calc_logZ(loglambda, nu)
    if (log) ans else exp(ans)
}
## internal (simref)
rcompois2 <- function(n, mean, nu) {
    logmean <- log(mean)
    loglambda <- getValues(compois_calc_loglambda(advector(logmean), advector(nu)))
    loglambda <- rep(loglambda, length.out=n)
    mapply(distr_rcompois, loglambda, nu) ## mapply re-cycles
}

################################################################################

##' @describeIn Distributions AD implementation of \link[stats]{pbinom}
setMethod("pbinom",
          signature(q = "ad", size = "ad", prob = "ad",
                    lower.tail = "missing", log.p = "missing"),
          function( q, size, prob ) {
              1 - pbeta(q = prob, shape1 = q + 1, shape2 = size - q)
          })
##' @describeIn Distributions Default method
setMethod("pbinom",
          signature(q = "num", size = "num", prob = "num",
                    lower.tail = "missing", log.p = "missing"),
          function( q, size, prob ) {
              stats:: pbinom ( q=q, size=size, prob=prob )
          })

################################################################################

## 'dmultinom' is already in R so we use S4 method to enhance it.
## ?dmultinom ---> 'dmultinom is currently _not vectorized_ at all'
## First we generate the version we want for AD types (dot signifies 'default argument')
##' @describeIn Distributions AD implementation of \link[stats]{dmultinom}
setMethod("dmultinom", signature("ad", "ad.", "ad", "logical."),
          function(x, size, prob, log) {
              ## Copy/paste from stats::dmultinom and remove if/else branching
              K <- length(prob)
              if (length(x) != K)
                  stop("x[] and prob[] must be equal length vectors.")
              s <- sum(prob)
              prob <- prob / s
              if (is.null(size))
                  size <- sum(x)
              r <- lgamma(size + 1) + sum(x * log(prob) - lgamma(x + 1))
              if (log)
                  r
              else exp(r)
          })
## This matches 'too much', so we fix by adding a specialization:
##' @describeIn Distributions Default method
setMethod("dmultinom", signature("num", "num.", "num", "logical."),
          function(x, size, prob, log) {
              stats::dmultinom(x, size, prob, log)
          })
## For S4 generics we add the OSA version like this:
##' @describeIn Distributions OSA implementation
setMethod("dmultinom", "osa", function(x, size, prob, log) {
    prob <- prob / sum(prob)
    if (is.null(size)) {
        size <- sum(x@x)
    }
    ## Factorize in succesive binomials
    perm <- order(attr(x@keep, "ord")) ## FIXME: Make extractor in osa.R ?
    x <- x[perm]
    prob <- prob[perm]
    ## Binomial parameters
    "c" <- ADoverload("c")
    cumsum0 <- function(x) c(0, cumsum(x[-length(x)]))
    size <- size - cumsum0(x@x)
    size <- getValues(advector(size)) ## Not variable
    prob <- prob / (1 - cumsum0(prob))
    sum(dbinom(x, size, prob, log=TRUE))
})
## For S4 generics we add the simref version like this:
##' @describeIn Distributions Simulation implementation. Modifies \code{x} and returns zero.
setMethod("dmultinom", "simref", function(x, size, prob, log) {
    if (is.null(size)) {
        if (is.null(x$getOrig))
            stop("Failed to determine 'size' parameter to use with simulation; please specify")
        obs <- x$getOrig(seq_len(length(x)))
        size <- sum(obs)
    }
    nrep <- 1 ## dmultinom is not vectorized
    x[] <- stats::rmultinom(nrep, size=size, prob=prob)
    rep(0, nrep)
})
## To prevent unintendend usage, we change the default method.
InvalidMethod <- function(which=-2) {
    fr <- sys.frame(which)
    cl <- match.call(sys.function(which), sys.call(which), FALSE)
    cls <- sapply(names(cl[-1]), function(nm)class(get(nm, envir=fr)))
    cl[-1] <- cls
    c("Unexpected combination of classes used in AD context:\n", deparse(cl))
}
##' @describeIn Distributions Default implementation that checks for invalid usage.
setMethod("dmultinom", signature("ANY", "ANY", "ANY", "ANY"),
          function (x, size, prob, log) {
              if (ad_context()) {
                  stop(InvalidMethod())
              }
              stats::dmultinom(x=x, size=size, prob=prob, log=log)
          })