distr.R

f <- system.file("include",package="TMB")
tmb <- unlist(lapply(dir(f,pattern=".hpp",full=TRUE),readLines))
tmb1 <- paste(tmb,collapse="") ## Everything on one line - warning: do not print!
d <- grep("^VECTORIZE.*\\(d", tmb, value=TRUE)
p <- grep("^VECTORIZE.*\\(p", tmb, value=TRUE)
q <- grep("^VECTORIZE.*\\(q", tmb, value=TRUE)
b <- grep("^VECTORIZE.*\\(bessel", tmb, value=TRUE)
cp <- grep("^VECTORIZE.*\\(compois", tmb, value=TRUE)
dp <- c(d,p,q,b,cp)
dp <- sub("VECTORIZE(.)_(.*)\\((.*)\\)", "\\3 \\1 \\2", dp)
## Manually add
more <- c("logspace_add 2 tt", "logspace_sub 2 tt")
dp <- c(dp, more)
df <- as.data.frame(t(do.call(cbind, sapply(dp, strsplit, " "))), stringsAsFactors=FALSE)
names(df) <- c("name","npar","code")
df <- subset(df,name!="pow") ## bogus
df <- subset(df, !(name=="pnorm" & npar==1) ) ## bogus
df <- subset(df, !(name=="qnorm" & npar==1) ) ## bogus
skip <- c("pSHASHo", "qSHASHo", "dnorm", "pnorm_approx", "qnorm_approx", "dzipois") ## RTMB uses a dnorm implementation in R
df <- subset(df, !(name %in% skip))

skip <- c("compois_calc_logZ", "compois_calc_loglambda")
df$export <- !(df$name %in% skip)

getsig <- function(name) {
    x <- sub(paste0(".*Type ",name,"\\((.*?)\\).*"),"\\1",tmb1)
    x <- gsub("[ ]*( )","\\1", x)
    x <- gsub("Type ", "", x)
    x <- gsub("const ", "", x)
    x <- gsub("\\&", "", x)
    x <- gsub("\t", "", x)
    x <- gsub("int give_log.*", "give_log", x)
    x[1]
}
df$signature_raw <- sapply(df$name,getsig)
sigtidy <- function(x) {
    x <- gsub("=.*?\\.","",x)
    x <- sub("^k","x",x)
    x <- sub("^y","x",x)
    ##x <- sub("give_log","log",x)
    x <- sub("[ ]*,",",",x)
    x
}
df$signature <- sigtidy(df$signature_raw)
## More tidy
i <- substring(df$name,1,1)=="p"
substring(df$signature[i],1,1) <- "q"
i <- substring(df$name,1,1)=="q"
substring(df$signature[i],1,1) <- "p"
df$stats <- sapply(df$name, exists, "package:stats")
getargs <- function(sig)gsub(" ","",strsplit(sig,",[ ]*")[[1]])
dblargs <- function(sig) { a <- getargs(sig); paste(a,a,sep="=",collapse=", ") }

newname <- function(x)paste0("distr_",x)
codegen <- function(i) {
    C <- df$signature[i]; npar <- as.numeric(df$npar[i])
    args <- strsplit(df$signature[i],",[ ]*")[[1]]
    fun <- df$name[i]
    type <- substring(fun,1,1)
    ## Codegen
    adind <- if (type=="d") (1:(npar-1)) else 1:npar
    nk <- paste0("n", adind)
    adargs <- args[adind]
    Xk <- paste0("X", adind)
    c(
    "// [[Rcpp::export]]",
    paste("ADrep", newname(fun),"(",paste(c(paste("ADrep", adargs ), paste("bool", tail(args,  1))  [type=="d"] ) , collapse=", "),")"),
    "{",
    paste0("int ", nk ,"=", adargs, ".size();"),
    paste0("int nmax = std::max({",paste(nk, collapse=", "),"});"),
    paste0("int nmin = std::min({",paste(nk, collapse=", "),"});"),
    "int n = (nmin == 0 ? 0 : nmax);",
    "ADrep ans(n);",
    paste0("const ad* ", Xk, " = adptr(", adargs ,");", collapse=" "),
    "ad* Y = adptr(ans);",
    paste0("for (int i=0; i<n; i++) Y[i] = ",fun, "(",  paste0(Xk, "[i % ", nk, "]", collapse=", "), ", give_log"[type=="d"] ,");"),
    "return ans;",
    "}")
}

header <- c(
    "// Autogenerated - do not edit",
    "#include \"RTMB.h\"")

xtra <- c(
    "// [[Rcpp::export]]",
    "double distr_rcompois(double loglambda, double nu) { return atomic::compois_utils::simulate(loglambda, nu); }")

## WriteLines **only if changed**
writeLines <- function(text, con) {
  if (file.exists(con)) {
    if(identical(readLines(con), text)) {
      cat("No change:", con, "\n")
      return (invisible(NULL))
    }
  }
  cat("Updating: ", con, "\n")
  base::writeLines(text, con)
}

