R/rcpp_tapply_generator.R

In kevinushey/Kmisc: Kevin Miscellaneous

##' Rcpp tapply Generator
##' 
##' Use this function as a generator for your own \code{tapply} functions;
##' that is, functions you would like to apply split over some grouping
##' variable. Note that we restrict ourselves to the case where we return
##' a scalar as output; rather than the more general output of \code{tapply}.
##' 
##' Note that we simplify differently than base \R \code{tapply}: when
##' \code{simplify=TRUE}, we call \code{unlist} on the output; hence,
##' a named vector is returned.
##' 
##' @importFrom Rcpp sourceCpp
##' @param fun A character string defining the C++ function. It must be
##' in terms of a variable \code{x}.
##' @param includes Other C++ libraries to include. For example, to include \code{boost/math.hpp},
##' you could pass \code{c("<boost/math.hpp>")}. Rcpp is included by default, unless RcppArmadillo
##' is included as well (since Rcpp is included as part of the RcppArmadillo include)
##' @param depends Other libraries to link to. Linking is done through Rcpp attributes.
##' @param returnType The return type of your function; note that we require that
##' the generator returns a scalar value of type \code{double}, \code{int}, or 
##' \code{bool}.
##' @param inline boolean; mark this function as inline? This may or may not
##' increase execution speed.
##' @param name An internal name for the function.
##' @param file A location to output the file. Defaults to a temporary
##' file as generated by \code{\link{tempfile}()}.
##' @param additional Other C++ code you want to include; e.g. helper functions.
##' This code will be inserted as-is above the code in \code{fun}.
##' @export
rcpp_tapply_generator <- function( fun, 
  includes=NULL,
  depends=NULL,
  inline=TRUE,
  returnType="double",
  name=NULL, 
  file=NULL,
  additional=NULL ) {
  
  if( getRversion() < "2.15.1" ||
      packageVersion("Rcpp") < "0.10.1"
  ) {
    message("Error: This function requires Rcpp > 0.10.1 and R > 2.15.1")
    return( invisible(NULL) )
  }
  
  ## generate a C++ source file based on 'fun'
  if( is.null(file) ) {
    
    cpp_source <- paste( sep="", tempfile(), ".cpp" )
    
  } else {
    
    if( length( grep( "\\.cpp", file ) ) == 0 ) {
      file <- paste( sep="", file, ".cpp" )
    }
    
    cpp_source <- file
    
  }
  
  cat("C++ source code will be written to", cpp_source, ".\n")
  
  ## an environment for the cppSource-exported functions to live in
  cpp_env <- new.env()
  
  ## internal name for the function
  if( is.null(name) ) {  
    name <- paste( collapse="", 
      sample( c(letters, LETTERS), size=20, replace=TRUE)
    )
  } 
  
  ## open a connection to the file, and ensure we close it after
  conn <- file( cpp_source, 'w' )
  on.exit( close(conn) )
  
  ## determine the Rcpp return type
  rcpp_returnType <- switch( returnType,
    double="NumericVector",
    int="IntegerVector",
    bool="LogicalVector",
    SEXP="CharacterVector",
    warning("unrecognized return type")
  )
  
  ## process the 'depends' argument
  include_Rcpp <- TRUE
  if( !is.null(depends) ) {
    if( "RcppArmadillo" %in% depends ) {
      include_Rcpp <- FALSE
    }
    depends <- paste("// [[Rcpp::depends(", paste(depends, collapse=", "), ")]]",
      sep="" )
    depends <- paste( sep="", depends, "\n" )
  }
  
  ## process the 'includes' argument
  if( !is.null(includes) ) {
    includes <- paste( paste("#include", includes), collapse="\n" )
    includes <- paste( includes, "\n", sep="" )
  }
  
  ## 'cat' our source code out to a file, to later be sourced
  cat( file=conn, sep="", paste( sep="", collapse="", 
    depends,
    if( include_Rcpp ) "#include <Rcpp.h>\n",
    includes, "\n",
    "using namespace Rcpp;\n",
    additional,
    "\n",
    "inline ", returnType, " do_", name, "( const ", rcpp_returnType, "& x ) {\n", fun, "\n }
                                 
// [[Rcpp::export]]
SEXP ", name, "( const List& X, bool simplify ) {
  SEXP output;
  if (simplify) {
    PROTECT(output = wrap(sapply(X, do_", name, ")));
  } else {
    PROTECT(output = wrap(lapply(X, do_", name, ")));
  }
  UNPROTECT(1);
  return output;
}
"))
  
  ## source the file into the 'cpp_env' environment
  cat("Compiling...\n")
  sourceCpp( cpp_source, env=cpp_env )
  cat("Done!")
  
  ## return a function that exposes the compiled code to the user
  return( function(X, gp, simplify=TRUE) {
    
    force( cpp_source )
    splat <- split_(X, gp)
    call <- call(name, splat, simplify)
    out <- eval(call, envir=cpp_env)
    names(out) <- names(splat)
    return(out)
    
  })
  
}

##' @rdname rcpp_tapply_generator
##' @export
Rcpp_tapply_generator <- rcpp_tapply_generator