Propagating multi-threading updates to python, temporarily removing enumeration unit tests after C++ refactor

Author: andrewherren

CMakeLists.txt

@@ -1,6 +1,8 @@
 # Build options
-option(USE_DEBUG "Set to ON for Debug mode" OFF)
+option(USE_DEBUG "Build with debug symbols and without optimization" OFF)
 option(USE_SANITIZER "Use santizer flags" OFF)
+option(USE_OPENMP "Use openMP" ON)
+option(USE_HOMEBREW_FALLBACK "(macOS-only) also look in 'brew --prefix' for libraries (e.g. OpenMP)" ON)
 option(BUILD_TEST "Build C++ tests with Google Test" OFF)
 option(BUILD_DEBUG_TARGETS "Build Standalone C++ Programs for Debugging" ON)
 option(BUILD_PYTHON "Build Shared Library for Python Package" OFF)
@@ -9,8 +11,8 @@ option(BUILD_PYTHON "Build Shared Library for Python Package" OFF)
 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 
-# Default to CMake 3.16
-cmake_minimum_required(VERSION 3.16)
+# Default to CMake 3.20
+cmake_minimum_required(VERSION 3.20)
 
 # Define the project
 project(stochtree LANGUAGES C CXX)
@@ -34,6 +36,13 @@ if(USE_DEBUG)
     add_definitions(-DDEBUG)
 endif()
 