## Make all TMB functions available via prefix 'distr_'
code <- unlist(lapply(1:nrow(df), codegen))
txt <- c(header, code, xtra)
writeLines(txt, "RTMB/src/distributions.cpp")

## ============================================================
## S4 R interface
df$signature <- sub("give_log","log",df$signature)
string <- function(x)paste0('"',x,'"')
getRmethod <- function(i) {
    name <- df$name[i]
    export <- df$export[i]
    is_density <- (substring(name,1,1) == "d")
    stats <- exists(name,"package:stats")
    if(stats) {
        Rsig <- head(as.list(args(get(name,"package:stats"))),-1)  ## R signature
        namespace <- paste0(environmentName(environment(get(name,"package:stats"))),"::") ## stats:: or base::
    } else {
        nm <- getargs(df$signature[i])
        Rsig <- structure(vector("list", length(nm)), names=nm )
        Rsig[] <- ""
        if ("log" %in% names(Rsig)) Rsig[["log"]] <- FALSE
    }
    a1 <- getargs(df$signature[i]) ## short (tmb)
    a2 <- names(Rsig) ## long (R)
    has.default <- nchar(sapply(Rsig,as.character))>0 ## long (R)
    if(any(is.na(match(a1,a2)))) {
        print(name)
    }
    a2.type <- ifelse(a2 %in% c("log","log.p","lower.tail") , "logical.", "ad")
    a2.type[a2.type=="ad" & has.default] <- "ad."
    ans <- ifelse(is.na(match(a2,a1)),"missing",a2.type)
    sig <- paste0("signature(",(paste(a2,"=",string(ans),collapse=", ")),")")
    cast <- ifelse(ans[match(a1,a2)] %in% c("ad","ad."),"advector","as.logical")
    ## Simple case: Not stats - create new simpel function
    if (!stats) {
        def <- c(
            paste(name,"<- function(", sub("log$","log=FALSE", df$signature[i]),") {"),
            paste('if (inherits(x,"osa")) return (dGenericOSA(',string(name),',',dblargs(df$signature[i]),'))')[is_density],
            paste('if (inherits(x,"simref")) return (dGenericSim(',string(name),',',dblargs(df$signature[i]),'))')[is_density],
            paste('dGenericEval(',string(name),',',dblargs(df$signature[i]),')'),
            "}")
        meth <- c(
            "##' @describeIn Distributions AD implementation"[export],
            def)
        return (meth)
    }
    def <- c(
        paste("function(", df$signature[i],") {"),
        paste(a1, "<- ",cast,"(",a1,")"),
        paste(newname(name),"(",df$signature[i],")"), "}")
    meth1 <-
        c(paste0("setMethod(", string(name), ","),
          paste0(sig, ","),
          def, ")" )
    ## Create specialization for numeric types
    sig2 <- gsub("ad", "num", sig)
    def2 <- c(
        paste("function(", df$signature[i],") {"),
        paste(namespace, name,"(",dblargs(df$signature[i]),")"), "}")
        ##paste("callNextMethod(",df$signature[i],")"), "}")
    meth2 <-
        c(paste0("setMethod(", string(name), ","),
          paste0(sig2, ","),
          def2, ")" )
    ## Create specialization for OSA
    if (is_density) {
        def3 <- c(
            paste("function(", df$signature[i],") {"),
            paste('  dGenericOSA(',string(name),',',dblargs(df$signature[i]),')'),
            "}")
        meth3 <- c(
            "##' @describeIn Distributions OSA implementation",
            paste0("setMethod(", string(name), ", ", string("osa"), ","),
            def3, ")" )
    } else {
        meth3 <- NULL
    }
    ## Create specialization for simref
    if (is_density) {
        def4 <- c(
            paste("function(", df$signature[i],") {"),
            paste('  dGenericSim(',string(name),',',dblargs(df$signature[i]),')'),
            "}")
        meth4 <- c(
            "##' @describeIn Distributions Simulation implementation. Modifies \\code{x} and returns zero.",
            paste0("setMethod(", string(name), ", ", string("simref"), ","),
            def4, ")" )
    } else {
        meth4 <- NULL
    }
    ## add roxygen
    meth1 <- c(
               paste0("##' @describeIn Distributions AD implementation of", " \\link[", gsub(":","",namespace), "]{", name, "}"),
               meth1)
    meth2 <- c(
               "##' @describeIn Distributions Default method",
               meth2)
    c(meth1,meth2,meth3,meth4)
}
## Start with stats methods only:
header <- c("## Autogenerated - do not edit","")
code.stats <- unlist(lapply(which(df$stats), getRmethod))
code.other <- unlist(lapply(which(!df$stats), getRmethod))
txt <- c(header, code.stats, code.other)
writeLines(txt, "RTMB/R/distributions.R")

## Argument documentation (partial)
## cat(paste("@param",unique(unlist(lapply(df$signature,strsplit,split="[ ]*,[ ]*"))), "parameter\n"))