+# Linker flags (empty by default, updated if using openmp)
+set(
+  STOCHTREE_LINK_FLAGS
+  ""
+)
+
+# Unix / MinGW compiler flags
 if(UNIX OR MINGW OR CYGWIN)
   set(
     CMAKE_CXX_FLAGS
@@ -42,11 +51,12 @@ if(UNIX OR MINGW OR CYGWIN)
   if(USE_DEBUG)
       set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -O0")
   else()
-      set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3")
+      set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -O3")
   endif()
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unknown-pragmas -Wno-unused-private-field")
 endif()
 
+# MSVC compiler flags
 if(MSVC)
     set(
       variables
@@ -72,6 +82,30 @@ else()
     endif()
 endif()
 
+# OpenMP
+if(USE_OPENMP)
+  add_definitions(-DSTOCHTREE_OPENMP_AVAILABLE)
+  if(APPLE)
+    find_package(OpenMP)
+    if(NOT OpenMP_FOUND)
+      if(USE_HOMEBREW_FALLBACK)
+        execute_process(COMMAND brew --prefix libomp
+                        OUTPUT_VARIABLE HOMEBREW_LIBOMP_PREFIX
+                        OUTPUT_STRIP_TRAILING_WHITESPACE)
+        set(OpenMP_CXX_FLAGS "-Xclang -fopenmp")
+        set(OpenMP_CXX_INCLUDE_DIR "-I${HOMEBREW_LIBOMP_PREFIX}/include")
+        set(OpenMP_CXX_LIB_NAMES omp)
+        set(OpenMP_libomp_LIBRARY ${HOMEBREW_LIBOMP_PREFIX}/lib/libomp.dylib)
+      endif()
+      find_package(OpenMP REQUIRED)
+    endif()
+  else()
+    find_package(OpenMP REQUIRED)
+  endif()
+  # Update flags with openmp 
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}")
+endif()
+
 # Header file directory
 set(StochTree_HEADER_DIR ${PROJECT_SOURCE_DIR}/include)
 
@@ -80,6 +114,8 @@ set(BOOSTMATH_HEADER_DIR ${PROJECT_SOURCE_DIR}/deps/boost_math/include)
 
 # Eigen header file directory
 set(EIGEN_HEADER_DIR ${PROJECT_SOURCE_DIR}/deps/eigen)
+add_definitions(-DEIGEN_MPL2_ONLY)
+add_definitions(-DEIGEN_DONT_PARALLELIZE)
 
 # fast_double_parser header file directory
 set(FAST_DOUBLE_PARSER_HEADER_DIR ${PROJECT_SOURCE_DIR}/deps/fast_double_parser/include)
@@ -109,10 +145,11 @@ file(
 add_library(stochtree_objs OBJECT ${SOURCES})
 
 # Include the headers in the source library
-target_include_directories(stochtree_objs PRIVATE ${StochTree_HEADER_DIR} ${BOOSTMATH_HEADER_DIR} ${EIGEN_HEADER_DIR} ${FAST_DOUBLE_PARSER_HEADER_DIR} ${FMT_HEADER_DIR})
-
-if(APPLE)
-  set(CMAKE_SHARED_LIBRARY_SUFFIX ".so")
+if(USE_OPENMP)
+  target_include_directories(stochtree_objs PRIVATE ${StochTree_HEADER_DIR} ${BOOSTMATH_HEADER_DIR} ${EIGEN_HEADER_DIR} ${FAST_DOUBLE_PARSER_HEADER_DIR} ${FMT_HEADER_DIR} ${OpenMP_CXX_INCLUDE_DIR})
+  target_link_libraries(stochtree_objs PRIVATE ${OpenMP_libomp_LIBRARY})
+else()
+  target_include_directories(stochtree_objs PRIVATE ${StochTree_HEADER_DIR} ${BOOSTMATH_HEADER_DIR} ${EIGEN_HEADER_DIR} ${FAST_DOUBLE_PARSER_HEADER_DIR} ${FMT_HEADER_DIR})
 endif()
 
 # Python shared library
@@ -122,8 +159,13 @@ if (BUILD_PYTHON)
   pybind11_add_module(stochtree_cpp src/py_stochtree.cpp)
 
   # Link to C++ source and headers
-  target_include_directories(stochtree_cpp PRIVATE ${StochTree_HEADER_DIR} ${BOOSTMATH_HEADER_DIR} ${EIGEN_HEADER_DIR} ${FAST_DOUBLE_PARSER_HEADER_DIR} ${FMT_HEADER_DIR})
-  target_link_libraries(stochtree_cpp PRIVATE stochtree_objs)
+  if(USE_OPENMP)
+    target_include_directories(stochtree_cpp PRIVATE ${StochTree_HEADER_DIR} ${BOOSTMATH_HEADER_DIR} ${EIGEN_HEADER_DIR} ${FAST_DOUBLE_PARSER_HEADER_DIR} ${FMT_HEADER_DIR} ${OpenMP_CXX_INCLUDE_DIR})
+    target_link_libraries(stochtree_cpp PRIVATE stochtree_objs ${OpenMP_libomp_LIBRARY})
+  else()
+    target_include_directories(stochtree_cpp PRIVATE ${StochTree_HEADER_DIR} ${BOOSTMATH_HEADER_DIR} ${EIGEN_HEADER_DIR} ${FAST_DOUBLE_PARSER_HEADER_DIR} ${FMT_HEADER_DIR})
+    target_link_libraries(stochtree_cpp PRIVATE stochtree_objs)
+  endif()
 
   # EXAMPLE_VERSION_INFO is defined by setup.py and passed into the C++ code as a
   # define (VERSION_INFO) here.
@@ -154,8 +196,13 @@ if(BUILD_TEST)
   file(GLOB CPP_TEST_SOURCES test/cpp/*.cpp)
   add_executable(teststochtree ${CPP_TEST_SOURCES})
   set(STOCHTREE_TEST_HEADER_DIR ${PROJECT_SOURCE_DIR}/test/cpp)
-  target_include_directories(teststochtree PRIVATE ${StochTree_HEADER_DIR} ${BOOSTMATH_HEADER_DIR} ${EIGEN_HEADER_DIR} ${STOCHTREE_TEST_HEADER_DIR} ${FAST_DOUBLE_PARSER_HEADER_DIR} ${FMT_HEADER_DIR})
-  target_link_libraries(teststochtree PRIVATE stochtree_objs GTest::gtest_main)
+  if(USE_OPENMP)
+    target_include_directories(teststochtree PRIVATE ${StochTree_HEADER_DIR} ${BOOSTMATH_HEADER_DIR} ${EIGEN_HEADER_DIR} ${STOCHTREE_TEST_HEADER_DIR} ${FAST_DOUBLE_PARSER_HEADER_DIR} ${FMT_HEADER_DIR} ${OpenMP_CXX_INCLUDE_DIR})
+    target_link_libraries(teststochtree PRIVATE stochtree_objs GTest::gtest_main ${OpenMP_libomp_LIBRARY})
+  else()
+    target_include_directories(teststochtree PRIVATE ${StochTree_HEADER_DIR} ${BOOSTMATH_HEADER_DIR} ${EIGEN_HEADER_DIR} ${STOCHTREE_TEST_HEADER_DIR} ${FAST_DOUBLE_PARSER_HEADER_DIR} ${FMT_HEADER_DIR})
+    target_link_libraries(teststochtree PRIVATE stochtree_objs GTest::gtest_main)
+  endif()
   gtest_discover_tests(teststochtree)
 endif()
 
@@ -164,7 +211,12 @@ if(BUILD_DEBUG_TARGETS)
   # Build test suite
   add_executable(debugstochtree debug/api_debug.cpp)
   set(StochTree_DEBUG_HEADER_DIR ${PROJECT_SOURCE_DIR}/debug)
-  target_include_directories(debugstochtree PRIVATE ${StochTree_HEADER_DIR} ${BOOSTMATH_HEADER_DIR} ${EIGEN_HEADER_DIR} ${StochTree_DEBUG_HEADER_DIR} ${FAST_DOUBLE_PARSER_HEADER_DIR} ${FMT_HEADER_DIR})
-  target_link_libraries(debugstochtree PRIVATE stochtree_objs)
+  if(USE_OPENMP)
+    target_include_directories(debugstochtree PRIVATE ${StochTree_HEADER_DIR} ${BOOSTMATH_HEADER_DIR} ${EIGEN_HEADER_DIR} ${StochTree_DEBUG_HEADER_DIR} ${FAST_DOUBLE_PARSER_HEADER_DIR} ${FMT_HEADER_DIR} ${OpenMP_CXX_INCLUDE_DIR})
+    target_link_libraries(debugstochtree PRIVATE stochtree_objs ${OpenMP_libomp_LIBRARY})
+  else()
+    target_include_directories(debugstochtree PRIVATE ${StochTree_HEADER_DIR} ${BOOSTMATH_HEADER_DIR} ${EIGEN_HEADER_DIR} ${StochTree_DEBUG_HEADER_DIR} ${FAST_DOUBLE_PARSER_HEADER_DIR} ${FMT_HEADER_DIR})
+    target_link_libraries(debugstochtree PRIVATE stochtree_objs)
+  endif()
 endif()
 

R/bcf.R

@@ -53,6 +53,7 @@
 #'   - `rfx_group_parameter_prior_cov` Prior covariance matrix for the random effects "group parameters." Default: `NULL`. Must be a square matrix whose dimension matches the number of random effects bases, or a scalar value that will be expanded to a diagonal matrix.
 #'   - `rfx_variance_prior_shape` Shape parameter for the inverse gamma prior on the variance of the random effects "group parameter." Default: `1`.
 #'   - `rfx_variance_prior_scale` Scale parameter for the inverse gamma prior on the variance of the random effects "group parameter." Default: `1`.
+#'   - `num_threads` Number of threads to use in the GFR and MCMC algorithms, as well as prediction. If OpenMP is not available on a user's setup, this will default to `1`, otherwise to the maximum number of available threads.
 #'
 #' @param prognostic_forest_params (Optional) A list of prognostic forest model parameters, each of which has a default value processed internally, so this argument list is optional.
 #'
@@ -174,7 +175,8 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
         rfx_working_parameter_prior_cov = NULL,
         rfx_group_parameter_prior_cov = NULL,
         rfx_variance_prior_shape = 1,
-        rfx_variance_prior_scale = 1
+        rfx_variance_prior_scale = 1, 
+        num_threads = -1
     )
     general_params_updated <- preprocessParams(
         general_params_default, general_params
@@ -248,6 +250,7 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
     rfx_group_parameter_prior_cov <- general_params_updated$rfx_group_parameter_prior_cov
     rfx_variance_prior_shape <- general_params_updated$rfx_variance_prior_shape
     rfx_variance_prior_scale <- general_params_updated$rfx_variance_prior_scale
+    num_threads <- general_params_updated$num_threads
 
     # 2. Mu forest parameters
     num_trees_mu <- prognostic_forest_params_updated$num_trees
@@ -1029,7 +1032,7 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
             forest_model_mu$sample_one_iteration(
                 forest_dataset = forest_dataset_train, residual = outcome_train, forest_samples = forest_samples_mu, 
                 active_forest = active_forest_mu, rng = rng, forest_model_config = forest_model_config_mu, 
-                global_model_config = global_model_config, keep_forest = keep_sample, gfr = TRUE
+                global_model_config = global_model_config, num_threads = num_threads, keep_forest = keep_sample, gfr = TRUE
             )
 
             # Cache train set predictions since they are already computed during sampling
@@ -1053,7 +1056,7 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
             forest_model_tau$sample_one_iteration(
                 forest_dataset = forest_dataset_train, residual = outcome_train, forest_samples = forest_samples_tau, 
                 active_forest = active_forest_tau, rng = rng, forest_model_config = forest_model_config_tau, 
-                global_model_config = global_model_config, keep_forest = keep_sample, gfr = TRUE
+                global_model_config = global_model_config, num_threads = num_threads, keep_forest = keep_sample, gfr = TRUE
             )
 
             # Cannot cache train set predictions for tau because the cached predictions in the 
@@ -1102,7 +1105,7 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
                 forest_model_variance$sample_one_iteration(
                     forest_dataset = forest_dataset_train, residual = outcome_train, forest_samples = forest_samples_variance, 
                     active_forest = active_forest_variance, rng = rng, forest_model_config = forest_model_config_variance, 
-                    global_model_config = global_model_config, keep_forest = keep_sample, gfr = TRUE
+                    global_model_config = global_model_config, num_threads = num_threads, keep_forest = keep_sample, gfr = TRUE
                 )
 
                 # Cache train set predictions since they are already computed during sampling
@@ -1309,7 +1312,7 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
                 forest_model_mu$sample_one_iteration(
                     forest_dataset = forest_dataset_train, residual = outcome_train, forest_samples = forest_samples_mu, 
                     active_forest = active_forest_mu, rng = rng, forest_model_config = forest_model_config_mu, 
-                    global_model_config = global_model_config, keep_forest = keep_sample, gfr = FALSE
+                    global_model_config = global_model_config, num_threads = num_threads, keep_forest = keep_sample, gfr = FALSE
                 )
 
                 # Cache train set predictions since they are already computed during sampling
@@ -1333,7 +1336,7 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
                 forest_model_tau$sample_one_iteration(
                     forest_dataset = forest_dataset_train, residual = outcome_train, forest_samples = forest_samples_tau, 
                     active_forest = active_forest_tau, rng = rng, forest_model_config = forest_model_config_tau, 
-                    global_model_config = global_model_config, keep_forest = keep_sample, gfr = FALSE
+                    global_model_config = global_model_config, num_threads = num_threads, keep_forest = keep_sample, gfr = FALSE
                 )
 
                 # Cannot cache train set predictions for tau because the cached predictions in the 
@@ -1382,7 +1385,7 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
                     forest_model_variance$sample_one_iteration(
                         forest_dataset = forest_dataset_train, residual = outcome_train, forest_samples = forest_samples_variance, 
                         active_forest = active_forest_variance, rng = rng, forest_model_config = forest_model_config_variance, 
-                        global_model_config = global_model_config, keep_forest = keep_sample, gfr = FALSE
+                        global_model_config = global_model_config, num_threads = num_threads, keep_forest = keep_sample, gfr = FALSE
                     )
 
                     # Cache train set predictions since they are already computed during sampling

R/model.R

@@ -67,7 +67,7 @@ ForestModel <- R6::R6Class(
         #' @param rng Wrapper around C++ random number generator
         #' @param forest_model_config ForestModelConfig object containing forest model parameters and settings
         #' @param global_model_config GlobalModelConfig object containing global model parameters and settings
-        #' @param num_threads Number of threads to use in the GFR and MCMC algorithms, as well as prediction. If OpenMP is not available on a user's setup, this will default to `1`, otherwise to the maximum number of available threads.
+        #' @param num_threads Number of threads to use in the GFR and MCMC algorithms, as well as prediction. If OpenMP is not available on a user's system, this will default to `1`, otherwise to the maximum number of available threads.
         #' @param keep_forest (Optional) Whether the updated forest sample should be saved to `forest_samples`. Default: `TRUE`.
         #' @param gfr (Optional) Whether or not the forest should be sampled using the "grow-from-root" (GFR) algorithm. Default: `TRUE`.
         sample_one_iteration = function(forest_dataset, residual, forest_samples, active_forest, 

debug/README.md

@@ -4,7 +4,7 @@ This subdirectory contains a debug program for the C++ codebase.
 The program takes several command line arguments (in order):
 
 1. Which data-generating process (DGP) to run (integer-coded, see below for a detailed description)
-1. Which leaf model to sample (integer-coded, see below for a detailed description)
+2. Which leaf model to sample (integer-coded, see below for a detailed description)
 3. Whether or not to include random effects (0 = no, 1 = yes)
 4. Number of grow-from-root (GFR) samples
 5. Number of MCMC samples
@@ -13,6 +13,7 @@ The program takes several command line arguments (in order):
 8. [Optional] index of outcome column in data file (leave this blank as `0`)
 9. [Optional] comma-delimited string of column indices of covariates (leave this blank as `""`)
 10. [Optional] comma-delimited string of column indices of leaf regression bases (leave this blank as `""`)
+11. [Optional] number of threads to use in the GFR sampler (leave this blank as `-1`)
 
 The DGPs are numbered as follows:
 
@@ -30,6 +31,6 @@ The models are numbered as follows:
 
 For an example of how to run this progam for DGP 0, leaf model 1, no random effects, 10 GFR samples, 100 MCMC samples and a default seed (`-1`), run
 
-`./build/debugstochtree 0 1 0 10 100 -1 "" 0 "" ""`
+`./build/debugstochtree 0 1 0 10 100 -1 "" 0 "" "" -1`
 
 from the main `stochtree` project directory after building with `BUILD_DEBUG_TARGETS` set to `ON`.