diff --git a/.clang-format b/.clang-format new file mode 100644 index 00000000..b6a55284 --- /dev/null +++ b/.clang-format @@ -0,0 +1,5 @@ +BasedOnStyle: Google +IndentWidth: 2 +AccessModifierOffset: -1 +ColumnLimit: 0 +SortIncludes: false diff --git a/.clang-tidy b/.clang-tidy new file mode 100644 index 00000000..1f3a66ea --- /dev/null +++ b/.clang-tidy @@ -0,0 +1,32 @@ +--- +# Conservative clang-tidy config for stochtree. +# Focuses on real bugs and performance issues; style is handled by clang-format. +# To run manually: clang-tidy -- (or via clangd in the editor) + +Checks: > + clang-analyzer-*, + bugprone-branch-clone, + bugprone-copy-constructor-init, + bugprone-dangling-handle, + bugprone-incorrect-roundings, + bugprone-infinite-loop, + bugprone-redundant-branch-condition, + bugprone-suspicious-include, + bugprone-use-after-move, + modernize-redundant-void-arg, + modernize-use-emplace, + modernize-use-nullptr, + modernize-use-override, + performance-for-range-copy, + performance-move-const-arg, + performance-unnecessary-copy-initialization, + -clang-analyzer-optin.performance.Padding + +# Leave empty to warn only, not fail the build. +WarningsAsErrors: "" + +# Only surface warnings for project headers, not dependencies. +HeaderFilterRegex: "include/stochtree/.*" + +# Respect .clang-format for any fixes clang-tidy applies. +FormatStyle: file diff --git a/.clangd b/.clangd new file mode 100644 index 00000000..c87746fa --- /dev/null +++ b/.clangd @@ -0,0 +1,14 @@ +--- +# Fallback compile flags for files not captured by compile_commands.json +# (e.g. R wrapper files when BUILD_PYTHON=OFF, or before cmake has been run). +# Files compiled by cmake use compile_commands.json and ignore these flags. +CompileFlags: + Add: + - "-std=c++17" + - "-Iinclude" + - "-Ideps/boost_math/include" + - "-Ideps/eigen" + - "-Ideps/fast_double_parser/include" + - "-Ideps/fmt/include" + - "-Ideps/pybind11/include" + Compiler: clang++ diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile index 8af2f042..33602460 100644 --- a/.devcontainer/Dockerfile +++ b/.devcontainer/Dockerfile @@ -5,7 +5,7 @@ RUN apt-get update -y && \ valgrind && \ rm -rf /var/lib/apt/lists/* -ARG REINSTALL_CMAKE_VERSION_FROM_SOURCE="3.22.2" +ARG REINSTALL_CMAKE_VERSION_FROM_SOURCE="3.29.3" # Optionally install the cmake for vcpkg COPY ./reinstall-cmake.sh /tmp/ diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index b519e257..f496e8ba 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -1,23 +1,34 @@ // For format details, see https://aka.ms/devcontainer.json. For config options, see the // README at: https://github.com/devcontainers/templates/tree/main/src/cpp { - "name": "C++", + "name": "stochtree C++ Dev", "build": { "dockerfile": "Dockerfile" - } + }, - // Features to add to the dev container. More info: https://containers.dev/features. - // "features": {}, + "features": { + "ghcr.io/devcontainers/features/git:1": {}, + "ghcr.io/devcontainers/features/github-cli:1": {} + }, - // Use 'forwardPorts' to make a list of ports inside the container available locally. - // "forwardPorts": [], + "customizations": { + "vscode": { + "extensions": [ + "ms-vscode.cmake-tools", + "llvm-vs-code-extensions.vscode-clangd", + "vadimcn.vscode-lldb" + ], + "settings": { + "clangd.path": "/usr/bin/clangd", + "clangd.arguments": ["--compile-commands-dir=${workspaceFolder}/build"], + "cmake.configureOnOpen": true, + "cmake.defaultConfigurePreset": "dev-quick", + "cmake.defaultBuildPreset": "dev-quick" + } + } + }, - // Use 'postCreateCommand' to run commands after the container is created. - // "postCreateCommand": "gcc -v", - - // Configure tool-specific properties. - // "customizations": {}, - - // Uncomment to connect as root instead. More info: https://aka.ms/dev-containers-non-root. - // "remoteUser": "root" + // Configure dev build (no test download) on container creation. + // Switch to the "dev" preset and re-run cmake when you need GoogleTest. + "postCreateCommand": "cmake --preset dev-quick && cmake --build --preset dev-quick" } diff --git a/.github/workflows/cross-language-parity.yml b/.github/workflows/cross-language-parity.yml new file mode 100644 index 00000000..945f8973 --- /dev/null +++ b/.github/workflows/cross-language-parity.yml @@ -0,0 +1,77 @@ +name: Cross-Language Parity Tests + +# Runs on every push/PR to main. Ubuntu only — this is a correctness check, +# not a platform-portability check. Python generates prediction fixtures; +# R reproduces the same model with the same seed and compares. + +on: + push: + branches: [main] + pull_request: + branches: [main] + workflow_dispatch: + +jobs: + parity: + name: cross-language-parity + runs-on: ubuntu-latest + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + submodules: 'recursive' + + # ----------------------------------------------------------------------- + # Python + # ----------------------------------------------------------------------- + - name: Setup Python 3.12 + uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: "pip" + + - name: Install Python package + run: | + pip install --upgrade pip + pip install -r requirements.txt + pip install . + + # ----------------------------------------------------------------------- + # R + # ----------------------------------------------------------------------- + - uses: r-lib/actions/setup-pandoc@v2 + + - uses: r-lib/actions/setup-r@v2 + with: + use-public-rspm: true + + - uses: r-lib/actions/setup-r-dependencies@v2 + with: + extra-packages: any::testthat, any::decor, any::rcmdcheck, any::jsonlite, any::devtools + needs: check + + - name: Build and install R package + run: | + Rscript cran-bootstrap.R 0 0 1 + R CMD INSTALL stochtree_cran + + # ----------------------------------------------------------------------- + # Parity tests + # ----------------------------------------------------------------------- + - name: Generate Python predictions + run: | + python test/cross_language/generate_predictions.py \ + --output-dir /tmp/parity_fixtures + + - name: Verify R predictions match Python + run: | + Rscript test/cross_language/verify_predictions.R /tmp/parity_fixtures + + # Upload fixtures on failure so they can be inspected + - name: Upload fixtures on failure + if: failure() + uses: actions/upload-artifact@v4 + with: + name: parity-fixtures + path: /tmp/parity_fixtures diff --git a/.gitignore b/.gitignore index 2298053c..feeb58ac 100644 --- a/.gitignore +++ b/.gitignore @@ -3,10 +3,15 @@ *.csv *.DS_Store lib/ -build*/ -.vscode/ +build/ +build-release*/ +.vscode/positron/ xcode/ *.json +!.vscode/extensions.json +!.vscode/tasks.json +!.vscode/launch.json +!CMakePresets.json !test/R/testthat/fixtures/*.json !test/python/fixtures/*.json .vs/ @@ -15,7 +20,7 @@ cpp_docs/doxyoutput/xml cpp_docs/doxyoutput/latex stochtree_cran *.trace -*.clangd +.cache/clangd/ *.claude *release_notes.md diff --git a/.vscode/extensions.json b/.vscode/extensions.json new file mode 100644 index 00000000..33e5b7db --- /dev/null +++ b/.vscode/extensions.json @@ -0,0 +1,8 @@ +{ + "recommendations": [ + "ms-vscode.cmake-tools", + "vadimcn.vscode-lldb", + "llvm-vs-code-extensions.vscode-clangd", + "Posit.air-vscode" + ] +} diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 00000000..7c8a35b3 --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,107 @@ +{ + "version": "0.2.0", + "configurations": [ + { + "name": "bart_debug (macOS)", + "type": "lldb", + "request": "launch", + "program": "${workspaceFolder}/build/bart_debug", + "args": ["--scenario", "1", "--n", "500", "--n_test", "100", "--p", "5", "--num_trees", "200", "--num_gfr", "10", "--num_mcmc", "100", "--seed", "-1"], + "cwd": "${workspaceFolder}", + "preLaunchTask": "CMake: Build (dev-quick)" + }, + { + "name": "bcf_debug (macOS)", + "type": "lldb", + "request": "launch", + "program": "${workspaceFolder}/build/bcf_debug", + "args": ["--scenario", "0", "--n", "500", "--n_test", "100", "--p", "5", "--num_trees_mu", "200", "--num_trees_tau", "50", "--num_gfr", "10", "--num_mcmc", "100", "--seed", "-1"], + "cwd": "${workspaceFolder}", + "preLaunchTask": "CMake: Build (dev-quick)" + }, + { + "name": "teststochtree (macOS)", + "type": "lldb", + "request": "launch", + "program": "${workspaceFolder}/build/teststochtree", + "args": [], + "cwd": "${workspaceFolder}", + "preLaunchTask": "CMake: Build (dev)" + }, + { + "name": "bart_debug (Linux/Container)", + "type": "cppdbg", + "request": "launch", + "program": "${workspaceFolder}/build/bart_debug", + "args": ["--scenario", "1", "--n", "500", "--n_test", "100", "--p", "5", "--num_trees", "200", "--num_gfr", "10", "--num_mcmc", "100", "--seed", "-1"], + "cwd": "${workspaceFolder}", + "MIMode": "gdb", + "preLaunchTask": "CMake: Build (dev-quick)" + }, + { + "name": "bcf_debug (Linux/Container)", + "type": "cppdbg", + "request": "launch", + "program": "${workspaceFolder}/build/bcf_debug", + "args": ["--scenario", "1", "--n", "200", "--n_test", "100", "--p", "5", "--num_trees_mu", "200", "--num_trees_tau", "50", "--num_gfr", "10", "--num_mcmc", "100", "--seed", "-1"], + "cwd": "${workspaceFolder}", + "MIMode": "gdb", + "preLaunchTask": "CMake: Build (dev-quick)" + }, + { + "name": "teststochtree (Linux/Container)", + "type": "cppdbg", + "request": "launch", + "program": "${workspaceFolder}/build/teststochtree", + "args": [], + "cwd": "${workspaceFolder}", + "MIMode": "gdb", + "preLaunchTask": "CMake: Build (dev)" + }, + { + "name": "bart_debug (macOS, Release)", + "type": "lldb", + "request": "launch", + "program": "${workspaceFolder}/build-release-drivers/bart_debug", + "args": ["--scenario", "1", "--n", "500", "--n_test", "100", "--p", "5", "--num_trees", "200", "--num_gfr", "10", "--num_mcmc", "100", "--seed", "-1"], + "cwd": "${workspaceFolder}", + "preLaunchTask": "CMake: Build (release-drivers)" + }, + { + "name": "bart_debug (Linux/Container, Release)", + "type": "cppdbg", + "request": "launch", + "program": "${workspaceFolder}/build-release-drivers/bart_debug", + "args": ["--scenario", "1", "--n", "500", "--n_test", "100", "--p", "5", "--num_trees", "200", "--num_gfr", "10", "--num_mcmc", "100", "--seed", "-1"], + "cwd": "${workspaceFolder}", + "MIMode": "gdb", + "preLaunchTask": "CMake: Build (release-drivers)" + }, + { + "name": "bcf_debug (macOS, Release)", + "type": "lldb", + "request": "launch", + "program": "${workspaceFolder}/build-release-drivers/bcf_debug", + "args": ["--scenario", "1", "--n", "200", "--n_test", "100", "--p", "5", "--num_trees_mu", "200", "--num_trees_tau", "50", "--num_gfr", "10", "--num_mcmc", "100", "--seed", "-1"], + "cwd": "${workspaceFolder}", + "preLaunchTask": "CMake: Build (release-drivers)" + }, + { + "name": "bcf_debug (Linux/Container, Release)", + "type": "cppdbg", + "request": "launch", + "program": "${workspaceFolder}/build-release-drivers/bcf_debug", + "args": ["--scenario", "1", "--n", "200", "--n_test", "100", "--p", "5", "--num_trees_mu", "200", "--num_trees_tau", "50", "--num_gfr", "10", "--num_mcmc", "100", "--seed", "-1"], + "cwd": "${workspaceFolder}", + "MIMode": "gdb", + "preLaunchTask": "CMake: Build (release-drivers)" + }, + { + "name": "Python Debugger: Current File", + "type": "debugpy", + "request": "launch", + "program": "${file}", + "console": "integratedTerminal" + } + ] +} diff --git a/.vscode/tasks.json b/.vscode/tasks.json new file mode 100644 index 00000000..e21400a0 --- /dev/null +++ b/.vscode/tasks.json @@ -0,0 +1,77 @@ +{ + "version": "2.0.0", + "tasks": [ + { + "label": "CMake: Configure (dev)", + "type": "shell", + "command": "cmake --preset dev", + "group": "build", + "problemMatcher": ["$gcc"], + "presentation": { "reveal": "always", "panel": "shared" } + }, + { + "label": "CMake: Configure (dev-quick)", + "type": "shell", + "command": "cmake --preset dev-quick", + "group": "build", + "problemMatcher": ["$gcc"], + "presentation": { "reveal": "always", "panel": "shared" } + }, + { + "label": "CMake: Build (dev)", + "type": "shell", + "command": "cmake --preset dev && cmake --build --preset dev", + "group": { "kind": "build", "isDefault": true }, + "problemMatcher": ["$gcc"], + "presentation": { "reveal": "always", "panel": "shared" } + }, + { + "label": "CMake: Build (dev-quick)", + "type": "shell", + "command": "cmake --preset dev-quick && cmake --build --preset dev-quick", + "group": "build", + "problemMatcher": ["$gcc"], + "presentation": { "reveal": "always", "panel": "shared" } + }, + { + "label": "CMake: Build (release)", + "type": "shell", + "command": "cmake --preset release && cmake --build --preset release", + "group": "build", + "problemMatcher": ["$gcc"], + "presentation": { "reveal": "always", "panel": "shared" } + }, + { + "label": "CMake: Build (release-drivers)", + "type": "shell", + "command": "cmake --preset release-drivers && cmake --build --preset release-drivers", + "group": "build", + "problemMatcher": ["$gcc"], + "presentation": { "reveal": "always", "panel": "shared" } + }, + { + "label": "CMake: Build (sanitizer)", + "type": "shell", + "command": "cmake --preset sanitizer && cmake --build --preset sanitizer", + "group": "build", + "problemMatcher": ["$gcc"], + "presentation": { "reveal": "always", "panel": "shared" } + }, + { + "label": "Test: Run All", + "type": "shell", + "command": "${workspaceFolder}/build/teststochtree", + "group": { "kind": "test", "isDefault": true }, + "problemMatcher": [], + "presentation": { "reveal": "always", "panel": "shared" } + }, + { + "label": "Test: Run All (sanitizer)", + "type": "shell", + "command": "${workspaceFolder}/build/teststochtree", + "group": "test", + "problemMatcher": [], + "presentation": { "reveal": "always", "panel": "shared" } + } + ] +} diff --git a/CMakeLists.txt b/CMakeLists.txt index b6471bd2..d825ba55 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -130,16 +130,18 @@ set(LIBRARY_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/build) # Aggregate the source files underpinning the implementation in the C++ library file( - GLOB + GLOB SOURCES + src/bart_sampler.cpp + src/bcf_sampler.cpp src/container.cpp src/cutpoint_candidates.cpp src/data.cpp src/io.cpp - src/json11.cpp src/leaf_model.cpp src/ordinal_sampler.cpp src/partition_tracker.cpp + src/prediction.cpp src/random_effects.cpp src/tree.cpp ) @@ -177,13 +179,15 @@ endif() # Build C++ test program if(BUILD_TEST) - # Check if user specified a local clone of the GoogleTest repo, use Github repo if not + include(FetchContent) + + # GoogleTest is fetched on first configure and cached in build/_deps/. + # Re-download is skipped automatically when FETCHCONTENT_UPDATES_DISCONNECTED=ON + # (set in the dev and sanitizer presets), so no submodule or internet access + # is needed after the initial cmake --preset dev. if (NOT DEFINED GOOGLETEST_GIT_REPO) set(GOOGLETEST_GIT_REPO https://github.com/google/googletest.git) endif() - - # Fetch and install GoogleTest dependency - include(FetchContent) FetchContent_Declare( googletest GIT_REPOSITORY ${GOOGLETEST_GIT_REPO} @@ -209,17 +213,28 @@ if(BUILD_TEST) gtest_discover_tests(teststochtree) endif() -# Standalone C++ Program for Debugging +# Standalone C++ Programs for Debugging if(BUILD_DEBUG_TARGETS) - # Build test suite - add_executable(debugstochtree debug/api_debug.cpp) set(StochTree_DEBUG_HEADER_DIR ${PROJECT_SOURCE_DIR}/debug) + + # BART debug driver + add_executable(bart_debug debug/bart_debug.cpp) + if(USE_OPENMP) + target_include_directories(bart_debug 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(bart_debug PRIVATE stochtree_objs ${OpenMP_libomp_LIBRARY}) + else() + target_include_directories(bart_debug 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(bart_debug PRIVATE stochtree_objs) + endif() + + # BCF debug driver + add_executable(bcf_debug debug/bcf_debug.cpp) 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}) + target_include_directories(bcf_debug 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(bcf_debug 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) + target_include_directories(bcf_debug 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(bcf_debug PRIVATE stochtree_objs) endif() endif() diff --git a/CMakePresets.json b/CMakePresets.json new file mode 100644 index 00000000..2a4730a7 --- /dev/null +++ b/CMakePresets.json @@ -0,0 +1,125 @@ +{ + "version": 6, + "configurePresets": [ + { + "name": "dev", + "displayName": "Dev (Debug + Tests)", + "description": "Debug build with GoogleTest and debug targets — primary development preset", + "binaryDir": "${sourceDir}/build", + "cacheVariables": { + "USE_DEBUG": "ON", + "BUILD_DEBUG_TARGETS": "ON", + "BUILD_TEST": "ON", + "BUILD_PYTHON": "OFF", + "CMAKE_EXPORT_COMPILE_COMMANDS": "ON", + "FETCHCONTENT_UPDATES_DISCONNECTED": "ON" + } + }, + { + "name": "dev-quick", + "displayName": "Dev (Debug, no tests)", + "description": "Debug build with debug targets, skips GoogleTest download — faster configure", + "binaryDir": "${sourceDir}/build", + "cacheVariables": { + "USE_DEBUG": "ON", + "BUILD_DEBUG_TARGETS": "ON", + "BUILD_TEST": "OFF", + "BUILD_PYTHON": "ON", + "CMAKE_EXPORT_COMPILE_COMMANDS": "ON" + } + }, + { + "name": "release", + "displayName": "Release", + "description": "Optimized release build", + "binaryDir": "${sourceDir}/build-release", + "cacheVariables": { + "USE_DEBUG": "OFF", + "BUILD_DEBUG_TARGETS": "OFF", + "BUILD_TEST": "OFF", + "BUILD_PYTHON": "OFF", + "CMAKE_EXPORT_COMPILE_COMMANDS": "OFF" + } + }, + { + "name": "release-drivers", + "displayName": "Release (with debug drivers)", + "description": "Optimized release build that also compiles the debug driver executables", + "binaryDir": "${sourceDir}/build-release-drivers", + "cacheVariables": { + "USE_DEBUG": "OFF", + "BUILD_DEBUG_TARGETS": "ON", + "BUILD_TEST": "OFF", + "BUILD_PYTHON": "OFF", + "CMAKE_EXPORT_COMPILE_COMMANDS": "OFF" + } + }, + { + "name": "sanitizer", + "displayName": "Sanitizer (ASAN + UBSAN)", + "description": "Debug build with address and undefined-behavior sanitizers", + "binaryDir": "${sourceDir}/build-sanitizer", + "cacheVariables": { + "USE_DEBUG": "ON", + "USE_SANITIZER": "ON", + "BUILD_DEBUG_TARGETS": "ON", + "BUILD_TEST": "ON", + "BUILD_PYTHON": "OFF", + "CMAKE_EXPORT_COMPILE_COMMANDS": "ON", + "FETCHCONTENT_UPDATES_DISCONNECTED": "ON" + } + } + ], + "buildPresets": [ + { + "name": "dev", + "displayName": "Dev (Debug + Tests)", + "configurePreset": "dev", + "jobs": 0 + }, + { + "name": "dev-quick", + "displayName": "Dev (Debug, no tests)", + "configurePreset": "dev-quick", + "jobs": 0 + }, + { + "name": "release", + "displayName": "Release", + "configurePreset": "release", + "jobs": 0 + }, + { + "name": "release-drivers", + "displayName": "Release (with debug drivers)", + "configurePreset": "release-drivers", + "jobs": 0 + }, + { + "name": "sanitizer", + "displayName": "Sanitizer", + "configurePreset": "sanitizer", + "jobs": 0 + } + ], + "testPresets": [ + { + "name": "dev", + "displayName": "Dev tests", + "configurePreset": "dev", + "output": { + "outputOnFailure": true, + "verbosity": "default" + } + }, + { + "name": "sanitizer", + "displayName": "Sanitizer tests", + "configurePreset": "sanitizer", + "output": { + "outputOnFailure": true, + "verbosity": "default" + } + } + ] +} diff --git a/MANIFEST.in b/MANIFEST.in index 82d41ade..e110685d 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -8,7 +8,7 @@ recursive-include include *.hpp # Vendor package source files (excluding R-specific) recursive-include src *.cpp recursive-include src *.h -exclude src/cpp11.cpp src/forest.cpp src/kernel.cpp src/R_data.cpp src/R_random_effects.cpp src/R_utils.cpp src/sampler.cpp src/serialization.cpp src/stochtree_types.h +exclude src/cpp11.cpp src/forest.cpp src/kernel.cpp src/R_data.cpp src/R_random_effects.cpp src/R_utils.cpp src/R_bart.cpp src/R_bcf.cpp src/sampler.cpp src/serialization.cpp src/stochtree_types.h # Remove the CRAN extensionless LICENSE file exclude LICENSE diff --git a/NAMESPACE b/NAMESPACE index 1029382e..54541abb 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -1,5 +1,7 @@ # Generated by roxygen2: do not edit by hand +S3method("$",bartmodel) +S3method("$",bcfmodel) S3method(extractParameter,bartmodel) S3method(extractParameter,bcfmodel) S3method(getRandomEffectSamples,bartmodel) diff --git a/R/bart.R b/R/bart.R index 53f747e6..ae9a180d 100644 --- a/R/bart.R +++ b/R/bart.R @@ -88,11 +88,12 @@ NULL #' We do not currently support (but plan to in the near future), test set evaluation for group labels #' that were not in the training set. #' @param rfx_basis_test (Optional) Test set basis for "random-slope" regression in additive random effects model. -#' @param observation_weights (Optional) Numeric vector of observation weights of length `nrow(X_train)`. Weights are +#' @param observation_weights_train (Optional) Numeric vector of observation weights of length `nrow(X_train)`. Weights are #' applied as `y_i | - ~ N(mu(X_i), sigma^2 / w_i)`, so larger weights increase an observation's influence on the fit. #' All weights must be non-negative. Default: `NULL` (all observations equally weighted). Compatible with Gaussian #' (continuous/identity) and probit outcome models; not compatible with cloglog link functions. Note: these are #' referred to internally in the C++ layer as "variance weights" (`var_weights`), since they scale the residual variance. +#' @param observation_weights Deprecated alias for `observation_weights_train`; will be removed in a future release. #' @param num_gfr Number of "warm-start" iterations run using the grow-from-root algorithm (He and Hahn, 2021). Default: 5. #' @param num_burnin Number of "burn-in" iterations of the MCMC sampler. Default: 0. #' @param num_mcmc Number of "retained" iterations of the MCMC sampler. Default: 100. @@ -157,6 +158,7 @@ NULL #' - `variance_prior_shape` Shape parameter for the inverse gamma prior on the variance of the random effects "group parameter." Default: `1`. #' - `variance_prior_scale` Scale parameter for the inverse gamma prior on the variance of the random effects "group parameter." Default: `1`. #' +#' #' @return List of sampling outputs and a wrapper around the sampled forests (which can be used for in-memory prediction on new data, or serialized to JSON on disk). #' @export #' @@ -194,6 +196,7 @@ bart <- function( leaf_basis_test = NULL, rfx_group_ids_test = NULL, rfx_basis_test = NULL, + observation_weights_train = NULL, observation_weights = NULL, num_gfr = 5, num_burnin = 0, @@ -388,19 +391,6 @@ bart <- function( )) } - # Set a function-scoped RNG if user provided a random seed - custom_rng <- random_seed >= 0 - has_existing_random_seed <- F - if (custom_rng) { - # Cache original global environment RNG state (if it exists) - if (exists(".Random.seed", envir = .GlobalEnv)) { - original_global_seed <- .Random.seed - has_existing_random_seed <- T - } - # Set new seed and store associated RNG state - set.seed(random_seed) - } - # Check if there are enough GFR samples to seed num_chains samplers if (num_gfr > 0) { if (num_chains > num_gfr) { @@ -450,23 +440,23 @@ bart <- function( previous_y_bar <- previous_bart_model$model_params$outcome_mean previous_y_scale <- previous_bart_model$model_params$outcome_scale if (previous_bart_model$model_params$include_mean_forest) { - previous_forest_samples_mean <- previous_bart_model$mean_forests + previous_forest_samples_mean <- previous_bart_model$samples$materialize_mean_forest() } else { previous_forest_samples_mean <- NULL } if (previous_bart_model$model_params$include_variance_forest) { - previous_forest_samples_variance <- previous_bart_model$variance_forests + previous_forest_samples_variance <- previous_bart_model$samples$materialize_variance_forest() } else { previous_forest_samples_variance <- NULL } if (previous_bart_model$model_params$sample_sigma2_global) { - previous_global_var_samples <- previous_bart_model$sigma2_global_samples / + previous_global_var_samples <- previous_bart_model$samples$global_var_samples() / (previous_y_scale * previous_y_scale) } else { previous_global_var_samples <- NULL } if (previous_bart_model$model_params$sample_sigma2_leaf) { - previous_leaf_var_samples <- previous_bart_model$sigma2_leaf_samples + previous_leaf_var_samples <- previous_bart_model$samples$leaf_scale_samples() } else { previous_leaf_var_samples <- NULL } @@ -476,7 +466,7 @@ bart <- function( previous_rfx_samples <- NULL } if (previous_bart_model$model_params$outcome_model$link == "cloglog") { - previous_cloglog_cutpoint_samples <- previous_bart_model$cloglog_cutpoint_samples + previous_cloglog_cutpoint_samples <- previous_bart_model$samples$cloglog_cutpoint_samples() previous_cloglog_num_categories <- previous_bart_model$cloglog_num_categories } else { previous_cloglog_cutpoint_samples <- NULL @@ -503,26 +493,38 @@ bart <- function( # Determine whether conditional mean, variance, or both will be modeled if (num_trees_variance > 0) { - include_variance_forest = TRUE + include_variance_forest <- TRUE } else { - include_variance_forest = FALSE + include_variance_forest <- FALSE } if (num_trees_mean > 0) { - include_mean_forest = TRUE + include_mean_forest <- TRUE } else { - include_mean_forest = FALSE + include_mean_forest <- FALSE } - # observation_weights compatibility checks + # `observation_weights` was renamed to `observation_weights_train`; honor the + # deprecated argument for one release cycle. if (!is.null(observation_weights)) { + warning( + "`observation_weights` is deprecated and will be removed in a future release; use `observation_weights_train` instead." + ) + if (is.null(observation_weights_train)) { + observation_weights_train <- observation_weights + } + } + + # observation_weights_train compatibility checks + if (!is.null(observation_weights_train)) { if (link_is_cloglog) { stop( - "observation_weights are not compatible with cloglog link functions." + "observation_weights_train are not compatible with cloglog link functions." ) } if (include_variance_forest) { - warning( - "Results may be unreliable when observation_weights are deployed alongside a variance forest model." + stop( + "observation_weights_train are not compatible with a variance forest model. ", + "Use either observation_weights_train or a variance forest, not both." ) } } @@ -545,28 +547,28 @@ bart <- function( # Variable weight preprocessing (and initialization if necessary) if (is.null(variable_weights)) { - variable_weights = rep(1 / ncol(X_train), ncol(X_train)) + variable_weights <- rep(1 / ncol(X_train), ncol(X_train)) } if (any(variable_weights < 0)) { stop("variable_weights cannot have any negative weights") } # Observation weight validation - if (!is.null(observation_weights)) { - if (!is.numeric(observation_weights)) { - stop("observation_weights must be a numeric vector") + if (!is.null(observation_weights_train)) { + if (!is.numeric(observation_weights_train)) { + stop("observation_weights_train must be a numeric vector") } - if (length(observation_weights) != nrow(X_train)) { - stop("length(observation_weights) must equal nrow(X_train)") + if (length(observation_weights_train) != nrow(X_train)) { + stop("length(observation_weights_train) must equal nrow(X_train)") } - if (any(observation_weights < 0)) { - stop("observation_weights cannot have any negative values") + if (any(observation_weights_train < 0)) { + stop("observation_weights_train cannot have any negative values") } - if (all(observation_weights == 0) && num_gfr > 0) { + if (all(observation_weights_train == 0) && num_gfr > 0) { stop( - "observation_weights are all zero (prior sampling mode) but num_gfr > 0. ", + "observation_weights_train are all zero (prior sampling mode) but num_gfr > 0. ", "GFR warm-start is data-dependent and ill-defined with zero weights. ", - "Set num_gfr = 0 when using all-zero observation_weights." + "Set num_gfr = 0 when using all-zero observation_weights_train." ) } } @@ -970,10 +972,10 @@ bart <- function( } # Determine whether a basis vector is provided - has_basis = !is.null(leaf_basis_train) + has_basis <- !is.null(leaf_basis_train) # Determine whether a test set is provided - has_test = !is.null(X_test) + has_test <- !is.null(X_test) # Preliminary runtime checks for probit link if (!include_mean_forest) { @@ -1055,163 +1057,6 @@ bart <- function( } } - # Handle standardization, prior calibration, and initialization of forest - # differently for binary and continuous outcomes - if (link_is_probit) { - # Probit-scale intercept: center the forest on the population-average latent mean. - # The forest predicts mu(X) and y_bar_train is added back at prediction time. - # The latent z sampling uses y_bar_train to set the correct truncated normal mean and to center z before the residual update. - y_bar_train <- qnorm(mean_cpp(as.numeric(y_train))) - y_std_train <- 1 - standardize <- FALSE - - # Set a pseudo outcome by subtracting mean_cpp(y_train) from y_train - resid_train <- y_train - mean_cpp(as.numeric(y_train)) - - # Set initial values of root nodes to 0.0 (in probit scale) - init_val_mean <- 0.0 - - # Calibrate priors for sigma^2 and tau - # Set sigma2_init to 1, ignoring default provided - sigma2_init <- 1.0 - # Skip variance_forest_init, since variance forests are not supported with probit link - if (is.null(b_leaf)) { - b_leaf <- 1 / (num_trees_mean) - } - if (has_basis) { - if (ncol(leaf_basis_train) > 1) { - if (is.null(sigma2_leaf_init)) { - sigma2_leaf_init <- diag( - 2 / (num_trees_mean), - ncol(leaf_basis_train) - ) - } - if (!is.matrix(sigma2_leaf_init)) { - current_leaf_scale <- as.matrix(diag( - sigma2_leaf_init, - ncol(leaf_basis_train) - )) - } else { - current_leaf_scale <- sigma2_leaf_init - } - } else { - if (is.null(sigma2_leaf_init)) { - sigma2_leaf_init <- as.matrix(2 / (num_trees_mean)) - } - if (!is.matrix(sigma2_leaf_init)) { - current_leaf_scale <- as.matrix(diag(sigma2_leaf_init, 1)) - } else { - current_leaf_scale <- sigma2_leaf_init - } - } - } else { - if (is.null(sigma2_leaf_init)) { - sigma2_leaf_init <- as.matrix(2 / (num_trees_mean)) - } - if (!is.matrix(sigma2_leaf_init)) { - current_leaf_scale <- as.matrix(diag(sigma2_leaf_init, 1)) - } else { - current_leaf_scale <- sigma2_leaf_init - } - } - current_sigma2 <- sigma2_init - } else if (link_is_cloglog) { - # Fix offset to 0 and scale to 1 - y_bar_train <- 0 - y_std_train <- 1 - standardize <- FALSE - - # Remap outcomes to start from 0 - resid_train <- as.numeric(y_train - min(unique_outcomes)) - cloglog_num_categories <- max(resid_train) + 1 - - # Set initial values of root nodes to 0.0 (in linear scale) - init_val_mean <- 0.0 - - # Calibrate priors for sigma^2 and tau - # Set sigma2_init to 1, ignoring default provided - sigma2_init <- 1.0 - if (is.null(sigma2_leaf_init)) { - sigma2_leaf_init <- as.matrix(2 / (num_trees_mean)) - } - current_sigma2 <- sigma2_init - current_leaf_scale <- sigma2_leaf_init - - # Set first cutpoint to 0 for identifiability - cloglog_cutpoint_0 <- 0 - - # Set shape and rate parameters for conditional gamma model - cloglog_forest_shape <- 2.0 - cloglog_forest_rate <- 2.0 - } else { - # Only standardize if user requested - if (standardize) { - y_bar_train <- mean_cpp(as.numeric(y_train)) - y_std_train <- sd_cpp(as.numeric(y_train)) - } else { - y_bar_train <- 0 - y_std_train <- 1 - } - - # Compute standardized outcome - resid_train <- (y_train - y_bar_train) / y_std_train - - # Compute initial value of root nodes in mean forest - init_val_mean <- mean_cpp(as.numeric(resid_train)) - - # Calibrate priors for sigma^2 and tau - if (is.null(sigma2_init)) { - sigma2_init <- 1.0 * var_cpp(as.numeric(resid_train)) - } - if (is.null(variance_forest_init)) { - variance_forest_init <- 1.0 * var_cpp(as.numeric(resid_train)) - } - if (is.null(b_leaf)) { - b_leaf <- var_cpp(as.numeric(resid_train)) / (2 * num_trees_mean) - } - if (has_basis) { - if (ncol(leaf_basis_train) > 1) { - if (is.null(sigma2_leaf_init)) { - sigma2_leaf_init <- diag( - 2 * var_cpp(as.numeric(resid_train)) / (num_trees_mean), - ncol(leaf_basis_train) - ) - } - if (!is.matrix(sigma2_leaf_init)) { - current_leaf_scale <- as.matrix(diag( - sigma2_leaf_init, - ncol(leaf_basis_train) - )) - } else { - current_leaf_scale <- sigma2_leaf_init - } - } else { - if (is.null(sigma2_leaf_init)) { - sigma2_leaf_init <- as.matrix( - 2 * var_cpp(as.numeric(resid_train)) / (num_trees_mean) - ) - } - if (!is.matrix(sigma2_leaf_init)) { - current_leaf_scale <- as.matrix(diag(sigma2_leaf_init, 1)) - } else { - current_leaf_scale <- sigma2_leaf_init - } - } - } else { - if (is.null(sigma2_leaf_init)) { - sigma2_leaf_init <- as.matrix( - 2 * var_cpp(as.numeric(resid_train)) / (num_trees_mean) - ) - } - if (!is.matrix(sigma2_leaf_init)) { - current_leaf_scale <- as.matrix(diag(sigma2_leaf_init, 1)) - } else { - current_leaf_scale <- sigma2_leaf_init - } - } - current_sigma2 <- sigma2_init - } - # Determine leaf model type if ((!has_basis) && (!link_is_cloglog)) { leaf_model_mean_forest <- 0 @@ -1230,21 +1075,21 @@ bart <- function( # Unpack model type info if (leaf_model_mean_forest == 0) { - leaf_dimension = 1 - is_leaf_constant = TRUE - leaf_regression = FALSE + leaf_dimension <- 1 + is_leaf_constant <- TRUE + leaf_regression <- FALSE } else if (leaf_model_mean_forest == 1) { stopifnot(has_basis) stopifnot(ncol(leaf_basis_train) == 1) - leaf_dimension = 1 - is_leaf_constant = FALSE - leaf_regression = TRUE + leaf_dimension <- 1 + is_leaf_constant <- FALSE + leaf_regression <- TRUE } else if (leaf_model_mean_forest == 2) { stopifnot(has_basis) stopifnot(ncol(leaf_basis_train) > 1) - leaf_dimension = ncol(leaf_basis_train) - is_leaf_constant = FALSE - leaf_regression = TRUE + leaf_dimension <- ncol(leaf_basis_train) + is_leaf_constant <- FALSE + leaf_regression <- TRUE if (sample_sigma2_leaf) { warning( "Sampling leaf scale not yet supported for multivariate leaf models, so the leaf scale parameter will not be sampled in this model." @@ -1252,1346 +1097,406 @@ bart <- function( sample_sigma2_leaf <- FALSE } } else if (leaf_model_mean_forest == 4) { - leaf_dimension = 1 - is_leaf_constant = TRUE - leaf_regression = FALSE - } - - # Data - if (leaf_regression) { - forest_dataset_train <- createForestDataset( - X_train, - leaf_basis_train, - observation_weights - ) - if (has_test) { - forest_dataset_test <- createForestDataset(X_test, leaf_basis_test) - } - requires_basis <- TRUE - } else { - forest_dataset_train <- createForestDataset( - X_train, - variance_weights = observation_weights - ) - if (has_test) { - forest_dataset_test <- createForestDataset(X_test) - } - requires_basis <- FALSE - } - outcome_train <- createOutcome(resid_train) - - # Random number generator (std::mt19937) - if (is.null(random_seed)) { - random_seed = sample(1:10000, 1, FALSE) - } - rng <- createCppRNG(random_seed) - - # Separate ordinal sampler object for cloglog - if (link_is_cloglog) { - ordinal_sampler <- ordinal_sampler_cpp() + leaf_dimension <- 1 + is_leaf_constant <- TRUE + leaf_regression <- FALSE } - # Sampling data structures - feature_types <- as.integer(feature_types) - global_model_config <- createGlobalModelConfig( - global_error_variance = current_sigma2 + cloglog_num_categories <- ifelse( + link_is_cloglog, + max(y_train - min(y_train)) + 1, + 0 ) - if (include_mean_forest) { - forest_model_config_mean <- createForestModelConfig( - feature_types = feature_types, - num_trees = num_trees_mean, - num_features = ncol(X_train), - num_observations = nrow(X_train), - variable_weights = variable_weights_mean, - leaf_dimension = leaf_dimension, - alpha = alpha_mean, - beta = beta_mean, - min_samples_leaf = min_samples_leaf_mean, - max_depth = max_depth_mean, - leaf_model_type = leaf_model_mean_forest, - leaf_model_scale = current_leaf_scale, - cutpoint_grid_size = cutpoint_grid_size, - num_features_subsample = num_features_subsample_mean - ) - if (link_is_cloglog) { - forest_model_config_mean$update_cloglog_forest_shape(cloglog_forest_shape) - forest_model_config_mean$update_cloglog_forest_rate(cloglog_forest_rate) - } - forest_model_mean <- createForestModel( - forest_dataset_train, - forest_model_config_mean, - global_model_config + model_params_r <- list( + "a_global" = a_global, + "b_global" = b_global, + "a_leaf" = a_leaf, + "standardize" = standardize, + "leaf_dimension" = leaf_dimension, + "is_leaf_constant" = is_leaf_constant, + "leaf_regression" = leaf_regression, + "requires_basis" = leaf_regression, + "num_covariates" = num_cov_orig, + "num_basis" = ifelse( + is.null(leaf_basis_train), + 0, + ncol(leaf_basis_train) + ), + "num_gfr" = num_gfr, + "num_burnin" = num_burnin, + "num_mcmc" = num_mcmc, + "keep_every" = keep_every, + "num_chains" = num_chains, + "has_basis" = !is.null(leaf_basis_train), + "has_rfx" = has_rfx, + "has_rfx_basis" = has_basis_rfx, + "num_rfx_basis" = num_basis_rfx, + "sample_sigma2_global" = sample_sigma2_global, + "sample_sigma2_leaf" = sample_sigma2_leaf, + "include_mean_forest" = include_mean_forest, + "include_variance_forest" = include_variance_forest, + "outcome_model" = outcome_model, + "probit_outcome_model" = probit_outcome_model, + "cloglog_num_categories" = cloglog_num_categories, + "rfx_model_spec" = rfx_model_spec + ) + + # Expand dimensions on RFX prior parameters if provided + # Working parameter (should be expanded to a vector if provided as a scalar) + if (!is.null(rfx_working_parameter_prior_mean)) { + rfx_working_parameter_prior_mean <- expand_dims_1d( + rfx_working_parameter_prior_mean, + num_rfx_components ) } - if (include_variance_forest) { - forest_model_config_variance <- createForestModelConfig( - feature_types = feature_types, - num_trees = num_trees_variance, - num_features = ncol(X_train), - num_observations = nrow(X_train), - variable_weights = variable_weights_variance, - leaf_dimension = 1, - alpha = alpha_variance, - beta = beta_variance, - min_samples_leaf = min_samples_leaf_variance, - max_depth = max_depth_variance, - leaf_model_type = leaf_model_variance_forest, - variance_forest_shape = a_forest, - variance_forest_scale = b_forest, - cutpoint_grid_size = cutpoint_grid_size, - num_features_subsample = num_features_subsample_variance - ) - forest_model_variance <- createForestModel( - forest_dataset_train, - forest_model_config_variance, - global_model_config + + # Group parameter (should be expanded to a matrix if provided as a scalar) + if (!is.null(rfx_group_parameter_prior_mean)) { + rfx_group_parameter_prior_mean <- expand_dims_2d( + rfx_group_parameter_prior_mean, + num_rfx_components, + num_rfx_groups ) } - # Container of forest samples - if (include_mean_forest) { - forest_samples_mean <- createForestSamples( - num_trees_mean, - leaf_dimension, - is_leaf_constant, - FALSE - ) - active_forest_mean <- createForest( - num_trees_mean, - leaf_dimension, - is_leaf_constant, - FALSE + # Working parameter (should be expanded to a diagonal matrix if provided as a scalar) + if (!is.null(rfx_working_parameter_prior_cov)) { + rfx_working_parameter_prior_cov <- expand_dims_2d_diag( + rfx_working_parameter_prior_cov, + num_rfx_components ) } - if (include_variance_forest) { - forest_samples_variance <- createForestSamples( - num_trees_variance, - 1, - TRUE, - TRUE - ) - active_forest_variance <- createForest( - num_trees_variance, - 1, - TRUE, - TRUE + + # Group parameter (should be expanded to a diagonal matrix if provided as a scalar) + if (!is.null(rfx_group_parameter_prior_cov)) { + rfx_group_parameter_prior_cov <- expand_dims_2d_diag( + rfx_group_parameter_prior_cov, + num_rfx_components ) } - # Random effects initialization - if (has_rfx) { - # Prior parameters - if (is.null(rfx_working_parameter_prior_mean)) { - if (num_rfx_components == 1) { - alpha_init <- c(0) - } else if (num_rfx_components > 1) { - alpha_init <- rep(0, num_rfx_components) - } else { - stop("There must be at least 1 random effect component") - } + # Specify the BART config + bart_config <- list( + "standardize_outcome" = standardize, + "num_threads" = num_threads, + "verbose" = verbose, + "cutpoint_grid_size" = cutpoint_grid_size, + "link_function" = ifelse( + outcome_model$link == "identity", + 0, + ifelse(outcome_model$link == "probit", 1, 2) + ), + "outcome_type" = ifelse( + outcome_model$outcome == "continuous", + 0, + ifelse(outcome_model$outcome == "binary", 1, 2) + ), + "random_seed" = random_seed, + "keep_gfr" = keep_gfr, + "keep_burnin" = keep_burnin, + "a_sigma2_global" = a_global, + "b_sigma2_global" = b_global, + "sigma2_global_init" = sigma2_init, + "sample_sigma2_global" = sample_sigma2_global, + "num_trees_mean" = num_trees_mean, + "alpha_mean" = alpha_mean, + "beta_mean" = beta_mean, + "min_samples_leaf_mean" = min_samples_leaf_mean, + "max_depth_mean" = max_depth_mean, + "leaf_constant_mean" = is_leaf_constant, + "leaf_dim_mean" = leaf_dimension, + "exponentiated_leaf_mean" = FALSE, + "num_features_subsample_mean" = num_features_subsample_mean, + "a_sigma2_mean" = a_leaf, + "b_sigma2_mean" = b_leaf, + "sigma2_mean_init" = if (is.matrix(sigma2_leaf_init)) { + NULL } else { - alpha_init <- expand_dims_1d( - rfx_working_parameter_prior_mean, - num_rfx_components - ) - } - - if (is.null(rfx_group_parameter_prior_mean)) { - xi_init <- matrix( - rep(alpha_init, num_rfx_groups), - num_rfx_components, - num_rfx_groups - ) + sigma2_leaf_init + }, + "sample_sigma2_leaf_mean" = sample_sigma2_leaf, + "mean_leaf_model_type" = leaf_model_mean_forest, + "sigma2_leaf_mean_matrix" = if (is.matrix(sigma2_leaf_init)) { + as.numeric(sigma2_leaf_init) } else { - xi_init <- expand_dims_2d( - rfx_group_parameter_prior_mean, - num_rfx_components, - num_rfx_groups - ) - } - - if (is.null(rfx_working_parameter_prior_cov)) { - sigma_alpha_init <- diag(1, num_rfx_components, num_rfx_components) + NULL + }, + "num_classes_cloglog" = cloglog_num_categories, + "cloglog_leaf_prior_shape" = cloglog_leaf_prior_shape, + "cloglog_leaf_prior_scale" = cloglog_leaf_prior_scale, + "cloglog_cutpoint_0" = 0, + "num_trees_variance" = num_trees_variance, + "leaf_prior_calibration_param" = a_0, + "shape_variance_forest" = a_forest, + "scale_variance_forest" = b_forest, + "variance_forest_leaf_init" = variance_forest_init, + "alpha_variance" = alpha_variance, + "beta_variance" = beta_variance, + "min_samples_leaf_variance" = min_samples_leaf_variance, + "max_depth_variance" = max_depth_variance, + "leaf_constant_variance" = TRUE, + "leaf_dim_variance" = 1, + "exponentiated_leaf_variance" = TRUE, + "num_features_subsample_variance" = num_features_subsample_variance, + "feature_types" = as.integer(feature_types), + "sweep_update_indices_mean" = if (num_trees_mean > 0) { + 0:(num_trees_mean - 1) + } else { + NULL + }, + "sweep_update_indices_variance" = if (num_trees_variance > 0) { + 0:(num_trees_variance - 1) } else { - sigma_alpha_init <- expand_dims_2d_diag( - rfx_working_parameter_prior_cov, - num_rfx_components + NULL + }, + "var_weights_mean" = variable_weights_mean, + "var_weights_variance" = variable_weights_variance, + "has_random_effects" = has_rfx, + "rfx_model_spec" = if (has_rfx) { + ifelse( + rfx_model_spec == "custom", + 0, + ifelse(rfx_model_spec == "intercept_only", 1, NULL) ) + } else { + NULL + }, + "rfx_working_parameter_mean_prior" = if (has_rfx) { + rfx_working_parameter_prior_mean + } else { + NULL + }, + "rfx_working_parameter_cov_prior" = if (has_rfx) { + rfx_working_parameter_prior_cov + } else { + NULL + }, + "rfx_group_parameter_mean_prior" = if (has_rfx) { + rfx_group_parameter_prior_mean + } else { + NULL + }, + "rfx_group_parameter_cov_prior" = if (has_rfx) { + rfx_group_parameter_prior_cov + } else { + NULL + }, + "rfx_variance_prior_shape" = if (has_rfx) { + rfx_variance_prior_shape + } else { + NULL + }, + "rfx_variance_prior_scale" = if (has_rfx) { + rfx_variance_prior_scale + } else { + NULL } + ) - if (is.null(rfx_group_parameter_prior_cov)) { - sigma_xi_init <- diag(1, num_rfx_components, num_rfx_components) + bart_samples <- BARTSamples$new() + bart_metadata <- bart_sample_cpp( + samples = bart_samples$samples_ptr, + X_train = X_train, + y_train = if (link_is_cloglog) { + as.numeric(y_train - min(y_train)) } else { - sigma_xi_init <- expand_dims_2d_diag( - rfx_group_parameter_prior_cov, - num_rfx_components - ) - } + y_train + }, + X_test = if (exists("X_test")) X_test else NULL, + n_train = nrow(X_train), + n_test = if (!is.null(X_test)) nrow(X_test) else 0L, + p = ncol(X_train), + basis_train = if (exists("leaf_basis_train")) leaf_basis_train else NULL, + basis_test = if (exists("leaf_basis_test")) leaf_basis_test else NULL, + basis_dim = if (!is.null(leaf_basis_train)) { + ncol(leaf_basis_train) + } else { + 0L + }, + obs_weights_train = observation_weights_train, + obs_weights_test = NULL, + rfx_group_ids_train = if (exists("rfx_group_ids_train")) { + rfx_group_ids_train + } else { + NULL + }, + rfx_group_ids_test = if (exists("rfx_group_ids_test")) { + rfx_group_ids_test + } else { + NULL + }, + rfx_basis_train = if (exists("rfx_basis_train")) { + rfx_basis_train + } else { + NULL + }, + rfx_basis_test = if (exists("rfx_basis_test")) rfx_basis_test else NULL, + rfx_num_groups = if (exists("num_rfx_groups")) { + as.integer(num_rfx_groups) + } else { + 0L + }, + rfx_basis_dim = as.integer(num_basis_rfx), + num_gfr = as.integer(num_gfr), + num_burnin = as.integer(num_burnin), + keep_every = as.integer(keep_every), + num_mcmc = as.integer(num_mcmc), + num_chains = as.integer(num_chains), + config_input = bart_config + ) + result <- list() + model_params_cpp <- list( + "sigma2_init" = bart_metadata[["sigma2_global_init"]], + "sigma2_leaf_init" = bart_metadata[["sigma2_mean_init"]], + "b_leaf" = bart_metadata[["b_sigma2_mean"]], + "a_forest" = bart_metadata[["shape_variance_forest"]], + "b_forest" = bart_metadata[["scale_variance_forest"]], + "outcome_mean" = bart_samples$y_bar(), + "outcome_scale" = bart_samples$y_std(), + "num_samples" = bart_samples$num_samples() + ) + model_params <- c(model_params_r, model_params_cpp) + result[["model_params"]] <- model_params + result[["train_set_metadata"]] <- X_train_metadata + result[["samples"]] <- bart_samples - sigma_xi_shape <- rfx_variance_prior_shape - sigma_xi_scale <- rfx_variance_prior_scale + # Unpack RFX samples + if (has_rfx) { + # Only need to store unique group IDs, everything else stored in BARTSamples + result[["rfx_unique_group_ids"]] <- levels(group_ids_factor) + } - # Random effects data structure and storage container - rfx_dataset_train <- createRandomEffectsDataset( - rfx_group_ids_train, - rfx_basis_train - ) - rfx_tracker_train <- createRandomEffectsTracker(rfx_group_ids_train) - rfx_model <- createRandomEffectsModel( - num_rfx_components, - num_rfx_groups + class(result) <- "bartmodel" + + return(result) +} + +#' Guard accessor for a `bartmodel`'s removed direct-access forest and parameter / prediction fields. +#' +#' The sampled forests / parameters / predictions are owned by a single `BARTSamples` object stored in `object$samples`. +#' They are no longer stored or accessible via `$mean_forests` / `$variance_forests`. +#' Similarly, sampled parameter vectors are no longer stored or accessible via `$sigma2_global_samples` / `$sigma2_leaf_samples` +#' and cached predictions are not accessible via `$y_hat_train` / `$y_hat_test` / `$sigma2_x_hat_train` / `$sigma2_x_hat_test`. +#' Accessing any of these model terms by name raises an error pointing at the supported extraction path. +#' @noRd +#' @export +`$.bartmodel` <- function(x, name) { + if (identical(name, "mean_forests") || identical(name, "variance_forests")) { + stop( + sprintf( + paste0( + "`bartmodel$%s` has been removed. The sampled forests are owned by `model$bart_samples`; ", + "you can extract a standalone copy with `extractForest()`." + ), + name + ), + call. = FALSE ) - rfx_model$set_working_parameter(alpha_init) - rfx_model$set_group_parameters(xi_init) - rfx_model$set_working_parameter_cov(sigma_alpha_init) - rfx_model$set_group_parameter_cov(sigma_xi_init) - rfx_model$set_variance_prior_shape(sigma_xi_shape) - rfx_model$set_variance_prior_scale(sigma_xi_scale) - rfx_samples <- createRandomEffectSamples( - num_rfx_components, - num_rfx_groups, - rfx_tracker_train + } else if ( + identical(name, "sigma2_global_samples") || + identical(name, "sigma2_leaf_samples") || + identical(name, "cloglog_cutpoint_samples") || + identical(name, "y_hat_train") || + identical(name, "y_hat_test") || + identical(name, "sigma2_x_hat_train") || + identical(name, "sigma2_x_hat_test") + ) { + stop( + sprintf( + paste0( + "`bartmodel$%s` has been removed. The parameters are stored in a C++ BARTSamples object; ", + "you can extract a standalone copy with `extractParameter()`." + ), + name + ), + call. = FALSE ) + } else { + .subset2(x, name) } +} - # Container of parameter samples - num_actual_mcmc_iter <- num_mcmc * keep_every - num_samples <- num_gfr + num_burnin + num_actual_mcmc_iter - # Delete GFR samples from these containers after the fact if desired - # num_retained_samples <- ifelse(keep_gfr, num_gfr, 0) + ifelse(keep_burnin, num_burnin, 0) + num_mcmc - num_retained_samples <- num_gfr + - ifelse(keep_burnin, num_burnin, 0) + - num_mcmc * num_chains - if (sample_sigma2_global) { - global_var_samples <- rep(NA, num_retained_samples) - } - if (sample_sigma2_leaf) { - leaf_scale_samples <- rep(NA, num_retained_samples) - } - if (link_is_cloglog) { - cloglog_cutpoint_samples <- matrix( - NA_real_, - cloglog_num_categories - 1, - num_retained_samples - ) +#' @title Predict from a BART Model +#' @description +#' Predict from a sampled BART model on new data +#' +#' @param object Object of type `bart` containing draws of a regression forest and associated sampling outputs. +#' @param X Covariates used to determine tree leaf predictions for each observation. Must be passed as a matrix or dataframe. +#' @param leaf_basis (Optional) Bases used for prediction (by e.g. dot product with leaf values). Default: `NULL`. +#' @param rfx_group_ids (Optional) Test set group labels used for an additive random effects model. +#' We do not currently support (but plan to in the near future), test set evaluation for group labels +#' that were not in the training set. +#' @param rfx_basis (Optional) Test set basis for "random-slope" regression in additive random effects model. +#' @param type (Optional) Type of prediction to return. Options are "mean", which averages the predictions from every draw of a BART model, and "posterior", which returns the entire matrix of posterior predictions. Default: "posterior". +#' @param terms (Optional) Which model terms to include in the prediction. This can be a single term or a list of model terms. Options include "y_hat", "mean_forest", "rfx", "variance_forest", or "all". If a model doesn't have mean forest, random effects, or variance forest predictions, but one of those terms is request, the request will simply be ignored. If none of the requested terms are present in a model, this function will return `NULL` along with a warning. Default: "all". +#' @param scale (Optional) Scale of mean function predictions. Options are "linear", which returns predictions on the original scale of the mean forest / RFX terms, "probability", which transforms predictions into class probabilities for models with discrete outcomes, and "class", which returns predicted outcome categories for discrete outcome models. "probability" is only valid for outcome models with `outcome == 'binary'` or `outcome == 'ordinal'`. For binary outcomes, this will return the probability that `y == 1`, and for ordinal outcomes, this will return probabilities for each outcome label. Default: "linear". +#' @param ... (Optional) Other prediction parameters. +#' +#' @return List of prediction matrices or single prediction matrix / vector, depending on the terms requested. +#' @export +#' +#' @examples +#' n <- 100 +#' p <- 5 +#' X <- matrix(runif(n*p), ncol = p) +#' f_XW <- ( +#' ((0 <= X[,1]) & (0.25 > X[,1])) * (-7.5) + +#' ((0.25 <= X[,1]) & (0.5 > X[,1])) * (-2.5) + +#' ((0.5 <= X[,1]) & (0.75 > X[,1])) * (2.5) + +#' ((0.75 <= X[,1]) & (1 > X[,1])) * (7.5) +#' ) +#' noise_sd <- 1 +#' y <- f_XW + rnorm(n, 0, noise_sd) +#' test_set_pct <- 0.2 +#' n_test <- round(test_set_pct*n) +#' n_train <- n - n_test +#' test_inds <- sort(sample(1:n, n_test, replace = FALSE)) +#' train_inds <- (1:n)[!((1:n) %in% test_inds)] +#' X_test <- X[test_inds,] +#' X_train <- X[train_inds,] +#' y_test <- y[test_inds] +#' y_train <- y[train_inds] +#' bart_model <- bart(X_train = X_train, y_train = y_train, +#' num_gfr = 10, num_burnin = 0, num_mcmc = 10) +#' y_hat_test <- predict(bart_model, X=X_test)$y_hat +predict.bartmodel <- function( + object, + X, + leaf_basis = NULL, + rfx_group_ids = NULL, + rfx_basis = NULL, + type = "posterior", + terms = "all", + scale = "linear", + ... +) { + # Handle mean function scale + if (!is.character(scale)) { + stop("scale must be a string or character vector") } - if (include_mean_forest) { - mean_forest_pred_train <- matrix( - NA_real_, - nrow(X_train), - num_retained_samples - ) + if (!(scale %in% c("linear", "probability", "class"))) { + stop("scale must either be 'linear', 'probability', or 'class'") } - if (include_variance_forest) { - variance_forest_pred_train <- matrix( - NA_real_, - nrow(X_train), - num_retained_samples + outcome_model <- object$model_params$outcome_model + is_probit <- (outcome_model$link == "probit" && + outcome_model$outcome == "binary") + is_binary_cloglog <- (outcome_model$link == "cloglog" && + outcome_model$outcome == "binary") + is_ordinal_cloglog <- (outcome_model$link == "cloglog" && + outcome_model$outcome == "ordinal") + is_cloglog <- is_binary_cloglog || is_ordinal_cloglog + if ((scale == "probability") && (!(is_probit || is_cloglog))) { + stop( + "scale cannot be 'probability' for models not fit with a probit or cloglog outcome model" ) } - sample_counter <- 0 - - # Initialize the leaves of each tree in the mean forest - if (include_mean_forest) { - if (requires_basis) { - # Handle the case in which we must initialize root values in a leaf basis regression - # when init_val_mean != 0. To do this, we regress rep(init_val_mean, nrow(y_train)) - # on leaf_basis_train and use (coefs / num_trees_mean) as initial values - if (abs(init_val_mean) > 0.00001) { - init_val_y <- rep(init_val_mean, nrow(y_train)) - init_val_model <- lm(init_val_y ~ 0 + leaf_basis_train) - init_values_mean_forest <- coef(init_val_model) - if (any(is.na(init_values_mean_forest))) { - init_values_mean_forest[which(is.na(init_values_mean_forest))] <- 0. - } - } else { - init_values_mean_forest <- rep(init_val_mean, ncol(leaf_basis_train)) - } - } else { - init_values_mean_forest <- init_val_mean - } - active_forest_mean$prepare_for_sampler( - forest_dataset_train, - outcome_train, - forest_model_mean, - leaf_model_mean_forest, - init_values_mean_forest + if ((scale == "class") && (!(is_probit || is_cloglog))) { + stop( + "scale cannot be 'class' for models not fit with a probit or cloglog outcome model" ) } - - # Initialize the leaves of each tree in the variance forest - if (include_variance_forest) { - active_forest_variance$prepare_for_sampler( - forest_dataset_train, - outcome_train, - forest_model_variance, - leaf_model_variance_forest, - variance_forest_init - ) - } - - # Initialize auxiliary data for cloglog - if (link_is_cloglog) { - ## Allocate auxiliary data - train_size <- nrow(X_train) - # Latent variable (Z in Alam et al (2025) notation) - forest_dataset_train$add_auxiliary_dimension(train_size) - # Forest predictions (eta in Alam et al (2025) notation) - forest_dataset_train$add_auxiliary_dimension(train_size) - # Log-scale non-cumulative cutpoint (gamma in Alam et al (2025) notation) - forest_dataset_train$add_auxiliary_dimension(cloglog_num_categories - 1) - # Exponentiated cumulative cutpoints (exp(c_k) in Alam et al (2025) notation) - # This auxiliary series is designed so that the element stored at position `i` - # corresponds to the sum of all exponentiated gamma_j values for j < i. - # It has cloglog_num_categories elements instead of cloglog_num_categories - 1 because - # even the largest categorical index has a valid value of sum_{j < i} exp(gamma_j) - forest_dataset_train$add_auxiliary_dimension(cloglog_num_categories) - - ## Set initial values for auxiliary data - # Initialize latent variables to zero (slot 0) - for (i in 1:train_size) { - forest_dataset_train$set_auxiliary_data_value(0, i - 1, 0.0) - } - # Initialize forest predictions to zero (slot 1) - for (i in 1:train_size) { - forest_dataset_train$set_auxiliary_data_value(1, i - 1, 0.0) - } - # Initialize log-scale cutpoints to 0 - initial_gamma <- rep(0.0, cloglog_num_categories - 1) - for (i in seq_along(initial_gamma)) { - forest_dataset_train$set_auxiliary_data_value(2, i - 1, initial_gamma[i]) - } - # Convert to cumulative exponentiated cutpoints directly in C++ - ordinal_sampler_update_cumsum_exp_cpp( - ordinal_sampler, - forest_dataset_train$data_ptr - ) - } - - # Run GFR (warm start) if specified - if (num_gfr > 0) { - for (i in 1:num_gfr) { - # Keep all GFR samples at this stage -- remove from ForestSamples after MCMC - # keep_sample <- ifelse(keep_gfr, TRUE, FALSE) - keep_sample <- TRUE - if (keep_sample) { - sample_counter <- sample_counter + 1 - } - # Print progress - if (verbose) { - if ((i %% 10 == 0) || (i == num_gfr)) { - cat( - "Sampling", - i, - "out of", - num_gfr, - "XBART (grow-from-root) draws\n" - ) - } - } - - if (include_mean_forest) { - if (link_is_probit) { - # Sample latent probit variable, z | - - # outcome_pred is the centered forest prediction (not including y_bar_train). - # The truncated normal mean is outcome_pred + y_bar_train (the full eta on the probit scale). - # The residual stored is z - y_bar_train - outcome_pred so the forest sees a - # zero-centered signal and the prior shrinkage toward 0 is well-calibrated. - outcome_pred <- active_forest_mean$predict( - forest_dataset_train - ) - if (has_rfx) { - rfx_pred <- rfx_model$predict( - rfx_dataset_train, - rfx_tracker_train - ) - outcome_pred <- outcome_pred + rfx_pred - } - eta_pred <- outcome_pred + y_bar_train - mu0 <- eta_pred[y_train == 0] - mu1 <- eta_pred[y_train == 1] - u0 <- runif(sum(y_train == 0), 0, pnorm(0 - mu0)) - u1 <- runif(sum(y_train == 1), pnorm(0 - mu1), 1) - resid_train[y_train == 0] <- mu0 + qnorm(u0) - resid_train[y_train == 1] <- mu1 + qnorm(u1) - - # Update outcome: center z by y_bar_train before passing to forest - outcome_train$update_data(resid_train - y_bar_train - outcome_pred) - } - - # Sample mean forest - forest_model_mean$sample_one_iteration( - forest_dataset = forest_dataset_train, - residual = outcome_train, - forest_samples = forest_samples_mean, - active_forest = active_forest_mean, - rng = rng, - forest_model_config = forest_model_config_mean, - 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 - if (keep_sample) { - mean_forest_pred_train[, - sample_counter - ] <- forest_model_mean$get_cached_forest_predictions() - } - - # Additional Gibbs updates needed for the cloglog model - if (link_is_cloglog) { - # Update auxiliary data to current forest predictions - forest_pred_current <- forest_model_mean$get_cached_forest_predictions() - for (i in 1:train_size) { - forest_dataset_train$set_auxiliary_data_value( - 1, - i - 1, - forest_pred_current[i] - ) - } - - # Sample latent z_i's using truncated exponential - ordinal_sampler_update_latent_variables_cpp( - ordinal_sampler, - forest_dataset_train$data_ptr, - outcome_train$data_ptr, - rng$rng_ptr - ) - - # Sample gamma parameters (cutpoints) - ordinal_sampler_update_gamma_params_cpp( - ordinal_sampler, - forest_dataset_train$data_ptr, - outcome_train$data_ptr, - cloglog_forest_shape, - cloglog_forest_rate, - cloglog_cutpoint_0, - rng$rng_ptr - ) - - # Update cumulative sum of exp(gamma) values - ordinal_sampler_update_cumsum_exp_cpp( - ordinal_sampler, - forest_dataset_train$data_ptr - ) - - # Retain cutpoint draw - if (keep_sample) { - cloglog_cutpoints <- forest_dataset_train$get_auxiliary_data_vector( - 2 - ) - cloglog_cutpoint_samples[, sample_counter] <- cloglog_cutpoints - } - } - } - if (include_variance_forest) { - 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, - num_threads = num_threads, - keep_forest = keep_sample, - gfr = TRUE - ) - - # Cache train set predictions since they are already computed during sampling - if (keep_sample) { - variance_forest_pred_train[, - sample_counter - ] <- forest_model_variance$get_cached_forest_predictions() - } - } - if (sample_sigma2_global) { - current_sigma2 <- sampleGlobalErrorVarianceOneIteration( - outcome_train, - forest_dataset_train, - rng, - a_global, - b_global - ) - if (keep_sample) { - global_var_samples[sample_counter] <- current_sigma2 - } - global_model_config$update_global_error_variance(current_sigma2) - } - if (sample_sigma2_leaf) { - leaf_scale_double <- sampleLeafVarianceOneIteration( - active_forest_mean, - rng, - a_leaf, - b_leaf - ) - current_leaf_scale <- as.matrix(leaf_scale_double) - if (keep_sample) { - leaf_scale_samples[sample_counter] <- leaf_scale_double - } - forest_model_config_mean$update_leaf_model_scale( - current_leaf_scale - ) - } - if (has_rfx) { - rfx_model$sample_random_effect( - rfx_dataset_train, - outcome_train, - rfx_tracker_train, - rfx_samples, - keep_sample, - current_sigma2, - rng - ) - } - } - } - - # Run MCMC - if (num_burnin + num_mcmc > 0) { - for (chain_num in 1:num_chains) { - if (verbose) { - cat("Sampling chain", chain_num, "of", num_chains, "\n") - } - if (num_gfr > 0) { - # Reset state of active_forest and forest_model based on a previous GFR sample - forest_ind <- num_gfr - chain_num - if (include_mean_forest) { - resetActiveForest( - active_forest_mean, - forest_samples_mean, - forest_ind - ) - resetForestModel( - forest_model_mean, - active_forest_mean, - forest_dataset_train, - outcome_train, - TRUE - ) - if (sample_sigma2_leaf) { - leaf_scale_double <- leaf_scale_samples[forest_ind + 1] - current_leaf_scale <- as.matrix(leaf_scale_double) - forest_model_config_mean$update_leaf_model_scale( - current_leaf_scale - ) - } - if (link_is_cloglog) { - # Restore ordinal labels corrupted by resetForestModel's - # residual adjustment (outcome stores category labels, not residuals) - outcome_train$update_data(resid_train) - # We can reset cutpoints from warm-start since cutpoints are retained - current_cutpoints <- cloglog_cutpoint_samples[, forest_ind + 1] - for (i in seq_along(current_cutpoints)) { - forest_dataset_train$set_auxiliary_data_value( - 2, - i - 1, - current_cutpoints[i] - ) - } - ordinal_sampler_update_cumsum_exp_cpp( - ordinal_sampler, - forest_dataset_train$data_ptr - ) - # Re-predict from the reconstituted active forest - active_forest_preds <- active_forest_mean$predict( - forest_dataset_train - ) - for (i in 1:train_size) { - forest_dataset_train$set_auxiliary_data_value( - 1, - i - 1, - active_forest_preds[i] - ) - # Latent variables must be reset to 0 and burnt in - forest_dataset_train$set_auxiliary_data_value(0, i - 1, 0.0) - } - } - } - if (include_variance_forest) { - resetActiveForest( - active_forest_variance, - forest_samples_variance, - forest_ind - ) - resetForestModel( - forest_model_variance, - active_forest_variance, - forest_dataset_train, - outcome_train, - FALSE - ) - } - if (has_rfx) { - resetRandomEffectsModel( - rfx_model, - rfx_samples, - forest_ind, - sigma_alpha_init - ) - resetRandomEffectsTracker( - rfx_tracker_train, - rfx_model, - rfx_dataset_train, - outcome_train, - rfx_samples - ) - } - if (sample_sigma2_global) { - current_sigma2 <- global_var_samples[forest_ind + 1] - global_model_config$update_global_error_variance( - current_sigma2 - ) - } - } else if (has_prev_model) { - warmstart_index <- ifelse( - previous_model_decrement, - previous_model_warmstart_sample_num - chain_num + 1, - previous_model_warmstart_sample_num - ) - if (include_mean_forest) { - resetActiveForest( - active_forest_mean, - previous_forest_samples_mean, - warmstart_index - 1 - ) - resetForestModel( - forest_model_mean, - active_forest_mean, - forest_dataset_train, - outcome_train, - TRUE - ) - if ( - sample_sigma2_leaf && - (!is.null(previous_leaf_var_samples)) - ) { - leaf_scale_double <- previous_leaf_var_samples[ - warmstart_index - ] - current_leaf_scale <- as.matrix(leaf_scale_double) - forest_model_config_mean$update_leaf_model_scale( - current_leaf_scale - ) - } - if (link_is_cloglog) { - # Restore ordinal labels corrupted by resetForestModel's - # residual adjustment (outcome stores category labels, not residuals) - outcome_train$update_data(resid_train) - # We can reset cutpoints from warm-start since cutpoints are retained - current_cutpoints <- previous_cloglog_cutpoint_samples[, - warmstart_index - ] - for (i in seq_along(current_cutpoints)) { - forest_dataset_train$set_auxiliary_data_value( - 2, - i - 1, - current_cutpoints[i] - ) - } - ordinal_sampler_update_cumsum_exp_cpp( - ordinal_sampler, - forest_dataset_train$data_ptr - ) - # Re-predict from the reconstituted active forest - active_forest_preds <- active_forest_mean$predict( - forest_dataset_train - ) - for (i in 1:train_size) { - forest_dataset_train$set_auxiliary_data_value( - 1, - i - 1, - active_forest_preds[i] - ) - # Latent variables must be reset to 0 and burnt in - forest_dataset_train$set_auxiliary_data_value(0, i - 1, 0.0) - } - } - } - if (include_variance_forest) { - resetActiveForest( - active_forest_variance, - previous_forest_samples_variance, - warmstart_index - 1 - ) - resetForestModel( - forest_model_variance, - active_forest_variance, - forest_dataset_train, - outcome_train, - FALSE - ) - } - if (has_rfx) { - if (is.null(previous_rfx_samples)) { - warning( - "`previous_model_json` did not have any random effects samples, so the RFX sampler will be run from scratch while the forests and any other parameters are warm started" - ) - rootResetRandomEffectsModel( - rfx_model, - alpha_init, - xi_init, - sigma_alpha_init, - sigma_xi_init, - sigma_xi_shape, - sigma_xi_scale - ) - rootResetRandomEffectsTracker( - rfx_tracker_train, - rfx_model, - rfx_dataset_train, - outcome_train - ) - } else { - resetRandomEffectsModel( - rfx_model, - previous_rfx_samples, - warmstart_index - 1, - sigma_alpha_init - ) - resetRandomEffectsTracker( - rfx_tracker_train, - rfx_model, - rfx_dataset_train, - outcome_train, - rfx_samples - ) - } - } - if (sample_sigma2_global) { - if (!is.null(previous_global_var_samples)) { - current_sigma2 <- previous_global_var_samples[ - warmstart_index - ] - global_model_config$update_global_error_variance( - current_sigma2 - ) - } - } - } else { - if (include_mean_forest) { - resetActiveForest(active_forest_mean) - active_forest_mean$set_root_leaves( - init_values_mean_forest / num_trees_mean - ) - resetForestModel( - forest_model_mean, - active_forest_mean, - forest_dataset_train, - outcome_train, - TRUE - ) - if (sample_sigma2_leaf) { - current_leaf_scale <- as.matrix(sigma2_leaf_init) - forest_model_config_mean$update_leaf_model_scale( - current_leaf_scale - ) - } - if (link_is_cloglog) { - # Restore ordinal labels corrupted by resetForestModel's - # residual adjustment (outcome stores category labels, not residuals) - outcome_train$update_data(resid_train) - # Reset all cloglog parameters to default values - for (i in 1:train_size) { - forest_dataset_train$set_auxiliary_data_value(0, i - 1, 0.0) - forest_dataset_train$set_auxiliary_data_value(1, i - 1, 0.0) - } - # Initialize log-scale cutpoints to 0 - initial_gamma <- rep(0.0, cloglog_num_categories - 1) - for (i in seq_along(initial_gamma)) { - forest_dataset_train$set_auxiliary_data_value( - 2, - i - 1, - initial_gamma[i] - ) - } - # Convert to cumulative exponentiated cutpoints directly in C++ - ordinal_sampler_update_cumsum_exp_cpp( - ordinal_sampler, - forest_dataset_train$data_ptr - ) - } - } - if (include_variance_forest) { - resetActiveForest(active_forest_variance) - active_forest_variance$set_root_leaves( - log(variance_forest_init) / num_trees_variance - ) - resetForestModel( - forest_model_variance, - active_forest_variance, - forest_dataset_train, - outcome_train, - FALSE - ) - } - if (has_rfx) { - rootResetRandomEffectsModel( - rfx_model, - alpha_init, - xi_init, - sigma_alpha_init, - sigma_xi_init, - sigma_xi_shape, - sigma_xi_scale - ) - rootResetRandomEffectsTracker( - rfx_tracker_train, - rfx_model, - rfx_dataset_train, - outcome_train - ) - } - if (sample_sigma2_global) { - current_sigma2 <- sigma2_init - global_model_config$update_global_error_variance( - current_sigma2 - ) - } - } - for (i in (num_gfr + 1):num_samples) { - is_mcmc <- i > (num_gfr + num_burnin) - if (is_mcmc) { - mcmc_counter <- i - (num_gfr + num_burnin) - if (mcmc_counter %% keep_every == 0) { - keep_sample <- TRUE - } else { - keep_sample <- FALSE - } - } else { - if (keep_burnin) { - keep_sample <- TRUE - } else { - keep_sample <- FALSE - } - } - if (keep_sample) { - sample_counter <- sample_counter + 1 - } - # Print progress - if (verbose) { - if (num_burnin > 0 && !is_mcmc) { - if ( - ((i - num_gfr) %% 100 == 0) || - ((i - num_gfr) == num_burnin) - ) { - cat( - "Sampling", - i - num_gfr, - "out of", - num_burnin, - "BART burn-in draws; Chain number ", - chain_num, - "\n" - ) - } - } - if (num_mcmc > 0 && is_mcmc) { - raw_iter <- i - num_gfr - num_burnin - if ((raw_iter %% 100 == 0) || (i == num_samples)) { - if (keep_every == 1) { - cat( - "Sampling", - raw_iter, - "out of", - num_mcmc, - "BART MCMC draws; Chain number ", - chain_num, - "\n" - ) - } else { - cat( - "Sampling raw draw", - raw_iter, - "of", - num_actual_mcmc_iter, - "BART MCMC draws (thinning by", - keep_every, - ":", - raw_iter %/% keep_every, - "of", - num_mcmc, - "retained); Chain number ", - chain_num, - "\n" - ) - } - } - } - } - - if (include_mean_forest) { - if (link_is_probit) { - # Sample latent probit variable, z | - - outcome_pred <- active_forest_mean$predict( - forest_dataset_train - ) - if (has_rfx) { - rfx_pred <- rfx_model$predict( - rfx_dataset_train, - rfx_tracker_train - ) - outcome_pred <- outcome_pred + rfx_pred - } - eta_pred <- outcome_pred + y_bar_train - mu0 <- eta_pred[y_train == 0] - mu1 <- eta_pred[y_train == 1] - u0 <- runif(sum(y_train == 0), 0, pnorm(0 - mu0)) - u1 <- runif(sum(y_train == 1), pnorm(0 - mu1), 1) - resid_train[y_train == 0] <- mu0 + qnorm(u0) - resid_train[y_train == 1] <- mu1 + qnorm(u1) - - # Update outcome: center z by y_bar_train before passing to forest - outcome_train$update_data( - resid_train - y_bar_train - outcome_pred - ) - } - - forest_model_mean$sample_one_iteration( - forest_dataset = forest_dataset_train, - residual = outcome_train, - forest_samples = forest_samples_mean, - active_forest = active_forest_mean, - rng = rng, - forest_model_config = forest_model_config_mean, - 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 - if (keep_sample) { - mean_forest_pred_train[, - sample_counter - ] <- forest_model_mean$get_cached_forest_predictions() - } - - # Additional Gibbs updates needed for the cloglog model - if (link_is_cloglog) { - # Update auxiliary data to current forest predictions - forest_pred_current <- forest_model_mean$get_cached_forest_predictions() - for (i in 1:train_size) { - forest_dataset_train$set_auxiliary_data_value( - 1, - i - 1, - forest_pred_current[i] - ) - } - - # Sample latent z_i's using truncated exponential - ordinal_sampler_update_latent_variables_cpp( - ordinal_sampler, - forest_dataset_train$data_ptr, - outcome_train$data_ptr, - rng$rng_ptr - ) - - # Sample gamma parameters (cutpoints) - ordinal_sampler_update_gamma_params_cpp( - ordinal_sampler, - forest_dataset_train$data_ptr, - outcome_train$data_ptr, - cloglog_forest_shape, - cloglog_forest_rate, - cloglog_cutpoint_0, - rng$rng_ptr - ) - - # Update cumulative sum of exp(gamma) values - ordinal_sampler_update_cumsum_exp_cpp( - ordinal_sampler, - forest_dataset_train$data_ptr - ) - - # Retain cutpoint draw - if (keep_sample) { - cloglog_cutpoints <- forest_dataset_train$get_auxiliary_data_vector( - 2 - ) - cloglog_cutpoint_samples[, sample_counter] <- cloglog_cutpoints - } - } - } - if (include_variance_forest) { - 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, - num_threads = num_threads, - keep_forest = keep_sample, - gfr = FALSE - ) - - # Cache train set predictions since they are already computed during sampling - if (keep_sample) { - variance_forest_pred_train[, - sample_counter - ] <- forest_model_variance$get_cached_forest_predictions() - } - } - if (sample_sigma2_global) { - current_sigma2 <- sampleGlobalErrorVarianceOneIteration( - outcome_train, - forest_dataset_train, - rng, - a_global, - b_global - ) - if (keep_sample) { - global_var_samples[sample_counter] <- current_sigma2 - } - global_model_config$update_global_error_variance( - current_sigma2 - ) - } - if (sample_sigma2_leaf) { - leaf_scale_double <- sampleLeafVarianceOneIteration( - active_forest_mean, - rng, - a_leaf, - b_leaf - ) - current_leaf_scale <- as.matrix(leaf_scale_double) - if (keep_sample) { - leaf_scale_samples[sample_counter] <- leaf_scale_double - } - forest_model_config_mean$update_leaf_model_scale( - current_leaf_scale - ) - } - if (has_rfx) { - rfx_model$sample_random_effect( - rfx_dataset_train, - outcome_train, - rfx_tracker_train, - rfx_samples, - keep_sample, - current_sigma2, - rng - ) - } - } - } - } - - # Remove GFR samples if they are not to be retained - if ((!keep_gfr) && (num_gfr > 0)) { - for (i in 1:num_gfr) { - if (include_mean_forest) { - forest_samples_mean$delete_sample(0) - } - if (include_variance_forest) { - forest_samples_variance$delete_sample(0) - } - if (has_rfx) { - rfx_samples$delete_sample(0) - } - } - if (include_mean_forest) { - mean_forest_pred_train <- mean_forest_pred_train[, - (num_gfr + 1):ncol(mean_forest_pred_train) - ] - if (link_is_cloglog) { - cloglog_cutpoint_samples <- cloglog_cutpoint_samples[, - (num_gfr + 1):ncol(cloglog_cutpoint_samples), - drop = FALSE - ] - } - } - if (include_variance_forest) { - variance_forest_pred_train <- variance_forest_pred_train[, - (num_gfr + 1):ncol(variance_forest_pred_train) - ] - } - if (sample_sigma2_global) { - global_var_samples <- global_var_samples[ - (num_gfr + 1):length(global_var_samples) - ] - } - if (sample_sigma2_leaf) { - leaf_scale_samples <- leaf_scale_samples[ - (num_gfr + 1):length(leaf_scale_samples) - ] - } - num_retained_samples <- num_retained_samples - num_gfr - } - - # Mean forest predictions - if (include_mean_forest) { - # y_hat_train <- forest_samples_mean$predict(forest_dataset_train)*y_std_train + y_bar_train - y_hat_train <- mean_forest_pred_train * y_std_train + y_bar_train - if (has_test) { - y_hat_test <- forest_samples_mean$predict(forest_dataset_test) * - y_std_train + - y_bar_train - } - } - - # Variance forest predictions - if (include_variance_forest) { - # sigma2_x_hat_train <- forest_samples_variance$predict(forest_dataset_train) - sigma2_x_hat_train <- exp(variance_forest_pred_train) - if (has_test) { - sigma2_x_hat_test <- forest_samples_variance$predict( - forest_dataset_test - ) - } - } - - # Random effects predictions - if (has_rfx) { - rfx_preds_train <- rfx_samples$predict( - rfx_group_ids_train, - rfx_basis_train - ) * - y_std_train - y_hat_train <- y_hat_train + rfx_preds_train - } - if ((has_rfx_test) && (has_test)) { - rfx_preds_test <- rfx_samples$predict( - rfx_group_ids_test, - rfx_basis_test - ) * - y_std_train - y_hat_test <- y_hat_test + rfx_preds_test - } - - # Global error variance - if (sample_sigma2_global) { - sigma2_global_samples <- global_var_samples * (y_std_train^2) - } - - # Leaf parameter variance - if (sample_sigma2_leaf) { - tau_samples <- leaf_scale_samples - } - - # Rescale variance forest prediction by global sigma2 (sampled or constant) - if (include_variance_forest) { - if (sample_sigma2_global) { - sigma2_x_hat_train <- sapply(1:num_retained_samples, function(i) { - sigma2_x_hat_train[, i] * sigma2_global_samples[i] - }) - if (has_test) { - sigma2_x_hat_test <- sapply( - 1:num_retained_samples, - function(i) { - sigma2_x_hat_test[, i] * sigma2_global_samples[i] - } - ) - } - } else { - sigma2_x_hat_train <- sigma2_x_hat_train * - sigma2_init * - y_std_train * - y_std_train - if (has_test) { - sigma2_x_hat_test <- sigma2_x_hat_test * - sigma2_init * - y_std_train * - y_std_train - } - } - } - - # Return results as a list - model_params <- list( - "sigma2_init" = sigma2_init, - "sigma2_leaf_init" = sigma2_leaf_init, - "a_global" = a_global, - "b_global" = b_global, - "a_leaf" = a_leaf, - "b_leaf" = b_leaf, - "a_forest" = a_forest, - "b_forest" = b_forest, - "outcome_mean" = y_bar_train, - "outcome_scale" = y_std_train, - "standardize" = standardize, - "leaf_dimension" = leaf_dimension, - "is_leaf_constant" = is_leaf_constant, - "leaf_regression" = leaf_regression, - "requires_basis" = requires_basis, - "num_covariates" = num_cov_orig, - "num_basis" = ifelse( - is.null(leaf_basis_train), - 0, - ncol(leaf_basis_train) - ), - "num_samples" = num_retained_samples, - "num_gfr" = num_gfr, - "num_burnin" = num_burnin, - "num_mcmc" = num_mcmc, - "keep_every" = keep_every, - "num_chains" = num_chains, - "has_basis" = !is.null(leaf_basis_train), - "has_rfx" = has_rfx, - "has_rfx_basis" = has_basis_rfx, - "num_rfx_basis" = num_basis_rfx, - "sample_sigma2_global" = sample_sigma2_global, - "sample_sigma2_leaf" = sample_sigma2_leaf, - "include_mean_forest" = include_mean_forest, - "include_variance_forest" = include_variance_forest, - "outcome_model" = outcome_model, - "probit_outcome_model" = probit_outcome_model, - "cloglog_num_categories" = ifelse( - link_is_cloglog, - cloglog_num_categories, - 0 - ), - "rfx_model_spec" = rfx_model_spec - ) - result <- list( - "model_params" = model_params, - "train_set_metadata" = X_train_metadata - ) - if (include_mean_forest) { - result[["mean_forests"]] = forest_samples_mean - result[["y_hat_train"]] = y_hat_train - if (has_test) { - result[["y_hat_test"]] = y_hat_test - } - if (link_is_cloglog && !outcome_is_binary) { - result[["cloglog_cutpoint_samples"]] = cloglog_cutpoint_samples - } - } - if (include_variance_forest) { - result[["variance_forests"]] = forest_samples_variance - result[["sigma2_x_hat_train"]] = sigma2_x_hat_train - if (has_test) result[["sigma2_x_hat_test"]] = sigma2_x_hat_test - } - if (sample_sigma2_global) { - result[["sigma2_global_samples"]] = sigma2_global_samples - } - if (sample_sigma2_leaf) { - result[["sigma2_leaf_samples"]] = tau_samples - } - if (has_rfx) { - result[["rfx_samples"]] = rfx_samples - result[["rfx_preds_train"]] = rfx_preds_train - result[["rfx_unique_group_ids"]] = levels(group_ids_factor) - } - if ((has_rfx_test) && (has_test)) { - result[["rfx_preds_test"]] = rfx_preds_test - } - class(result) <- "bartmodel" - - # Clean up classes with external pointers to C++ data structures - if (include_mean_forest) { - rm(forest_model_mean) - } - if (include_variance_forest) { - rm(forest_model_variance) - } - rm(forest_dataset_train) - if (has_test) { - rm(forest_dataset_test) - } - if (has_rfx) { - rm(rfx_dataset_train, rfx_tracker_train, rfx_model) - } - rm(outcome_train) - rm(rng) - - # Restore global RNG state if user provided a random seed - if (custom_rng) { - if (has_existing_random_seed) { - .Random.seed <- original_global_seed - } else { - rm(".Random.seed", envir = .GlobalEnv) - } - } - - return(result) -} - -#' @title Predict from a BART Model -#' @description -#' Predict from a sampled BART model on new data -#' -#' @param object Object of type `bart` containing draws of a regression forest and associated sampling outputs. -#' @param X Covariates used to determine tree leaf predictions for each observation. Must be passed as a matrix or dataframe. -#' @param leaf_basis (Optional) Bases used for prediction (by e.g. dot product with leaf values). Default: `NULL`. -#' @param rfx_group_ids (Optional) Test set group labels used for an additive random effects model. -#' We do not currently support (but plan to in the near future), test set evaluation for group labels -#' that were not in the training set. -#' @param rfx_basis (Optional) Test set basis for "random-slope" regression in additive random effects model. -#' @param type (Optional) Type of prediction to return. Options are "mean", which averages the predictions from every draw of a BART model, and "posterior", which returns the entire matrix of posterior predictions. Default: "posterior". -#' @param terms (Optional) Which model terms to include in the prediction. This can be a single term or a list of model terms. Options include "y_hat", "mean_forest", "rfx", "variance_forest", or "all". If a model doesn't have mean forest, random effects, or variance forest predictions, but one of those terms is request, the request will simply be ignored. If none of the requested terms are present in a model, this function will return `NULL` along with a warning. Default: "all". -#' @param scale (Optional) Scale of mean function predictions. Options are "linear", which returns predictions on the original scale of the mean forest / RFX terms, "probability", which transforms predictions into class probabilities for models with discrete outcomes, and "class", which returns predicted outcome categories for discrete outcome models. "probability" is only valid for outcome models with `outcome == 'binary'` or `outcome == 'ordinal'`. For binary outcomes, this will return the probability that `y == 1`, and for ordinal outcomes, this will return probabilities for each outcome label. Default: "linear". -#' @param ... (Optional) Other prediction parameters. -#' -#' @return List of prediction matrices or single prediction matrix / vector, depending on the terms requested. -#' @export -#' -#' @examples -#' n <- 100 -#' p <- 5 -#' X <- matrix(runif(n*p), ncol = p) -#' f_XW <- ( -#' ((0 <= X[,1]) & (0.25 > X[,1])) * (-7.5) + -#' ((0.25 <= X[,1]) & (0.5 > X[,1])) * (-2.5) + -#' ((0.5 <= X[,1]) & (0.75 > X[,1])) * (2.5) + -#' ((0.75 <= X[,1]) & (1 > X[,1])) * (7.5) -#' ) -#' noise_sd <- 1 -#' y <- f_XW + rnorm(n, 0, noise_sd) -#' test_set_pct <- 0.2 -#' n_test <- round(test_set_pct*n) -#' n_train <- n - n_test -#' test_inds <- sort(sample(1:n, n_test, replace = FALSE)) -#' train_inds <- (1:n)[!((1:n) %in% test_inds)] -#' X_test <- X[test_inds,] -#' X_train <- X[train_inds,] -#' y_test <- y[test_inds] -#' y_train <- y[train_inds] -#' bart_model <- bart(X_train = X_train, y_train = y_train, -#' num_gfr = 10, num_burnin = 0, num_mcmc = 10) -#' y_hat_test <- predict(bart_model, X=X_test)$y_hat -predict.bartmodel <- function( - object, - X, - leaf_basis = NULL, - rfx_group_ids = NULL, - rfx_basis = NULL, - type = "posterior", - terms = "all", - scale = "linear", - ... -) { - # Handle mean function scale - if (!is.character(scale)) { - stop("scale must be a string or character vector") - } - if (!(scale %in% c("linear", "probability", "class"))) { - stop("scale must either be 'linear', 'probability', or 'class'") - } - outcome_model <- object$model_params$outcome_model - is_probit <- (outcome_model$link == "probit" && - outcome_model$outcome == "binary") - is_binary_cloglog <- (outcome_model$link == "cloglog" && - outcome_model$outcome == "binary") - is_ordinal_cloglog <- (outcome_model$link == "cloglog" && - outcome_model$outcome == "ordinal") - is_cloglog <- is_binary_cloglog || is_ordinal_cloglog - if ((scale == "probability") && (!(is_probit || is_cloglog))) { - stop( - "scale cannot be 'probability' for models not fit with a probit or cloglog outcome model" - ) - } - if ((scale == "class") && (!(is_probit || is_cloglog))) { - stop( - "scale cannot be 'class' for models not fit with a probit or cloglog outcome model" - ) - } - probability_scale <- scale == "probability" - class_scale <- scale == "class" + probability_scale <- scale == "probability" + class_scale <- scale == "class" # Handle prediction type if (!is.character(type)) { @@ -2736,226 +1641,153 @@ predict.bartmodel <- function( } } - # Create prediction dataset - if (!is.null(leaf_basis)) { - prediction_dataset <- createForestDataset(X, leaf_basis) - } else { - prediction_dataset <- createForestDataset(X) - } - - # Compute variance forest predictions - if (predict_variance_forest) { - s_x_raw <- object$variance_forests$predict(prediction_dataset) - } - - # Scale variance forest predictions - num_samples <- object$model_params$num_samples - y_std <- object$model_params$outcome_scale - y_bar <- object$model_params$outcome_mean - sigma2_init <- object$model_params$sigma2_init - if (predict_variance_forest) { - if (object$model_params$sample_sigma2_global) { - sigma2_global_samples <- object$sigma2_global_samples - variance_forest_predictions <- sapply(1:num_samples, function(i) { - s_x_raw[, i] * sigma2_global_samples[i] - }) + # Read the forests through borrowed (non-owning) pointers into the single-owner + # samples object -- no deep copy, and avoids tripping the deprecated $mean_forests + # accessor's warning on internal prediction. + bart_samples <- object$samples + bart_metadata_list <- list( + num_samples = as.integer(object$model_params$num_samples), + y_bar = as.double(object$model_params$outcome_mean), + y_std = as.double(object$model_params$outcome_scale), + include_variance_forest = has_variance_forest, + has_rfx = has_rfx, + rfx_model_spec = if (has_rfx) { + object$model_params$rfx_model_spec } else { - variance_forest_predictions <- s_x_raw * sigma2_init * y_std * y_std - } - if (predict_mean) { - variance_forest_predictions <- rowMeans(variance_forest_predictions) - } - } - - # Compute mean forest predictions - if (predict_mean_forest || predict_mean_forest_intermediate) { - mean_forest_predictions <- object$mean_forests$predict( - prediction_dataset - ) * - y_std + - y_bar - } - - # Compute rfx predictions (if needed) - if (predict_rfx || predict_rfx_intermediate) { - if (!is.null(rfx_basis)) { - rfx_predictions <- object$rfx_samples$predict( - rfx_group_ids, - rfx_basis - ) * - y_std + "" + }, + link_function = object$model_params$outcome_model$link, + outcome_type = object$model_params$outcome_model$outcome, + cloglog_num_classes = if ( + !is.null(object$model_params$num_classes_cloglog) + ) { + as.integer(object$model_params$num_classes_cloglog) + } else if (!is.null(object$model_params$cloglog_num_categories)) { + as.integer(object$model_params$cloglog_num_categories) } else { - # Sanity check -- this branch should only occur if rfx_model_spec == "intercept_only" - if (!rfx_intercept) { - stop( - "rfx_basis must be provided for random effects models with random slopes" - ) - } - - # Extract the raw RFX samples and scale by train set outcome standard deviation - rfx_param_list <- object$rfx_samples$extract_parameter_samples() - rfx_beta_draws <- rfx_param_list$beta_samples * y_std - - # Promote to an array with consistent dimensions when there's one rfx term - if (length(dim(rfx_beta_draws)) == 2) { - dim(rfx_beta_draws) <- c(1, dim(rfx_beta_draws)) - } - - # Construct a matrix with the appropriate group random effects arranged for each observation - rfx_predictions_raw <- array( - NA, - dim = c( - nrow(X), - ncol(rfx_basis), - object$model_params$num_samples - ) - ) - for (i in 1:nrow(X)) { - rfx_predictions_raw[i, , ] <- - rfx_beta_draws[, rfx_group_ids[i], ] - } - - # Intercept-only model, so the random effect prediction is simply the - # value of the respective group's intercept coefficient for each observation - rfx_predictions = rfx_predictions_raw[, 1, ] + 0L } - } + ) - # Combine into y hat predictions - if (probability_scale || class_scale) { - if (is_probit) { - if (predict_y_hat) { - if (has_mean_forest && has_rfx) { - y_hat <- pnorm(mean_forest_predictions + rfx_predictions) - mean_forest_predictions <- pnorm(mean_forest_predictions) - rfx_predictions <- pnorm(rfx_predictions) - } else if (has_mean_forest) { - y_hat <- pnorm(mean_forest_predictions) - mean_forest_predictions <- pnorm(mean_forest_predictions) - } else if (has_rfx) { - y_hat <- pnorm(rfx_predictions) - rfx_predictions <- pnorm(rfx_predictions) - } - } else { - if (has_mean_forest && has_rfx) { - mean_forest_predictions <- pnorm(mean_forest_predictions) - rfx_predictions <- pnorm(rfx_predictions) - } else if (has_mean_forest) { - mean_forest_predictions <- pnorm(mean_forest_predictions) - } else if (has_rfx) { - rfx_predictions <- pnorm(rfx_predictions) - } - } - } else if (is_binary_cloglog) { - mean_forest_predictions <- exp(-exp(mean_forest_predictions)) - if (predict_y_hat) { - y_hat <- mean_forest_predictions - } - } else if (is_ordinal_cloglog) { - cloglog_num_categories <- object$model_params$cloglog_num_categories - cloglog_cutpoint_samples <- object$cloglog_cutpoint_samples - n_obs_pred <- nrow(X) - n_samp_pred <- object$model_params$num_samples - mean_forest_probabilities <- array( - NA_real_, - dim = c(n_obs_pred, cloglog_num_categories, n_samp_pred) - ) - # Sequential ordinal cloglog: P(Y=k) = prod_{j X[,1])) * (-7.5) + +#' ((0.25 <= X[,1]) & (0.5 > X[,1])) * (-2.5) + +#' ((0.5 <= X[,1]) & (0.75 > X[,1])) * (2.5) + +#' ((0.75 <= X[,1]) & (1 > X[,1])) * (7.5) +#' ) +#' snr <- 3 +#' group_ids <- rep(c(1,2), n %/% 2) +#' rfx_coefs <- matrix(c(-1, -1, 1, 1), nrow=2, byrow=TRUE) +#' rfx_basis <- cbind(1, runif(n, -1, 1)) +#' rfx_term <- rowSums(rfx_coefs[group_ids,] * rfx_basis) +#' E_y <- f_XW + rfx_term +#' y <- E_y + rnorm(n, 0, 1)*(sd(E_y)/snr) +#' test_set_pct <- 0.2 +#' n_test <- round(test_set_pct*n) +#' n_train <- n - n_test +#' test_inds <- sort(sample(1:n, n_test, replace = FALSE)) +#' train_inds <- (1:n)[!((1:n) %in% test_inds)] +#' X_test <- X[test_inds,] +#' X_train <- X[train_inds,] +#' y_test <- y[test_inds] +#' y_train <- y[train_inds] +#' rfx_group_ids_test <- group_ids[test_inds] +#' rfx_group_ids_train <- group_ids[train_inds] +#' rfx_basis_test <- rfx_basis[test_inds,] +#' rfx_basis_train <- rfx_basis[train_inds,] +#' rfx_term_test <- rfx_term[test_inds] +#' rfx_term_train <- rfx_term[train_inds] +#' bart_model <- bart(X_train = X_train, y_train = y_train, X_test = X_test, +#' rfx_group_ids_train = rfx_group_ids_train, +#' rfx_group_ids_test = rfx_group_ids_test, +#' rfx_basis_train = rfx_basis_train, +#' rfx_basis_test = rfx_basis_test, +#' num_gfr = 10, num_burnin = 0, num_mcmc = 10) +#' mean_forest <- extractForest(bart_model, "mean") +extractForest.bartmodel <- function(object, term) { + if (term %in% c("mean", "mean_forest")) { + if (object$samples$has_mean_forest()) { + return(object$samples$materialize_mean_forest()) + } else { + stop("This model does not have a mean forest") + } + } + + if (term %in% c("variance", "variance_forest")) { + if (object$samples$has_variance_forest()) { + return(object$samples$materialize_variance_forest()) + } else { + stop("This model does not have a variance forest") + } + } + + stop(paste0("term ", term, " is not a valid BART forest term")) +} + #' @title Extract BART Parameter Samples #' @description Extract a vector, matrix or array of parameter samples from a BART model by name. #' Random effects are handled by a separate `getRandomEffectSamples` function due to the complexity of the random effects parameters. @@ -3457,24 +2361,24 @@ getRandomEffectSamples.bartmodel <- function(object, ...) { #' sigma2_samples <- extractParameter(bart_model, "sigma2") extractParameter.bartmodel <- function(object, term) { if (term %in% c("sigma2", "global_error_scale", "sigma2_global")) { - if (!is.null(object$sigma2_global_samples)) { - return(object$sigma2_global_samples) + if (object$samples$has_global_var_samples()) { + return(object$samples$global_var_samples()) } else { stop("This model does not have global variance parameter samples") } } if (term %in% c("sigma2_leaf", "leaf_scale")) { - if (!is.null(object$sigma2_leaf_samples)) { - return(object$sigma2_leaf_samples) + if (object$samples$has_leaf_scale_samples()) { + return(object$samples$leaf_scale_samples()) } else { stop("This model does not have leaf variance parameter samples") } } if (term %in% c("y_hat_train")) { - if (!is.null(object$y_hat_train)) { - return(object$y_hat_train) + if (object$samples$has_mean_forest_predictions_train()) { + return(object$samples$mean_forest_predictions_train()) } else { stop( "This model does not have in-sample mean function prediction samples" @@ -3483,32 +2387,32 @@ extractParameter.bartmodel <- function(object, term) { } if (term %in% c("y_hat_test")) { - if (!is.null(object$y_hat_test)) { - return(object$y_hat_test) + if (object$samples$has_mean_forest_predictions_test()) { + return(object$samples$mean_forest_predictions_test()) } else { stop("This model does not have test set mean function prediction samples") } } if (term %in% c("sigma2_x_train", "var_x_train")) { - if (!is.null(object$sigma2_x_hat_train)) { - return(object$sigma2_x_hat_train) + if (object$samples$has_variance_forest_predictions_train()) { + return(object$samples$variance_forest_predictions_train()) } else { stop("This model does not have in-sample variance forest predictions") } } if (term %in% c("sigma2_x_test", "var_x_test")) { - if (!is.null(object$sigma2_x_hat_test)) { - return(object$sigma2_x_hat_test) + if (object$samples$has_variance_forest_predictions_test()) { + return(object$samples$variance_forest_predictions_test()) } else { stop("This model does not have test set variance forest predictions") } } if (term %in% c("cloglog_cutpoints", "cutpoints")) { - if (!is.null(object$cloglog_cutpoint_samples)) { - return(object$cloglog_cutpoint_samples) + if (object$samples$has_cloglog_cutpoint_samples()) { + return(object$samples$cloglog_cutpoint_samples()) } else { stop("This model does not have ordinal cutpoint samples") } @@ -3532,56 +2436,14 @@ saveBARTModelToJson <- function(object) { stop("This BCF model has not yet been sampled") } - # Add the forests - if (object$model_params$include_mean_forest) { - jsonobj$add_forest(object$mean_forests) - } - if (object$model_params$include_variance_forest) { - jsonobj$add_forest(object$variance_forests) - } - - # Add metadata - jsonobj$add_scalar( - "num_numeric_vars", - object$train_set_metadata$num_numeric_vars - ) - jsonobj$add_scalar( - "num_ordered_cat_vars", - object$train_set_metadata$num_ordered_cat_vars - ) - jsonobj$add_scalar( - "num_unordered_cat_vars", - object$train_set_metadata$num_unordered_cat_vars - ) - if (object$train_set_metadata$num_numeric_vars > 0) { - jsonobj$add_string_vector( - "numeric_vars", - object$train_set_metadata$numeric_vars - ) - } - if (object$train_set_metadata$num_ordered_cat_vars > 0) { - jsonobj$add_string_vector( - "ordered_cat_vars", - object$train_set_metadata$ordered_cat_vars - ) - jsonobj$add_string_list( - "ordered_unique_levels", - object$train_set_metadata$ordered_unique_levels - ) - } - if (object$train_set_metadata$num_unordered_cat_vars > 0) { - jsonobj$add_string_vector( - "unordered_cat_vars", - object$train_set_metadata$unordered_cat_vars - ) - jsonobj$add_string_list( - "unordered_unique_levels", - object$train_set_metadata$unordered_unique_levels - ) - } + # Add the samples C++ object to the JSON object, which will handle serialization of all the sampled parameters and predictions in a cross-platform way. + bart_samples <- object$samples + bart_samples$append_to_json(jsonobj) # Add version stamp and global parameters jsonobj$add_string("stochtree_version", getStochtreeVersion()) + jsonobj$add_string("platform", "R") + jsonobj$add_integer("schema_version", STOCHTREE_SCHEMA_VERSION) jsonobj$add_scalar("outcome_scale", object$model_params$outcome_scale) jsonobj$add_scalar("outcome_mean", object$model_params$outcome_mean) jsonobj$add_boolean("standardize", object$model_params$standardize) @@ -3637,37 +2499,22 @@ saveBARTModelToJson <- function(object) { "cloglog_num_categories", object$model_params$cloglog_num_categories ) - if (object$model_params$outcome_model$outcome == "ordinal") { - for (i in 1:(object$model_params$cloglog_num_categories - 1)) { - jsonobj$add_vector( - paste0("cloglog_cutpoint_samples_", i), - object$cloglog_cutpoint_samples[i, ], - "parameters" - ) - } - } - } - if (object$model_params$sample_sigma2_global) { - jsonobj$add_vector( - "sigma2_global_samples", - object$sigma2_global_samples, - "parameters" - ) - } - if (object$model_params$sample_sigma2_leaf) { - jsonobj$add_vector( - "sigma2_leaf_samples", - object$sigma2_leaf_samples, - "parameters" - ) } # Add random effects (if present) if (object$model_params$has_rfx) { - jsonobj$add_random_effects(object$rfx_samples) jsonobj$add_string_vector( "rfx_unique_group_ids", - object$rfx_unique_group_ids + object$rfx_unique_group_ids, + subfolder_name = "random_effects" + ) + # Cross-platform compatible on the rfx axis iff the group id levels are + # integer-valued (Python supports only integer group ids). + rfx_compatible <- all(grepl("^-?[0-9]+$", object$rfx_unique_group_ids)) + jsonobj$add_boolean( + "cross_platform_compatible", + rfx_compatible, + subfolder_name = "random_effects" ) } @@ -3675,7 +2522,7 @@ saveBARTModelToJson <- function(object) { preprocessor_metadata_string <- savePreprocessorToJsonString( object$train_set_metadata ) - jsonobj$add_string("preprocessor_metadata", preprocessor_metadata_string) + jsonobj$add_string("covariate_preprocessor", preprocessor_metadata_string) return(jsonobj) } @@ -3721,415 +2568,107 @@ saveBARTModelToJsonString <- function(object) { model_params } -#' @title Convert JSON to BART Model -#' @rdname BARTSerialization -#' @param json_object Object of type `CppJson` containing Json representation of a BART model -#' @export -createBARTModelFromJson <- function(json_object) { - # Initialize the BCF model - output <- list() - - # Helpers for optional-field presence checks - .ver <- inferStochtreeJsonVersion(json_object) - has_field <- function(name) { - json_contains_field_cpp(json_object$json_ptr, name) - } - has_subfolder_field <- function(subfolder, name) { - json_contains_field_subfolder_cpp(json_object$json_ptr, subfolder, name) - } - - # Unpack the forests - include_mean_forest <- json_object$get_boolean("include_mean_forest") - include_variance_forest <- json_object$get_boolean( - "include_variance_forest" - ) - if (include_mean_forest) { - output[["mean_forests"]] <- loadForestContainerJson( - json_object, - "forest_0" - ) - if (include_variance_forest) { - output[["variance_forests"]] <- loadForestContainerJson( - json_object, - "forest_1" - ) - } - } else { - output[["variance_forests"]] <- loadForestContainerJson( - json_object, - "forest_0" - ) - } - - # Unpack metadata - train_set_metadata = list() - train_set_metadata[["num_numeric_vars"]] <- json_object$get_scalar( - "num_numeric_vars" - ) - train_set_metadata[["num_ordered_cat_vars"]] <- json_object$get_scalar( - "num_ordered_cat_vars" - ) - train_set_metadata[["num_unordered_cat_vars"]] <- json_object$get_scalar( - "num_unordered_cat_vars" - ) - if (train_set_metadata[["num_numeric_vars"]] > 0) { - train_set_metadata[["numeric_vars"]] <- json_object$get_string_vector( - "numeric_vars" - ) - } - if (train_set_metadata[["num_ordered_cat_vars"]] > 0) { - train_set_metadata[[ - "ordered_cat_vars" - ]] <- json_object$get_string_vector("ordered_cat_vars") - train_set_metadata[[ - "ordered_unique_levels" - ]] <- json_object$get_string_list( - "ordered_unique_levels", - train_set_metadata[["ordered_cat_vars"]] - ) - } - if (train_set_metadata[["num_unordered_cat_vars"]] > 0) { - train_set_metadata[[ - "unordered_cat_vars" - ]] <- json_object$get_string_vector("unordered_cat_vars") - train_set_metadata[[ - "unordered_unique_levels" - ]] <- json_object$get_string_list( - "unordered_unique_levels", - train_set_metadata[["unordered_cat_vars"]] - ) - } - output[["train_set_metadata"]] <- train_set_metadata - - # Unpack model params - model_params = list() - model_params[["outcome_scale"]] <- json_object$get_scalar("outcome_scale") - model_params[["outcome_mean"]] <- json_object$get_scalar("outcome_mean") - model_params[["standardize"]] <- json_object$get_boolean("standardize") - model_params[["sigma2_init"]] <- json_object$get_scalar("sigma2_init") - model_params[["sample_sigma2_global"]] <- json_object$get_boolean( - "sample_sigma2_global" - ) - model_params[["sample_sigma2_leaf"]] <- json_object$get_boolean( - "sample_sigma2_leaf" +# In-place v0 -> v1 migration for a BART model envelope: positional forest keys +# (forests/forest_0, ...) -> named keys (mean_forest / variance_forest), driven by +# the include_*_forest flags (unchanged across v0/v1). +.migrateBartJsonV0ToV1 <- function(json_object, loaded_version) { + json_object$add_string("platform", inferPlatformV0(json_object, "R")) + include_mean <- json_object$get_boolean_or_default( + "include_mean_forest", + FALSE ) - model_params[["include_mean_forest"]] <- include_mean_forest - model_params[["include_variance_forest"]] <- include_variance_forest - model_params[["has_rfx"]] <- json_object$get_boolean("has_rfx") - - if (has_field("has_rfx_basis")) { - model_params[["has_rfx_basis"]] <- json_object$get_boolean("has_rfx_basis") - model_params[["num_rfx_basis"]] <- json_object$get_scalar("num_rfx_basis") - } else { - model_params[["has_rfx_basis"]] <- FALSE - model_params[["num_rfx_basis"]] <- 1 - warning(paste0( - "Fields 'has_rfx_basis' and 'num_rfx_basis' not found in JSON (model appears to have been ", - "serialized under stochtree ", - .ver, - "). Defaulting to FALSE / 1. ", - "Re-save your model to suppress this warning." - )) - } - - model_params[["num_gfr"]] <- json_object$get_scalar("num_gfr") - model_params[["num_burnin"]] <- json_object$get_scalar("num_burnin") - model_params[["num_mcmc"]] <- json_object$get_scalar("num_mcmc") - model_params[["num_samples"]] <- json_object$get_scalar("num_samples") - model_params[["num_covariates"]] <- if (has_field("num_covariates")) { - json_object$get_scalar("num_covariates") - } else { - NA_real_ - } - model_params[["num_basis"]] <- json_object$get_scalar("num_basis") - model_params <- .reconstructBartLeafModelFields(model_params) - model_params[["requires_basis"]] <- json_object$get_boolean("requires_basis") - - if (has_field("num_chains")) { - model_params[["num_chains"]] <- json_object$get_scalar("num_chains") - } else { - model_params[["num_chains"]] <- 1 - warning(paste0( - "Field 'num_chains' not found in JSON (model appears to have been serialized under stochtree ", - .ver, - "). Defaulting to 1. Re-save your model to suppress this warning." - )) - } - - if (has_field("keep_every")) { - model_params[["keep_every"]] <- json_object$get_scalar("keep_every") - } else { - model_params[["keep_every"]] <- 1 - warning(paste0( - "Field 'keep_every' not found in JSON (model appears to have been serialized under stochtree ", - .ver, - "). Defaulting to 1. Re-save your model to suppress this warning." - )) - } - - model_params[["probit_outcome_model"]] <- if ( - has_field("probit_outcome_model") - ) { - json_object$get_boolean("probit_outcome_model") - } else { + include_variance <- json_object$get_boolean_or_default( + "include_variance_forest", FALSE - } - - if ( - has_subfolder_field("outcome_model", "outcome") && - has_subfolder_field("outcome_model", "link") - ) { - outcome_model_outcome <- json_object$get_string("outcome", "outcome_model") - outcome_model_link <- json_object$get_string("link", "outcome_model") - } else { - outcome_model_outcome <- "continuous" - outcome_model_link <- "identity" - warning(paste0( - "Fields 'outcome' and 'link' not found under 'outcome_model' in JSON (model appears to have ", - "been serialized under stochtree ", - .ver, - "). Defaulting to outcome='continuous', ", - "link='identity'. Re-save your model to suppress this warning." - )) - } - model_params[["outcome_model"]] <- OutcomeModel( - outcome = outcome_model_outcome, - link = outcome_model_link ) - - if (has_field("rfx_model_spec")) { - model_params[["rfx_model_spec"]] <- json_object$get_string("rfx_model_spec") - } else { - model_params[["rfx_model_spec"]] <- "" - if (model_params[["has_rfx"]]) { - warning(paste0( - "Field 'rfx_model_spec' not found in JSON (model appears to have been serialized under ", - "stochtree ", - .ver, - "). Defaulting to ''. Re-save your model to suppress this warning." - )) - } - } - if (model_params[["outcome_model"]]$link == "cloglog") { - cloglog_num_categories <- json_object$get_scalar("cloglog_num_categories") - model_params[["cloglog_num_categories"]] <- cloglog_num_categories - } else { - model_params[["cloglog_num_categories"]] <- 0 - } - - output[["model_params"]] <- model_params - - # Unpack sampled parameters - if (model_params[["sample_sigma2_global"]]) { - output[["sigma2_global_samples"]] <- json_object$get_vector( - "sigma2_global_samples", - "parameters" - ) - } - if (model_params[["sample_sigma2_leaf"]]) { - output[["sigma2_leaf_samples"]] <- json_object$get_vector( - "sigma2_leaf_samples", - "parameters" - ) - } - if ( - model_params[["outcome_model"]]$link == "cloglog" && - model_params[["outcome_model"]]$outcome == "ordinal" - ) { - cloglog_cutpoint_samples <- matrix( - NA_real_, - model_params[["cloglog_num_categories"]] - 1, - model_params[["num_samples"]] - ) - for (i in 1:(model_params[["cloglog_num_categories"]] - 1)) { - cloglog_cutpoint_samples[i, ] <- json_object$get_vector( - paste0("cloglog_cutpoint_samples_", i), - "parameters" + if (include_mean) { + json_object$rename_field( + "forest_0", + "mean_forest", + subfolder_name = "forests" + ) + if (include_variance) { + json_object$rename_field( + "forest_1", + "variance_forest", + subfolder_name = "forests" ) } - output[["cloglog_cutpoint_samples"]] <- cloglog_cutpoint_samples - } - - # Unpack random effects - if (model_params[["has_rfx"]]) { - output[["rfx_unique_group_ids"]] <- json_object$get_string_vector( - "rfx_unique_group_ids" - ) - output[["rfx_samples"]] <- loadRandomEffectSamplesJson(json_object, 0) - } - - # Unpack covariate preprocessor - if (has_field("preprocessor_metadata")) { - preprocessor_metadata_string <- json_object$get_string( - "preprocessor_metadata" - ) - output[["train_set_metadata"]] <- createPreprocessorFromJsonString( - preprocessor_metadata_string + } else if (include_variance) { + json_object$rename_field( + "forest_0", + "variance_forest", + subfolder_name = "forests" + ) + } + # R's legacy preprocessor key -> unified v1 key (no-op for Python v0 JSON, + # which already uses `covariate_preprocessor`). + json_object$rename_field("preprocessor_metadata", "covariate_preprocessor") + # Relocate R's top-level rfx unique group ids into the random_effects subfolder + # (no-op for Python v0 JSON, which never wrote this field). + if (json_object$contains("rfx_unique_group_ids")) { + .rfx_uids <- json_object$get_string_vector("rfx_unique_group_ids") + json_object$add_string_vector( + "rfx_unique_group_ids", + .rfx_uids, + subfolder_name = "random_effects" ) - } else { - output[["train_set_metadata"]] <- NULL - warning(paste0( - "Field 'preprocessor_metadata' not found in JSON (model appears to have been serialized ", - "under stochtree ", - .ver, - "). DataFrame covariates will not be supported for prediction. ", - "Re-save your model to suppress this warning." - )) + json_object$erase_field("rfx_unique_group_ids") } - - class(output) <- "bartmodel" - return(output) -} - -#' @title Convert JSON File to BART Model -#' @rdname BARTSerialization -#' @param json_filename String of filepath, must end in ".json" -#' @export -createBARTModelFromJsonFile <- function(json_filename) { - # Load a `CppJson` object from file - bart_json <- createCppJsonFile(json_filename) - - # Create and return the BART object - bart_object <- createBARTModelFromJson(bart_json) - - return(bart_object) } -#' @title Convert JSON String to BART Model -#' @rdname BARTSerialization -#' @param json_string JSON string dump -#' @export -createBARTModelFromJsonString <- function(json_string) { - # Load a `CppJson` object from string - bart_json <- createCppJsonString(json_string) - - # Create and return the BART object - bart_object <- createBARTModelFromJson(bart_json) - - return(bart_object) +#' @description Create a BARTSamples object from JSON +#' @noRd +createBARTSamplesFromJson <- function(json) { + bart_samples <- BARTSamples$new() + bart_samples$from_json(json) + bart_samples } -#' @title Convert JSON List to Single BART Model +#' @title Convert JSON to BART Model #' @rdname BARTSerialization -#' @param json_object_list List of objects of type `CppJson` containing Json representation of a BART model +#' @param json_object Object of type `CppJson` containing Json representation of a BART model #' @export -createBARTModelFromCombinedJson <- function(json_object_list) { +createBARTModelFromJson <- function(json_object) { # Initialize the BCF model output <- list() - - # For scalar / preprocessing details which aren't sample-dependent, - # defer to the first json - json_object_default <- json_object_list[[1]] - - # Helpers for optional-field presence checks - .ver <- inferStochtreeJsonVersion(json_object_default) - has_field <- function(name) { - json_contains_field_cpp(json_object_default$json_ptr, name) - } - has_subfolder_field <- function(subfolder, name) { - json_contains_field_subfolder_cpp( - json_object_default$json_ptr, - subfolder, - name - ) - } - - # Unpack the forests - include_mean_forest <- json_object_default$get_boolean( - "include_mean_forest" - ) - include_variance_forest <- json_object_default$get_boolean( - "include_variance_forest" - ) - if (include_mean_forest) { - output[["mean_forests"]] <- loadForestContainerCombinedJson( - json_object_list, - "forest_0" - ) - if (include_variance_forest) { - output[["variance_forests"]] <- loadForestContainerCombinedJson( - json_object_list, - "forest_1" - ) - } - } else { - output[["variance_forests"]] <- loadForestContainerCombinedJson( - json_object_list, - "forest_0" - ) - } - - # Unpack metadata - train_set_metadata = list() - train_set_metadata[["num_numeric_vars"]] <- json_object_default$get_scalar( - "num_numeric_vars" - ) - train_set_metadata[[ - "num_ordered_cat_vars" - ]] <- json_object_default$get_scalar("num_ordered_cat_vars") - train_set_metadata[[ - "num_unordered_cat_vars" - ]] <- json_object_default$get_scalar("num_unordered_cat_vars") - if (train_set_metadata[["num_numeric_vars"]] > 0) { - train_set_metadata[[ - "numeric_vars" - ]] <- json_object_default$get_string_vector("numeric_vars") - } - if (train_set_metadata[["num_ordered_cat_vars"]] > 0) { - train_set_metadata[[ - "ordered_cat_vars" - ]] <- json_object_default$get_string_vector("ordered_cat_vars") - train_set_metadata[[ - "ordered_unique_levels" - ]] <- json_object_default$get_string_list( - "ordered_unique_levels", - train_set_metadata[["ordered_cat_vars"]] - ) + + # Helpers for optional-field presence checks + .ver <- inferStochtreeJsonVersion(json_object) + resolveSchemaVersion(json_object, migrate = .migrateBartJsonV0ToV1) + cross_platform <- enforceCrossPlatformGate(json_object, "R") + has_field <- function(name) { + json_contains_field_cpp(json_object$json_ptr, name) } - if (train_set_metadata[["num_unordered_cat_vars"]] > 0) { - train_set_metadata[[ - "unordered_cat_vars" - ]] <- json_object_default$get_string_vector("unordered_cat_vars") - train_set_metadata[[ - "unordered_unique_levels" - ]] <- json_object_default$get_string_list( - "unordered_unique_levels", - train_set_metadata[["unordered_cat_vars"]] - ) + has_subfolder_field <- function(subfolder, name) { + json_contains_field_subfolder_cpp(json_object$json_ptr, subfolder, name) } - output[["train_set_metadata"]] <- train_set_metadata # Unpack model params - model_params = list() - model_params[["outcome_scale"]] <- json_object_default$get_scalar( - "outcome_scale" - ) - model_params[["outcome_mean"]] <- json_object_default$get_scalar( - "outcome_mean" - ) - model_params[["standardize"]] <- json_object_default$get_boolean( - "standardize" + include_mean_forest <- json_object$get_boolean( + "include_mean_forest" ) - model_params[["sigma2_init"]] <- json_object_default$get_scalar( - "sigma2_init" + include_variance_forest <- json_object$get_boolean( + "include_variance_forest" ) - model_params[["sample_sigma2_global"]] <- json_object_default$get_boolean( + model_params <- list() + model_params[["outcome_scale"]] <- json_object$get_scalar("outcome_scale") + model_params[["outcome_mean"]] <- json_object$get_scalar("outcome_mean") + model_params[["standardize"]] <- json_object$get_boolean("standardize") + model_params[["sigma2_init"]] <- json_object$get_scalar("sigma2_init") + model_params[["sample_sigma2_global"]] <- json_object$get_boolean( "sample_sigma2_global" ) - model_params[["sample_sigma2_leaf"]] <- json_object_default$get_boolean( + model_params[["sample_sigma2_leaf"]] <- json_object$get_boolean( "sample_sigma2_leaf" ) model_params[["include_mean_forest"]] <- include_mean_forest model_params[["include_variance_forest"]] <- include_variance_forest - model_params[["has_rfx"]] <- json_object_default$get_boolean("has_rfx") + model_params[["has_rfx"]] <- json_object$get_boolean("has_rfx") if (has_field("has_rfx_basis")) { - model_params[["has_rfx_basis"]] <- json_object_default$get_boolean( - "has_rfx_basis" - ) - model_params[["num_rfx_basis"]] <- json_object_default$get_scalar( - "num_rfx_basis" - ) + model_params[["has_rfx_basis"]] <- json_object$get_boolean("has_rfx_basis") + model_params[["num_rfx_basis"]] <- json_object$get_scalar("num_rfx_basis") } else { model_params[["has_rfx_basis"]] <- FALSE model_params[["num_rfx_basis"]] <- 1 @@ -4142,21 +2681,45 @@ createBARTModelFromCombinedJson <- function(json_object_list) { )) } + model_params[["num_gfr"]] <- json_object$get_scalar("num_gfr") + model_params[["num_burnin"]] <- json_object$get_scalar("num_burnin") + model_params[["num_mcmc"]] <- json_object$get_scalar("num_mcmc") + model_params[["num_samples"]] <- json_object$get_scalar("num_samples") model_params[["num_covariates"]] <- if (has_field("num_covariates")) { - json_object_default$get_scalar("num_covariates") + json_object$get_scalar("num_covariates") } else { NA_real_ } - model_params[["num_basis"]] <- json_object_default$get_scalar("num_basis") + model_params[["num_basis"]] <- json_object$get_scalar("num_basis") model_params <- .reconstructBartLeafModelFields(model_params) - model_params[["requires_basis"]] <- json_object_default$get_boolean( - "requires_basis" - ) + model_params[["requires_basis"]] <- json_object$get_boolean("requires_basis") + + if (has_field("num_chains")) { + model_params[["num_chains"]] <- json_object$get_scalar("num_chains") + } else { + model_params[["num_chains"]] <- 1 + warning(paste0( + "Field 'num_chains' not found in JSON (model appears to have been serialized under stochtree ", + .ver, + "). Defaulting to 1. Re-save your model to suppress this warning." + )) + } + + if (has_field("keep_every")) { + model_params[["keep_every"]] <- json_object$get_scalar("keep_every") + } else { + model_params[["keep_every"]] <- 1 + warning(paste0( + "Field 'keep_every' not found in JSON (model appears to have been serialized under stochtree ", + .ver, + "). Defaulting to 1. Re-save your model to suppress this warning." + )) + } model_params[["probit_outcome_model"]] <- if ( has_field("probit_outcome_model") ) { - json_object_default$get_boolean("probit_outcome_model") + json_object$get_boolean("probit_outcome_model") } else { FALSE } @@ -4165,14 +2728,8 @@ createBARTModelFromCombinedJson <- function(json_object_list) { has_subfolder_field("outcome_model", "outcome") && has_subfolder_field("outcome_model", "link") ) { - outcome_model_outcome <- json_object_default$get_string( - "outcome", - "outcome_model" - ) - outcome_model_link <- json_object_default$get_string( - "link", - "outcome_model" - ) + outcome_model_outcome <- json_object$get_string("outcome", "outcome_model") + outcome_model_link <- json_object$get_string("link", "outcome_model") } else { outcome_model_outcome <- "continuous" outcome_model_link <- "identity" @@ -4190,9 +2747,7 @@ createBARTModelFromCombinedJson <- function(json_object_list) { ) if (has_field("rfx_model_spec")) { - model_params[["rfx_model_spec"]] <- json_object_default$get_string( - "rfx_model_spec" - ) + model_params[["rfx_model_spec"]] <- json_object$get_string("rfx_model_spec") } else { model_params[["rfx_model_spec"]] <- "" if (model_params[["has_rfx"]]) { @@ -4204,136 +2759,36 @@ createBARTModelFromCombinedJson <- function(json_object_list) { )) } } - - if (has_field("num_chains")) { - model_params[["num_chains"]] <- json_object_default$get_scalar("num_chains") - } else { - model_params[["num_chains"]] <- 1 - warning(paste0( - "Field 'num_chains' not found in JSON (model appears to have been serialized under stochtree ", - .ver, - "). Defaulting to 1. Re-save your model to suppress this warning." - )) - } - - if (has_field("keep_every")) { - model_params[["keep_every"]] <- json_object_default$get_scalar("keep_every") - } else { - model_params[["keep_every"]] <- 1 - warning(paste0( - "Field 'keep_every' not found in JSON (model appears to have been serialized under stochtree ", - .ver, - "). Defaulting to 1. Re-save your model to suppress this warning." - )) - } if (model_params[["outcome_model"]]$link == "cloglog") { - cloglog_num_categories <- json_object_default$get_scalar( - "cloglog_num_categories" - ) + cloglog_num_categories <- json_object$get_scalar("cloglog_num_categories") model_params[["cloglog_num_categories"]] <- cloglog_num_categories } else { model_params[["cloglog_num_categories"]] <- 0 } - # Combine values that are sample-specific - for (i in 1:length(json_object_list)) { - json_object <- json_object_list[[i]] - if (i == 1) { - model_params[["num_gfr"]] <- json_object$get_scalar("num_gfr") - model_params[["num_burnin"]] <- json_object$get_scalar("num_burnin") - model_params[["num_mcmc"]] <- json_object$get_scalar("num_mcmc") - model_params[["num_samples"]] <- json_object$get_scalar( - "num_samples" - ) - } else { - prev_json <- json_object_list[[i - 1]] - model_params[["num_gfr"]] <- model_params[["num_gfr"]] + - json_object$get_scalar("num_gfr") - model_params[["num_burnin"]] <- model_params[["num_burnin"]] + - json_object$get_scalar("num_burnin") - model_params[["num_mcmc"]] <- model_params[["num_mcmc"]] + - json_object$get_scalar("num_mcmc") - model_params[["num_samples"]] <- model_params[["num_samples"]] + - json_object$get_scalar("num_samples") - } - } output[["model_params"]] <- model_params - # Unpack sampled parameters - if (model_params[["sample_sigma2_global"]]) { - for (i in 1:length(json_object_list)) { - json_object <- json_object_list[[i]] - if (i == 1) { - output[["sigma2_global_samples"]] <- json_object$get_vector( - "sigma2_global_samples", - "parameters" - ) - } else { - output[["sigma2_global_samples"]] <- c( - output[["sigma2_global_samples"]], - json_object$get_vector( - "sigma2_global_samples", - "parameters" - ) - ) - } - } - } - if (model_params[["sample_sigma2_leaf"]]) { - for (i in 1:length(json_object_list)) { - json_object <- json_object_list[[i]] - if (i == 1) { - output[["sigma2_leaf_samples"]] <- json_object$get_vector( - "sigma2_leaf_samples", - "parameters" - ) - } else { - output[["sigma2_leaf_samples"]] <- c( - output[["sigma2_leaf_samples"]], - json_object$get_vector("sigma2_leaf_samples", "parameters") - ) - } - } - } - if ( - model_params[["outcome_model"]]$link == "cloglog" && - model_params[["outcome_model"]]$outcome == "ordinal" - ) { - cloglog_cutpoint_samples <- matrix( - NA_real_, - model_params[["cloglog_num_categories"]] - 1, - model_params[["num_samples"]] - ) - index_start <- 1 - for (i in 1:length(json_object_list)) { - json_object <- json_object_list[[i]] - num_samples <- json_object$get_scalar("num_samples") - subset_inds <- index_start:(index_start + num_samples - 1) - for (j in 1:(model_params[["cloglog_num_categories"]] - 1)) { - cloglog_cutpoint_samples[j, subset_inds] <- json_object$get_vector( - paste0("cloglog_cutpoint_samples_", j), - "parameters" - ) - } - } - output[["cloglog_cutpoint_samples"]] <- cloglog_cutpoint_samples - } + # Unpack samples + output[["samples"]] <- createBARTSamplesFromJson(json_object) - # Unpack random effects + # Unpack random effects group IDs if (model_params[["has_rfx"]]) { - output[[ - "rfx_unique_group_ids" - ]] <- json_object_default$get_string_vector("rfx_unique_group_ids") - output[["rfx_samples"]] <- loadRandomEffectSamplesCombinedJson( - json_object_list, - 0 + output[["rfx_unique_group_ids"]] <- resolveRfxUniqueGroupIds( + json_object, + output[["rfx_samples"]]$materialize_rfx() ) } # Unpack covariate preprocessor - if (has_field("preprocessor_metadata")) { - preprocessor_metadata_string <- json_object_default$get_string( - "preprocessor_metadata" + if (cross_platform) { + # Identity metadata for the cross-platform all-numeric path (gate enforced); + # the foreign native preprocessor is not reconstructed. + output[["train_set_metadata"]] <- buildIdentityPreprocessorMetadata( + json_object + ) + } else if (has_field("covariate_preprocessor")) { + preprocessor_metadata_string <- json_object$get_string( + "covariate_preprocessor" ) output[["train_set_metadata"]] <- createPreprocessorFromJsonString( preprocessor_metadata_string @@ -4341,7 +2796,7 @@ createBARTModelFromCombinedJson <- function(json_object_list) { } else { output[["train_set_metadata"]] <- NULL warning(paste0( - "Field 'preprocessor_metadata' not found in JSON (model appears to have been serialized ", + "Field 'covariate_preprocessor' not found in JSON (model appears to have been serialized ", "under stochtree ", .ver, "). DataFrame covariates will not be supported for prediction. ", @@ -4353,27 +2808,52 @@ createBARTModelFromCombinedJson <- function(json_object_list) { return(output) } -#' @title Convert JSON String List to Single BART Model +#' @title Convert JSON File to BART Model #' @rdname BARTSerialization -#' @param json_string_list List of JSON strings which can be parsed to objects of type `CppJson` containing Json representation of a BART model +#' @param json_filename String of filepath, must end in ".json" #' @export -createBARTModelFromCombinedJsonString <- function(json_string_list) { +createBARTModelFromJsonFile <- function(json_filename) { + # Load a `CppJson` object from file + bart_json <- createCppJsonFile(json_filename) + + # Create and return the BART object + bart_object <- createBARTModelFromJson(bart_json) + + return(bart_object) +} + +#' @title Convert JSON String to BART Model +#' @rdname BARTSerialization +#' @param json_string JSON string dump +#' @export +createBARTModelFromJsonString <- function(json_string) { + # Load a `CppJson` object from string + bart_json <- createCppJsonString(json_string) + + # Create and return the BART object + bart_object <- createBARTModelFromJson(bart_json) + + return(bart_object) +} + +#' @title Convert JSON List to Single BART Model +#' @rdname BARTSerialization +#' @param json_object_list List of objects of type `CppJson` containing Json representation of a BART model +#' @export +createBARTModelFromCombinedJson <- function(json_object_list) { # Initialize the BCF model output <- list() - # Convert JSON strings - json_object_list <- list() - for (i in 1:length(json_string_list)) { - json_string <- json_string_list[[i]] - json_object_list[[i]] <- createCppJsonString(json_string) - } - # For scalar / preprocessing details which aren't sample-dependent, # defer to the first json json_object_default <- json_object_list[[1]] # Helpers for optional-field presence checks .ver <- inferStochtreeJsonVersion(json_object_default) + for (.jo in json_object_list) { + resolveSchemaVersion(.jo, migrate = .migrateBartJsonV0ToV1) + } + cross_platform <- enforceCrossPlatformGate(json_object_default, "R") has_field <- function(name) { json_contains_field_cpp(json_object_default$json_ptr, name) } @@ -4385,73 +2865,14 @@ createBARTModelFromCombinedJsonString <- function(json_string_list) { ) } - # Unpack the forests + # Unpack model params include_mean_forest <- json_object_default$get_boolean( "include_mean_forest" ) include_variance_forest <- json_object_default$get_boolean( "include_variance_forest" ) - if (include_mean_forest) { - output[["mean_forests"]] <- loadForestContainerCombinedJson( - json_object_list, - "forest_0" - ) - if (include_variance_forest) { - output[["variance_forests"]] <- loadForestContainerCombinedJson( - json_object_list, - "forest_1" - ) - } - } else { - output[["variance_forests"]] <- loadForestContainerCombinedJson( - json_object_list, - "forest_0" - ) - } - - # Unpack metadata - train_set_metadata = list() - train_set_metadata[["num_numeric_vars"]] <- json_object_default$get_scalar( - "num_numeric_vars" - ) - train_set_metadata[[ - "num_ordered_cat_vars" - ]] <- json_object_default$get_scalar("num_ordered_cat_vars") - train_set_metadata[[ - "num_unordered_cat_vars" - ]] <- json_object_default$get_scalar("num_unordered_cat_vars") - if (train_set_metadata[["num_numeric_vars"]] > 0) { - train_set_metadata[[ - "numeric_vars" - ]] <- json_object_default$get_string_vector("numeric_vars") - } - if (train_set_metadata[["num_ordered_cat_vars"]] > 0) { - train_set_metadata[[ - "ordered_cat_vars" - ]] <- json_object_default$get_string_vector("ordered_cat_vars") - train_set_metadata[[ - "ordered_unique_levels" - ]] <- json_object_default$get_string_list( - "ordered_unique_levels", - train_set_metadata[["ordered_cat_vars"]] - ) - } - if (train_set_metadata[["num_unordered_cat_vars"]] > 0) { - train_set_metadata[[ - "unordered_cat_vars" - ]] <- json_object_default$get_string_vector("unordered_cat_vars") - train_set_metadata[[ - "unordered_unique_levels" - ]] <- json_object_default$get_string_list( - "unordered_unique_levels", - train_set_metadata[["unordered_cat_vars"]] - ) - } - output[["train_set_metadata"]] <- train_set_metadata - - # Unpack model params - model_params = list() + model_params <- list() model_params[["outcome_scale"]] <- json_object_default$get_scalar( "outcome_scale" ) @@ -4577,7 +2998,6 @@ createBARTModelFromCombinedJsonString <- function(json_string_list) { "). Defaulting to 1. Re-save your model to suppress this warning." )) } - if (model_params[["outcome_model"]]$link == "cloglog") { cloglog_num_categories <- json_object_default$get_scalar( "cloglog_num_categories" @@ -4611,81 +3031,36 @@ createBARTModelFromCombinedJsonString <- function(json_string_list) { } output[["model_params"]] <- model_params - # Unpack sampled parameters - if (model_params[["sample_sigma2_global"]]) { - for (i in 1:length(json_object_list)) { - json_object <- json_object_list[[i]] - if (i == 1) { - output[["sigma2_global_samples"]] <- json_object$get_vector( - "sigma2_global_samples", - "parameters" - ) - } else { - output[["sigma2_global_samples"]] <- c( - output[["sigma2_global_samples"]], - json_object$get_vector( - "sigma2_global_samples", - "parameters" - ) - ) - } - } - } - if (model_params[["sample_sigma2_leaf"]]) { - for (i in 1:length(json_object_list)) { - json_object <- json_object_list[[i]] - if (i == 1) { - output[["sigma2_leaf_samples"]] <- json_object$get_vector( - "sigma2_leaf_samples", - "parameters" - ) - } else { - output[["sigma2_leaf_samples"]] <- c( - output[["sigma2_leaf_samples"]], - json_object$get_vector("sigma2_leaf_samples", "parameters") - ) - } - } - } - if ( - model_params[["outcome_model"]]$link == "cloglog" && - model_params[["outcome_model"]]$outcome == "ordinal" - ) { - cloglog_cutpoint_samples <- matrix( - NA_real_, - model_params[["cloglog_num_categories"]] - 1, - model_params[["num_samples"]] - ) - index_start <- 1 - for (i in 1:length(json_object_list)) { - json_object <- json_object_list[[i]] - num_samples <- json_object$get_scalar("num_samples") - subset_inds <- index_start:(index_start + num_samples - 1) - for (j in 1:(model_params[["cloglog_num_categories"]] - 1)) { - cloglog_cutpoint_samples[j, subset_inds] <- json_object$get_vector( - paste0("cloglog_cutpoint_samples_", j), - "parameters" - ) - } + # Unpack samples + for (i in 1:length(json_object_list)) { + json_object <- json_object_list[[i]] + if (i == 1) { + combined_samples <- createBARTSamplesFromJson(json_object) + } else { + additional_samples <- createBARTSamplesFromJson(json_object) + combined_samples$merge(additional_samples) } - output[["cloglog_cutpoint_samples"]] <- cloglog_cutpoint_samples } + output[["samples"]] <- combined_samples - # Unpack random effects + # Unpack random effects group IDs if (model_params[["has_rfx"]]) { - output[[ - "rfx_unique_group_ids" - ]] <- json_object_default$get_string_vector("rfx_unique_group_ids") - output[["rfx_samples"]] <- loadRandomEffectSamplesCombinedJson( - json_object_list, - 0 + output[["rfx_unique_group_ids"]] <- resolveRfxUniqueGroupIds( + json_object_default, + output[["rfx_samples"]] ## TODO: write materialize wrapper for RFX ) } # Unpack covariate preprocessor - if (has_field("preprocessor_metadata")) { + if (cross_platform) { + # Identity metadata for the cross-platform all-numeric path (gate enforced); + # the foreign native preprocessor is not reconstructed. + output[["train_set_metadata"]] <- buildIdentityPreprocessorMetadata( + json_object_default + ) + } else if (has_field("covariate_preprocessor")) { preprocessor_metadata_string <- json_object_default$get_string( - "preprocessor_metadata" + "covariate_preprocessor" ) output[["train_set_metadata"]] <- createPreprocessorFromJsonString( preprocessor_metadata_string @@ -4693,7 +3068,7 @@ createBARTModelFromCombinedJsonString <- function(json_string_list) { } else { output[["train_set_metadata"]] <- NULL warning(paste0( - "Field 'preprocessor_metadata' not found in JSON (model appears to have been serialized ", + "Field 'covariate_preprocessor' not found in JSON (model appears to have been serialized ", "under stochtree ", .ver, "). DataFrame covariates will not be supported for prediction. ", @@ -4704,3 +3079,19 @@ createBARTModelFromCombinedJsonString <- function(json_string_list) { class(output) <- "bartmodel" return(output) } + +#' @title Convert JSON String List to Single BART Model +#' @rdname BARTSerialization +#' @param json_string_list List of JSON strings which can be parsed to objects of type `CppJson` containing Json representation of a BART model +#' @export +createBARTModelFromCombinedJsonString <- function(json_string_list) { + # Convert JSON strings + json_object_list <- list() + for (i in 1:length(json_string_list)) { + json_string <- json_string_list[[i]] + json_object_list[[i]] <- createCppJsonString(json_string) + } + + # Create BART model from list of JSON objects + createBARTModelFromCombinedJson(json_object_list) +} diff --git a/R/bcf.R b/R/bcf.R index 431611ae..c75f06f7 100644 --- a/R/bcf.R +++ b/R/bcf.R @@ -89,11 +89,12 @@ NULL #' We do not currently support (but plan to in the near future), test set evaluation for group labels #' that were not in the training set. #' @param rfx_basis_test (Optional) Test set basis for "random-slope" regression in additive random effects model. -#' @param observation_weights (Optional) Numeric vector of observation weights of length `nrow(X_train)`. Weights are +#' @param observation_weights_train (Optional) Numeric vector of observation weights of length `nrow(X_train)`. Weights are #' applied as `y_i | - ~ N(mu(X_i), sigma^2 / w_i)`, so larger weights increase an observation's influence on the fit. #' All weights must be non-negative. Default: `NULL` (all observations equally weighted). Applied to both the #' prognostic and treatment effect forests. Compatible with Gaussian (continuous/identity) and probit outcome models; #' not compatible with cloglog link functions. +#' @param observation_weights Deprecated alias for `observation_weights_train`; will be removed in a future release. #' @param num_gfr Number of "warm-start" iterations run using the grow-from-root algorithm (He and Hahn, 2021). Default: 5. #' @param num_burnin Number of "burn-in" iterations of the MCMC sampler. Default: 0. #' @param num_mcmc Number of "retained" iterations of the MCMC sampler. Default: 100. @@ -113,7 +114,7 @@ NULL #' - `sigma2_global_shape` Shape parameter in the `IG(sigma2_global_shape, sigma2_global_scale)` global error variance model. Default: `0`. #' - `sigma2_global_scale` Scale parameter in the `IG(sigma2_global_shape, sigma2_global_scale)` global error variance model. Default: `0`. #' - `variable_weights` Numeric weights reflecting the relative probability of splitting on each variable. Does not need to sum to 1 but cannot be negative. Defaults to `rep(1/ncol(X_train), ncol(X_train))` if not set here. Note that if the propensity score is included as a covariate in either forest, its weight will default to `1/ncol(X_train)`. A workaround if you wish to provide a custom weight for the propensity score is to include it as a column in `X_train` and then set `propensity_covariate` to `'none'` adjust `keep_vars` accordingly for the `prognostic` or `treatment_effect` forests. -#' - `propensity_covariate` Whether to include the propensity score as a covariate in either or both of the forests. Enter `"none"` for neither, `"prognostic"` for the prognostic forest, `"treatment_effect"` for the treatment forest, and `"both"` for both forests. If this is not `"none"` and a propensity score is not provided, it will be estimated from (`X_train`, `Z_train`) using `stochtree::bart()`. Default: `"mu"`. +#' - `propensity_covariate` Whether to include the propensity score as a covariate in either or both of the forests. Enter `"none"` for neither, `"prognostic"` for the prognostic forest, `"treatment_effect"` for the treatment forest, and `"both"` for both forests. If this is not `"none"` and a propensity score is not provided, it will be estimated from (`X_train`, `Z_train`) using `stochtree::bart()`. Default: `"prognostic"`. #' - `adaptive_coding` Whether or not to use an "adaptive coding" scheme in which a binary treatment variable is not coded manually as (0,1) or (-1,1) but learned via parameters `b_0` and `b_1` that attach to the outcome model `[b_0 (1-Z) + b_1 Z] tau(X)`. This is ignored when Z is not binary. Default: `FALSE`. #' - `control_coding_init` Initial value of the "control" group coding parameter. This is ignored when Z is not binary. Default: `-0.5`. #' - `treated_coding_init` Initial value of the "treatment" group coding parameter. This is ignored when Z is not binary. Default: `0.5`. @@ -188,6 +189,7 @@ NULL #' - `variance_prior_shape` Shape parameter for the inverse gamma prior on the variance of the random effects "group parameter." Default: `1`. #' - `variance_prior_scale` Scale parameter for the inverse gamma prior on the variance of the random effects "group parameter." Default: `1`. #' +#' #' @return List of sampling outputs and a wrapper around the sampled forests (which can be used for in-memory prediction on new data, or serialized to JSON on disk). #' @export #' @@ -249,6 +251,7 @@ bcf <- function( propensity_test = NULL, rfx_group_ids_test = NULL, rfx_basis_test = NULL, + observation_weights_train = NULL, observation_weights = NULL, num_gfr = 5, num_burnin = 0, @@ -430,7 +433,7 @@ bcf <- function( min_samples_leaf_variance <- variance_forest_params_updated$min_samples_leaf max_depth_variance <- variance_forest_params_updated$max_depth a_0 <- variance_forest_params_updated$leaf_prior_calibration_param - variance_forest_init <- variance_forest_params_updated$init_root_val + variance_forest_init <- variance_forest_params_updated$variance_forest_init a_forest <- variance_forest_params_updated$var_forest_prior_shape b_forest <- variance_forest_params_updated$var_forest_prior_scale keep_vars_variance <- variance_forest_params_updated$keep_vars @@ -502,19 +505,6 @@ bcf <- function( ) } - # Set a function-scoped RNG if user provided a random seed - custom_rng <- random_seed >= 0 - has_existing_random_seed <- F - if (custom_rng) { - # Cache original global environment RNG state (if it exists) - if (exists(".Random.seed", envir = .GlobalEnv)) { - original_global_seed <- .Random.seed - has_existing_random_seed <- T - } - # Set new seed and store associated RNG state - set.seed(random_seed) - } - # Check if there are enough GFR samples to seed num_chains samplers if (num_gfr > 0) { if (num_chains > num_gfr) { @@ -561,43 +551,43 @@ bcf <- function( } previous_y_bar <- previous_bcf_model$model_params$outcome_mean previous_y_scale <- previous_bcf_model$model_params$outcome_scale - previous_forest_samples_mu <- previous_bcf_model$forests_mu - previous_forest_samples_tau <- previous_bcf_model$forests_tau + previous_forest_samples_mu <- previous_bcf_model$samples$materialize_mu_forest() + previous_forest_samples_tau <- previous_bcf_model$samples$materialize_tau_forest() if (previous_bcf_model$model_params$include_variance_forest) { - previous_forest_samples_variance <- previous_bcf_model$forests_variance + previous_forest_samples_variance <- previous_bcf_model$samples$materialize_variance_forest() } else { previous_forest_samples_variance <- NULL } if (previous_bcf_model$model_params$sample_sigma2_global) { - previous_global_var_samples <- previous_bcf_model$sigma2_global_samples / + previous_global_var_samples <- previous_bcf_model$samples$global_var_samples() / (previous_y_scale * previous_y_scale) } else { previous_global_var_samples <- NULL } if (previous_bcf_model$model_params$sample_sigma2_leaf_mu) { - previous_leaf_var_mu_samples <- previous_bcf_model$sigma2_leaf_mu_samples + previous_leaf_var_mu_samples <- previous_bcf_model$samples$leaf_scale_mu_samples() } else { previous_leaf_var_mu_samples <- NULL } if (previous_bcf_model$model_params$sample_sigma2_leaf_tau) { - previous_leaf_var_tau_samples <- previous_bcf_model$sigma2_leaf_tau_samples + previous_leaf_var_tau_samples <- previous_bcf_model$samples$leaf_scale_tau_samples() } else { previous_leaf_var_tau_samples <- NULL } if (previous_bcf_model$model_params$has_rfx) { - previous_rfx_samples <- previous_bcf_model$rfx_samples + previous_rfx_samples <- previous_bcf_model$samples$materialize_rfx() } else { previous_rfx_samples <- NULL } if (previous_bcf_model$model_params$adaptive_coding) { - previous_b_1_samples <- previous_bcf_model$b_1_samples - previous_b_0_samples <- previous_bcf_model$b_0_samples + previous_b_1_samples <- previous_bcf_model$samples$b1_samples() + previous_b_0_samples <- previous_bcf_model$samples$b0_samples() } else { previous_b_1_samples <- NULL previous_b_0_samples <- NULL } if (previous_bcf_model$model_params$sample_tau_0) { - previous_tau_0_samples <- previous_bcf_model$tau_0_samples + previous_tau_0_samples <- previous_bcf_model$samples$tau_0_samples() } else { previous_tau_0_samples <- NULL } @@ -624,37 +614,49 @@ bcf <- function( # Determine whether conditional variance will be modeled if (num_trees_variance > 0) { - include_variance_forest = TRUE + include_variance_forest <- TRUE } else { - include_variance_forest = FALSE + include_variance_forest <- FALSE } - # observation_weights validation and compatibility checks + # `observation_weights` was renamed to `observation_weights_train`; honor the + # deprecated argument for one release cycle. if (!is.null(observation_weights)) { - if (!is.numeric(observation_weights)) { - stop("observation_weights must be a numeric vector") + warning( + "`observation_weights` is deprecated and will be removed in a future release; use `observation_weights_train` instead." + ) + if (is.null(observation_weights_train)) { + observation_weights_train <- observation_weights + } + } + + # observation_weights_train validation and compatibility checks + if (!is.null(observation_weights_train)) { + if (!is.numeric(observation_weights_train)) { + stop("observation_weights_train must be a numeric vector") } - if (length(observation_weights) != nrow(X_train)) { - stop("length(observation_weights) must equal nrow(X_train)") + if (length(observation_weights_train) != nrow(X_train)) { + stop("length(observation_weights_train) must equal nrow(X_train)") } - if (any(observation_weights < 0)) { - stop("observation_weights cannot have any negative values") + if (any(observation_weights_train < 0)) { + stop("observation_weights_train cannot have any negative values") } - if (all(observation_weights == 0) && num_gfr > 0) { + if (all(observation_weights_train == 0) && num_gfr > 0) { stop( - "observation_weights are all zero (prior sampling mode) but num_gfr > 0. ", + "observation_weights_train are all zero (prior sampling mode) but num_gfr > 0. ", "GFR warm-start is data-dependent and ill-defined with zero weights. ", - "Set num_gfr = 0 when using all-zero observation_weights." + "Set num_gfr = 0 when using all-zero observation_weights_train." ) } if (link_is_cloglog) { stop( - "observation_weights are not compatible with cloglog link functions." + "observation_weights_train are not compatible with cloglog link functions." ) } if (include_variance_forest) { - warning( - "Results may be unreliable when observation_weights are deployed alongside a variance forest model." + stop( + "observation_weights_train are not compatible with a variance forest model. ", + "Use either observation_weights_train or a variance forest, not both." ) } } @@ -672,7 +674,7 @@ bcf <- function( # Variable weight preprocessing (and initialization if necessary) if (is.null(variable_weights)) { - variable_weights = rep(1 / ncol(X_train), ncol(X_train)) + variable_weights <- rep(1 / ncol(X_train), ncol(X_train)) } if (any(variable_weights < 0)) { stop("variable_weights cannot have any negative weights") @@ -1239,7 +1241,7 @@ bcf <- function( stopifnot(num_mcmc >= 0) # Determine whether a test set is provided - has_test = !is.null(X_test) + has_test <- !is.null(X_test) # Convert y_train to numeric vector if not already converted if (!is.null(dim(y_train))) { @@ -1329,7 +1331,8 @@ bcf <- function( X_test = X_test, num_gfr = num_gfr_propensity, num_burnin = num_burnin_propensity, - num_mcmc = num_mcmc_propensity + num_mcmc = num_mcmc_propensity, + general_params = list(random_seed = random_seed) ) propensity_train <- predict( bart_model_propensity, @@ -1460,6 +1463,20 @@ bcf <- function( ) sample_sigma2_global <- F } + # Calibrate the treatment effect leaf scale from delta_max when the user has not + # set sigma2_leaf_init directly for the treatment forest. + if (is.null(sigma2_leaf_tau)) { + # Prior calibrated so that P(abs(tau(X)) < delta_max / dnorm(0)) = p (p = 0.6827) + p <- 0.6827 + q_quantile <- qnorm((p + 1) / 2) + sigma2_leaf_tau_scalar <- ((delta_max / (q_quantile * dnorm(0)))^2) / + num_trees_tau + if (has_multivariate_treatment) { + sigma2_leaf_tau <- diag(sigma2_leaf_tau_scalar, ncol(Z_train)) + } else { + sigma2_leaf_tau <- sigma2_leaf_tau_scalar + } + } } # Runtime checks for variance forest @@ -1472,146 +1489,6 @@ bcf <- function( } } - # Handle standardization, prior calibration, and initialization of forest - # differently for binary and continuous outcomes - if (link_is_probit) { - # Probit-scale intercept: center the forest on the population-average latent mean. - # The forest predicts mu(X) and y_bar_train is added back at prediction time. - # The latent z sampling uses y_bar_train to set the correct truncated normal mean and to center z before the residual update. - y_bar_train <- qnorm(mean_cpp(as.numeric(y_train))) - y_std_train <- 1 - - # Set a pseudo outcome by subtracting mean_cpp(y_train) from y_train - resid_train <- y_train - mean_cpp(as.numeric(y_train)) - - # Set initial value for the mu forest - init_mu <- 0.0 - - # Calibrate priors for global sigma^2 and sigma2_leaf_mu / sigma2_leaf_tau - # Set sigma2_init to 1, ignoring any defaults provided - sigma2_init <- 1.0 - # Skip variance_forest_init, since variance forests are not supported with probit link - if (is.null(b_leaf_mu)) { - b_leaf_mu <- 1 / num_trees_mu - } - if (is.null(b_leaf_tau)) { - b_leaf_tau <- 1 / (2 * num_trees_tau) - } - if (is.null(sigma2_leaf_mu)) { - sigma2_leaf_mu <- 2 / (num_trees_mu) - current_leaf_scale_mu <- as.matrix(sigma2_leaf_mu) - } else { - if (!is.matrix(sigma2_leaf_mu)) { - current_leaf_scale_mu <- as.matrix(sigma2_leaf_mu) - } else { - current_leaf_scale_mu <- sigma2_leaf_mu - } - } - if (is.null(sigma2_leaf_tau)) { - # Calibrate prior so that P(abs(tau(X)) < delta_max / dnorm(0)) = p - # Use p = 0.9 as an internal default rather than adding another - # user-facing "parameter" of the binary outcome BCF prior. - # Can be overriden by specifying `sigma2_leaf_init` in - # treatment_effect_forest_params. - p <- 0.6827 - q_quantile <- qnorm((p + 1) / 2) - sigma2_leaf_tau <- ((delta_max / (q_quantile * dnorm(0)))^2) / - num_trees_tau - current_leaf_scale_tau <- as.matrix(diag( - sigma2_leaf_tau, - ncol(Z_train) - )) - } else { - if (!is.matrix(sigma2_leaf_tau)) { - current_leaf_scale_tau <- as.matrix(diag( - sigma2_leaf_tau, - ncol(Z_train) - )) - } else { - if (ncol(sigma2_leaf_tau) != ncol(Z_train)) { - stop( - "sigma2_leaf_init for the tau forest must have the same number of columns / rows as columns in the Z_train matrix" - ) - } - if (nrow(sigma2_leaf_tau) != ncol(Z_train)) { - stop( - "sigma2_leaf_init for the tau forest must have the same number of columns / rows as columns in the Z_train matrix" - ) - } - current_leaf_scale_tau <- sigma2_leaf_tau - } - } - current_sigma2 <- sigma2_init - } else { - # Only standardize if user requested - if (standardize) { - y_bar_train <- mean_cpp(as.numeric(y_train)) - y_std_train <- sd_cpp(as.numeric(y_train)) - } else { - y_bar_train <- 0 - y_std_train <- 1 - } - - # Compute standardized outcome - resid_train <- (y_train - y_bar_train) / y_std_train - - # Set initial value for the mu forest - init_mu <- mean_cpp(as.numeric(resid_train)) - - # Calibrate priors for global sigma^2 and sigma2_leaf_mu / sigma2_leaf_tau - if (is.null(sigma2_init)) { - sigma2_init <- 1.0 * var_cpp(as.numeric(resid_train)) - } - if (is.null(variance_forest_init)) { - variance_forest_init <- 1.0 * var_cpp(as.numeric(resid_train)) - } - if (is.null(b_leaf_mu)) { - b_leaf_mu <- var_cpp(as.numeric(resid_train)) / (num_trees_mu) - } - if (is.null(b_leaf_tau)) { - b_leaf_tau <- var_cpp(as.numeric(resid_train)) / (2 * num_trees_tau) - } - if (is.null(sigma2_leaf_mu)) { - sigma2_leaf_mu <- 2.0 * var_cpp(as.numeric(resid_train)) / (num_trees_mu) - current_leaf_scale_mu <- as.matrix(sigma2_leaf_mu) - } else { - if (!is.matrix(sigma2_leaf_mu)) { - current_leaf_scale_mu <- as.matrix(sigma2_leaf_mu) - } else { - current_leaf_scale_mu <- sigma2_leaf_mu - } - } - if (is.null(sigma2_leaf_tau)) { - sigma2_leaf_tau <- 0.5 * - var_cpp(as.numeric(resid_train)) / - (num_trees_tau) - current_leaf_scale_tau <- as.matrix(diag( - sigma2_leaf_tau, - ncol(Z_train) - )) - } else { - if (!is.matrix(sigma2_leaf_tau)) { - current_leaf_scale_tau <- as.matrix(diag( - sigma2_leaf_tau, - ncol(Z_train) - )) - } else { - if (ncol(sigma2_leaf_tau) != ncol(Z_train)) { - stop( - "sigma2_leaf_init for the tau forest must have the same number of columns / rows as columns in the Z_train matrix" - ) - } - if (nrow(sigma2_leaf_tau) != ncol(Z_train)) { - stop( - "sigma2_leaf_init for the tau forest must have the same number of columns / rows as columns in the Z_train matrix" - ) - } - current_leaf_scale_tau <- sigma2_leaf_tau - } - } - current_sigma2 <- sigma2_init - } - # Set mu and tau leaf models / dimensions leaf_model_mu_forest <- 0 leaf_dimension_mu_forest <- 1 @@ -1627,2857 +1504,406 @@ bcf <- function( leaf_model_variance_forest <- 3 leaf_dimension_variance_forest <- 1 - # Random effects prior parameters - if (has_rfx) { - # Prior parameters - if (is.null(rfx_working_parameter_prior_mean)) { - if (num_rfx_components == 1) { - alpha_init <- c(0) - } else if (num_rfx_components > 1) { - alpha_init <- rep(0, num_rfx_components) - } else { - stop("There must be at least 1 random effect component") - } - } else { - alpha_init <- expand_dims_1d( - rfx_working_parameter_prior_mean, - num_rfx_components - ) - } - - if (is.null(rfx_group_parameter_prior_mean)) { - xi_init <- matrix( - rep(alpha_init, num_rfx_groups), - num_rfx_components, - num_rfx_groups - ) - } else { - xi_init <- expand_dims_2d( - rfx_group_parameter_prior_mean, - num_rfx_components, - num_rfx_groups - ) - } - - if (is.null(rfx_working_parameter_prior_cov)) { - sigma_alpha_init <- diag(1, num_rfx_components, num_rfx_components) - } else { - sigma_alpha_init <- expand_dims_2d_diag( - rfx_working_parameter_prior_cov, - num_rfx_components - ) - } - - if (is.null(rfx_group_parameter_prior_cov)) { - sigma_xi_init <- diag(1, num_rfx_components, num_rfx_components) - } else { - sigma_xi_init <- expand_dims_2d_diag( - rfx_group_parameter_prior_cov, - num_rfx_components - ) - } - - sigma_xi_shape <- rfx_variance_prior_shape - sigma_xi_scale <- rfx_variance_prior_scale - } - - # Random effects data structure and storage container - if (has_rfx) { - rfx_dataset_train <- createRandomEffectsDataset( - rfx_group_ids_train, - rfx_basis_train - ) - rfx_tracker_train <- createRandomEffectsTracker(rfx_group_ids_train) - rfx_model <- createRandomEffectsModel( - num_rfx_components, - num_rfx_groups - ) - rfx_model$set_working_parameter(alpha_init) - rfx_model$set_group_parameters(xi_init) - rfx_model$set_working_parameter_cov(sigma_alpha_init) - rfx_model$set_group_parameter_cov(sigma_xi_init) - rfx_model$set_variance_prior_shape(sigma_xi_shape) - rfx_model$set_variance_prior_scale(sigma_xi_scale) - rfx_samples <- createRandomEffectSamples( - num_rfx_components, - num_rfx_groups, - rfx_tracker_train - ) - } - - # Container of variance parameter samples - num_actual_mcmc_iter <- num_mcmc * keep_every - num_samples <- num_gfr + num_burnin + num_actual_mcmc_iter - # Delete GFR samples from these containers after the fact if desired - # num_retained_samples <- ifelse(keep_gfr, num_gfr, 0) + ifelse(keep_burnin, num_burnin, 0) + num_mcmc - num_retained_samples <- num_gfr + - ifelse(keep_burnin, num_burnin, 0) + - num_mcmc * num_chains - if (sample_sigma2_global) { - global_var_samples <- rep(NA, num_retained_samples) - } - if (sample_sigma2_leaf_mu) { - leaf_scale_mu_samples <- rep(NA, num_retained_samples) - } - if (sample_sigma2_leaf_tau) { - leaf_scale_tau_samples <- rep(NA, num_retained_samples) - } - if (sample_tau_0) { - p_tau0 <- ncol(as.matrix(Z_train)) - tau_0_samples <- matrix(NA_real_, p_tau0, num_retained_samples) - } - muhat_train_raw <- matrix(NA_real_, nrow(X_train), num_retained_samples) + # Model params set without calibration / initialization if (include_variance_forest) { - sigma2_x_train_raw <- matrix( - NA_real_, - nrow(X_train), - num_retained_samples - ) - } - sample_counter <- 0 - - # Prepare adaptive coding structure - if ( - (!is.numeric(b_0)) || - (!is.numeric(b_1)) || - (length(b_0) > 1) || - (length(b_1) > 1) - ) { - stop("b_0 and b_1 must be single numeric values") - } - if (adaptive_coding) { - b_0_samples <- rep(NA, num_retained_samples) - b_1_samples <- rep(NA, num_retained_samples) - current_b_0 <- b_0 - current_b_1 <- b_1 - tau_basis_train <- (1 - Z_train) * current_b_0 + Z_train * current_b_1 - if (has_test) { - tau_basis_test <- (1 - Z_test) * current_b_0 + Z_test * current_b_1 - } + num_variance_covariates <- sum(variable_weights_variance > 0) } else { - tau_basis_train <- Z_train - if (has_test) tau_basis_test <- Z_test - } - - # Prepare tau_0 (global treatment effect intercept) structure - if (sample_tau_0) { - if (!exists("p_tau0")) { - p_tau0 <- ncol(as.matrix(Z_train)) - } - tau_0 <- rep(0.0, p_tau0) - # Auto-calibrate prior variance if not provided - if (is.null(tau_0_prior_var)) { - tau_0_prior_var <- var_cpp(as.numeric(resid_train)) - } - prior_var_tau0 <- diag(p_tau0) * tau_0_prior_var + num_variance_covariates <- 0 } - - # Data - forest_dataset_train <- createForestDataset( - X_train, - tau_basis_train, - observation_weights + model_params_r <- list( + "initial_b_0" = b_0, + "initial_b_1" = b_1, + "a_global" = a_global, + "b_global" = b_global, + "a_leaf_mu" = a_leaf_mu, + "a_leaf_tau" = a_leaf_tau, + "standardize" = standardize, + "num_covariates" = num_cov_orig, + "num_prognostic_covariates" = sum(variable_weights_mu > 0), + "num_treatment_covariates" = sum(variable_weights_tau > 0), + "num_variance_covariates" = num_variance_covariates, + "treatment_dim" = ncol(Z_train), + "propensity_covariate" = propensity_covariate, + "binary_treatment" = binary_treatment, + "multivariate_treatment" = has_multivariate_treatment, + "adaptive_coding" = adaptive_coding, + "internal_propensity_model" = internal_propensity_model, + "num_gfr" = num_gfr, + "num_burnin" = num_burnin, + "num_mcmc" = num_mcmc, + "keep_every" = keep_every, + "num_chains" = num_chains, + "has_rfx" = has_rfx, + "has_rfx_basis" = has_basis_rfx, + "num_rfx_basis" = num_basis_rfx, + "include_variance_forest" = include_variance_forest, + "sample_sigma2_global" = sample_sigma2_global, + "sample_sigma2_leaf_mu" = sample_sigma2_leaf_mu, + "sample_sigma2_leaf_tau" = sample_sigma2_leaf_tau, + "probit_outcome_model" = probit_outcome_model, + "outcome_model" = outcome_model, + "rfx_model_spec" = rfx_model_spec ) - if (has_test) { - forest_dataset_test <- createForestDataset(X_test, tau_basis_test) - } - outcome_train <- createOutcome(resid_train) - # Random number generator (std::mt19937) - if (is.null(random_seed)) { - random_seed = sample(1:10000, 1, FALSE) + # Expand dimensions on RFX prior parameters if provided + # Working parameter (should be expanded to a vector if provided as a scalar) + if (!is.null(rfx_working_parameter_prior_mean)) { + rfx_working_parameter_prior_mean <- expand_dims_1d( + rfx_working_parameter_prior_mean, + num_rfx_components + ) } - rng <- createCppRNG(random_seed) - # Sampling data structures - global_model_config <- createGlobalModelConfig( - global_error_variance = current_sigma2 - ) - forest_model_config_mu <- createForestModelConfig( - feature_types = feature_types, - num_trees = num_trees_mu, - num_features = ncol(X_train), - num_observations = nrow(X_train), - variable_weights = variable_weights_mu, - leaf_dimension = leaf_dimension_mu_forest, - alpha = alpha_mu, - beta = beta_mu, - min_samples_leaf = min_samples_leaf_mu, - max_depth = max_depth_mu, - leaf_model_type = leaf_model_mu_forest, - leaf_model_scale = current_leaf_scale_mu, - cutpoint_grid_size = cutpoint_grid_size, - num_features_subsample = num_features_subsample_mu - ) - forest_model_config_tau <- createForestModelConfig( - feature_types = feature_types, - num_trees = num_trees_tau, - num_features = ncol(X_train), - num_observations = nrow(X_train), - variable_weights = variable_weights_tau, - leaf_dimension = leaf_dimension_tau_forest, - alpha = alpha_tau, - beta = beta_tau, - min_samples_leaf = min_samples_leaf_tau, - max_depth = max_depth_tau, - leaf_model_type = leaf_model_tau_forest, - leaf_model_scale = current_leaf_scale_tau, - cutpoint_grid_size = cutpoint_grid_size, - num_features_subsample = num_features_subsample_tau - ) - forest_model_mu <- createForestModel( - forest_dataset_train, - forest_model_config_mu, - global_model_config - ) - forest_model_tau <- createForestModel( - forest_dataset_train, - forest_model_config_tau, - global_model_config - ) - if (include_variance_forest) { - forest_model_config_variance <- createForestModelConfig( - feature_types = feature_types, - num_trees = num_trees_variance, - num_features = ncol(X_train), - num_observations = nrow(X_train), - variable_weights = variable_weights_variance, - leaf_dimension = leaf_dimension_variance_forest, - alpha = alpha_variance, - beta = beta_variance, - min_samples_leaf = min_samples_leaf_variance, - max_depth = max_depth_variance, - leaf_model_type = leaf_model_variance_forest, - cutpoint_grid_size = cutpoint_grid_size, - num_features_subsample = num_features_subsample_variance - ) - forest_model_variance <- createForestModel( - forest_dataset_train, - forest_model_config_variance, - global_model_config + # Group parameter (should be expanded to a matrix if provided as a scalar) + if (!is.null(rfx_group_parameter_prior_mean)) { + rfx_group_parameter_prior_mean <- expand_dims_2d( + rfx_group_parameter_prior_mean, + num_rfx_components, + num_rfx_groups ) } - # Container of forest samples - forest_samples_mu <- createForestSamples(num_trees_mu, 1, TRUE) - forest_samples_tau <- createForestSamples( - num_trees_tau, - ncol(Z_train), - FALSE - ) - active_forest_mu <- createForest(num_trees_mu, 1, TRUE) - active_forest_tau <- createForest(num_trees_tau, ncol(Z_train), FALSE) - if (include_variance_forest) { - forest_samples_variance <- createForestSamples( - num_trees_variance, - 1, - TRUE, - TRUE - ) - active_forest_variance <- createForest( - num_trees_variance, - 1, - TRUE, - TRUE + # Working parameter (should be expanded to a diagonal matrix if provided as a scalar) + if (!is.null(rfx_working_parameter_prior_cov)) { + rfx_working_parameter_prior_cov <- expand_dims_2d_diag( + rfx_working_parameter_prior_cov, + num_rfx_components ) } - # Initialize the leaves of each tree in the prognostic forest - active_forest_mu$prepare_for_sampler( - forest_dataset_train, - outcome_train, - forest_model_mu, - leaf_model_mu_forest, - init_mu - ) - active_forest_mu$adjust_residual( - forest_dataset_train, - outcome_train, - forest_model_mu, - FALSE, - FALSE - ) - - # Initialize the leaves of each tree in the treatment effect forest - init_tau <- rep(0., ncol(Z_train)) - active_forest_tau$prepare_for_sampler( - forest_dataset_train, - outcome_train, - forest_model_tau, - leaf_model_tau_forest, - init_tau - ) - active_forest_tau$adjust_residual( - forest_dataset_train, - outcome_train, - forest_model_tau, - TRUE, - FALSE - ) - - # Initialize the leaves of each tree in the variance forest - if (include_variance_forest) { - active_forest_variance$prepare_for_sampler( - forest_dataset_train, - outcome_train, - forest_model_variance, - leaf_model_variance_forest, - variance_forest_init + # Group parameter (should be expanded to a diagonal matrix if provided as a scalar) + if (!is.null(rfx_group_parameter_prior_cov)) { + rfx_group_parameter_prior_cov <- expand_dims_2d_diag( + rfx_group_parameter_prior_cov, + num_rfx_components ) } - # Run GFR (warm start) if specified - if (num_gfr > 0) { - for (i in 1:num_gfr) { - # Keep all GFR samples at this stage -- remove from ForestSamples after MCMC - # keep_sample <- ifelse(keep_gfr, TRUE, FALSE) - keep_sample <- TRUE - if (keep_sample) { - sample_counter <- sample_counter + 1 - } - # Print progress - if (verbose) { - if ((i %% 10 == 0) || (i == num_gfr)) { - cat( - "Sampling", - i, - "out of", - num_gfr, - "XBCF (grow-from-root) draws\n" - ) - } - } - - if (link_is_probit) { - # Sample latent probit variable, z | - - # outcome_pred is the centered forest prediction (not including y_bar_train). - # The truncated normal mean is outcome_pred + y_bar_train (the full eta on the probit scale). - # The residual stored is z - y_bar_train - outcome_pred so the forest sees a - # zero-centered signal and the prior shrinkage toward 0 is well-calibrated. - mu_forest_pred <- active_forest_mu$predict(forest_dataset_train) - tau_forest_pred <- active_forest_tau$predict( - forest_dataset_train + # Specify the BCF config + bcf_config <- list( + "standardize_outcome" = standardize, + "num_threads" = num_threads, + "verbose" = verbose, + "cutpoint_grid_size" = cutpoint_grid_size, + "link_function" = ifelse( + outcome_model$link == "identity", + 0, + ifelse(outcome_model$link == "probit", 1, 2) + ), + "outcome_type" = ifelse( + outcome_model$outcome == "continuous", + 0, + ifelse(outcome_model$outcome == "binary", 1, 2) + ), + "random_seed" = random_seed, + "keep_gfr" = keep_gfr, + "keep_burnin" = keep_burnin, + "adaptive_coding" = adaptive_coding, + "b_0_init" = b_0, + "b_1_init" = b_1, + "a_sigma2_global" = a_global, + "b_sigma2_global" = b_global, + "sigma2_global_init" = sigma2_init, + "sample_sigma2_global" = sample_sigma2_global, + "num_trees_mu" = num_trees_mu, + "alpha_mu" = alpha_mu, + "beta_mu" = beta_mu, + "min_samples_leaf_mu" = min_samples_leaf_mu, + "max_depth_mu" = max_depth_mu, + "leaf_constant_mu" = TRUE, + "leaf_dim_mu" = leaf_dimension_mu_forest, + "exponentiated_leaf_mu" = FALSE, + "num_features_subsample_mu" = num_features_subsample_mu, + "a_sigma2_mu" = a_leaf_mu, + "b_sigma2_mu" = b_leaf_mu, + "sigma2_mu_init" = if (is.matrix(sigma2_leaf_mu)) { + NULL + } else { + sigma2_leaf_mu + }, + "sample_sigma2_leaf_mu" = sample_sigma2_leaf_mu, + "mean_leaf_model_type" = leaf_model_mu_forest, + "sigma2_leaf_mu_matrix" = if (is.matrix(sigma2_leaf_mu)) { + as.numeric(sigma2_leaf_mu) + } else { + NULL + }, + "num_trees_tau" = num_trees_tau, + "alpha_tau" = alpha_tau, + "beta_tau" = beta_tau, + "min_samples_leaf_tau" = min_samples_leaf_tau, + "max_depth_tau" = max_depth_tau, + "leaf_constant_tau" = FALSE, + "leaf_dim_tau" = leaf_dimension_tau_forest, + "exponentiated_leaf_tau" = FALSE, + "num_features_subsample_tau" = num_features_subsample_tau, + "a_sigma2_tau" = a_leaf_tau, + "b_sigma2_tau" = b_leaf_tau, + "sigma2_tau_init" = if (is.matrix(sigma2_leaf_tau)) { + NULL + } else { + sigma2_leaf_tau + }, + "sample_sigma2_leaf_tau" = sample_sigma2_leaf_tau, + "tau_leaf_model_type" = leaf_model_tau_forest, + "sigma2_leaf_tau_matrix" = if (is.matrix(sigma2_leaf_tau)) { + as.numeric(sigma2_leaf_tau) + } else { + NULL + }, + "sample_tau_0" = sample_tau_0, + "tau_0_prior_var_scalar" = if (is.matrix(tau_0_prior_var)) { + NULL + } else { + tau_0_prior_var + }, + "tau_0_prior_var_multivariate" = if (is.matrix(tau_0_prior_var)) { + as.numeric(tau_0_prior_var) + } else { + NULL + }, + "num_trees_variance" = num_trees_variance, + "leaf_prior_calibration_param" = a_0, + "shape_variance_forest" = a_forest, + "scale_variance_forest" = b_forest, + "variance_forest_leaf_init" = variance_forest_init, + "alpha_variance" = alpha_variance, + "beta_variance" = beta_variance, + "min_samples_leaf_variance" = min_samples_leaf_variance, + "max_depth_variance" = max_depth_variance, + "leaf_constant_variance" = TRUE, + "leaf_dim_variance" = leaf_dimension_variance_forest, + "exponentiated_leaf_variance" = TRUE, + "num_features_subsample_variance" = num_features_subsample_variance, + "feature_types" = as.integer(feature_types), + "sweep_update_indices_mu" = if (num_trees_mu > 0) { + 0:(num_trees_mu - 1) + } else { + NULL + }, + "sweep_update_indices_tau" = if (num_trees_tau > 0) { + 0:(num_trees_tau - 1) + } else { + NULL + }, + "sweep_update_indices_variance" = if (num_trees_variance > 0) { + 0:(num_trees_variance - 1) + } else { + NULL + }, + "var_weights_mu" = variable_weights_mu, + "var_weights_tau" = variable_weights_tau, + "var_weights_variance" = variable_weights_variance, + "has_random_effects" = has_rfx, + "rfx_model_spec" = if (has_rfx) { + ifelse( + rfx_model_spec == "custom", + 0, + ifelse( + rfx_model_spec == "intercept_only", + 1, + ifelse(rfx_model_spec == "intercept_plus_treatment", 2, NULL) ) - outcome_pred <- mu_forest_pred + tau_forest_pred - if (has_rfx) { - rfx_pred <- rfx_model$predict( - rfx_dataset_train, - rfx_tracker_train - ) - outcome_pred <- outcome_pred + rfx_pred - } - eta_pred <- outcome_pred + y_bar_train - mu0 <- eta_pred[y_train == 0] - mu1 <- eta_pred[y_train == 1] - u0 <- runif(sum(y_train == 0), 0, pnorm(0 - mu0)) - u1 <- runif(sum(y_train == 1), pnorm(0 - mu1), 1) - resid_train[y_train == 0] <- mu0 + qnorm(u0) - resid_train[y_train == 1] <- mu1 + qnorm(u1) - - # Update outcome: center z by y_bar_train before passing to forests - outcome_train$update_data(resid_train - y_bar_train - outcome_pred) - } - - # Sample the prognostic forest - 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, - num_threads = num_threads, - keep_forest = keep_sample, - gfr = TRUE ) + } else { + NULL + }, + "rfx_working_parameter_mean_prior" = if (has_rfx) { + rfx_working_parameter_prior_mean + } else { + NULL + }, + "rfx_working_parameter_cov_prior" = if (has_rfx) { + rfx_working_parameter_prior_cov + } else { + NULL + }, + "rfx_group_parameter_mean_prior" = if (has_rfx) { + rfx_group_parameter_prior_mean + } else { + NULL + }, + "rfx_group_parameter_cov_prior" = if (has_rfx) { + rfx_group_parameter_prior_cov + } else { + NULL + }, + "rfx_variance_prior_shape" = if (has_rfx) { + rfx_variance_prior_shape + } else { + NULL + }, + "rfx_variance_prior_scale" = if (has_rfx) { + rfx_variance_prior_scale + } else { + NULL + } + ) + + bcf_samples <- BCFSamples$new() + bcf_metadata <- bcf_sample_cpp( + samples = bcf_samples$samples_ptr, + X_train = X_train, + Z_train = Z_train, + y_train = y_train, + X_test = if (exists("X_test")) X_test else NULL, + Z_test = if (exists("Z_test")) Z_test else NULL, + n_train = nrow(X_train), + n_test = if (!is.null(X_test)) nrow(X_test) else 0L, + p = ncol(X_train), + treatment_dim = ncol(Z_train), + obs_weights_train = observation_weights_train, + obs_weights_test = NULL, + rfx_group_ids_train = if (exists("rfx_group_ids_train")) { + rfx_group_ids_train + } else { + NULL + }, + rfx_group_ids_test = if (exists("rfx_group_ids_test")) { + rfx_group_ids_test + } else { + NULL + }, + rfx_basis_train = if (exists("rfx_basis_train")) { + rfx_basis_train + } else { + NULL + }, + rfx_basis_test = if (exists("rfx_basis_test")) rfx_basis_test else NULL, + rfx_num_groups = if (exists("num_rfx_groups")) { + as.integer(num_rfx_groups) + } else { + 0L + }, + rfx_basis_dim = as.integer(num_basis_rfx), + num_gfr = as.integer(num_gfr), + num_burnin = as.integer(num_burnin), + keep_every = as.integer(keep_every), + num_mcmc = as.integer(num_mcmc), + num_chains = as.integer(num_chains), + adaptive_coding = adaptive_coding, + config_input = bcf_config + ) + result <- list() + model_params_cpp <- list( + "initial_sigma2" = bcf_metadata[["sigma2_global_init"]], + "sigma2_leaf_mu" = bcf_metadata[["sigma2_mu_init"]], + "sigma2_leaf_tau" = bcf_metadata[["sigma2_tau_init"]], + "b_leaf_mu" = bcf_metadata[["b_sigma2_mu"]], + "b_leaf_tau" = bcf_metadata[["b_sigma2_tau"]], + "a_forest" = bcf_metadata[["shape_variance_forest"]], + "b_forest" = bcf_metadata[["scale_variance_forest"]], + "outcome_mean" = bcf_samples$y_bar(), + "outcome_scale" = bcf_samples$y_std(), + "num_samples" = bcf_samples$num_samples(), + "sample_tau_0" = sample_tau_0, + "tau_0_prior_var" = if (sample_tau_0) tau_0_prior_var else NULL + ) + model_params <- c(model_params_r, model_params_cpp) + result[["model_params"]] <- model_params + result[["train_set_metadata"]] <- X_train_metadata + result[["samples"]] <- bcf_samples - # Cache train set predictions since they are already computed during sampling - if (keep_sample) { - muhat_train_raw[, - sample_counter - ] <- forest_model_mu$get_cached_forest_predictions() - } - - # Sample variance parameters (if requested) - if (sample_sigma2_global) { - current_sigma2 <- sampleGlobalErrorVarianceOneIteration( - outcome_train, - forest_dataset_train, - rng, - a_global, - b_global - ) - global_model_config$update_global_error_variance(current_sigma2) - } - if (sample_sigma2_leaf_mu) { - leaf_scale_mu_double <- sampleLeafVarianceOneIteration( - active_forest_mu, - rng, - a_leaf_mu, - b_leaf_mu - ) - current_leaf_scale_mu <- as.matrix(leaf_scale_mu_double) - if (keep_sample) { - leaf_scale_mu_samples[ - sample_counter - ] <- leaf_scale_mu_double - } - forest_model_config_mu$update_leaf_model_scale( - current_leaf_scale_mu - ) - } - - # Sample tau_0 (global treatment effect intercept, if requested) - if (sample_tau_0) { - mu_x_raw_tau0 <- active_forest_mu$predict_raw(forest_dataset_train) - tau_x_raw_tau0 <- active_forest_tau$predict_raw(forest_dataset_train) - Z_basis_mat <- as.matrix(tau_basis_train) - # tau(X) * basis contribution per observation - tau_x_full <- rowSums(Z_basis_mat * as.matrix(tau_x_raw_tau0)) - # For probit, resid_train holds the full-scale latent z; center it so that - # tau_0 does not absorb the probit intercept y_bar_train. - resid_for_tau0 <- if (link_is_probit) { - resid_train - y_bar_train - } else { - resid_train - } - partial_resid_tau0 <- resid_for_tau0 - - as.numeric(mu_x_raw_tau0) - - tau_x_full - if (has_rfx) { - partial_resid_tau0 <- partial_resid_tau0 - - as.numeric( - rfx_model$predict(rfx_dataset_train, rfx_tracker_train) - ) - } - Ztr_tau0 <- t(Z_basis_mat) %*% as.matrix(partial_resid_tau0) - ZtZ_current <- crossprod(Z_basis_mat) - Sigma_post <- solve( - ZtZ_current / current_sigma2 + diag(p_tau0) / tau_0_prior_var - ) - mu_post_tau0 <- as.numeric(Sigma_post %*% Ztr_tau0 / current_sigma2) - if (p_tau0 == 1) { - tau_0_new <- rnorm(1, mu_post_tau0, sqrt(as.numeric(Sigma_post))) - } else { - tau_0_new <- as.numeric( - mu_post_tau0 + t(chol(Sigma_post)) %*% rnorm(p_tau0) - ) - } - resid_delta <- as.numeric( - Z_basis_mat %*% matrix(tau_0_new - tau_0, ncol = 1) - ) - outcome_train$subtract_vector(resid_delta) - tau_0 <- tau_0_new - if (keep_sample) { - tau_0_samples[, sample_counter] <- tau_0 - } - } + # Unpack RFX samples + if (has_rfx) { + # Only need to store unique group IDs, everything else stored in BARTSamples + result[["rfx_unique_group_ids"]] <- levels(group_ids_factor) + } - # Sample the treatment forest - 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, - num_threads = num_threads, - keep_forest = keep_sample, - gfr = TRUE - ) + if (internal_propensity_model) { + result[["bart_propensity_model"]] <- bart_model_propensity + } - # Cannot cache train set predictions for tau because the cached predictions in the - # tracking data structures are pre-multiplied by the basis (treatment) - # ... + class(result) <- "bcfmodel" - # Sample coding parameters (if requested) - if (adaptive_coding) { - # Estimate mu(X) and tau(X) and compute y - mu(X) - mu_x_raw_train <- active_forest_mu$predict_raw( - forest_dataset_train - ) - tau_x_raw_train <- active_forest_tau$predict_raw( - forest_dataset_train - ) - # For probit, resid_train holds full-scale z; center it so b_0/b_1 do not - # absorb the probit intercept y_bar_train. - resid_for_coding <- if (link_is_probit) { - resid_train - y_bar_train - } else { - resid_train - } - partial_resid_mu_train <- resid_for_coding - mu_x_raw_train - if (has_rfx) { - rfx_preds_train <- rfx_model$predict( - rfx_dataset_train, - rfx_tracker_train - ) - partial_resid_mu_train <- partial_resid_mu_train - - rfx_preds_train - } + return(result) +} - # Compute sufficient statistics for regression of y - mu(X) on [tau_total(X)(1-Z), tau_total(X)Z] - # where tau_total(X) = tau_0 + tau(X) when sample_tau_0, else tau(X) - tau_x_for_coding <- if (sample_tau_0) { - tau_x_raw_train + tau_0[1] - } else { - tau_x_raw_train - } - s_tt0 <- sum(tau_x_for_coding * tau_x_for_coding * (Z_train == 0)) - s_tt1 <- sum(tau_x_for_coding * tau_x_for_coding * (Z_train == 1)) - s_ty0 <- sum( - tau_x_for_coding * partial_resid_mu_train * (Z_train == 0) - ) - s_ty1 <- sum( - tau_x_for_coding * partial_resid_mu_train * (Z_train == 1) - ) +#' Guard accessor for a `bcfmodel`'s removed direct-forest fields. +#' +#' The sampled forests / parameters / predictions are owned by a single `BCFSamples` object stored in `object$samples`. +#' They are no longer stored or accessible via `$forests_mu` / `$forests_tau` / `$forests_variance`. +#' Similarly, sampled parameter vectors are no longer stored or accessible via `$sigma2_global_samples` / `$sigma2_leaf_mu_samples` / `$sigma2_leaf_tau_samples` +#' and cached predictions are no longer accessible via `$y_hat_train` / `$y_hat_test` / `$mu_hat_train` / `$mu_hat_test` / `$tau_hat_train` / `$tau_hat_test` / +#' `$sigma2_x_hat_train` / `$sigma2_x_hat_test`. Accessing any of these model terms by name raises an error pointing at the supported extraction path. +#' @noRd +#' @export +`$.bcfmodel` <- function(x, name) { + if ( + identical(name, "forests_mu") || + identical(name, "forests_tau") || + identical(name, "forests_variance") + ) { + stop( + sprintf( + paste0( + "`bcfmodel$%s` has been removed. The sampled forests are owned by `model$samples`; ", + "extract a standalone copy with `extractForest()`." + ), + name + ), + call. = FALSE + ) + } else if ( + identical(name, "sigma2_global_samples") || + identical(name, "sigma2_leaf_mu_samples") || + identical(name, "sigma2_leaf_tau_samples") || + identical(name, "b_0_samples") || + identical(name, "b_1_samples") || + identical(name, "tau_0_samples") || + identical(name, "y_hat_train") || + identical(name, "y_hat_test") || + identical(name, "mu_hat_train") || + identical(name, "mu_hat_test") || + identical(name, "tau_hat_train") || + identical(name, "tau_hat_test") || + identical(name, "sigma2_x_hat_train") || + identical(name, "sigma2_x_hat_test") + ) { + stop( + sprintf( + paste0( + "`bcfmodel$%s` has been removed. The parameters are stored in a C++ BCFSamples object; ", + "you can extract a standalone copy with `extractParameter()`." + ), + name + ), + call. = FALSE + ) + } else { + .subset2(x, name) + } +} - # Sample b0 (coefficient on tau_total(X)(1-Z)) and b1 (coefficient on tau_total(X)Z) - current_b_0 <- rnorm( - 1, - (s_ty0 / (s_tt0 + 2 * current_sigma2)), - sqrt(current_sigma2 / (s_tt0 + 2 * current_sigma2)) - ) - current_b_1 <- rnorm( - 1, - (s_ty1 / (s_tt1 + 2 * current_sigma2)), - sqrt(current_sigma2 / (s_tt1 + 2 * current_sigma2)) - ) - - # Update basis for the leaf regression - if (sample_tau_0) { - tau_basis_old <- tau_basis_train - } - tau_basis_train <- (1 - Z_train) * - current_b_0 + - Z_train * current_b_1 - forest_dataset_train$update_basis(tau_basis_train) - if (keep_sample) { - b_0_samples[sample_counter] <- current_b_0 - b_1_samples[sample_counter] <- current_b_1 - } - if (has_test) { - tau_basis_test <- (1 - Z_test) * - current_b_0 + - Z_test * current_b_1 - forest_dataset_test$update_basis(tau_basis_test) - } - - # Update leaf predictions and residual - forest_model_tau$propagate_basis_update( - forest_dataset_train, - outcome_train, - active_forest_tau - ) - - # Fix tau_0 component of residual after basis change - if (sample_tau_0) { - outcome_train$subtract_vector( - as.numeric(tau_basis_train - tau_basis_old) * tau_0[1] - ) - } - } - - # Sample variance parameters (if requested) - if (include_variance_forest) { - 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, - num_threads = num_threads, - keep_forest = keep_sample, - gfr = TRUE - ) - - # Cache train set predictions since they are already computed during sampling - if (keep_sample) { - sigma2_x_train_raw[, - sample_counter - ] <- forest_model_variance$get_cached_forest_predictions() - } - } - if (sample_sigma2_global) { - current_sigma2 <- sampleGlobalErrorVarianceOneIteration( - outcome_train, - forest_dataset_train, - rng, - a_global, - b_global - ) - if (keep_sample) { - global_var_samples[sample_counter] <- current_sigma2 - } - global_model_config$update_global_error_variance(current_sigma2) - } - if (sample_sigma2_leaf_tau) { - leaf_scale_tau_double <- sampleLeafVarianceOneIteration( - active_forest_tau, - rng, - a_leaf_tau, - b_leaf_tau - ) - current_leaf_scale_tau <- as.matrix(leaf_scale_tau_double) - if (keep_sample) { - leaf_scale_tau_samples[ - sample_counter - ] <- leaf_scale_tau_double - } - forest_model_config_mu$update_leaf_model_scale( - current_leaf_scale_mu - ) - } - - # Sample random effects parameters (if requested) - if (has_rfx) { - rfx_model$sample_random_effect( - rfx_dataset_train, - outcome_train, - rfx_tracker_train, - rfx_samples, - keep_sample, - current_sigma2, - rng - ) - } - } - } - - # Run MCMC - if (num_burnin + num_mcmc > 0) { - for (chain_num in 1:num_chains) { - if (verbose) { - cat("Sampling chain", chain_num, "of", num_chains, "\n") - } - if (num_gfr > 0) { - # Reset state of active_forest and forest_model based on a previous GFR sample - forest_ind <- num_gfr - chain_num - resetActiveForest( - active_forest_mu, - forest_samples_mu, - forest_ind - ) - resetForestModel( - forest_model_mu, - active_forest_mu, - forest_dataset_train, - outcome_train, - TRUE - ) - resetActiveForest( - active_forest_tau, - forest_samples_tau, - forest_ind - ) - resetForestModel( - forest_model_tau, - active_forest_tau, - forest_dataset_train, - outcome_train, - TRUE - ) - if (sample_sigma2_leaf_mu) { - leaf_scale_mu_double <- leaf_scale_mu_samples[ - forest_ind + 1 - ] - current_leaf_scale_mu <- as.matrix(leaf_scale_mu_double) - forest_model_config_mu$update_leaf_model_scale( - current_leaf_scale_mu - ) - } - if (sample_sigma2_leaf_tau) { - leaf_scale_tau_double <- leaf_scale_tau_samples[ - forest_ind + 1 - ] - current_leaf_scale_tau <- as.matrix(leaf_scale_tau_double) - forest_model_config_tau$update_leaf_model_scale( - current_leaf_scale_tau - ) - } - if (include_variance_forest) { - resetActiveForest( - active_forest_variance, - forest_samples_variance, - forest_ind - ) - resetForestModel( - forest_model_variance, - active_forest_variance, - forest_dataset_train, - outcome_train, - FALSE - ) - } - if (has_rfx) { - resetRandomEffectsModel( - rfx_model, - rfx_samples, - forest_ind, - sigma_alpha_init - ) - resetRandomEffectsTracker( - rfx_tracker_train, - rfx_model, - rfx_dataset_train, - outcome_train, - rfx_samples - ) - } - if (adaptive_coding) { - tau_basis_train_old <- tau_basis_train - current_b_1 <- b_1_samples[forest_ind + 1] - current_b_0 <- b_0_samples[forest_ind + 1] - tau_basis_train <- (1 - Z_train) * - current_b_0 + - Z_train * current_b_1 - forest_dataset_train$update_basis(tau_basis_train) - if (has_test) { - tau_basis_test <- (1 - Z_test) * - current_b_0 + - Z_test * current_b_1 - forest_dataset_test$update_basis(tau_basis_test) - } - forest_model_tau$propagate_basis_update( - forest_dataset_train, - outcome_train, - active_forest_tau - ) - # Correct residual for tau_0 component of the basis change - if (sample_tau_0) { - outcome_train$subtract_vector( - as.numeric((tau_basis_train - tau_basis_train_old) * tau_0[1]) - ) - } - } - # Reset tau_0 intercept and correct the running residual - if (sample_tau_0) { - tau_0_old <- tau_0 - tau_0 <- tau_0_samples[, forest_ind + 1] - Z_basis_gfr <- as.matrix(tau_basis_train) - outcome_train$subtract_vector( - as.numeric(Z_basis_gfr %*% matrix(tau_0 - tau_0_old, ncol = 1)) - ) - } - if (sample_sigma2_global) { - current_sigma2 <- global_var_samples[forest_ind + 1] - global_model_config$update_global_error_variance( - current_sigma2 - ) - } - } else if (has_prev_model) { - warmstart_index <- ifelse( - previous_model_decrement, - previous_model_warmstart_sample_num - chain_num + 1, - previous_model_warmstart_sample_num - ) - resetActiveForest( - active_forest_mu, - previous_forest_samples_mu, - warmstart_index - 1 - ) - resetForestModel( - forest_model_mu, - active_forest_mu, - forest_dataset_train, - outcome_train, - TRUE - ) - resetActiveForest( - active_forest_tau, - previous_forest_samples_tau, - warmstart_index - 1 - ) - resetForestModel( - forest_model_tau, - active_forest_tau, - forest_dataset_train, - outcome_train, - TRUE - ) - if (include_variance_forest) { - resetActiveForest( - active_forest_variance, - previous_forest_samples_variance, - warmstart_index - 1 - ) - resetForestModel( - forest_model_variance, - active_forest_variance, - forest_dataset_train, - outcome_train, - FALSE - ) - } - if ( - sample_sigma2_leaf_mu && - (!is.null(previous_leaf_var_mu_samples)) - ) { - leaf_scale_mu_double <- previous_leaf_var_mu_samples[ - warmstart_index - ] - current_leaf_scale_mu <- as.matrix(leaf_scale_mu_double) - forest_model_config_mu$update_leaf_model_scale( - current_leaf_scale_mu - ) - } - if ( - sample_sigma2_leaf_tau && - (!is.null(previous_leaf_var_tau_samples)) - ) { - leaf_scale_tau_double <- previous_leaf_var_tau_samples[ - warmstart_index - ] - current_leaf_scale_tau <- as.matrix(leaf_scale_tau_double) - forest_model_config_tau$update_leaf_model_scale( - current_leaf_scale_tau - ) - } - if (adaptive_coding) { - tau_basis_train_old <- tau_basis_train - if (!is.null(previous_b_1_samples)) { - current_b_1 <- previous_b_1_samples[ - warmstart_index - ] - } - if (!is.null(previous_b_0_samples)) { - current_b_0 <- previous_b_0_samples[ - warmstart_index - ] - } - tau_basis_train <- (1 - Z_train) * - current_b_0 + - Z_train * current_b_1 - forest_dataset_train$update_basis(tau_basis_train) - if (has_test) { - tau_basis_test <- (1 - Z_test) * - current_b_0 + - Z_test * current_b_1 - forest_dataset_test$update_basis(tau_basis_test) - } - forest_model_tau$propagate_basis_update( - forest_dataset_train, - outcome_train, - active_forest_tau - ) - # Correct residual for tau_0 component of the basis change - if (sample_tau_0) { - outcome_train$subtract_vector( - as.numeric((tau_basis_train - tau_basis_train_old) * tau_0[1]) - ) - } - } - # Reset tau_0 intercept and correct the running residual - if (sample_tau_0 && !is.null(previous_tau_0_samples)) { - tau_0_old <- tau_0 - # previous model stores tau_0 in original scale; convert to standardized scale - tau_0 <- as.numeric( - previous_tau_0_samples[, warmstart_index] / previous_y_scale - ) - Z_basis_ws <- as.matrix(tau_basis_train) - outcome_train$subtract_vector( - as.numeric(Z_basis_ws %*% matrix(tau_0 - tau_0_old, ncol = 1)) - ) - } - if (has_rfx) { - if (is.null(previous_rfx_samples)) { - warning( - "`previous_model_json` did not have any random effects samples, so the RFX sampler will be run from scratch while the forests and any other parameters are warm started" - ) - rootResetRandomEffectsModel( - rfx_model, - alpha_init, - xi_init, - sigma_alpha_init, - sigma_xi_init, - sigma_xi_shape, - sigma_xi_scale - ) - rootResetRandomEffectsTracker( - rfx_tracker_train, - rfx_model, - rfx_dataset_train, - outcome_train - ) - } else { - resetRandomEffectsModel( - rfx_model, - previous_rfx_samples, - warmstart_index - 1, - sigma_alpha_init - ) - resetRandomEffectsTracker( - rfx_tracker_train, - rfx_model, - rfx_dataset_train, - outcome_train, - rfx_samples - ) - } - } - if (sample_sigma2_global) { - if (!is.null(previous_global_var_samples)) { - current_sigma2 <- previous_global_var_samples[ - warmstart_index - ] - } - global_model_config$update_global_error_variance( - current_sigma2 - ) - } - } else { - resetActiveForest(active_forest_mu) - active_forest_mu$set_root_leaves(init_mu / num_trees_mu) - resetForestModel( - forest_model_mu, - active_forest_mu, - forest_dataset_train, - outcome_train, - TRUE - ) - resetActiveForest(active_forest_tau) - active_forest_tau$set_root_leaves(init_tau / num_trees_tau) - resetForestModel( - forest_model_tau, - active_forest_tau, - forest_dataset_train, - outcome_train, - TRUE - ) - if (sample_sigma2_leaf_mu) { - current_leaf_scale_mu <- as.matrix(sigma2_leaf_mu) - forest_model_config_mu$update_leaf_model_scale( - current_leaf_scale_mu - ) - } - if (sample_sigma2_leaf_tau) { - current_leaf_scale_tau <- as.matrix(sigma2_leaf_tau) - forest_model_config_tau$update_leaf_model_scale( - current_leaf_scale_tau - ) - } - if (include_variance_forest) { - resetActiveForest(active_forest_variance) - active_forest_variance$set_root_leaves( - log(variance_forest_init) / num_trees_variance - ) - resetForestModel( - forest_model_variance, - active_forest_variance, - forest_dataset_train, - outcome_train, - FALSE - ) - } - if (has_rfx) { - rootResetRandomEffectsModel( - rfx_model, - alpha_init, - xi_init, - sigma_alpha_init, - sigma_xi_init, - sigma_xi_shape, - sigma_xi_scale - ) - rootResetRandomEffectsTracker( - rfx_tracker_train, - rfx_model, - rfx_dataset_train, - outcome_train - ) - } - if (adaptive_coding) { - tau_basis_train_old <- tau_basis_train - current_b_1 <- b_1 - current_b_0 <- b_0 - tau_basis_train <- (1 - Z_train) * - current_b_0 + - Z_train * current_b_1 - forest_dataset_train$update_basis(tau_basis_train) - if (has_test) { - tau_basis_test <- (1 - Z_test) * - current_b_0 + - Z_test * current_b_1 - forest_dataset_test$update_basis(tau_basis_test) - } - forest_model_tau$propagate_basis_update( - forest_dataset_train, - outcome_train, - active_forest_tau - ) - # Correct residual for tau_0 component of the basis change - if (sample_tau_0) { - outcome_train$subtract_vector( - as.numeric((tau_basis_train - tau_basis_train_old) * tau_0[1]) - ) - } - } - # Reset tau_0 to initial value (0) and correct the running residual - if (sample_tau_0) { - tau_0_old <- tau_0 - tau_0 <- rep(0.0, p_tau0) - Z_basis_reset <- as.matrix(tau_basis_train) - outcome_train$subtract_vector( - as.numeric(Z_basis_reset %*% matrix(tau_0 - tau_0_old, ncol = 1)) - ) - } - if (sample_sigma2_global) { - current_sigma2 <- sigma2_init - global_model_config$update_global_error_variance( - current_sigma2 - ) - } - } - for (i in (num_gfr + 1):num_samples) { - is_mcmc <- i > (num_gfr + num_burnin) - if (is_mcmc) { - mcmc_counter <- i - (num_gfr + num_burnin) - if (mcmc_counter %% keep_every == 0) { - keep_sample <- TRUE - } else { - keep_sample <- FALSE - } - } else { - if (keep_burnin) { - keep_sample <- TRUE - } else { - keep_sample <- FALSE - } - } - if (keep_sample) { - sample_counter <- sample_counter + 1 - } - # Print progress - if (verbose) { - if (num_burnin > 0 && !is_mcmc) { - if ( - ((i - num_gfr) %% 100 == 0) || - ((i - num_gfr) == num_burnin) - ) { - cat( - "Sampling", - i - num_gfr, - "out of", - num_burnin, - "BCF burn-in draws; Chain number ", - chain_num, - "\n" - ) - } - } - if (num_mcmc > 0 && is_mcmc) { - raw_iter <- i - num_gfr - num_burnin - if ((raw_iter %% 100 == 0) || (i == num_samples)) { - if (keep_every == 1) { - cat( - "Sampling", - raw_iter, - "out of", - num_mcmc, - "BCF MCMC draws; Chain number ", - chain_num, - "\n" - ) - } else { - cat( - "Sampling raw draw", - raw_iter, - "of", - num_actual_mcmc_iter, - "BCF MCMC draws (thinning by", - keep_every, - ":", - raw_iter %/% keep_every, - "of", - num_mcmc, - "retained); Chain number ", - chain_num, - "\n" - ) - } - } - } - } - - if (link_is_probit) { - # Sample latent probit variable, z | - - mu_forest_pred <- active_forest_mu$predict( - forest_dataset_train - ) - tau_forest_pred <- active_forest_tau$predict( - forest_dataset_train - ) - outcome_pred <- mu_forest_pred + tau_forest_pred - if (has_rfx) { - rfx_pred <- rfx_model$predict( - rfx_dataset_train, - rfx_tracker_train - ) - outcome_pred <- outcome_pred + rfx_pred - } - eta_pred <- outcome_pred + y_bar_train - mu0 <- eta_pred[y_train == 0] - mu1 <- eta_pred[y_train == 1] - u0 <- runif(sum(y_train == 0), 0, pnorm(0 - mu0)) - u1 <- runif(sum(y_train == 1), pnorm(0 - mu1), 1) - resid_train[y_train == 0] <- mu0 + qnorm(u0) - resid_train[y_train == 1] <- mu1 + qnorm(u1) - - # Update outcome: center z by y_bar_train before passing to forests - outcome_train$update_data(resid_train - y_bar_train - outcome_pred) - } - - # Sample the prognostic forest - 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, - num_threads = num_threads, - keep_forest = keep_sample, - gfr = FALSE - ) - - # Cache train set predictions since they are already computed during sampling - if (keep_sample) { - muhat_train_raw[, - sample_counter - ] <- forest_model_mu$get_cached_forest_predictions() - } - - # Sample variance parameters (if requested) - if (sample_sigma2_global) { - current_sigma2 <- sampleGlobalErrorVarianceOneIteration( - outcome_train, - forest_dataset_train, - rng, - a_global, - b_global - ) - global_model_config$update_global_error_variance( - current_sigma2 - ) - } - if (sample_sigma2_leaf_mu) { - leaf_scale_mu_double <- sampleLeafVarianceOneIteration( - active_forest_mu, - rng, - a_leaf_mu, - b_leaf_mu - ) - current_leaf_scale_mu <- as.matrix(leaf_scale_mu_double) - if (keep_sample) { - leaf_scale_mu_samples[ - sample_counter - ] <- leaf_scale_mu_double - } - forest_model_config_mu$update_leaf_model_scale( - current_leaf_scale_mu - ) - } - - # Sample tau_0 (global treatment effect intercept, if requested) - if (sample_tau_0) { - mu_x_raw_tau0 <- active_forest_mu$predict_raw(forest_dataset_train) - tau_x_raw_tau0 <- active_forest_tau$predict_raw(forest_dataset_train) - Z_basis_mat <- as.matrix(tau_basis_train) - # tau(X) * basis contribution per observation - tau_x_full <- rowSums(Z_basis_mat * as.matrix(tau_x_raw_tau0)) - # For probit, resid_train holds the full-scale latent z; center it so that - # tau_0 does not absorb the probit intercept y_bar_train. - resid_for_tau0 <- if (link_is_probit) { - resid_train - y_bar_train - } else { - resid_train - } - partial_resid_tau0 <- resid_for_tau0 - - as.numeric(mu_x_raw_tau0) - - tau_x_full - if (has_rfx) { - partial_resid_tau0 <- partial_resid_tau0 - - as.numeric( - rfx_model$predict(rfx_dataset_train, rfx_tracker_train) - ) - } - Ztr_tau0 <- t(Z_basis_mat) %*% as.matrix(partial_resid_tau0) - ZtZ_current <- crossprod(Z_basis_mat) - Sigma_post <- solve( - ZtZ_current / current_sigma2 + diag(p_tau0) / tau_0_prior_var - ) - mu_post_tau0 <- as.numeric(Sigma_post %*% Ztr_tau0 / current_sigma2) - if (p_tau0 == 1) { - tau_0_new <- rnorm(1, mu_post_tau0, sqrt(as.numeric(Sigma_post))) - } else { - tau_0_new <- as.numeric( - mu_post_tau0 + t(chol(Sigma_post)) %*% rnorm(p_tau0) - ) - } - resid_delta <- as.numeric( - Z_basis_mat %*% matrix(tau_0_new - tau_0, ncol = 1) - ) - outcome_train$subtract_vector(resid_delta) - tau_0 <- tau_0_new - if (keep_sample) { - tau_0_samples[, sample_counter] <- tau_0 - } - } - - # Sample the treatment forest - 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, - num_threads = num_threads, - keep_forest = keep_sample, - gfr = FALSE - ) - - # Cannot cache train set predictions for tau because the cached predictions in the - # tracking data structures are pre-multiplied by the basis (treatment) - # ... - - # Sample coding parameters (if requested) - if (adaptive_coding) { - # Estimate mu(X) and tau(X) and compute y - mu(X) - mu_x_raw_train <- active_forest_mu$predict_raw( - forest_dataset_train - ) - tau_x_raw_train <- active_forest_tau$predict_raw( - forest_dataset_train - ) - # For probit, resid_train holds full-scale z; center it so b_0/b_1 do not - # absorb the probit intercept y_bar_train. - resid_for_coding <- if (link_is_probit) { - resid_train - y_bar_train - } else { - resid_train - } - partial_resid_mu_train <- resid_for_coding - mu_x_raw_train - if (has_rfx) { - rfx_preds_train <- rfx_model$predict( - rfx_dataset_train, - rfx_tracker_train - ) - partial_resid_mu_train <- partial_resid_mu_train - - rfx_preds_train - } - - # Compute sufficient statistics for regression of y - mu(X) on [tau_total(X)(1-Z), tau_total(X)Z] - # where tau_total(X) = tau_0 + tau(X) when sample_tau_0, else tau(X) - tau_x_for_coding <- if (sample_tau_0) { - tau_x_raw_train + tau_0[1] - } else { - tau_x_raw_train - } - s_tt0 <- sum( - tau_x_for_coding * tau_x_for_coding * (Z_train == 0) - ) - s_tt1 <- sum( - tau_x_for_coding * tau_x_for_coding * (Z_train == 1) - ) - s_ty0 <- sum( - tau_x_for_coding * - partial_resid_mu_train * - (Z_train == 0) - ) - s_ty1 <- sum( - tau_x_for_coding * - partial_resid_mu_train * - (Z_train == 1) - ) - - # Sample b0 (coefficient on tau_total(X)(1-Z)) and b1 (coefficient on tau_total(X)Z) - current_b_0 <- rnorm( - 1, - (s_ty0 / (s_tt0 + 2 * current_sigma2)), - sqrt(current_sigma2 / (s_tt0 + 2 * current_sigma2)) - ) - current_b_1 <- rnorm( - 1, - (s_ty1 / (s_tt1 + 2 * current_sigma2)), - sqrt(current_sigma2 / (s_tt1 + 2 * current_sigma2)) - ) - - # Update basis for the leaf regression - if (sample_tau_0) { - tau_basis_old <- tau_basis_train - } - tau_basis_train <- (1 - Z_train) * - current_b_0 + - Z_train * current_b_1 - forest_dataset_train$update_basis(tau_basis_train) - if (keep_sample) { - b_0_samples[sample_counter] <- current_b_0 - b_1_samples[sample_counter] <- current_b_1 - } - if (has_test) { - tau_basis_test <- (1 - Z_test) * - current_b_0 + - Z_test * current_b_1 - forest_dataset_test$update_basis(tau_basis_test) - } - - # Update leaf predictions and residual - forest_model_tau$propagate_basis_update( - forest_dataset_train, - outcome_train, - active_forest_tau - ) - - # Fix tau_0 component of residual after basis change - if (sample_tau_0) { - outcome_train$subtract_vector( - as.numeric(tau_basis_train - tau_basis_old) * tau_0[1] - ) - } - } - - # Sample variance parameters (if requested) - if (include_variance_forest) { - 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, - num_threads = num_threads, - keep_forest = keep_sample, - gfr = FALSE - ) - - # Cache train set predictions since they are already computed during sampling - if (keep_sample) { - sigma2_x_train_raw[, - sample_counter - ] <- forest_model_variance$get_cached_forest_predictions() - } - } - if (sample_sigma2_global) { - current_sigma2 <- sampleGlobalErrorVarianceOneIteration( - outcome_train, - forest_dataset_train, - rng, - a_global, - b_global - ) - if (keep_sample) { - global_var_samples[sample_counter] <- current_sigma2 - } - global_model_config$update_global_error_variance( - current_sigma2 - ) - } - if (sample_sigma2_leaf_tau) { - leaf_scale_tau_double <- sampleLeafVarianceOneIteration( - active_forest_tau, - rng, - a_leaf_tau, - b_leaf_tau - ) - current_leaf_scale_tau <- as.matrix(leaf_scale_tau_double) - if (keep_sample) { - leaf_scale_tau_samples[ - sample_counter - ] <- leaf_scale_tau_double - } - forest_model_config_tau$update_leaf_model_scale( - current_leaf_scale_tau - ) - } - - # Sample random effects parameters (if requested) - if (has_rfx) { - rfx_model$sample_random_effect( - rfx_dataset_train, - outcome_train, - rfx_tracker_train, - rfx_samples, - keep_sample, - current_sigma2, - rng - ) - } - } - } - } - - # Remove GFR samples if they are not to be retained - if ((!keep_gfr) && (num_gfr > 0)) { - for (i in 1:num_gfr) { - forest_samples_mu$delete_sample(0) - forest_samples_tau$delete_sample(0) - if (include_variance_forest) { - forest_samples_variance$delete_sample(0) - } - if (has_rfx) { - rfx_samples$delete_sample(0) - } - } - if (sample_sigma2_global) { - global_var_samples <- global_var_samples[ - (num_gfr + 1):length(global_var_samples) - ] - } - if (sample_sigma2_leaf_mu) { - leaf_scale_mu_samples <- leaf_scale_mu_samples[ - (num_gfr + 1):length(leaf_scale_mu_samples) - ] - } - if (sample_sigma2_leaf_tau) { - leaf_scale_tau_samples <- leaf_scale_tau_samples[ - (num_gfr + 1):length(leaf_scale_tau_samples) - ] - } - if (adaptive_coding) { - b_1_samples <- b_1_samples[(num_gfr + 1):length(b_1_samples)] - b_0_samples <- b_0_samples[(num_gfr + 1):length(b_0_samples)] - } - if (sample_tau_0) { - tau_0_samples <- tau_0_samples[, - (num_gfr + 1):ncol(tau_0_samples), - drop = FALSE - ] - } - muhat_train_raw <- muhat_train_raw[, - (num_gfr + 1):ncol(muhat_train_raw) - ] - if (include_variance_forest) { - sigma2_x_train_raw <- sigma2_x_train_raw[, - (num_gfr + 1):ncol(sigma2_x_train_raw) - ] - } - num_retained_samples <- num_retained_samples - num_gfr - } - - # Forest predictions - mu_hat_train <- muhat_train_raw * y_std_train + y_bar_train - if (adaptive_coding) { - tau_hat_train_raw <- forest_samples_tau$predict_raw( - forest_dataset_train - ) - tau_hat_train <- t(t(tau_hat_train_raw) * (b_1_samples - b_0_samples)) * - y_std_train - control_adj_train <- t(t(tau_hat_train_raw) * b_0_samples) * y_std_train - mu_hat_train <- mu_hat_train + control_adj_train - } else { - tau_hat_train <- forest_samples_tau$predict_raw(forest_dataset_train) * - y_std_train - } - # Fold tau_0 into tau_hat_train so it holds the full CATE (tau_0 + tau(X)) - if (sample_tau_0) { - tau_0_vec <- as.numeric(tau_0_samples) # num_retained_samples vector (scalar treatment) - if (adaptive_coding) { - # CATE = (b_1 - b_0) * (tau_0 + tau(X)); control adj to mu = b_0 * (tau_0 + tau(X)) - tau_hat_train <- sweep( - tau_hat_train, - 2, - (b_1_samples - b_0_samples) * tau_0_vec * y_std_train, - "+" - ) - mu_hat_train <- sweep( - mu_hat_train, - 2, - b_0_samples * tau_0_vec * y_std_train, - "+" - ) - } else if (!has_multivariate_treatment) { - tau_hat_train <- sweep(tau_hat_train, 2, tau_0_vec * y_std_train, "+") - } else { - # tau_hat_train: n x p x num_retained_samples; tau_0_samples: p x num_retained_samples - for (j in seq_len(p_tau0)) { - tau_hat_train[, j, ] <- tau_hat_train[, j, ] + - outer(rep(1, nrow(X_train)), tau_0_samples[j, ] * y_std_train) - } - } - } - if (has_multivariate_treatment) { - tau_train_dim <- dim(tau_hat_train) - tau_num_obs <- tau_train_dim[1] - tau_num_samples <- tau_train_dim[3] - treatment_term_train <- matrix( - NA_real_, - nrow = tau_num_obs, - tau_num_samples - ) - for (i in 1:nrow(Z_train)) { - treatment_term_train[i, ] <- colSums( - tau_hat_train[i, , ] * Z_train[i, ] - ) - } - } else { - treatment_term_train <- tau_hat_train * as.numeric(Z_train) - } - y_hat_train <- mu_hat_train + treatment_term_train - if (has_test) { - mu_hat_test <- forest_samples_mu$predict(forest_dataset_test) * - y_std_train + - y_bar_train - if (adaptive_coding) { - tau_hat_test_raw <- forest_samples_tau$predict_raw( - forest_dataset_test - ) - tau_hat_test <- t( - t(tau_hat_test_raw) * (b_1_samples - b_0_samples) - ) * - y_std_train - control_adj_test <- t(t(tau_hat_test_raw) * b_0_samples) * y_std_train - mu_hat_test <- mu_hat_test + control_adj_test - } else { - tau_hat_test <- forest_samples_tau$predict_raw( - forest_dataset_test - ) * - y_std_train - } - # Fold tau_0 into tau_hat_test so it holds the full CATE (tau_0 + tau(X)) - if (sample_tau_0) { - if (adaptive_coding) { - tau_hat_test <- sweep( - tau_hat_test, - 2, - (b_1_samples - b_0_samples) * tau_0_vec * y_std_train, - "+" - ) - mu_hat_test <- sweep( - mu_hat_test, - 2, - b_0_samples * tau_0_vec * y_std_train, - "+" - ) - } else if (!has_multivariate_treatment) { - tau_hat_test <- sweep(tau_hat_test, 2, tau_0_vec * y_std_train, "+") - } else { - for (j in seq_len(p_tau0)) { - tau_hat_test[, j, ] <- tau_hat_test[, j, ] + - outer(rep(1, nrow(X_test)), tau_0_samples[j, ] * y_std_train) - } - } - } - if (has_multivariate_treatment) { - tau_test_dim <- dim(tau_hat_test) - tau_num_obs <- tau_test_dim[1] - tau_num_samples <- tau_test_dim[3] - treatment_term_test <- matrix( - NA_real_, - nrow = tau_num_obs, - tau_num_samples - ) - for (i in 1:nrow(Z_test)) { - treatment_term_test[i, ] <- colSums( - tau_hat_test[i, , ] * Z_test[i, ] - ) - } - } else { - treatment_term_test <- tau_hat_test * as.numeric(Z_test) - } - y_hat_test <- mu_hat_test + treatment_term_test - } - if (include_variance_forest) { - sigma2_x_hat_train <- exp(sigma2_x_train_raw) - if (has_test) { - sigma2_x_hat_test <- forest_samples_variance$predict( - forest_dataset_test - ) - } - } - - # Random effects predictions - if (has_rfx) { - rfx_preds_train <- rfx_samples$predict( - rfx_group_ids_train, - rfx_basis_train - ) * - y_std_train - y_hat_train <- y_hat_train + rfx_preds_train - } - if ((has_rfx_test) && (has_test)) { - rfx_preds_test <- rfx_samples$predict( - rfx_group_ids_test, - rfx_basis_test - ) * - y_std_train - y_hat_test <- y_hat_test + rfx_preds_test - } - - # Global error variance - if (sample_sigma2_global) { - sigma2_global_samples <- global_var_samples * (y_std_train^2) - } - - # Leaf parameter variance for prognostic forest - if (sample_sigma2_leaf_mu) { - sigma2_leaf_mu_samples <- leaf_scale_mu_samples - } - - # Leaf parameter variance for treatment effect forest - if (sample_sigma2_leaf_tau) { - sigma2_leaf_tau_samples <- leaf_scale_tau_samples - } - - # Rescale variance forest prediction by global sigma2 (sampled or constant) - if (include_variance_forest) { - if (sample_sigma2_global) { - sigma2_x_hat_train <- sapply(1:num_retained_samples, function(i) { - sigma2_x_hat_train[, i] * sigma2_global_samples[i] - }) - if (has_test) { - sigma2_x_hat_test <- sapply( - 1:num_retained_samples, - function(i) { - sigma2_x_hat_test[, i] * sigma2_global_samples[i] - } - ) - } - } else { - sigma2_x_hat_train <- sigma2_x_hat_train * - sigma2_init * - y_std_train * - y_std_train - if (has_test) { - sigma2_x_hat_test <- sigma2_x_hat_test * - sigma2_init * - y_std_train * - y_std_train - } - } - } - - # Return results as a list - if (include_variance_forest) { - num_variance_covariates <- sum(variable_weights_variance > 0) - } else { - num_variance_covariates <- 0 - } - model_params <- list( - "initial_sigma2" = sigma2_init, - "initial_sigma2_leaf_mu" = sigma2_leaf_mu, - "initial_sigma2_leaf_tau" = sigma2_leaf_tau, - "initial_b_0" = b_0, - "initial_b_1" = b_1, - "a_global" = a_global, - "b_global" = b_global, - "a_leaf_mu" = a_leaf_mu, - "b_leaf_mu" = b_leaf_mu, - "a_leaf_tau" = a_leaf_tau, - "b_leaf_tau" = b_leaf_tau, - "a_forest" = a_forest, - "b_forest" = b_forest, - "outcome_mean" = y_bar_train, - "outcome_scale" = y_std_train, - "standardize" = standardize, - "num_covariates" = num_cov_orig, - "num_prognostic_covariates" = sum(variable_weights_mu > 0), - "num_treatment_covariates" = sum(variable_weights_tau > 0), - "num_variance_covariates" = num_variance_covariates, - "treatment_dim" = ncol(Z_train), - "propensity_covariate" = propensity_covariate, - "binary_treatment" = binary_treatment, - "multivariate_treatment" = has_multivariate_treatment, - "adaptive_coding" = adaptive_coding, - "sample_tau_0" = sample_tau_0, - "tau_0_prior_var" = if (sample_tau_0) tau_0_prior_var else NULL, - "internal_propensity_model" = internal_propensity_model, - "num_samples" = num_retained_samples, - "num_gfr" = num_gfr, - "num_burnin" = num_burnin, - "num_mcmc" = num_mcmc, - "keep_every" = keep_every, - "num_chains" = num_chains, - "has_rfx" = has_rfx, - "has_rfx_basis" = has_basis_rfx, - "num_rfx_basis" = num_basis_rfx, - "include_variance_forest" = include_variance_forest, - "sample_sigma2_global" = sample_sigma2_global, - "sample_sigma2_leaf_mu" = sample_sigma2_leaf_mu, - "sample_sigma2_leaf_tau" = sample_sigma2_leaf_tau, - "probit_outcome_model" = probit_outcome_model, - "outcome_model" = outcome_model, - "rfx_model_spec" = rfx_model_spec - ) - result <- list( - "forests_mu" = forest_samples_mu, - "forests_tau" = forest_samples_tau, - "model_params" = model_params, - "mu_hat_train" = mu_hat_train, - "tau_hat_train" = tau_hat_train, - "y_hat_train" = y_hat_train, - "train_set_metadata" = X_train_metadata - ) - if (has_test) { - result[["mu_hat_test"]] = mu_hat_test - } - if (has_test) { - result[["tau_hat_test"]] = tau_hat_test - } - if (has_test) { - result[["y_hat_test"]] = y_hat_test - } - if (include_variance_forest) { - result[["forests_variance"]] = forest_samples_variance - result[["sigma2_x_hat_train"]] = sigma2_x_hat_train - if (has_test) result[["sigma2_x_hat_test"]] = sigma2_x_hat_test - } - if (sample_sigma2_global) { - result[["sigma2_global_samples"]] = sigma2_global_samples - } - if (sample_sigma2_leaf_mu) { - result[["sigma2_leaf_mu_samples"]] = sigma2_leaf_mu_samples - } - if (sample_sigma2_leaf_tau) { - result[["sigma2_leaf_tau_samples"]] = sigma2_leaf_tau_samples - } - if (adaptive_coding) { - result[["b_0_samples"]] = b_0_samples - result[["b_1_samples"]] = b_1_samples - } - if (sample_tau_0) { - result[["tau_0_samples"]] = tau_0_samples * y_std_train - } - if (has_rfx) { - result[["rfx_samples"]] = rfx_samples - result[["rfx_preds_train"]] = rfx_preds_train - result[["rfx_unique_group_ids"]] = levels(group_ids_factor) - } - if ((has_rfx_test) && (has_test)) { - result[["rfx_preds_test"]] = rfx_preds_test - } - if (internal_propensity_model) { - result[["bart_propensity_model"]] = bart_model_propensity - } - class(result) <- "bcfmodel" - - # Restore global RNG state if user provided a random seed - if (custom_rng) { - if (has_existing_random_seed) { - .Random.seed <- original_global_seed - } else { - rm(".Random.seed", envir = .GlobalEnv) - } - } - - return(result) -} - -#' @title Predict from BCF Model -#' @description -#' Predict from a sampled BCF model on new data -#' -#' @param object Object of type `bcfmodel` containing draws of a Bayesian causal forest model and associated sampling outputs. -#' @param X Covariates used to determine tree leaf predictions for each observation. Must be passed as a matrix or dataframe. -#' @param Z Treatments used for prediction. -#' @param propensity (Optional) Propensities used for prediction. -#' @param rfx_group_ids (Optional) Test set group labels used for an additive random effects model. -#' We do not currently support (but plan to in the near future), test set evaluation for group labels -#' that were not in the training set. -#' @param rfx_basis (Optional) Test set basis for "random-slope" regression in additive random effects model. If the model was sampled with a random effects `model_spec` of "intercept_only" or "intercept_plus_treatment", this is optional, but if it is provided, it will be used. -#' @param type (Optional) Type of prediction to return. Options are "mean", which averages the predictions from every draw of a BCF model, and "posterior", which returns the entire matrix of posterior predictions. Default: "posterior". -#' @param terms (Optional) Which model terms to include in the prediction. Options include `"y_hat"`, `"prognostic_function"`, `"mu"`, `"cate"`, `"tau"`, `"rfx"`, `"variance_forest"`, or `"all"`. -#' -#' The treatment effect terms follow a three-level hierarchy: -#' \itemize{ -#' \item `"tau"` returns `tau_0 + tau(X)`: the parametric treatment intercept (if sampled) plus the treatment forest. This matches `model$tau_hat_train` / `model$tau_hat_test`. -#' \item `"cate"` additionally folds in the random slope on treatment when random effects are fit with `rfx_model_spec = "intercept_plus_treatment"`; otherwise it is identical to `"tau"`. -#' \item The raw forest-only component (without `tau_0`) is not directly returned by this method; use `model$forests_tau` to access it. -#' } -#' -#' Similarly for the prognostic term: `"mu"` returns the prognostic forest only, while `"prognostic_function"` additionally folds in the random intercept when `rfx_model_spec` is `"intercept_only"` or `"intercept_plus_treatment"`; otherwise the two are identical. -#' -#' If a model doesn't have random effects or variance forest predictions but one of those terms is requested, the request will simply be ignored. If none of the requested terms are present, this function will return `NULL` along with a warning. Default: `"all"`. -#' @param scale (Optional) Scale of mean function predictions. Options are "linear", which returns predictions on the original scale of the mean forest / RFX terms, and "probability", which transforms predictions into a probability of observing `y == 1`. "probability" is only valid for models fit with a probit outcome model. Default: "linear". -#' @param ... (Optional) Other prediction parameters. -#' -#' @return List of prediction matrices or single prediction matrix / vector, depending on the terms requested. -#' @export -#' -#' @examples -#' n <- 500 -#' p <- 5 -#' X <- matrix(runif(n*p), ncol = p) -#' mu_x <- ( -#' ((0 <= X[,1]) & (0.25 > X[,1])) * (-7.5) + -#' ((0.25 <= X[,1]) & (0.5 > X[,1])) * (-2.5) + -#' ((0.5 <= X[,1]) & (0.75 > X[,1])) * (2.5) + -#' ((0.75 <= X[,1]) & (1 > X[,1])) * (7.5) -#' ) -#' pi_x <- ( -#' ((0 <= X[,1]) & (0.25 > X[,1])) * (0.2) + -#' ((0.25 <= X[,1]) & (0.5 > X[,1])) * (0.4) + -#' ((0.5 <= X[,1]) & (0.75 > X[,1])) * (0.6) + -#' ((0.75 <= X[,1]) & (1 > X[,1])) * (0.8) -#' ) -#' tau_x <- ( -#' ((0 <= X[,2]) & (0.25 > X[,2])) * (0.5) + -#' ((0.25 <= X[,2]) & (0.5 > X[,2])) * (1.0) + -#' ((0.5 <= X[,2]) & (0.75 > X[,2])) * (1.5) + -#' ((0.75 <= X[,2]) & (1 > X[,2])) * (2.0) -#' ) -#' Z <- rbinom(n, 1, pi_x) -#' noise_sd <- 1 -#' y <- mu_x + tau_x*Z + rnorm(n, 0, noise_sd) -#' test_set_pct <- 0.2 -#' n_test <- round(test_set_pct*n) -#' n_train <- n - n_test -#' test_inds <- sort(sample(1:n, n_test, replace = FALSE)) -#' train_inds <- (1:n)[!((1:n) %in% test_inds)] -#' X_test <- X[test_inds,] -#' X_train <- X[train_inds,] -#' pi_test <- pi_x[test_inds] -#' pi_train <- pi_x[train_inds] -#' Z_test <- Z[test_inds] -#' Z_train <- Z[train_inds] -#' y_test <- y[test_inds] -#' y_train <- y[train_inds] -#' mu_test <- mu_x[test_inds] -#' mu_train <- mu_x[train_inds] -#' tau_test <- tau_x[test_inds] -#' tau_train <- tau_x[train_inds] -#' bcf_model <- bcf(X_train = X_train, Z_train = Z_train, y_train = y_train, -#' propensity_train = pi_train, num_gfr = 10, -#' num_burnin = 0, num_mcmc = 10) -#' preds <- predict(bcf_model, X_test, Z_test, pi_test) -predict.bcfmodel <- function( - object, - X, - Z, - propensity = NULL, - rfx_group_ids = NULL, - rfx_basis = NULL, - type = "posterior", - terms = "all", - scale = "linear", - ... -) { - # Handle mean function scale - if (!is.character(scale)) { - stop("scale must be a string or character vector") - } - if (!(scale %in% c("linear", "probability"))) { - stop("scale must either be 'linear' or 'probability'") - } - is_probit <- object$model_params$outcome_model$link == "probit" - if ((scale == "probability") && (!is_probit)) { - stop( - "scale cannot be 'probability' for models not fit with a probit outcome model" - ) - } - probability_scale <- scale == "probability" - - # Handle prediction type - if (!is.character(type)) { - stop("type must be a string or character vector") - } - if (!(type %in% c("mean", "posterior"))) { - stop("type must either be 'mean' or 'posterior") - } - predict_mean <- type == "mean" - - # Warn users about CATE / prognostic function when rfx_model_spec is "custom" - if (object$model_params$has_rfx) { - if (object$model_params$rfx_model_spec == "custom") { - if (("prognostic_function" %in% terms) || ("cate" %in% terms)) { - warning(paste0( - "This BCF model was fit with a custom random effects model specification (i.e. a user-provided basis). ", - "As a result, 'prognostic_function' and 'cate' refer only to the prognostic ('mu') ", - "and treatment effect 'tau' forests, respectively, and do not include any random ", - "effects contributions. If your user-provided random effects basis includes a random intercept or a ", - "random slope on the treatment variable, you will need to compute the prognostic or CATE functions manually by predicting ", - "'yhat' for different covariate and rfx_basis values." - )) - } - } - } - - # Handle prediction terms - rfx_model_spec = object$model_params$rfx_model_spec - rfx_intercept_only <- rfx_model_spec == "intercept_only" - rfx_intercept_plus_treatment <- rfx_model_spec == "intercept_plus_treatment" - rfx_intercept <- rfx_intercept_only || rfx_intercept_plus_treatment - mu_prog_separate <- ifelse(rfx_intercept, TRUE, FALSE) - tau_cate_separate <- ifelse(rfx_intercept_plus_treatment, TRUE, FALSE) - if (!is.character(terms)) { - stop("type must be a string or character vector") - } - for (term in terms) { - if ( - !(term %in% - c( - "y_hat", - "prognostic_function", - "mu", - "cate", - "tau", - "rfx", - "variance_forest", - "all" - )) - ) { - warning(paste0( - "Term '", - term, - "' was requested. Valid terms are 'y_hat', 'prognostic_function', 'mu', 'cate', 'tau', 'rfx', 'variance_forest', and 'all'.", - " This term will be ignored and prediction will only proceed if other requested terms are available in the model." - )) - } - } - num_terms <- length(terms) - has_mu_forest <- T - has_tau_forest <- T - has_variance_forest <- object$model_params$include_variance_forest - has_rfx <- object$model_params$has_rfx - has_y_hat <- T - predict_y_hat <- (((has_y_hat) && ("y_hat" %in% terms)) || - ((has_y_hat) && ("all" %in% terms))) - predict_mu_forest <- (((has_mu_forest) && ("all" %in% terms)) || - ((has_mu_forest) && ("mu" %in% terms))) - predict_tau_forest <- (((has_tau_forest) && ("tau" %in% terms)) || - ((has_tau_forest) && ("all" %in% terms))) - predict_prog_function <- (((has_mu_forest) && - ("prognostic_function" %in% terms)) || - ((has_mu_forest) && ("all" %in% terms))) - predict_cate_function <- (((has_tau_forest) && ("cate" %in% terms)) || - ((has_tau_forest) && ("all" %in% terms))) - predict_rfx <- (((has_rfx) && ("rfx" %in% terms)) || - ((has_rfx) && ("all" %in% terms))) - predict_variance_forest <- (((has_variance_forest) && - ("variance_forest" %in% terms)) || - ((has_variance_forest) && ("all" %in% terms))) - predict_count <- sum(c( - predict_y_hat, - predict_mu_forest, - predict_prog_function, - predict_tau_forest, - predict_cate_function, - predict_rfx, - predict_variance_forest - )) - if (predict_count == 0) { - warning(paste0( - "None of the requested model terms, ", - paste(terms, collapse = ", "), - ", were fit in this model" - )) - return(NULL) - } - predict_rfx_intermediate <- (predict_y_hat && has_rfx) - predict_rfx_raw <- ((predict_prog_function && has_rfx && rfx_intercept) || - (predict_cate_function && has_rfx && rfx_intercept_plus_treatment)) - predict_mu_forest_intermediate <- ((predict_y_hat || predict_prog_function) && - has_mu_forest) - predict_tau_forest_intermediate <- ((predict_y_hat || - predict_cate_function || - (object$model_params$adaptive_coding && - (predict_mu_forest || predict_prog_function))) && - has_tau_forest) - - # Make sure covariates are matrix or data frame - if ((!is.data.frame(X)) && (!is.matrix(X))) { - stop("X must be a matrix or dataframe") - } - - # Handle factor-valued treatment before numeric conversion - if (is.factor(Z)) { - lvls <- levels(Z) - if (length(lvls) != 2) { - stop("Factor Z must have exactly 2 levels for binary treatment") - } - warning( - "Z is a factor; recoding to 0/1 using level order: ", - lvls[1], - " = 0, ", - lvls[2], - " = 1" - ) - Z <- as.integer(Z) - 1L - } - - # Convert all input data to matrices if not already converted - if ((is.null(dim(Z))) && (!is.null(Z))) { - Z <- as.matrix(as.numeric(Z)) - } - if ((is.null(dim(propensity))) && (!is.null(propensity))) { - propensity <- as.matrix(propensity) - } - if ((is.null(dim(rfx_basis))) && (!is.null(rfx_basis))) { - rfx_basis <- as.matrix(rfx_basis) - } - - # Data checks - if (nrow(X) != nrow(Z)) { - stop("X and Z must have the same number of rows") - } - if (object$model_params$num_covariates != ncol(X)) { - stop( - "X and must have the same number of columns as the covariates used to train the model" - ) - } - if ((object$model_params$has_rfx) && (is.null(rfx_group_ids))) { - stop( - "Random effect group labels (rfx_group_ids) must be provided for this model" - ) - } - if ((object$model_params$has_rfx_basis) && (is.null(rfx_basis))) { - if (object$model_params$rfx_model_spec == "custom") { - stop("Random effects basis (rfx_basis) must be provided for this model") - } - } - if ((object$model_params$num_rfx_basis > 0) && (!is.null(rfx_basis))) { - if (ncol(rfx_basis) != object$model_params$num_rfx_basis) { - stop( - "Random effects basis has a different dimension than the basis used to train this model" - ) - } - } - - # Preprocess covariates before any prediction calls that depend on X being - # a numeric matrix (e.g. the internal propensity BART model was trained on a - # preprocessed matrix, so it expects a matrix, not the raw data frame) - train_set_metadata <- object$train_set_metadata - X <- preprocessPredictionData(X, train_set_metadata) - - # Compute propensity score using the internal bart model - if ( - (object$model_params$propensity_covariate != "none") && - (is.null(propensity)) - ) { - if (!object$model_params$internal_propensity_model) { - stop("propensity must be provided for this model") - } - propensity <- rowMeans(predict(object$bart_propensity_model, X)$y_hat) - } - - # Recode group IDs to integer vector (if passed as, for example, a vector of county names, etc...) - if (!is.null(rfx_group_ids)) { - rfx_unique_group_ids <- object$rfx_unique_group_ids - group_ids_factor <- factor(rfx_group_ids, levels = rfx_unique_group_ids) - if (sum(is.na(group_ids_factor)) > 0) { - stop( - "All random effect group labels provided in rfx_group_ids must have been present at sampling time" - ) - } - rfx_group_ids <- as.integer(group_ids_factor) - } - - # Handle RFX model specification - if (has_rfx) { - if (object$model_params$rfx_model_spec == "custom") { - if (is.null(rfx_basis)) { - stop( - "A user-provided basis (`rfx_basis`) must be provided when the model was sampled with a random effects model spec set to 'custom'" - ) - } - } else if (object$model_params$rfx_model_spec == "intercept_only") { - # Only construct a basis if user-provided basis missing - if (is.null(rfx_basis)) { - rfx_basis <- matrix( - rep(1, nrow(X)), - nrow = nrow(X), - ncol = 1 - ) - } - } else if ( - object$model_params$rfx_model_spec == "intercept_plus_treatment" - ) { - # Only construct a basis if user-provided basis missing - if (is.null(rfx_basis)) { - rfx_basis <- cbind( - rep(1, nrow(X)), - Z - ) - } - } - } - - # Add propensities to covariate set if necessary - X_combined <- X - if (object$model_params$propensity_covariate != "none") { - X_combined <- cbind(X, propensity) - } - - # Create prediction datasets - forest_dataset_pred <- createForestDataset(X_combined, Z) - - # Compute variance forest predictions - if (predict_variance_forest) { - s_x_raw <- object$forests_variance$predict(forest_dataset_pred) - } - - # Scale variance forest predictions - num_samples <- object$model_params$num_samples - y_std <- object$model_params$outcome_scale - y_bar <- object$model_params$outcome_mean - initial_sigma2 <- object$model_params$initial_sigma2 - if (predict_variance_forest) { - if (object$model_params$sample_sigma2_global) { - sigma2_global_samples <- object$sigma2_global_samples - variance_forest_predictions <- sapply(1:num_samples, function(i) { - s_x_raw[, i] * sigma2_global_samples[i] - }) - } else { - variance_forest_predictions <- s_x_raw * - initial_sigma2 * - y_std * - y_std - } - if (predict_mean) { - variance_forest_predictions <- rowMeans(variance_forest_predictions) - } - } - - # Compute mu forest predictions - if (predict_mu_forest || predict_mu_forest_intermediate) { - mu_hat_forest <- object$forests_mu$predict(forest_dataset_pred) * - y_std + - y_bar - } - - # Compute CATE forest predictions - if (predict_tau_forest || predict_tau_forest_intermediate) { - if (object$model_params$adaptive_coding) { - tau_hat_raw <- object$forests_tau$predict_raw(forest_dataset_pred) - tau_hat_forest <- t( - t(tau_hat_raw) * (object$b_1_samples - object$b_0_samples) - ) * - y_std - if (predict_mu_forest || predict_mu_forest_intermediate) { - control_adj <- t(t(tau_hat_raw) * object$b_0_samples) * y_std - mu_hat_forest <- mu_hat_forest + control_adj - } - } else { - tau_hat_forest <- object$forests_tau$predict_raw(forest_dataset_pred) * - y_std - } - # tau_hat_forest is the forest-only component tau(X); compute cate_hat_forest - # (tau_0 + tau(X)) for the "cate" term and treatment_term used in y_hat - if (object$model_params$sample_tau_0 && !is.null(object$tau_0_samples)) { - tau_0_samp <- object$tau_0_samples # p_tau0 x num_samples (already in original scale) - if (object$model_params$adaptive_coding) { - cate_hat_forest <- sweep( - tau_hat_forest, - 2, - (object$b_1_samples - object$b_0_samples) * as.numeric(tau_0_samp), - "+" - ) - if (predict_mu_forest || predict_mu_forest_intermediate) { - mu_hat_forest <- sweep( - mu_hat_forest, - 2, - object$b_0_samples * as.numeric(tau_0_samp), - "+" - ) - } - } else if (!object$model_params$multivariate_treatment) { - cate_hat_forest <- sweep(tau_hat_forest, 2, as.numeric(tau_0_samp), "+") - } else { - p_tau0 <- nrow(tau_0_samp) - cate_hat_forest <- tau_hat_forest - for (j in seq_len(p_tau0)) { - cate_hat_forest[, j, ] <- cate_hat_forest[, j, ] + - outer(rep(1, nrow(X)), tau_0_samp[j, ]) - } - } - } else { - cate_hat_forest <- tau_hat_forest - } - if (object$model_params$multivariate_treatment) { - tau_dim <- dim(cate_hat_forest) - tau_num_obs <- tau_dim[1] - tau_num_samples <- tau_dim[3] - treatment_term <- matrix(NA_real_, nrow = tau_num_obs, tau_num_samples) - for (i in 1:nrow(Z)) { - treatment_term[i, ] <- colSums(cate_hat_forest[i, , ] * Z[i, ]) - } - } else { - treatment_term <- cate_hat_forest * as.numeric(Z) - } - } - - # Compute rfx predictions - if (predict_rfx || predict_rfx_intermediate) { - rfx_predictions <- object$rfx_samples$predict( - rfx_group_ids, - rfx_basis - ) * - y_std - } - - # Extract "raw" rfx coefficients for each rfx basis term if needed - if (predict_rfx_raw) { - # Extract the raw RFX samples and scale by train set outcome standard deviation - rfx_param_list <- object$rfx_samples$extract_parameter_samples() - rfx_beta_draws <- rfx_param_list$beta_samples * - object$model_params$outcome_scale - - # Promote to an array with consistent dimensions when there's one rfx term - if (length(dim(rfx_beta_draws)) == 2) { - dim(rfx_beta_draws) <- c(1, dim(rfx_beta_draws)) - } - - # Construct a matrix with the appropriate group random effects arranged for each observation - rfx_predictions_raw <- array( - NA, - dim = c( - nrow(X), - ncol(rfx_basis), - object$model_params$num_samples - ) - ) - for (i in 1:nrow(X)) { - rfx_predictions_raw[i, , ] <- - rfx_beta_draws[, rfx_group_ids[i], ] - } - } - - # Add raw RFX predictions to mu and tau if warranted by the RFX model spec - if (predict_prog_function) { - if (mu_prog_separate) { - prognostic_function <- mu_hat_forest + rfx_predictions_raw[, 1, ] - } else { - prognostic_function <- mu_hat_forest - } - } - if (predict_cate_function) { - if (tau_cate_separate) { - cate <- (cate_hat_forest + - rfx_predictions_raw[, 2:ncol(rfx_basis), ]) - } else { - cate <- cate_hat_forest - } - } - - # Combine into y hat predictions - needs_mean_term_preds <- predict_y_hat || - predict_mu_forest || - predict_tau_forest || - predict_prog_function || - predict_cate_function || - predict_rfx - if (needs_mean_term_preds) { - if (probability_scale) { - if (has_rfx) { - if (predict_y_hat) { - y_hat <- pnorm(mu_hat_forest + treatment_term + rfx_predictions) - } - if (predict_rfx) { - rfx_predictions <- pnorm(rfx_predictions) - } - } else { - if (predict_y_hat) { - y_hat <- pnorm(mu_hat_forest + treatment_term) - } - } - if (predict_mu_forest) { - mu_hat <- pnorm(mu_hat_forest) - } - if (predict_tau_forest) { - tau_hat <- pnorm(cate_hat_forest) - } - if (predict_prog_function) { - prognostic_function <- pnorm(prognostic_function) - } - if (predict_cate_function) { - cate <- pnorm(cate) - } - } else { - if (has_rfx) { - if (predict_y_hat) { - y_hat <- mu_hat_forest + treatment_term + rfx_predictions - } - } else { - if (predict_y_hat) { - y_hat <- mu_hat_forest + treatment_term - } - } - if (predict_mu_forest) { - mu_hat <- mu_hat_forest - } - if (predict_tau_forest) { - tau_hat <- cate_hat_forest - } - if (predict_prog_function) { - prognostic_function <- prognostic_function - } - if (predict_cate_function) { - cate <- cate - } - } - } - - # Collapse to posterior mean predictions if requested - if (predict_mean) { - if (predict_mu_forest) { - mu_hat <- rowMeans(mu_hat) - } - if (predict_tau_forest) { - if (object$model_params$multivariate_treatment) { - tau_hat <- apply(tau_hat, c(1, 2), mean) - } else { - tau_hat <- rowMeans(tau_hat) - } - } - if (predict_prog_function) { - prognostic_function <- rowMeans(prognostic_function) - } - if (predict_cate_function) { - if (object$model_params$multivariate_treatment) { - cate <- apply(cate, c(1, 2), mean) - } else { - cate <- rowMeans(cate) - } - } - if (predict_rfx) { - rfx_predictions <- rowMeans(rfx_predictions) - } - if (predict_y_hat) { - y_hat <- rowMeans(y_hat) - } - } - - # Return results - if (predict_count == 1) { - if (predict_y_hat) { - return(y_hat) - } else if (predict_mu_forest) { - return(mu_hat) - } else if (predict_tau_forest) { - return(tau_hat) - } else if (predict_prog_function) { - return(prognostic_function) - } else if (predict_cate_function) { - return(cate) - } else if (predict_rfx) { - return(rfx_predictions) - } else if (predict_variance_forest) { - return(variance_forest_predictions) - } - } else { - result <- list() - if (predict_y_hat) { - result[["y_hat"]] = y_hat - } else { - result[["y_hat"]] <- NULL - } - if (predict_mu_forest) { - result[["mu_hat"]] = mu_hat - } else { - result[["mu_hat"]] <- NULL - } - if (predict_tau_forest) { - result[["tau_hat"]] = tau_hat - } else { - result[["tau_hat"]] <- NULL - } - if (predict_prog_function) { - result[["prognostic_function"]] = prognostic_function - } else { - result[["prognostic_function"]] <- NULL - } - if (predict_cate_function) { - result[["cate"]] = cate - } else { - result[["cate"]] <- NULL - } - if (predict_rfx) { - result[["rfx_predictions"]] = rfx_predictions - } else { - result[["rfx_predictions"]] <- NULL - } - if (predict_variance_forest) { - result[["variance_forest_predictions"]] = variance_forest_predictions - } else { - result[["variance_forest_predictions"]] <- NULL - } - } - return(result) -} - -#' @title Print Summary of BCF Model -#' @description Prints a summary of the BCF model, including the model terms and their specifications. -#' @param x The BCF model object -#' @param ... Additional arguments (currently unused) -#' @export -#' @return BCF model object unchanged after printing summary -print.bcfmodel <- function(x, ...) { - # What type of model was run - model_terms <- c("prognostic forest", "treatment effect forest") - if (x$model_params$include_variance_forest) { - model_terms <- c(model_terms, "variance forest") - } - if (x$model_params$has_rfx) { - model_terms <- c(model_terms, "additive random effects") - } - if (x$model_params$sample_sigma2_global) { - model_terms <- c(model_terms, "global error variance model") - } - if (x$model_params$sample_sigma2_leaf_mu) { - model_terms <- c(model_terms, "prognostic forest leaf scale model") - } - if (x$model_params$sample_sigma2_leaf_tau) { - model_terms <- c(model_terms, "treatment effect forest leaf scale model") - } - if (x$model_params$sample_tau_0) { - model_terms <- c(model_terms, "treatment effect intercept model") - } - if (length(model_terms) > 2) { - summary_message <- paste0( - "stochtree::bcf() run with ", - paste0( - paste0(model_terms[1:(length(model_terms) - 1)], collapse = ", "), - ", and ", - model_terms[length(model_terms)] - ) - ) - } else if (length(model_terms) == 2) { - summary_message <- paste0( - "stochtree::bcf() run with ", - paste0(model_terms, collapse = " and ") - ) - } else { - summary_message <- paste0("stochtree::bcf() run with ", model_terms) - } - - # Outcome details - is_probit <- x$model_params$outcome_model$link == "probit" - summary_message <- paste0( - summary_message, - "\n", - "Outcome was modeled ", - ifelse( - is_probit, - "with a probit link", - "as gaussian" - ) - ) - - # Treatment details - if (x$model_params$binary_treatment) { - summary_message <- paste0( - summary_message, - "\n", - "Treatment was binary and ", - ifelse( - x$model_params$adaptive_coding, - "its effect was estimated with adaptive coding", - "its effect was estimated with default coding" - ) - ) - } else if (x$model_params$multivariate_treatment) { - summary_message <- paste0( - summary_message, - "\n", - "Treatment was multivariate with ", - x$model_params$treatment_dim, - " dimensions" - ) - } else { - summary_message <- paste0( - summary_message, - "\n", - "Treatment was univariate but not binary" - ) - } - - # Standardization - if (x$model_params$standardize) { - summary_message <- paste0( - summary_message, - "\n", - "outcome was standardized" - ) - } - - # Internal propensity model - if (x$model_params$propensity_covariate == "none") { - summary_message <- paste0( - summary_message, - "\n", - "Propensity scores were not used in either forest of the model" - ) - } else { - if (x$model_params$internal_propensity_model) { - summary_message <- paste0( - summary_message, - "\n", - "An internal propensity model was fit using stochtree::bart() in lieu of user-provided propensity scores" - ) - } else { - summary_message <- paste0( - summary_message, - "\n", - "User-provided propensity scores were included in the model" - ) - } - } - - # Random effects details - if (x$model_params$has_rfx) { - if (x$model_params$rfx_model_spec == "custom") { - summary_message <- paste0( - summary_message, - "\n", - "Random effects were fit with a user-supplied basis" - ) - } else if (x$model_params$rfx_model_spec == "intercept_only") { - summary_message <- paste0( - summary_message, - "\n", - "Random effects were fit with an 'intercept-only' parameterization" - ) - } else if (x$model_params$rfx_model_spec == "intercept_plus_treatment") { - summary_message <- paste0( - summary_message, - "\n", - "Random effects were fit with an 'intercept-plus-treatment' parameterization" - ) - } - } - - # Sampler details - summary_message <- paste0( - summary_message, - "\n", - "The sampler was run for ", - x$model_params$num_gfr, - " GFR iterations, with ", - x$model_params$num_chains, - ifelse(x$model_params$num_chains == 1, " chain of ", " chains of "), - x$model_params$num_burnin, - " burn-in iterations and ", - x$model_params$num_mcmc, - " MCMC iterations, ", - ifelse( - x$model_params$keep_every == 1, - "retaining every iteration (i.e. no thinning)", - paste0( - "retaining every ", - x$model_params$keep_every, - "th iteration (i.e. thinning)" - ) - ) - ) - - # Print the model details - cat(summary_message, "\n") - - # Return bcf model invisibly - invisible(x) -} - -#' @title Summarize BCF Model -#' @description Summarize a BCF fit with a description of the model that was fit and numeric summaries of any sampled quantities. -#' @param object The BCF model object -#' @param ... Additional arguments -#' @export -#' @return BCF model object unchanged after summarizing -summary.bcfmodel <- function(object, ...) { - # First, print the BCF model - tmp <- print(object) - - # Summarize any sampled quantities - - # Global error scale - if (object$model_params$sample_sigma2_global) { - sigma2_samples <- object$sigma2_global_samples - n_samples <- length(sigma2_samples) - mean_sigma2 <- mean(sigma2_samples) - sd_sigma2 <- sd(sigma2_samples) - quantiles_sigma2 <- quantile( - sigma2_samples, - probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) - ) - cat(sprintf( - "Summary of sigma^2 posterior: \n%d samples, mean = %.3f, standard deviation = %.3f, quantiles:\n", - n_samples, - mean_sigma2, - sd_sigma2 - )) - print(quantiles_sigma2) - } - - # Leaf scale for the prognostic forest - if (object$model_params$sample_sigma2_leaf_mu) { - sigma2_leaf_samples <- object$sigma2_leaf_mu_samples - n_samples <- length(sigma2_leaf_samples) - mean_sigma2 <- mean(sigma2_leaf_samples) - sd_sigma2 <- sd(sigma2_leaf_samples) - quantiles_sigma2 <- quantile( - sigma2_leaf_samples, - probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) - ) - cat(sprintf( - "Summary of prognostic forest leaf scale posterior: \n%d samples, mean = %.3f, standard deviation = %.3f, quantiles:\n", - n_samples, - mean_sigma2, - sd_sigma2 - )) - print(quantiles_sigma2) - } - - # Leaf scale for the treatment effect forest - if (object$model_params$sample_sigma2_leaf_tau) { - sigma2_leaf_samples <- object$sigma2_leaf_tau_samples - n_samples <- length(sigma2_leaf_samples) - mean_sigma2 <- mean(sigma2_leaf_samples) - sd_sigma2 <- sd(sigma2_leaf_samples) - quantiles_sigma2 <- quantile( - sigma2_leaf_samples, - probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) - ) - cat(sprintf( - "Summary of treatment effect forest leaf scale posterior: \n%d samples, mean = %.3f, standard deviation = %.3f, quantiles:\n", - n_samples, - mean_sigma2, - sd_sigma2 - )) - print(quantiles_sigma2) - } - - # Adaptive coding parameters - if (object$model_params$adaptive_coding) { - b0_samples <- object$b_0_samples - b1_samples <- object$b_1_samples - n_samples <- length(b0_samples) - mean_b0 <- mean(b0_samples) - mean_b1 <- mean(b1_samples) - sd_b0 <- sd(b0_samples) - sd_b1 <- sd(b1_samples) - quantiles_b0 <- quantile( - b0_samples, - probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) - ) - quantiles_b1 <- quantile( - b1_samples, - probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) - ) - cat(sprintf( - "Summary of adaptive coding parameters: \n%d samples, mean (control) = %.3f, mean (treated) = %.3f, standard deviation (control) = %.3f, standard deviation (treated) = %.3f\n", - n_samples, - mean_b0, - mean_b1, - sd_b0, - sd_b1 - )) - cat("quantiles (control):\n") - print(quantiles_b0) - cat("quantiles (treated):\n") - print(quantiles_b1) - } - - # Treatment effect intercept (tau_0) - if (object$model_params$sample_tau_0 && !is.null(object$tau_0_samples)) { - tau_0_vec <- as.numeric(object$tau_0_samples) - n_samples <- ncol(object$tau_0_samples) - mean_tau_0 <- mean(tau_0_vec) - sd_tau_0 <- sd(tau_0_vec) - quantiles_tau_0 <- quantile( - tau_0_vec, - probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) - ) - cat(sprintf( - "Summary of treatment effect intercept (tau_0) posterior: \n%d samples, mean = %.3f, standard deviation = %.3f, quantiles:\n", - n_samples, - mean_tau_0, - sd_tau_0 - )) - print(quantiles_tau_0) - } - - # In-sample predictions - if (!is.null(object$y_hat_train)) { - y_hat_train_mean <- rowMeans(object$y_hat_train) - n_y_hat_train <- length(y_hat_train_mean) - mean_y_hat_train <- mean(y_hat_train_mean) - sd_y_hat_train <- sd(y_hat_train_mean) - quantiles_y_hat_train <- quantile( - y_hat_train_mean, - probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) - ) - cat(sprintf( - "Summary of in-sample posterior mean predictions: \n%d observations, mean = %.3f, standard deviation = %.3f, quantiles:\n", - n_y_hat_train, - mean_y_hat_train, - sd_y_hat_train - )) - print(quantiles_y_hat_train) - } - - # Test-set predictions - if (!is.null(object$y_hat_test)) { - y_hat_test_mean <- rowMeans(object$y_hat_test) - n_y_hat_test <- length(y_hat_test_mean) - mean_y_hat_test <- mean(y_hat_test_mean) - sd_y_hat_test <- sd(y_hat_test_mean) - quantiles_y_hat_test <- quantile( - y_hat_test_mean, - probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) - ) - cat(sprintf( - "Summary of test-set posterior mean predictions: \n%d observations, mean = %.3f, standard deviation = %.3f, quantiles:\n", - n_y_hat_test, - mean_y_hat_test, - sd_y_hat_test - )) - print(quantiles_y_hat_test) - } - - # In-sample treatment effect function estimates - if (!is.null(object$tau_hat_train)) { - if (!object$model_params$multivariate_treatment) { - tau_hat_train_mean <- rowMeans(object$tau_hat_train) - n_tau_hat_train <- length(tau_hat_train_mean) - mean_tau_hat_train <- mean(tau_hat_train_mean) - sd_tau_hat_train <- sd(tau_hat_train_mean) - quantiles_tau_hat_train <- quantile( - tau_hat_train_mean, - probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) - ) - cat(sprintf( - "Summary of in-sample posterior mean CATEs: \n%d observations, mean = %.3f, standard deviation = %.3f, quantiles:\n", - n_tau_hat_train, - mean_tau_hat_train, - sd_tau_hat_train - )) - print(quantiles_tau_hat_train) - } - } - - # Test set treatment effect function estimates - if (!is.null(object$tau_hat_test)) { - if (!object$model_params$multivariate_treatment) { - tau_hat_test_mean <- rowMeans(object$tau_hat_test) - n_tau_hat_test <- length(tau_hat_test_mean) - mean_tau_hat_test <- mean(tau_hat_test_mean) - sd_tau_hat_test <- sd(tau_hat_test_mean) - quantiles_tau_hat_test <- quantile( - tau_hat_test_mean, - probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) - ) - cat(sprintf( - "Summary of test-set posterior mean CATEs: \n%d observations, mean = %.3f, standard deviation = %.3f, quantiles:\n", - n_tau_hat_test, - mean_tau_hat_test, - sd_tau_hat_test - )) - print(quantiles_tau_hat_test) - } - } - - # Random effects - if (object$model_params$has_rfx) { - rfx_samples <- getRandomEffectSamples(object) - rfx_beta_samples <- rfx_samples$beta_samples - if (length(dim(rfx_beta_samples)) > 2) { - cat( - "Random effects summary of variance components across groups and posterior draws:\n" - ) - rfx_component_means <- apply(rfx_beta_samples, 1, mean) - rfx_component_sds <- apply(rfx_beta_samples, 1, sd) - cat("Variance component means: ", rfx_component_means, "\n") - cat("Variance component standard deviations: ", rfx_component_sds, "\n") - quantile_summary <- t(apply( - rfx_beta_samples, - 1, - quantile, - probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) - )) - cat("Variance component quantiles:\n") - print(quantile_summary) - } else { - cat( - "Random effects summary of variance components across groups and posterior draws:\n" - ) - rfx_component_means <- mean(rfx_beta_samples) - rfx_component_sds <- sd(rfx_beta_samples) - cat("Random effects overall mean: ", rfx_component_means, "\n") - cat( - "Random effects overall standard deviation: ", - rfx_component_sds, - "\n" - ) - quantile_summary <- quantile( - rfx_beta_samples, - probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) - ) - cat("Random effects overall quantiles:\n") - print(quantile_summary) - } - } - - # Return bcf model invisibly - invisible(object) -} - -#' @title Plot BCF Model -#' @description Plot the BCF model fit and any relevant sampled quantities. This will default to a traceplot of the global error scale and the in-sample mean forest predictions for the first train set observation. Since `stochtree::bcf()` is flexible and it's possible to sample a model with a fixed global error scale and no mean forest, this procedure will throw an error if these two default terms are omitted. -#' @param x The BCF model object -#' @param ... Additional arguments -#' @export -#' @return BCF model object unchanged after summarizing -plot.bcfmodel <- function(x, ...) { - # Check if model has global error scale samples - has_sigma2_samples <- x$model_params$sample_sigma2_global - has_mean_forest_preds <- !is.null(x$y_hat_train) - - # First try combinations of sigma2 and mean forest predictions - if (has_sigma2_samples || has_mean_forest_preds) { - if (has_sigma2_samples) { - plot( - x$sigma2_global_samples, - type = "l", - ylab = "Sigma^2", - main = "Global error scale traceplot" - ) - } else if (has_mean_forest_preds) { - plot( - x$y_hat_train[1, ], - type = "l", - ylab = "Predictions", - main = "In-sample mean function trace for the first train set observation" - ) - } - } else { - stop( - "This model does not have enough model terms / parameter traces to produce stochtree's default plots. See `predict.bcfmodel()` for examples of how to further investigate your model." - ) - } - - # Return bcf model invisibly - invisible(x) -} - -#' @title Extract BCF Parameter Samples -#' @description Extract a vector, matrix or array of parameter samples from a BCF model by name. -#' Random effects are handled by a separate `getRandomEffectSamples` function due to the complexity of the random effects parameters. -#' If the requested model term is not found, an error is thrown. -#' The following conventions are used for parameter names: -#' - Global error variance: `"sigma2"`, `"global_error_scale"`, `"sigma2_global"` -#' - Prognostic forest leaf scale: `"sigma2_leaf_mu"`, `"leaf_scale_mu"`, `"mu_leaf_scale"` -#' - Treatment effect forest leaf scale: `"sigma2_leaf_tau"`, `"leaf_scale_tau"`, `"tau_leaf_scale"` -#' - Adaptive coding parameters: `"adaptive_coding"` (returns both the control and treated parameters jointly, with control in the first row and treated in the second row) -#' - In-sample mean function predictions: `"y_hat_train"` -#' - Test set mean function predictions: `"y_hat_test"` -#' - In-sample treatment effect forest predictions: `"tau_hat_train"` -#' - Test set treatment effect forest predictions: `"tau_hat_test"` -#' - Treatment effect intercept: `"tau_0"`, `"treatment_intercept"`, `"tau_intercept"` -#' - In-sample variance forest predictions: `"sigma2_x_train"`, `"var_x_train"` -#' - Test set variance forest predictions: `"sigma2_x_test"`, `"var_x_test"` +#' @title Predict from BCF Model +#' @description +#' Predict from a sampled BCF model on new data #' -#' @param object Object of type `bcfmodel` containing draws of a BCF model and associated sampling outputs. -#' @param term Name of the parameter to extract (e.g., `"sigma2"`, `"y_hat_train"`, etc.) -#' @return Array of parameter samples. If the underlying parameter is a scalar, this will be a vector of length `num_samples`. -#' If the underlying parameter is vector-valued, this will be (`parameter_dimension` x `num_samples`) matrix, and if the underlying -#' parameter is multidimensional, this will be an array of dimension (`parameter_dimension_1` x `parameter_dimension_2` x ... x `num_samples`). +#' @param object Object of type `bcfmodel` containing draws of a Bayesian causal forest model and associated sampling outputs. +#' @param X Covariates used to determine tree leaf predictions for each observation. Must be passed as a matrix or dataframe. +#' @param Z Treatments used for prediction. +#' @param propensity (Optional) Propensities used for prediction. +#' @param rfx_group_ids (Optional) Test set group labels used for an additive random effects model. +#' We do not currently support (but plan to in the near future), test set evaluation for group labels +#' that were not in the training set. +#' @param rfx_basis (Optional) Test set basis for "random-slope" regression in additive random effects model. If the model was sampled with a random effects `model_spec` of "intercept_only" or "intercept_plus_treatment", this is optional, but if it is provided, it will be used. +#' @param type (Optional) Type of prediction to return. Options are "mean", which averages the predictions from every draw of a BCF model, and "posterior", which returns the entire matrix of posterior predictions. Default: "posterior". +#' @param terms (Optional) Which model terms to include in the prediction. Options include `"y_hat"`, `"prognostic_function"`, `"mu"`, `"cate"`, `"tau"`, `"rfx"`, `"variance_forest"`, or `"all"`. +#' +#' The treatment effect terms follow a three-level hierarchy: +#' \itemize{ +#' \item `"tau"` returns `tau_0 + tau(X)`: the parametric treatment intercept (if sampled) plus the treatment forest. This matches `model$tau_hat_train` / `model$tau_hat_test`. +#' \item `"cate"` additionally folds in the random slope on treatment when random effects are fit with `rfx_model_spec = "intercept_plus_treatment"`; otherwise it is identical to `"tau"`. +#' \item The raw forest-only component (without `tau_0`) is not directly returned by this method; extract the treatment forest with `model$samples$materialize_tau_forest()` to access it. +#' } +#' +#' Similarly for the prognostic term: `"mu"` returns the prognostic forest only, while `"prognostic_function"` additionally folds in the random intercept when `rfx_model_spec` is `"intercept_only"` or `"intercept_plus_treatment"`; otherwise the two are identical. +#' +#' If a model doesn't have random effects or variance forest predictions but one of those terms is requested, the request will simply be ignored. If none of the requested terms are present, this function will return `NULL` along with a warning. Default: `"all"`. +#' @param scale (Optional) Scale of mean function predictions. Options are "linear", which returns predictions on the original scale of the mean forest / RFX terms, and "probability", which transforms predictions into a probability of observing `y == 1`. "probability" is only valid for models fit with a probit outcome model. Default: "linear". +#' @param ... (Optional) Other prediction parameters. +#' +#' @return List of prediction matrices or single prediction matrix / vector, depending on the terms requested. #' @export #' #' @examples @@ -4503,969 +1929,1497 @@ plot.bcfmodel <- function(x, ...) { #' ((0.75 <= X[,2]) & (1 > X[,2])) * (2.0) #' ) #' Z <- rbinom(n, 1, pi_x) -#' E_XZ <- mu_x + Z*tau_x -#' snr <- 3 -#' rfx_group_ids <- rep(c(1,2), n %/% 2) -#' rfx_coefs <- matrix(c(-1, -1, 1, 1), nrow=2, byrow=TRUE) -#' rfx_basis <- cbind(1, runif(n, -1, 1)) -#' rfx_term <- rowSums(rfx_coefs[rfx_group_ids,] * rfx_basis) -#' y <- E_XZ + rfx_term + rnorm(n, 0, 1)*(sd(E_XZ)/snr) -#' bcf_model <- bcf(X_train = X, y_train = y, Z_train = Z, -#' rfx_group_ids_train = rfx_group_ids, -#' rfx_basis_train = rfx_basis, -#' num_gfr = 10, num_burnin = 0, num_mcmc = 10) -#' sigma2_samples <- extractParameter(bcf_model, "sigma2") -extractParameter.bcfmodel <- function(object, term) { - if (term %in% c("sigma2", "global_error_scale", "sigma2_global")) { - if (!is.null(object$sigma2_global_samples)) { - return(object$sigma2_global_samples) - } else { - stop("This model does not have global variance parameter samples") +#' noise_sd <- 1 +#' y <- mu_x + tau_x*Z + rnorm(n, 0, noise_sd) +#' test_set_pct <- 0.2 +#' n_test <- round(test_set_pct*n) +#' n_train <- n - n_test +#' test_inds <- sort(sample(1:n, n_test, replace = FALSE)) +#' train_inds <- (1:n)[!((1:n) %in% test_inds)] +#' X_test <- X[test_inds,] +#' X_train <- X[train_inds,] +#' pi_test <- pi_x[test_inds] +#' pi_train <- pi_x[train_inds] +#' Z_test <- Z[test_inds] +#' Z_train <- Z[train_inds] +#' y_test <- y[test_inds] +#' y_train <- y[train_inds] +#' mu_test <- mu_x[test_inds] +#' mu_train <- mu_x[train_inds] +#' tau_test <- tau_x[test_inds] +#' tau_train <- tau_x[train_inds] +#' bcf_model <- bcf(X_train = X_train, Z_train = Z_train, y_train = y_train, +#' propensity_train = pi_train, num_gfr = 10, +#' num_burnin = 0, num_mcmc = 10) +#' preds <- predict(bcf_model, X_test, Z_test, pi_test) +predict.bcfmodel <- function( + object, + X, + Z, + propensity = NULL, + rfx_group_ids = NULL, + rfx_basis = NULL, + type = "posterior", + terms = "all", + scale = "linear", + ... +) { + # Handle mean function scale + if (!is.character(scale)) { + stop("scale must be a string or character vector") + } + if (!(scale %in% c("linear", "probability"))) { + stop("scale must either be 'linear' or 'probability'") + } + is_probit <- object$model_params$outcome_model$link == "probit" + if ((scale == "probability") && (!is_probit)) { + stop( + "scale cannot be 'probability' for models not fit with a probit outcome model" + ) + } + probability_scale <- scale == "probability" + + # Handle prediction type + if (!is.character(type)) { + stop("type must be a string or character vector") + } + if (!(type %in% c("mean", "posterior"))) { + stop("type must either be 'mean' or 'posterior") + } + predict_mean <- type == "mean" + + # Warn users about CATE / prognostic function when rfx_model_spec is "custom" + if (object$model_params$has_rfx) { + if (object$model_params$rfx_model_spec == "custom") { + if (("prognostic_function" %in% terms) || ("cate" %in% terms)) { + warning(paste0( + "This BCF model was fit with a custom random effects model specification (i.e. a user-provided basis). ", + "As a result, 'prognostic_function' and 'cate' refer only to the prognostic ('mu') ", + "and treatment effect 'tau' forests, respectively, and do not include any random ", + "effects contributions. If your user-provided random effects basis includes a random intercept or a ", + "random slope on the treatment variable, you will need to compute the prognostic or CATE functions manually by predicting ", + "'yhat' for different covariate and rfx_basis values." + )) + } } } - if (term %in% c("sigma2_leaf_mu", "leaf_scale_mu", "mu_leaf_scale")) { - if (!is.null(object$sigma2_leaf_mu_samples)) { - return(object$sigma2_leaf_mu_samples) - } else { - stop( - "This model does not have prognostic forest leaf variance parameter samples" - ) + # Handle prediction terms + rfx_model_spec <- object$model_params$rfx_model_spec + rfx_intercept_only <- rfx_model_spec == "intercept_only" + rfx_intercept_plus_treatment <- rfx_model_spec == "intercept_plus_treatment" + rfx_intercept <- rfx_intercept_only || rfx_intercept_plus_treatment + mu_prog_separate <- ifelse(rfx_intercept, TRUE, FALSE) + tau_cate_separate <- ifelse(rfx_intercept_plus_treatment, TRUE, FALSE) + if (!is.character(terms)) { + stop("type must be a string or character vector") + } + for (term in terms) { + if ( + !(term %in% + c( + "y_hat", + "prognostic_function", + "mu", + "cate", + "tau", + "rfx", + "variance_forest", + "all" + )) + ) { + warning(paste0( + "Term '", + term, + "' was requested. Valid terms are 'y_hat', 'prognostic_function', 'mu', 'cate', 'tau', 'rfx', 'variance_forest', and 'all'.", + " This term will be ignored and prediction will only proceed if other requested terms are available in the model." + )) } } - if (term %in% c("sigma2_leaf_tau", "leaf_scale_tau", "tau_leaf_scale")) { - if (!is.null(object$sigma2_leaf_tau_samples)) { - return(object$sigma2_leaf_tau_samples) - } else { - stop( - "This model does not have treatment effect forest leaf variance parameter samples" - ) + num_terms <- length(terms) + has_mu_forest <- T + has_tau_forest <- T + has_variance_forest <- object$model_params$include_variance_forest + has_rfx <- object$model_params$has_rfx + has_y_hat <- T + predict_y_hat <- (((has_y_hat) && ("y_hat" %in% terms)) || + ((has_y_hat) && ("all" %in% terms))) + predict_mu_forest <- (((has_mu_forest) && ("all" %in% terms)) || + ((has_mu_forest) && ("mu" %in% terms))) + predict_tau_forest <- (((has_tau_forest) && ("tau" %in% terms)) || + ((has_tau_forest) && ("all" %in% terms))) + predict_prog_function <- (((has_mu_forest) && + ("prognostic_function" %in% terms)) || + ((has_mu_forest) && ("all" %in% terms))) + predict_cate_function <- (((has_tau_forest) && ("cate" %in% terms)) || + ((has_tau_forest) && ("all" %in% terms))) + predict_rfx <- (((has_rfx) && ("rfx" %in% terms)) || + ((has_rfx) && ("all" %in% terms))) + predict_variance_forest <- (((has_variance_forest) && + ("variance_forest" %in% terms)) || + ((has_variance_forest) && ("all" %in% terms))) + predict_count <- sum(c( + predict_y_hat, + predict_mu_forest, + predict_prog_function, + predict_tau_forest, + predict_cate_function, + predict_rfx, + predict_variance_forest + )) + if (predict_count == 0) { + warning(paste0( + "None of the requested model terms, ", + paste(terms, collapse = ", "), + ", were fit in this model" + )) + return(NULL) + } + predict_rfx_intermediate <- (predict_y_hat && has_rfx) + predict_rfx_raw <- ((predict_prog_function && has_rfx && rfx_intercept) || + (predict_cate_function && has_rfx && rfx_intercept_plus_treatment)) + predict_mu_forest_intermediate <- ((predict_y_hat || predict_prog_function) && + has_mu_forest) + predict_tau_forest_intermediate <- ((predict_y_hat || + predict_cate_function || + (object$model_params$adaptive_coding && + (predict_mu_forest || predict_prog_function))) && + has_tau_forest) + + # Make sure covariates are matrix or data frame + if ((!is.data.frame(X)) && (!is.matrix(X))) { + stop("X must be a matrix or dataframe") + } + + # Handle factor-valued treatment before numeric conversion + if (is.factor(Z)) { + lvls <- levels(Z) + if (length(lvls) != 2) { + stop("Factor Z must have exactly 2 levels for binary treatment") } + warning( + "Z is a factor; recoding to 0/1 using level order: ", + lvls[1], + " = 0, ", + lvls[2], + " = 1" + ) + Z <- as.integer(Z) - 1L } - if (term %in% c("adaptive_coding")) { - if (object$model_params$adaptive_coding) { - b0_samples <- object$b_0_samples - b1_samples <- object$b_1_samples - return(unname(rbind(b0_samples, b1_samples))) - } else { - stop("This model does not have adaptive coding parameter samples") - } + # Convert all input data to matrices if not already converted + if ((is.null(dim(Z))) && (!is.null(Z))) { + Z <- as.matrix(as.numeric(Z)) } - - if (term %in% c("y_hat_train")) { - if (!is.null(object$y_hat_train)) { - return(object$y_hat_train) - } else { - stop( - "This model does not have in-sample mean function prediction samples" - ) - } + if ((is.null(dim(propensity))) && (!is.null(propensity))) { + propensity <- as.matrix(propensity) } - - if (term %in% c("y_hat_test")) { - if (!is.null(object$y_hat_test)) { - return(object$y_hat_test) - } else { - stop("This model does not have test set mean function prediction samples") - } + if ((is.null(dim(rfx_basis))) && (!is.null(rfx_basis))) { + rfx_basis <- as.matrix(rfx_basis) } - if (term %in% c("mu_hat_train", "prognostic_function_train")) { - if (!is.null(object$mu_hat_train)) { - return(object$mu_hat_train) - } else { - stop( - "This model does not have in-sample prognostic function predictions" - ) + # Data checks + if (nrow(X) != nrow(Z)) { + stop("X and Z must have the same number of rows") + } + if (object$model_params$num_covariates != ncol(X)) { + stop( + "X and must have the same number of columns as the covariates used to train the model" + ) + } + if ((object$model_params$has_rfx) && (is.null(rfx_group_ids))) { + stop( + "Random effect group labels (rfx_group_ids) must be provided for this model" + ) + } + if ((object$model_params$has_rfx_basis) && (is.null(rfx_basis))) { + if (object$model_params$rfx_model_spec == "custom") { + stop("Random effects basis (rfx_basis) must be provided for this model") } } - - if (term %in% c("mu_hat_test", "prognostic_function_test")) { - if (!is.null(object$mu_hat_test)) { - return(object$mu_hat_test) - } else { + if ((object$model_params$num_rfx_basis > 0) && (!is.null(rfx_basis))) { + if (ncol(rfx_basis) != object$model_params$num_rfx_basis) { stop( - "This model does not have test set prognostic function predictions" + "Random effects basis has a different dimension than the basis used to train this model" ) } } - if (term %in% c("tau_hat_train", "cate_train")) { - if (!is.null(object$tau_hat_train)) { - return(object$tau_hat_train) - } else { - stop( - "This model does not have in-sample treatment effect forest predictions" - ) + # Preprocess covariates before any prediction calls that depend on X being + # a numeric matrix (e.g. the internal propensity BART model was trained on a + # preprocessed matrix, so it expects a matrix, not the raw data frame) + train_set_metadata <- object$train_set_metadata + X <- preprocessPredictionData(X, train_set_metadata) + + # Compute propensity score using the internal bart model + if ( + (object$model_params$propensity_covariate != "none") && + (is.null(propensity)) + ) { + if (!object$model_params$internal_propensity_model) { + stop("propensity must be provided for this model") } + propensity <- rowMeans(predict(object$bart_propensity_model, X)$y_hat) } - if (term %in% c("tau_hat_test", "cate_test")) { - if (!is.null(object$tau_hat_test)) { - return(object$tau_hat_test) - } else { + # Recode group IDs to integer vector (if passed as, for example, a vector of county names, etc...) + if (!is.null(rfx_group_ids)) { + rfx_unique_group_ids <- object$rfx_unique_group_ids + group_ids_factor <- factor(rfx_group_ids, levels = rfx_unique_group_ids) + if (sum(is.na(group_ids_factor)) > 0) { stop( - "This model does not have test set treatment effect forest predictions" + "All random effect group labels provided in rfx_group_ids must have been present at sampling time" ) } + rfx_group_ids <- as.integer(group_ids_factor) } - if (term %in% c("sigma2_x_train", "var_x_train")) { - if (!is.null(object$sigma2_x_hat_train)) { - return(object$sigma2_x_hat_train) - } else { - stop("This model does not have in-sample variance forest predictions") + # Handle RFX model specification + if (has_rfx) { + if (object$model_params$rfx_model_spec == "custom") { + if (is.null(rfx_basis)) { + stop( + "A user-provided basis (`rfx_basis`) must be provided when the model was sampled with a random effects model spec set to 'custom'" + ) + } + } else if (object$model_params$rfx_model_spec == "intercept_only") { + # Only construct a basis if user-provided basis missing + if (is.null(rfx_basis)) { + rfx_basis <- matrix( + rep(1, nrow(X)), + nrow = nrow(X), + ncol = 1 + ) + } + } else if ( + object$model_params$rfx_model_spec == "intercept_plus_treatment" + ) { + # Only construct a basis if user-provided basis missing + if (is.null(rfx_basis)) { + rfx_basis <- cbind( + rep(1, nrow(X)), + Z + ) + } } } - if (term %in% c("sigma2_x_test", "var_x_test")) { - if (!is.null(object$sigma2_x_hat_test)) { - return(object$sigma2_x_hat_test) - } else { - stop("This model does not have test set variance forest predictions") - } + # Add propensities to covariate set if necessary + X_combined <- X + if (object$model_params$propensity_covariate != "none") { + X_combined <- cbind(X, propensity) } - if (term %in% c("tau_0", "treatment_intercept", "tau_intercept")) { - if (!is.null(object$tau_0_samples)) { - return(object$tau_0_samples) - } else { - stop( - "This model does not have treatment effect intercept (tau_0) samples" - ) - } + # Dimensions needed by both the C++ and R predict paths + n <- nrow(X_combined) + p <- ncol(X_combined) + treatment_dim <- ncol(Z) + obs_weights <- NULL + rfx_num_groups <- if (!is.null(rfx_group_ids)) { + length(unique(rfx_group_ids)) + } else { + 0L } + rfx_basis_dim <- if (!is.null(rfx_basis)) ncol(rfx_basis) else 0L - stop(paste0("term ", term, " is not a valid BCF model term")) -} - -#' @title Extract Random Effect Samples from BCF Model -#' @description -#' Extract raw sample values for each of the random effect parameter terms. -#' -#' @param object Object of type `bcfmodel` containing draws of a Bayesian causal forest model and associated sampling outputs. -#' @param ... Other parameters to be used in random effects extraction -#' @return List of arrays. The alpha array has dimension (`num_components`, `num_samples`) and is simply a vector if `num_components = 1`. -#' The xi and beta arrays have dimension (`num_components`, `num_groups`, `num_samples`) and is simply a matrix if `num_components = 1`. -#' The sigma array has dimension (`num_components`, `num_samples`) and is simply a vector if `num_components = 1`. -#' @export -#' -#' @examples -#' n <- 500 -#' p <- 5 -#' X <- matrix(runif(n*p), ncol = p) -#' mu_x <- ( -#' ((0 <= X[,1]) & (0.25 > X[,1])) * (-7.5) + -#' ((0.25 <= X[,1]) & (0.5 > X[,1])) * (-2.5) + -#' ((0.5 <= X[,1]) & (0.75 > X[,1])) * (2.5) + -#' ((0.75 <= X[,1]) & (1 > X[,1])) * (7.5) -#' ) -#' pi_x <- ( -#' ((0 <= X[,1]) & (0.25 > X[,1])) * (0.2) + -#' ((0.25 <= X[,1]) & (0.5 > X[,1])) * (0.4) + -#' ((0.5 <= X[,1]) & (0.75 > X[,1])) * (0.6) + -#' ((0.75 <= X[,1]) & (1 > X[,1])) * (0.8) -#' ) -#' tau_x <- ( -#' ((0 <= X[,2]) & (0.25 > X[,2])) * (0.5) + -#' ((0.25 <= X[,2]) & (0.5 > X[,2])) * (1.0) + -#' ((0.5 <= X[,2]) & (0.75 > X[,2])) * (1.5) + -#' ((0.75 <= X[,2]) & (1 > X[,2])) * (2.0) -#' ) -#' Z <- rbinom(n, 1, pi_x) -#' E_XZ <- mu_x + Z*tau_x -#' snr <- 3 -#' rfx_group_ids <- rep(c(1,2), n %/% 2) -#' rfx_coefs <- matrix(c(-1, -1, 1, 1), nrow=2, byrow=TRUE) -#' rfx_basis <- cbind(1, runif(n, -1, 1)) -#' rfx_term <- rowSums(rfx_coefs[rfx_group_ids,] * rfx_basis) -#' y <- E_XZ + rfx_term + rnorm(n, 0, 1)*(sd(E_XZ)/snr) -#' test_set_pct <- 0.2 -#' n_test <- round(test_set_pct*n) -#' n_train <- n - n_test -#' test_inds <- sort(sample(1:n, n_test, replace = FALSE)) -#' train_inds <- (1:n)[!((1:n) %in% test_inds)] -#' X_test <- X[test_inds,] -#' X_train <- X[train_inds,] -#' pi_test <- pi_x[test_inds] -#' pi_train <- pi_x[train_inds] -#' Z_test <- Z[test_inds] -#' Z_train <- Z[train_inds] -#' y_test <- y[test_inds] -#' y_train <- y[train_inds] -#' mu_test <- mu_x[test_inds] -#' mu_train <- mu_x[train_inds] -#' tau_test <- tau_x[test_inds] -#' tau_train <- tau_x[train_inds] -#' rfx_group_ids_test <- rfx_group_ids[test_inds] -#' rfx_group_ids_train <- rfx_group_ids[train_inds] -#' rfx_basis_test <- rfx_basis[test_inds,] -#' rfx_basis_train <- rfx_basis[train_inds,] -#' rfx_term_test <- rfx_term[test_inds] -#' rfx_term_train <- rfx_term[train_inds] -#' mu_params <- list(sample_sigma2_leaf = TRUE) -#' tau_params <- list(sample_sigma2_leaf = FALSE) -#' bcf_model <- bcf(X_train = X_train, Z_train = Z_train, y_train = y_train, -#' propensity_train = pi_train, -#' rfx_group_ids_train = rfx_group_ids_train, -#' rfx_basis_train = rfx_basis_train, X_test = X_test, -#' Z_test = Z_test, propensity_test = pi_test, -#' rfx_group_ids_test = rfx_group_ids_test, -#' rfx_basis_test = rfx_basis_test, -#' num_gfr = 10, num_burnin = 0, num_mcmc = 10, -#' prognostic_forest_params = mu_params, -#' treatment_effect_forest_params = tau_params) -#' rfx_samples <- getRandomEffectSamples(bcf_model) -getRandomEffectSamples.bcfmodel <- function(object, ...) { - result = list() + scale_int <- switch( + scale, + "linear" = 0L, + "probability" = 1L, + "class" = 2L, + 0L + ) - if (!object$model_params$has_rfx) { - warning("This model has no RFX terms, returning an empty list") - return(result) + # Build a flat list of model components for bcf_predict_cpp, since the bcfmodel + # object uses R6 wrappers and nested model_params that C++ cannot navigate directly. + has_variance_forest_model <- isTRUE( + object$model_params$include_variance_forest + ) + # Read forests through borrowed (non-owning) pointers into the single-owner + # samples object -- no deep copy, no deprecated-accessor error. + bcf_samples <- object$samples + has_rfx_model <- isTRUE(object$model_params$has_rfx) + bcf_metadata_list <- list( + num_samples = as.integer(object$model_params$num_samples), + y_bar = as.double(object$model_params$outcome_mean), + y_std = as.double(object$model_params$outcome_scale), + include_variance_forest = has_variance_forest_model, + has_rfx = has_rfx_model, + rfx_model_spec = if (has_rfx_model) { + object$model_params$rfx_model_spec + } else { + "" + }, + adaptive_coding = isTRUE(object$model_params$adaptive_coding), + sample_tau_0 = isTRUE(object$model_params$sample_tau_0) + ) + + output <- bcf_predict_cpp( + bcf_samples_ptr = bcf_samples$samples_ptr, + bcf_model_metadata = bcf_metadata_list, + X = X_combined, + Z = Z, + n = n, + p = p, + treatment_dim = treatment_dim, + obs_weights = obs_weights, + rfx_group_ids = rfx_group_ids, + rfx_basis = rfx_basis, + rfx_num_groups = rfx_num_groups, + rfx_basis_dim = rfx_basis_dim, + posterior = type == "posterior", + scale = scale_int, + predict_y_hat = predict_y_hat, + predict_mu_x = predict_mu_forest, + predict_tau_x = predict_tau_forest, + predict_prognostic_function = predict_prog_function, + predict_cate = predict_cate_function, + predict_conditional_variance = predict_variance_forest, + predict_random_effects = predict_rfx + ) + # Reshape flat C++ output vectors to matrices (n x num_samples) and rename + # fields to match the R predict path. For type="mean", num_samples_output=1 + # so we drop the trailing singleton to return a plain vector. + num_samples_raw <- as.integer(object$model_params$num_samples) + num_samples_output <- if (type == "posterior") num_samples_raw else 1L + reshape_cpp_pred_2d <- function(v, dim1, dim2) { + if (is.null(v)) { + return(NULL) + } + if (dim2 == 1L) { + return(as.vector(v)) + } + m <- v + dim(m) <- c(dim1, dim2) + m + } + reshape_cpp_pred_3d <- function(v, dim1, dim2, dim3) { + if (is.null(v)) { + return(NULL) + } + if (dim2 == 1L && dim3 == 1L) { + return(as.vector(v)) + } + if (dim3 == 1L) { + m <- v + dim(m) <- c(dim1, dim2) + return(m) + } + if (dim2 == 1L) { + m <- v + dim(m) <- c(dim1, dim3) + return(m) + } + a <- v + dim(a) <- c(dim1, dim2, dim3) + a + } + result <- list( + y_hat = reshape_cpp_pred_2d(output$y_hat, n, num_samples_output), + mu_hat = reshape_cpp_pred_2d(output$mu_x, n, num_samples_output), + tau_hat = reshape_cpp_pred_3d( + output$tau_x, + n, + treatment_dim, + num_samples_output + ), + prognostic_function = reshape_cpp_pred_2d( + output$prognostic_function, + n, + num_samples_output + ), + cate = reshape_cpp_pred_3d( + output$cate, + n, + treatment_dim, + num_samples_output + ), + rfx_predictions = reshape_cpp_pred_2d( + output$random_effects, + n, + num_samples_output + ), + variance_forest_predictions = reshape_cpp_pred_2d( + output$conditional_variance, + n, + num_samples_output + ) + ) + if (predict_count == 1L) { + if (predict_y_hat) { + return(result[["y_hat"]]) + } + if (predict_mu_forest) { + return(result[["mu_hat"]]) + } + if (predict_prog_function) { + return(result[["prognostic_function"]]) + } + if (predict_tau_forest) { + return(result[["tau_hat"]]) + } + if (predict_cate_function) { + return(result[["cate"]]) + } + if (predict_rfx) { + return(result[["rfx_predictions"]]) + } + if (predict_variance_forest) { + return(result[["variance_forest_predictions"]]) + } } - - # Extract the samples - result <- object$rfx_samples$extract_parameter_samples() - - # Scale by sd(y_train) - result$beta_samples <- result$beta_samples * - object$model_params$outcome_scale - result$xi_samples <- result$xi_samples * object$model_params$outcome_scale - result$alpha_samples <- result$alpha_samples * - object$model_params$outcome_scale - result$sigma_samples <- result$sigma_samples * - (object$model_params$outcome_scale^2) - return(result) } -#' @title Convert BCF Model to JSON -#' @param object Object of type `bcfmodel` containing draws of a Bayesian causal forest model and associated sampling outputs. +#' @title Print Summary of BCF Model +#' @description Prints a summary of the BCF model, including the model terms and their specifications. +#' @param x The BCF model object +#' @param ... Additional arguments (currently unused) #' @export -#' @rdname BCFSerialization -saveBCFModelToJson <- function(object) { - jsonobj <- createCppJson() - - if (!inherits(object, "bcfmodel")) { - stop("`object` must be a BCF model") +#' @return BCF model object unchanged after printing summary +print.bcfmodel <- function(x, ...) { + # What type of model was run + model_terms <- c("prognostic forest", "treatment effect forest") + if (x$model_params$include_variance_forest) { + model_terms <- c(model_terms, "variance forest") } - - if (is.null(object$model_params)) { - stop("This BCF model has not yet been sampled") + if (x$model_params$has_rfx) { + model_terms <- c(model_terms, "additive random effects") } - - # Add the forests - jsonobj$add_forest(object$forests_mu) - jsonobj$add_forest(object$forests_tau) - if (object$model_params$include_variance_forest) { - jsonobj$add_forest(object$forests_variance) + if (x$model_params$sample_sigma2_global) { + model_terms <- c(model_terms, "global error variance model") } - - # Add metadata - jsonobj$add_scalar( - "num_numeric_vars", - object$train_set_metadata$num_numeric_vars - ) - jsonobj$add_scalar( - "num_ordered_cat_vars", - object$train_set_metadata$num_ordered_cat_vars - ) - jsonobj$add_scalar( - "num_unordered_cat_vars", - object$train_set_metadata$num_unordered_cat_vars - ) - if (object$train_set_metadata$num_numeric_vars > 0) { - jsonobj$add_string_vector( - "numeric_vars", - object$train_set_metadata$numeric_vars - ) + if (x$model_params$sample_sigma2_leaf_mu) { + model_terms <- c(model_terms, "prognostic forest leaf scale model") } - if (object$train_set_metadata$num_ordered_cat_vars > 0) { - jsonobj$add_string_vector( - "ordered_cat_vars", - object$train_set_metadata$ordered_cat_vars - ) - jsonobj$add_string_list( - "ordered_unique_levels", - object$train_set_metadata$ordered_unique_levels - ) + if (x$model_params$sample_sigma2_leaf_tau) { + model_terms <- c(model_terms, "treatment effect forest leaf scale model") } - if (object$train_set_metadata$num_unordered_cat_vars > 0) { - jsonobj$add_string_vector( - "unordered_cat_vars", - object$train_set_metadata$unordered_cat_vars + if (x$model_params$sample_tau_0) { + model_terms <- c(model_terms, "treatment effect intercept model") + } + if (length(model_terms) > 2) { + summary_message <- paste0( + "stochtree::bcf() run with ", + paste0( + paste0(model_terms[1:(length(model_terms) - 1)], collapse = ", "), + ", and ", + model_terms[length(model_terms)] + ) ) - jsonobj$add_string_list( - "unordered_unique_levels", - object$train_set_metadata$unordered_unique_levels + } else if (length(model_terms) == 2) { + summary_message <- paste0( + "stochtree::bcf() run with ", + paste0(model_terms, collapse = " and ") ) + } else { + summary_message <- paste0("stochtree::bcf() run with ", model_terms) } - # Add version stamp and global parameters - jsonobj$add_string("stochtree_version", getStochtreeVersion()) - jsonobj$add_scalar("outcome_scale", object$model_params$outcome_scale) - jsonobj$add_scalar("outcome_mean", object$model_params$outcome_mean) - jsonobj$add_boolean("standardize", object$model_params$standardize) - jsonobj$add_scalar("sigma2_init", object$model_params$initial_sigma2) - jsonobj$add_boolean( - "sample_sigma2_global", - object$model_params$sample_sigma2_global - ) - jsonobj$add_boolean( - "sample_sigma2_leaf_mu", - object$model_params$sample_sigma2_leaf_mu - ) - jsonobj$add_boolean( - "sample_sigma2_leaf_tau", - object$model_params$sample_sigma2_leaf_tau - ) - jsonobj$add_boolean( - "include_variance_forest", - object$model_params$include_variance_forest - ) - jsonobj$add_string( - "propensity_covariate", - object$model_params$propensity_covariate - ) - jsonobj$add_boolean("has_rfx", object$model_params$has_rfx) - jsonobj$add_boolean("has_rfx_basis", object$model_params$has_rfx_basis) - jsonobj$add_scalar("num_rfx_basis", object$model_params$num_rfx_basis) - jsonobj$add_boolean( - "multivariate_treatment", - object$model_params$multivariate_treatment - ) - jsonobj$add_boolean("adaptive_coding", object$model_params$adaptive_coding) - jsonobj$add_boolean( - "binary_treatment", - object$model_params$binary_treatment - ) - jsonobj$add_scalar("treatment_dim", object$model_params$treatment_dim) - jsonobj$add_boolean("sample_tau_0", object$model_params$sample_tau_0) - jsonobj$add_boolean( - "internal_propensity_model", - object$model_params$internal_propensity_model - ) - jsonobj$add_scalar("num_gfr", object$model_params$num_gfr) - jsonobj$add_scalar("num_burnin", object$model_params$num_burnin) - jsonobj$add_scalar("num_mcmc", object$model_params$num_mcmc) - jsonobj$add_scalar("num_samples", object$model_params$num_samples) - jsonobj$add_scalar("keep_every", object$model_params$keep_every) - jsonobj$add_scalar("num_chains", object$model_params$num_chains) - jsonobj$add_scalar("num_covariates", object$model_params$num_covariates) - jsonobj$add_boolean( - "probit_outcome_model", - object$model_params$probit_outcome_model - ) - outcome_model_outcome <- object$model_params$outcome_model$outcome - outcome_model_link <- object$model_params$outcome_model$link - jsonobj$add_string( - "outcome", - outcome_model_outcome, - "outcome_model" - ) - jsonobj$add_string( - "link", - outcome_model_link, - "outcome_model" - ) - if (object$model_params$sample_sigma2_global) { - jsonobj$add_vector( - "sigma2_global_samples", - object$sigma2_global_samples, - "parameters" + # Outcome details + is_probit <- x$model_params$outcome_model$link == "probit" + summary_message <- paste0( + summary_message, + "\n", + "Outcome was modeled ", + ifelse( + is_probit, + "with a probit link", + "as gaussian" ) - } - if (object$model_params$sample_sigma2_leaf_mu) { - jsonobj$add_vector( - "sigma2_leaf_mu_samples", - object$sigma2_leaf_mu_samples, - "parameters" + ) + + # Treatment details + if (x$model_params$binary_treatment) { + summary_message <- paste0( + summary_message, + "\n", + "Treatment was binary and ", + ifelse( + x$model_params$adaptive_coding, + "its effect was estimated with adaptive coding", + "its effect was estimated with default coding" + ) ) - } - if (object$model_params$sample_sigma2_leaf_tau) { - jsonobj$add_vector( - "sigma2_leaf_tau_samples", - object$sigma2_leaf_tau_samples, - "parameters" + } else if (x$model_params$multivariate_treatment) { + summary_message <- paste0( + summary_message, + "\n", + "Treatment was multivariate with ", + x$model_params$treatment_dim, + " dimensions" ) - } - if (object$model_params$adaptive_coding) { - jsonobj$add_vector("b1_samples", object$b_1_samples, "parameters") - jsonobj$add_vector("b0_samples", object$b_0_samples, "parameters") - } - if (object$model_params$sample_tau_0 && !is.null(object$tau_0_samples)) { - jsonobj$add_scalar("tau_0_dim", nrow(object$tau_0_samples)) - jsonobj$add_vector( - "tau_0_samples", - as.numeric(object$tau_0_samples), - "parameters" + } else { + summary_message <- paste0( + summary_message, + "\n", + "Treatment was univariate but not binary" ) } - # Add random effects (if present) - if (object$model_params$has_rfx) { - jsonobj$add_random_effects(object$rfx_samples) - jsonobj$add_string_vector( - "rfx_unique_group_ids", - object$rfx_unique_group_ids + # Standardization + if (x$model_params$standardize) { + summary_message <- paste0( + summary_message, + "\n", + "outcome was standardized" ) } - jsonobj$add_string( - "rfx_model_spec", - object$model_params$rfx_model_spec - ) - # Add propensity model (if it exists) - if (object$model_params$internal_propensity_model) { - bart_propensity_string <- saveBARTModelToJsonString( - object$bart_propensity_model + # Internal propensity model + if (x$model_params$propensity_covariate == "none") { + summary_message <- paste0( + summary_message, + "\n", + "Propensity scores were not used in either forest of the model" ) - jsonobj$add_string("bart_propensity_model", bart_propensity_string) + } else { + if (x$model_params$internal_propensity_model) { + summary_message <- paste0( + summary_message, + "\n", + "An internal propensity model was fit using stochtree::bart() in lieu of user-provided propensity scores" + ) + } else { + summary_message <- paste0( + summary_message, + "\n", + "User-provided propensity scores were included in the model" + ) + } } - # Add covariate preprocessor metadata - preprocessor_metadata_string <- savePreprocessorToJsonString( - object$train_set_metadata - ) - jsonobj$add_string("preprocessor_metadata", preprocessor_metadata_string) - - return(jsonobj) -} - -#' @title Save BCF Model to JSON File -#' @param object Object of type `bcfmodel` containing draws of a Bayesian causal forest model and associated sampling outputs. -#' @param filename String of filepath, must end in ".json" -#' @export -#' @rdname BCFSerialization -saveBCFModelToJsonFile <- function(object, filename) { - # Convert to Json - jsonobj <- saveBCFModelToJson(object) - - # Save to file - jsonobj$save_file(filename) -} + # Random effects details + if (x$model_params$has_rfx) { + if (x$model_params$rfx_model_spec == "custom") { + summary_message <- paste0( + summary_message, + "\n", + "Random effects were fit with a user-supplied basis" + ) + } else if (x$model_params$rfx_model_spec == "intercept_only") { + summary_message <- paste0( + summary_message, + "\n", + "Random effects were fit with an 'intercept-only' parameterization" + ) + } else if (x$model_params$rfx_model_spec == "intercept_plus_treatment") { + summary_message <- paste0( + summary_message, + "\n", + "Random effects were fit with an 'intercept-plus-treatment' parameterization" + ) + } + } -#' @title Convert BCF Model to JSON String -#' @param object Object of type `bcfmodel` containing draws of a Bayesian causal forest model and associated sampling outputs. -#' @export -#' @rdname BCFSerialization -saveBCFModelToJsonString <- function(object) { - # Convert to Json - jsonobj <- saveBCFModelToJson(object) + # Sampler details + summary_message <- paste0( + summary_message, + "\n", + "The sampler was run for ", + x$model_params$num_gfr, + " GFR iterations, with ", + x$model_params$num_chains, + ifelse(x$model_params$num_chains == 1, " chain of ", " chains of "), + x$model_params$num_burnin, + " burn-in iterations and ", + x$model_params$num_mcmc, + " MCMC iterations, ", + ifelse( + x$model_params$keep_every == 1, + "retaining every iteration (i.e. no thinning)", + paste0( + "retaining every ", + x$model_params$keep_every, + "th iteration (i.e. thinning)" + ) + ) + ) - # Dump to string - return(jsonobj$return_json_string()) -} + # Print the model details + cat(summary_message, "\n") -# Recover `binary_treatment` for legacy BCF JSON written before the field was -# serialized. Two sample-time invariants make this exact in most cases: -# * multivariate treatment is never binary, and -# * adaptive_coding is forced FALSE unless the treatment is binary, so -# adaptive_coding == TRUE implies binary_treatment == TRUE. -# A univariate, default-coded treatment is genuinely ambiguous from JSON alone -# (it would require the original Z_train), so we conservatively return FALSE. -# Reads directly from the JSON object so it does not depend on the order in -# which the calling load path populates model_params. -.inferBinaryTreatmentFromJson <- function(json_obj, has_field_fn) { - multivariate <- if (has_field_fn("multivariate_treatment")) { - json_obj$get_boolean("multivariate_treatment") - } else { - FALSE - } - if (isTRUE(multivariate)) { - return(FALSE) - } - adaptive <- if (has_field_fn("adaptive_coding")) { - json_obj$get_boolean("adaptive_coding") - } else { - FALSE - } - isTRUE(adaptive) + # Return bcf model invisibly + invisible(x) } -#' @title Convert JSON to BCF Model -#' @param json_object Object of type `CppJson` containing Json representation of a BCF model +#' @title Summarize BCF Model +#' @description Summarize a BCF fit with a description of the model that was fit and numeric summaries of any sampled quantities. +#' @param object The BCF model object +#' @param ... Additional arguments #' @export -#' @rdname BCFSerialization -createBCFModelFromJson <- function(json_object) { - # Initialize the BCF model - output <- list() +#' @return BCF model object unchanged after summarizing +summary.bcfmodel <- function(object, ...) { + # First, print the BCF model + tmp <- print(object) - # Version inference and presence-check helpers - .ver <- inferStochtreeJsonVersion(json_object) - has_field <- function(name) { - json_contains_field_cpp(json_object$json_ptr, name) - } - has_subfolder_field <- function(subfolder, name) { - json_contains_field_subfolder_cpp(json_object$json_ptr, subfolder, name) - } + # Summarize any sampled quantities - # Unpack the forests - output[["forests_mu"]] <- loadForestContainerJson(json_object, "forest_0") - output[["forests_tau"]] <- loadForestContainerJson(json_object, "forest_1") - include_variance_forest <- json_object$get_boolean( - "include_variance_forest" - ) - if (include_variance_forest) { - output[["forests_variance"]] <- loadForestContainerJson( - json_object, - "forest_2" + # Global error scale + if (object$model_params$sample_sigma2_global) { + sigma2_samples <- object$samples$global_var_samples() + n_samples <- length(sigma2_samples) + mean_sigma2 <- mean(sigma2_samples) + sd_sigma2 <- sd(sigma2_samples) + quantiles_sigma2 <- quantile( + sigma2_samples, + probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) ) + cat(sprintf( + "Summary of sigma^2 posterior: \n%d samples, mean = %.3f, standard deviation = %.3f, quantiles:\n", + n_samples, + mean_sigma2, + sd_sigma2 + )) + print(quantiles_sigma2) } - # Unpack metadata - train_set_metadata = list() - train_set_metadata[["num_numeric_vars"]] <- json_object$get_scalar( - "num_numeric_vars" - ) - train_set_metadata[["num_ordered_cat_vars"]] <- json_object$get_scalar( - "num_ordered_cat_vars" - ) - train_set_metadata[["num_unordered_cat_vars"]] <- json_object$get_scalar( - "num_unordered_cat_vars" - ) - if (train_set_metadata[["num_numeric_vars"]] > 0) { - train_set_metadata[["numeric_vars"]] <- json_object$get_string_vector( - "numeric_vars" - ) - } - if (train_set_metadata[["num_ordered_cat_vars"]] > 0) { - train_set_metadata[[ - "ordered_cat_vars" - ]] <- json_object$get_string_vector("ordered_cat_vars") - train_set_metadata[[ - "ordered_unique_levels" - ]] <- json_object$get_string_list( - "ordered_unique_levels", - train_set_metadata[["ordered_cat_vars"]] - ) - } - if (train_set_metadata[["num_unordered_cat_vars"]] > 0) { - train_set_metadata[[ - "unordered_cat_vars" - ]] <- json_object$get_string_vector("unordered_cat_vars") - train_set_metadata[[ - "unordered_unique_levels" - ]] <- json_object$get_string_list( - "unordered_unique_levels", - train_set_metadata[["unordered_cat_vars"]] + # Leaf scale for the prognostic forest + if (object$model_params$sample_sigma2_leaf_mu) { + sigma2_leaf_samples <- object$samples$leaf_scale_mu_samples() + n_samples <- length(sigma2_leaf_samples) + mean_sigma2 <- mean(sigma2_leaf_samples) + sd_sigma2 <- sd(sigma2_leaf_samples) + quantiles_sigma2 <- quantile( + sigma2_leaf_samples, + probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) ) - } - output[["train_set_metadata"]] <- train_set_metadata - - # Unpack model params - model_params = list() - model_params[["outcome_scale"]] <- json_object$get_scalar("outcome_scale") - model_params[["outcome_mean"]] <- json_object$get_scalar("outcome_mean") - model_params[["standardize"]] <- json_object$get_boolean("standardize") - if (has_field("sigma2_init")) { - model_params[["initial_sigma2"]] <- json_object$get_scalar("sigma2_init") - } else { - model_params[["initial_sigma2"]] <- json_object$get_scalar("initial_sigma2") - warning(sprintf( - "JSON field 'initial_sigma2' is deprecated; please re-save the model to use 'sigma2_init' (inferred version: %s).", - .ver + cat(sprintf( + "Summary of prognostic forest leaf scale posterior: \n%d samples, mean = %.3f, standard deviation = %.3f, quantiles:\n", + n_samples, + mean_sigma2, + sd_sigma2 )) + print(quantiles_sigma2) } - model_params[["sample_sigma2_global"]] <- json_object$get_boolean( - "sample_sigma2_global" - ) - model_params[["sample_sigma2_leaf_mu"]] <- json_object$get_boolean( - "sample_sigma2_leaf_mu" - ) - model_params[["sample_sigma2_leaf_tau"]] <- json_object$get_boolean( - "sample_sigma2_leaf_tau" - ) - model_params[["include_variance_forest"]] <- include_variance_forest - model_params[["propensity_covariate"]] <- json_object$get_string( - "propensity_covariate" - ) - model_params[["has_rfx"]] <- json_object$get_boolean("has_rfx") - if (has_field("has_rfx_basis")) { - model_params[["has_rfx_basis"]] <- json_object$get_boolean("has_rfx_basis") - model_params[["num_rfx_basis"]] <- json_object$get_scalar("num_rfx_basis") - } else { - model_params[["has_rfx_basis"]] <- FALSE - model_params[["num_rfx_basis"]] <- 1 - warning(sprintf( - "Fields 'has_rfx_basis' and 'num_rfx_basis' not found in BCF JSON (inferred version: %s). Defaulting to has_rfx_basis=FALSE, num_rfx_basis=1.", - .ver + + # Leaf scale for the treatment effect forest + if (object$model_params$sample_sigma2_leaf_tau) { + sigma2_leaf_samples <- object$samples$leaf_scale_tau_samples() + n_samples <- length(sigma2_leaf_samples) + mean_sigma2 <- mean(sigma2_leaf_samples) + sd_sigma2 <- sd(sigma2_leaf_samples) + quantiles_sigma2 <- quantile( + sigma2_leaf_samples, + probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) + ) + cat(sprintf( + "Summary of treatment effect forest leaf scale posterior: \n%d samples, mean = %.3f, standard deviation = %.3f, quantiles:\n", + n_samples, + mean_sigma2, + sd_sigma2 )) + print(quantiles_sigma2) } - model_params[["adaptive_coding"]] <- json_object$get_boolean( - "adaptive_coding" - ) - if (has_field("binary_treatment")) { - model_params[["binary_treatment"]] <- json_object$get_boolean( - "binary_treatment" + + # Adaptive coding parameters + if (object$model_params$adaptive_coding) { + b0_samples <- object$samples$b0_samples() + b1_samples <- object$samples$b1_samples() + n_samples <- length(b0_samples) + mean_b0 <- mean(b0_samples) + mean_b1 <- mean(b1_samples) + sd_b0 <- sd(b0_samples) + sd_b1 <- sd(b1_samples) + quantiles_b0 <- quantile( + b0_samples, + probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) ) - } else { - model_params[["binary_treatment"]] <- .inferBinaryTreatmentFromJson( - json_object, - has_field + quantiles_b1 <- quantile( + b1_samples, + probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) ) - warning(sprintf( - "Field 'binary_treatment' not found in BCF JSON (inferred version: %s). Inferred binary_treatment=%s from other JSON fields.", - .ver, - model_params[["binary_treatment"]] + cat(sprintf( + "Summary of adaptive coding parameters: \n%d samples, mean (control) = %.3f, mean (treated) = %.3f, standard deviation (control) = %.3f, standard deviation (treated) = %.3f\n", + n_samples, + mean_b0, + mean_b1, + sd_b0, + sd_b1 )) + cat("quantiles (control):\n") + print(quantiles_b0) + cat("quantiles (treated):\n") + print(quantiles_b1) } - if (has_field("treatment_dim")) { - model_params[["treatment_dim"]] <- json_object$get_scalar("treatment_dim") - } else { - model_params[["treatment_dim"]] <- 1 - if ( - has_field("multivariate_treatment") && - isTRUE(json_object$get_boolean("multivariate_treatment")) - ) { - warning(sprintf( - "Field 'treatment_dim' not found in BCF JSON (inferred version: %s) for a multivariate-treatment model. Defaulting to 1.", - .ver - )) + + # Treatment effect intercept (tau_0) + tau_0_samples <- object$samples$tau_0_samples() + if (object$model_params$sample_tau_0 && length(tau_0_samples) > 0) { + tau_0_vec <- as.numeric(tau_0_samples) + n_samples <- if (is.matrix(tau_0_samples)) { + ncol(tau_0_samples) + } else { + length(tau_0_samples) } - } - if (has_field("sample_tau_0")) { - model_params[["sample_tau_0"]] <- json_object$get_boolean("sample_tau_0") - } else { - model_params[["sample_tau_0"]] <- FALSE - warning(sprintf( - "Field 'sample_tau_0' not found in BCF JSON (inferred version: %s). Defaulting to FALSE.", - .ver - )) - } - if (has_field("multivariate_treatment")) { - model_params[["multivariate_treatment"]] <- json_object$get_boolean( - "multivariate_treatment" - ) - } else { - model_params[["multivariate_treatment"]] <- FALSE - warning(sprintf( - "Field 'multivariate_treatment' not found in BCF JSON (inferred version: %s). Defaulting to FALSE.", - .ver - )) - } - if (has_field("internal_propensity_model")) { - model_params[["internal_propensity_model"]] <- json_object$get_boolean( - "internal_propensity_model" + mean_tau_0 <- mean(tau_0_vec) + sd_tau_0 <- sd(tau_0_vec) + quantiles_tau_0 <- quantile( + tau_0_vec, + probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) ) - } else { - model_params[["internal_propensity_model"]] <- FALSE - warning(sprintf( - "Field 'internal_propensity_model' not found in BCF JSON (inferred version: %s). Defaulting to FALSE.", - .ver - )) - } - model_params[["num_gfr"]] <- json_object$get_scalar("num_gfr") - model_params[["num_burnin"]] <- json_object$get_scalar("num_burnin") - model_params[["num_mcmc"]] <- json_object$get_scalar("num_mcmc") - model_params[["num_samples"]] <- json_object$get_scalar("num_samples") - model_params[["num_covariates"]] <- json_object$get_scalar("num_covariates") - if (has_field("num_chains")) { - model_params[["num_chains"]] <- json_object$get_scalar("num_chains") - } else { - model_params[["num_chains"]] <- 1 - warning(sprintf( - "Field 'num_chains' not found in BCF JSON (inferred version: %s). Defaulting to 1.", - .ver + cat(sprintf( + "Summary of treatment effect intercept (tau_0) posterior: \n%d samples, mean = %.3f, standard deviation = %.3f, quantiles:\n", + n_samples, + mean_tau_0, + sd_tau_0 )) + print(quantiles_tau_0) } - if (has_field("keep_every")) { - model_params[["keep_every"]] <- json_object$get_scalar("keep_every") - } else { - model_params[["keep_every"]] <- 1 - warning(sprintf( - "Field 'keep_every' not found in BCF JSON (inferred version: %s). Defaulting to 1.", - .ver + + # In-sample predictions + y_hat_train <- object$samples$y_hat_train() + if (length(y_hat_train) > 0) { + y_hat_train_mean <- rowMeans(y_hat_train) + n_y_hat_train <- length(y_hat_train_mean) + mean_y_hat_train <- mean(y_hat_train_mean) + sd_y_hat_train <- sd(y_hat_train_mean) + quantiles_y_hat_train <- quantile( + y_hat_train_mean, + probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) + ) + cat(sprintf( + "Summary of in-sample posterior mean predictions: \n%d observations, mean = %.3f, standard deviation = %.3f, quantiles:\n", + n_y_hat_train, + mean_y_hat_train, + sd_y_hat_train )) + print(quantiles_y_hat_train) } - if (has_field("probit_outcome_model")) { - model_params[["probit_outcome_model"]] <- json_object$get_boolean( - "probit_outcome_model" + + # Test-set predictions + y_hat_test <- object$samples$y_hat_test() + if (length(y_hat_test) > 0) { + y_hat_test_mean <- rowMeans(y_hat_test) + n_y_hat_test <- length(y_hat_test_mean) + mean_y_hat_test <- mean(y_hat_test_mean) + sd_y_hat_test <- sd(y_hat_test_mean) + quantiles_y_hat_test <- quantile( + y_hat_test_mean, + probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) ) - } else { - model_params[["probit_outcome_model"]] <- FALSE - warning(sprintf( - "Field 'probit_outcome_model' not found in BCF JSON (inferred version: %s). Defaulting to FALSE.", - .ver + cat(sprintf( + "Summary of test-set posterior mean predictions: \n%d observations, mean = %.3f, standard deviation = %.3f, quantiles:\n", + n_y_hat_test, + mean_y_hat_test, + sd_y_hat_test )) + print(quantiles_y_hat_test) } - if (has_subfolder_field("outcome_model", "outcome")) { - outcome_model_outcome <- json_object$get_string("outcome", "outcome_model") - outcome_model_link <- json_object$get_string("link", "outcome_model") - } else { - outcome_model_outcome <- "continuous" - outcome_model_link <- "identity" - warning(sprintf( - "Subfolder 'outcome_model' not found in BCF JSON (inferred version: %s). Defaulting to outcome='continuous', link='identity'.", - .ver - )) + + # In-sample treatment effect function estimates + tau_hat_train <- object$samples$tau_forest_predictions_train() + if (length(tau_hat_train) > 0) { + if (!object$model_params$multivariate_treatment) { + tau_hat_train_mean <- rowMeans(tau_hat_train) + n_tau_hat_train <- length(tau_hat_train_mean) + mean_tau_hat_train <- mean(tau_hat_train_mean) + sd_tau_hat_train <- sd(tau_hat_train_mean) + quantiles_tau_hat_train <- quantile( + tau_hat_train_mean, + probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) + ) + cat(sprintf( + "Summary of in-sample posterior mean CATEs: \n%d observations, mean = %.3f, standard deviation = %.3f, quantiles:\n", + n_tau_hat_train, + mean_tau_hat_train, + sd_tau_hat_train + )) + print(quantiles_tau_hat_train) + } } - model_params[["outcome_model"]] <- OutcomeModel( - outcome = outcome_model_outcome, - link = outcome_model_link - ) - if (has_field("rfx_model_spec")) { - model_params[["rfx_model_spec"]] <- json_object$get_string( - "rfx_model_spec" - ) - } else { - model_params[["rfx_model_spec"]] <- "" - if (model_params[["has_rfx"]]) { - warning(sprintf( - "Field 'rfx_model_spec' not found in BCF JSON (inferred version: %s) but has_rfx=TRUE.", - .ver + + # Test set treatment effect function estimates + tau_hat_test <- object$samples$tau_forest_predictions_test() + if (length(tau_hat_test) > 0) { + if (!object$model_params$multivariate_treatment) { + tau_hat_test_mean <- rowMeans(tau_hat_test) + n_tau_hat_test <- length(tau_hat_test_mean) + mean_tau_hat_test <- mean(tau_hat_test_mean) + sd_tau_hat_test <- sd(tau_hat_test_mean) + quantiles_tau_hat_test <- quantile( + tau_hat_test_mean, + probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) + ) + cat(sprintf( + "Summary of test-set posterior mean CATEs: \n%d observations, mean = %.3f, standard deviation = %.3f, quantiles:\n", + n_tau_hat_test, + mean_tau_hat_test, + sd_tau_hat_test + )) + print(quantiles_tau_hat_test) + } + } + + # Random effects + if (object$model_params$has_rfx) { + rfx_samples <- getRandomEffectSamples(object) + rfx_beta_samples <- rfx_samples$beta_samples + if (length(dim(rfx_beta_samples)) > 2) { + cat( + "Random effects summary of variance components across groups and posterior draws:\n" + ) + rfx_component_means <- apply(rfx_beta_samples, 1, mean) + rfx_component_sds <- apply(rfx_beta_samples, 1, sd) + cat("Variance component means: ", rfx_component_means, "\n") + cat("Variance component standard deviations: ", rfx_component_sds, "\n") + quantile_summary <- t(apply( + rfx_beta_samples, + 1, + quantile, + probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) )) + cat("Variance component quantiles:\n") + print(quantile_summary) + } else { + cat( + "Random effects summary of variance components across groups and posterior draws:\n" + ) + rfx_component_means <- mean(rfx_beta_samples) + rfx_component_sds <- sd(rfx_beta_samples) + cat("Random effects overall mean: ", rfx_component_means, "\n") + cat( + "Random effects overall standard deviation: ", + rfx_component_sds, + "\n" + ) + quantile_summary <- quantile( + rfx_beta_samples, + probs = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975) + ) + cat("Random effects overall quantiles:\n") + print(quantile_summary) + } + } + + # Return bcf model invisibly + invisible(object) +} + +#' @title Plot BCF Model +#' @description Plot the BCF model fit and any relevant sampled quantities. This will default to a traceplot of the global error scale and the in-sample mean forest predictions for the first train set observation. Since `stochtree::bcf()` is flexible and it's possible to sample a model with a fixed global error scale and no mean forest, this procedure will throw an error if these two default terms are omitted. +#' @param x The BCF model object +#' @param ... Additional arguments +#' @export +#' @return BCF model object unchanged after summarizing +plot.bcfmodel <- function(x, ...) { + # Check if model has global error scale samples + has_sigma2_samples <- x$model_params$sample_sigma2_global + has_mean_forest_preds <- length(x$samples$y_hat_train()) > 0 + + # First try combinations of sigma2 and mean forest predictions + if (has_sigma2_samples || has_mean_forest_preds) { + if (has_sigma2_samples) { + plot( + x$samples$global_var_samples(), + type = "l", + ylab = "Sigma^2", + main = "Global error scale traceplot" + ) + } else if (has_mean_forest_preds) { + plot( + x$samples$y_hat_train()[1, ], + type = "l", + ylab = "Predictions", + main = "In-sample mean function trace for the first train set observation" + ) + } + } else { + stop( + "This model does not have enough model terms / parameter traces to produce stochtree's default plots. See `predict.bcfmodel()` for examples of how to further investigate your model." + ) + } + + # Return bcf model invisibly + invisible(x) +} + +#' @title Extract BCF Parameter Samples +#' @description Extract a vector, matrix or array of parameter samples from a BCF model by name. +#' Random effects are handled by a separate `getRandomEffectSamples` function due to the complexity of the random effects parameters. +#' If the requested model term is not found, an error is thrown. +#' The following conventions are used for parameter names: +#' - Global error variance: `"sigma2"`, `"global_error_scale"`, `"sigma2_global"` +#' - Prognostic forest leaf scale: `"sigma2_leaf_mu"`, `"leaf_scale_mu"`, `"mu_leaf_scale"` +#' - Treatment effect forest leaf scale: `"sigma2_leaf_tau"`, `"leaf_scale_tau"`, `"tau_leaf_scale"` +#' - Adaptive coding parameters: `"adaptive_coding"` (returns both the control and treated parameters jointly, with control in the first row and treated in the second row) +#' - In-sample mean function predictions: `"y_hat_train"` +#' - Test set mean function predictions: `"y_hat_test"` +#' - In-sample treatment effect forest predictions: `"tau_hat_train"` +#' - Test set treatment effect forest predictions: `"tau_hat_test"` +#' - Treatment effect intercept: `"tau_0"`, `"treatment_intercept"`, `"tau_intercept"` +#' - In-sample variance forest predictions: `"sigma2_x_train"`, `"var_x_train"` +#' - Test set variance forest predictions: `"sigma2_x_test"`, `"var_x_test"` +#' +#' @param object Object of type `bcfmodel` containing draws of a BCF model and associated sampling outputs. +#' @param term Name of the parameter to extract (e.g., `"sigma2"`, `"y_hat_train"`, etc.) +#' @return Array of parameter samples. If the underlying parameter is a scalar, this will be a vector of length `num_samples`. +#' If the underlying parameter is vector-valued, this will be (`parameter_dimension` x `num_samples`) matrix, and if the underlying +#' parameter is multidimensional, this will be an array of dimension (`parameter_dimension_1` x `parameter_dimension_2` x ... x `num_samples`). +#' @export +#' +#' @examples +#' n <- 500 +#' p <- 5 +#' X <- matrix(runif(n*p), ncol = p) +#' mu_x <- ( +#' ((0 <= X[,1]) & (0.25 > X[,1])) * (-7.5) + +#' ((0.25 <= X[,1]) & (0.5 > X[,1])) * (-2.5) + +#' ((0.5 <= X[,1]) & (0.75 > X[,1])) * (2.5) + +#' ((0.75 <= X[,1]) & (1 > X[,1])) * (7.5) +#' ) +#' pi_x <- ( +#' ((0 <= X[,1]) & (0.25 > X[,1])) * (0.2) + +#' ((0.25 <= X[,1]) & (0.5 > X[,1])) * (0.4) + +#' ((0.5 <= X[,1]) & (0.75 > X[,1])) * (0.6) + +#' ((0.75 <= X[,1]) & (1 > X[,1])) * (0.8) +#' ) +#' tau_x <- ( +#' ((0 <= X[,2]) & (0.25 > X[,2])) * (0.5) + +#' ((0.25 <= X[,2]) & (0.5 > X[,2])) * (1.0) + +#' ((0.5 <= X[,2]) & (0.75 > X[,2])) * (1.5) + +#' ((0.75 <= X[,2]) & (1 > X[,2])) * (2.0) +#' ) +#' Z <- rbinom(n, 1, pi_x) +#' E_XZ <- mu_x + Z*tau_x +#' snr <- 3 +#' rfx_group_ids <- rep(c(1,2), n %/% 2) +#' rfx_coefs <- matrix(c(-1, -1, 1, 1), nrow=2, byrow=TRUE) +#' rfx_basis <- cbind(1, runif(n, -1, 1)) +#' rfx_term <- rowSums(rfx_coefs[rfx_group_ids,] * rfx_basis) +#' y <- E_XZ + rfx_term + rnorm(n, 0, 1)*(sd(E_XZ)/snr) +#' bcf_model <- bcf(X_train = X, y_train = y, Z_train = Z, +#' rfx_group_ids_train = rfx_group_ids, +#' rfx_basis_train = rfx_basis, +#' num_gfr = 10, num_burnin = 0, num_mcmc = 10) +#' sigma2_samples <- extractParameter(bcf_model, "sigma2") +extractParameter.bcfmodel <- function(object, term) { + if (term %in% c("sigma2", "global_error_scale", "sigma2_global")) { + s <- object$samples$global_var_samples() + if (length(s) > 0) { + return(s) + } else { + stop("This model does not have global variance parameter samples") + } + } + + if (term %in% c("sigma2_leaf_mu", "leaf_scale_mu", "mu_leaf_scale")) { + s <- object$samples$leaf_scale_mu_samples() + if (length(s) > 0) { + return(s) + } else { + stop( + "This model does not have prognostic forest leaf variance parameter samples" + ) } } - output[["model_params"]] <- model_params - # Unpack sampled parameters - if (model_params[["sample_sigma2_global"]]) { - output[["sigma2_global_samples"]] <- json_object$get_vector( - "sigma2_global_samples", - "parameters" - ) + if (term %in% c("sigma2_leaf_tau", "leaf_scale_tau", "tau_leaf_scale")) { + s <- object$samples$leaf_scale_tau_samples() + if (length(s) > 0) { + return(s) + } else { + stop( + "This model does not have treatment effect forest leaf variance parameter samples" + ) + } } - if (model_params[["sample_sigma2_leaf_mu"]]) { - output[["sigma2_leaf_mu_samples"]] <- json_object$get_vector( - "sigma2_leaf_mu_samples", - "parameters" - ) + + if (term %in% c("adaptive_coding")) { + if (object$model_params$adaptive_coding) { + b0_samples <- object$samples$b0_samples() + b1_samples <- object$samples$b1_samples() + return(unname(rbind(b0_samples, b1_samples))) + } else { + stop("This model does not have adaptive coding parameter samples") + } } - if (model_params[["sample_sigma2_leaf_tau"]]) { - output[["sigma2_leaf_tau_samples"]] <- json_object$get_vector( - "sigma2_leaf_tau_samples", - "parameters" - ) + + if (term %in% c("y_hat_train")) { + preds <- object$samples$y_hat_train() + if (length(preds) > 0) { + return(preds) + } else { + stop( + "This model does not have in-sample mean function prediction samples" + ) + } + } + + if (term %in% c("y_hat_test")) { + preds <- object$samples$y_hat_test() + if (length(preds) > 0) { + return(preds) + } else { + stop("This model does not have test set mean function prediction samples") + } } - if (model_params[["adaptive_coding"]]) { - if (has_subfolder_field("parameters", "b1_samples")) { - output[["b_1_samples"]] <- json_object$get_vector( - "b1_samples", - "parameters" + + if (term %in% c("mu_hat_train", "prognostic_function_train")) { + preds <- object$samples$mu_forest_predictions_train() + if (length(preds) > 0) { + return(preds) + } else { + stop( + "This model does not have in-sample prognostic function predictions" ) - output[["b_0_samples"]] <- json_object$get_vector( - "b0_samples", - "parameters" + } + } + + if (term %in% c("mu_hat_test", "prognostic_function_test")) { + preds <- object$samples$mu_forest_predictions_test() + if (length(preds) > 0) { + return(preds) + } else { + stop( + "This model does not have test set prognostic function predictions" ) + } + } + + # tau_hat / cate is the FULL CATE: the sampler folds tau_0 (and the adaptive-coding + # (b1 - b0) scaling) into tau_forest_predictions before caching, so these accessors already + # return tau_0 + tau(x) -- do NOT add tau_0 again here (see bcf_sampler.cpp postprocess / GH #376). + if (term %in% c("tau_hat_train", "cate_train")) { + preds <- object$samples$tau_forest_predictions_train() + if (length(preds) > 0) { + return(preds) } else { - output[["b_1_samples"]] <- json_object$get_vector( - "b_1_samples", - "parameters" + stop( + "This model does not have in-sample treatment effect forest predictions" ) - output[["b_0_samples"]] <- json_object$get_vector( - "b_0_samples", - "parameters" + } + } + + if (term %in% c("tau_hat_test", "cate_test")) { + preds <- object$samples$tau_forest_predictions_test() + if (length(preds) > 0) { + return(preds) + } else { + stop( + "This model does not have test set treatment effect forest predictions" ) - warning(sprintf( - "JSON fields 'b_1_samples'/'b_0_samples' are deprecated; please re-save the model to use 'b1_samples'/'b0_samples' (inferred version: %s).", - .ver - )) } } - if (model_params[["sample_tau_0"]]) { - tau_0_dim <- as.integer(json_object$get_scalar("tau_0_dim")) - tau_0_vec <- json_object$get_vector("tau_0_samples", "parameters") - output[["tau_0_samples"]] <- matrix(tau_0_vec, nrow = tau_0_dim) + + if (term %in% c("sigma2_x_train", "sigma2_x_hat_train", "var_x_train")) { + preds <- object$samples$variance_forest_predictions_train() + if (length(preds) > 0) { + return(preds) + } else { + stop("This model does not have in-sample variance forest predictions") + } + } + + if (term %in% c("sigma2_x_test", "sigma2_x_hat_test", "var_x_test")) { + preds <- object$samples$variance_forest_predictions_test() + if (length(preds) > 0) { + return(preds) + } else { + stop("This model does not have test set variance forest predictions") + } } - # Unpack random effects - if (model_params[["has_rfx"]]) { - output[["rfx_unique_group_ids"]] <- json_object$get_string_vector( - "rfx_unique_group_ids" - ) - output[["rfx_samples"]] <- loadRandomEffectSamplesJson(json_object, 0) + if (term %in% c("tau_0", "treatment_intercept", "tau_intercept")) { + s <- object$samples$tau_0_samples() + if (length(s) > 0) { + return(s) + } else { + stop( + "This model does not have treatment effect intercept (tau_0) samples" + ) + } + } + + stop(paste0("term ", term, " is not a valid BCF model term")) +} + +#' @title Extract Random Effect Samples from BCF Model +#' @description +#' Extract raw sample values for each of the random effect parameter terms. +#' +#' @param object Object of type `bcfmodel` containing draws of a Bayesian causal forest model and associated sampling outputs. +#' @param ... Other parameters to be used in random effects extraction +#' @return List of arrays. The alpha array has dimension (`num_components`, `num_samples`) and is simply a vector if `num_components = 1`. +#' The xi and beta arrays have dimension (`num_components`, `num_groups`, `num_samples`) and is simply a matrix if `num_components = 1`. +#' The sigma array has dimension (`num_components`, `num_samples`) and is simply a vector if `num_components = 1`. +#' @export +#' +#' @examples +#' n <- 500 +#' p <- 5 +#' X <- matrix(runif(n*p), ncol = p) +#' mu_x <- ( +#' ((0 <= X[,1]) & (0.25 > X[,1])) * (-7.5) + +#' ((0.25 <= X[,1]) & (0.5 > X[,1])) * (-2.5) + +#' ((0.5 <= X[,1]) & (0.75 > X[,1])) * (2.5) + +#' ((0.75 <= X[,1]) & (1 > X[,1])) * (7.5) +#' ) +#' pi_x <- ( +#' ((0 <= X[,1]) & (0.25 > X[,1])) * (0.2) + +#' ((0.25 <= X[,1]) & (0.5 > X[,1])) * (0.4) + +#' ((0.5 <= X[,1]) & (0.75 > X[,1])) * (0.6) + +#' ((0.75 <= X[,1]) & (1 > X[,1])) * (0.8) +#' ) +#' tau_x <- ( +#' ((0 <= X[,2]) & (0.25 > X[,2])) * (0.5) + +#' ((0.25 <= X[,2]) & (0.5 > X[,2])) * (1.0) + +#' ((0.5 <= X[,2]) & (0.75 > X[,2])) * (1.5) + +#' ((0.75 <= X[,2]) & (1 > X[,2])) * (2.0) +#' ) +#' Z <- rbinom(n, 1, pi_x) +#' E_XZ <- mu_x + Z*tau_x +#' snr <- 3 +#' rfx_group_ids <- rep(c(1,2), n %/% 2) +#' rfx_coefs <- matrix(c(-1, -1, 1, 1), nrow=2, byrow=TRUE) +#' rfx_basis <- cbind(1, runif(n, -1, 1)) +#' rfx_term <- rowSums(rfx_coefs[rfx_group_ids,] * rfx_basis) +#' y <- E_XZ + rfx_term + rnorm(n, 0, 1)*(sd(E_XZ)/snr) +#' test_set_pct <- 0.2 +#' n_test <- round(test_set_pct*n) +#' n_train <- n - n_test +#' test_inds <- sort(sample(1:n, n_test, replace = FALSE)) +#' train_inds <- (1:n)[!((1:n) %in% test_inds)] +#' X_test <- X[test_inds,] +#' X_train <- X[train_inds,] +#' pi_test <- pi_x[test_inds] +#' pi_train <- pi_x[train_inds] +#' Z_test <- Z[test_inds] +#' Z_train <- Z[train_inds] +#' y_test <- y[test_inds] +#' y_train <- y[train_inds] +#' mu_test <- mu_x[test_inds] +#' mu_train <- mu_x[train_inds] +#' tau_test <- tau_x[test_inds] +#' tau_train <- tau_x[train_inds] +#' rfx_group_ids_test <- rfx_group_ids[test_inds] +#' rfx_group_ids_train <- rfx_group_ids[train_inds] +#' rfx_basis_test <- rfx_basis[test_inds,] +#' rfx_basis_train <- rfx_basis[train_inds,] +#' rfx_term_test <- rfx_term[test_inds] +#' rfx_term_train <- rfx_term[train_inds] +#' mu_params <- list(sample_sigma2_leaf = TRUE) +#' tau_params <- list(sample_sigma2_leaf = FALSE) +#' bcf_model <- bcf(X_train = X_train, Z_train = Z_train, y_train = y_train, +#' propensity_train = pi_train, +#' rfx_group_ids_train = rfx_group_ids_train, +#' rfx_basis_train = rfx_basis_train, X_test = X_test, +#' Z_test = Z_test, propensity_test = pi_test, +#' rfx_group_ids_test = rfx_group_ids_test, +#' rfx_basis_test = rfx_basis_test, +#' num_gfr = 10, num_burnin = 0, num_mcmc = 10, +#' prognostic_forest_params = mu_params, +#' treatment_effect_forest_params = tau_params) +#' rfx_samples <- getRandomEffectSamples(bcf_model) +getRandomEffectSamples.bcfmodel <- function(object, ...) { + result <- list() + + if (!object$model_params$has_rfx) { + warning("This model has no RFX terms, returning an empty list") + return(result) + } + + # Extract the samples + rfx_samples <- object$samples$materialize_rfx() + result <- rfx_samples$extract_parameter_samples() + + # Scale by sd(y_train) + result$beta_samples <- result$beta_samples * + object$model_params$outcome_scale + result$xi_samples <- result$xi_samples * object$model_params$outcome_scale + result$alpha_samples <- result$alpha_samples * + object$model_params$outcome_scale + result$sigma_samples <- result$sigma_samples * + (object$model_params$outcome_scale^2) + + return(result) +} + +#' @title Convert BCF Model to JSON +#' @param object Object of type `bcfmodel` containing draws of a Bayesian causal forest model and associated sampling outputs. +#' @export +#' @rdname BCFSerialization +saveBCFModelToJson <- function(object) { + jsonobj <- createCppJson() + + if (!inherits(object, "bcfmodel")) { + stop("`object` must be a BCF model") + } + + if (is.null(object$model_params)) { + stop("This BCF model has not yet been sampled") } - # Unpack propensity model (if it exists) - if (model_params[["internal_propensity_model"]]) { - bart_propensity_string <- json_object$get_string( - "bart_propensity_model" + # Add the samples to the JSON object + bcf_samples <- object$samples + bcf_samples$append_to_json(jsonobj) + + # Add version stamp and global parameters + jsonobj$add_string("stochtree_version", getStochtreeVersion()) + jsonobj$add_string("platform", "R") + jsonobj$add_integer("schema_version", STOCHTREE_SCHEMA_VERSION) + jsonobj$add_scalar("outcome_scale", object$model_params$outcome_scale) + jsonobj$add_scalar("outcome_mean", object$model_params$outcome_mean) + jsonobj$add_boolean("standardize", object$model_params$standardize) + jsonobj$add_scalar("sigma2_init", object$model_params$initial_sigma2) + jsonobj$add_boolean( + "sample_sigma2_global", + object$model_params$sample_sigma2_global + ) + jsonobj$add_boolean( + "sample_sigma2_leaf_mu", + object$model_params$sample_sigma2_leaf_mu + ) + jsonobj$add_boolean( + "sample_sigma2_leaf_tau", + object$model_params$sample_sigma2_leaf_tau + ) + jsonobj$add_boolean( + "include_variance_forest", + object$model_params$include_variance_forest + ) + jsonobj$add_string( + "propensity_covariate", + object$model_params$propensity_covariate + ) + jsonobj$add_boolean("has_rfx", object$model_params$has_rfx) + jsonobj$add_boolean("has_rfx_basis", object$model_params$has_rfx_basis) + jsonobj$add_scalar("num_rfx_basis", object$model_params$num_rfx_basis) + jsonobj$add_boolean( + "multivariate_treatment", + object$model_params$multivariate_treatment + ) + jsonobj$add_boolean("adaptive_coding", object$model_params$adaptive_coding) + jsonobj$add_boolean( + "binary_treatment", + object$model_params$binary_treatment + ) + jsonobj$add_scalar("treatment_dim", object$model_params$treatment_dim) + jsonobj$add_boolean("sample_tau_0", object$model_params$sample_tau_0) + jsonobj$add_boolean( + "internal_propensity_model", + object$model_params$internal_propensity_model + ) + jsonobj$add_scalar("num_gfr", object$model_params$num_gfr) + jsonobj$add_scalar("num_burnin", object$model_params$num_burnin) + jsonobj$add_scalar("num_mcmc", object$model_params$num_mcmc) + jsonobj$add_scalar("num_samples", object$model_params$num_samples) + jsonobj$add_scalar("keep_every", object$model_params$keep_every) + jsonobj$add_scalar("num_chains", object$model_params$num_chains) + jsonobj$add_scalar("num_covariates", object$model_params$num_covariates) + jsonobj$add_boolean( + "probit_outcome_model", + object$model_params$probit_outcome_model + ) + outcome_model_outcome <- object$model_params$outcome_model$outcome + outcome_model_link <- object$model_params$outcome_model$link + jsonobj$add_string( + "outcome", + outcome_model_outcome, + "outcome_model" + ) + jsonobj$add_string( + "link", + outcome_model_link, + "outcome_model" + ) + + # Add random effects group IDs + if (object$model_params$has_rfx) { + jsonobj$add_string_vector( + "rfx_unique_group_ids", + object$rfx_unique_group_ids, + subfolder_name = "random_effects" ) - output[["bart_propensity_model"]] <- createBARTModelFromJsonString( - bart_propensity_string + # Cross-platform compatible on the rfx axis iff the group id levels are + # integer-valued (Python supports only integer group ids). + rfx_compatible <- all(grepl("^-?[0-9]+$", object$rfx_unique_group_ids)) + jsonobj$add_boolean( + "cross_platform_compatible", + rfx_compatible, + subfolder_name = "random_effects" ) } + jsonobj$add_string( + "rfx_model_spec", + object$model_params$rfx_model_spec + ) - # Unpack covariate preprocessor - if (has_field("preprocessor_metadata")) { - preprocessor_metadata_string <- json_object$get_string( - "preprocessor_metadata" - ) - output[["train_set_metadata"]] <- createPreprocessorFromJsonString( - preprocessor_metadata_string + # Add propensity model (if it exists) + if (object$model_params$internal_propensity_model) { + bart_propensity_string <- saveBARTModelToJsonString( + object$bart_propensity_model ) - } else { - output[["train_set_metadata"]] <- NULL - warning(sprintf( - "Field 'preprocessor_metadata' not found in BCF JSON (inferred version: %s). Preprocessor is unavailable; prediction may fail.", - .ver - )) + jsonobj$add_string("bart_propensity_model", bart_propensity_string) } - class(output) <- "bcfmodel" - return(output) + # Add covariate preprocessor metadata + preprocessor_metadata_string <- savePreprocessorToJsonString( + object$train_set_metadata + ) + jsonobj$add_string("covariate_preprocessor", preprocessor_metadata_string) + + return(jsonobj) } -#' @title Convert JSON File to BCF Model -#' @param json_filename String of filepath, must end in ".json" +#' @title Save BCF Model to JSON File +#' @param object Object of type `bcfmodel` containing draws of a Bayesian causal forest model and associated sampling outputs. +#' @param filename String of filepath, must end in ".json" #' @export #' @rdname BCFSerialization -createBCFModelFromJsonFile <- function(json_filename) { - # Load a `CppJson` object from file - bcf_json <- createCppJsonFile(json_filename) - - # Create and return the BCF object - bcf_object <- createBCFModelFromJson(bcf_json) +saveBCFModelToJsonFile <- function(object, filename) { + # Convert to Json + jsonobj <- saveBCFModelToJson(object) - return(bcf_object) + # Save to file + jsonobj$save_file(filename) } -#' @title Convert JSON String to BCF Model -#' @param json_string JSON string dump +#' @title Convert BCF Model to JSON String +#' @param object Object of type `bcfmodel` containing draws of a Bayesian causal forest model and associated sampling outputs. #' @export #' @rdname BCFSerialization -createBCFModelFromJsonString <- function(json_string) { - # Load a `CppJson` object from string - bcf_json <- createCppJsonString(json_string) +saveBCFModelToJsonString <- function(object) { + # Convert to Json + jsonobj <- saveBCFModelToJson(object) - # Create and return the BCF object - bcf_object <- createBCFModelFromJson(bcf_json) + # Dump to string + return(jsonobj$return_json_string()) +} - return(bcf_object) +# Recover `binary_treatment` for legacy BCF JSON written before the field was +# serialized. Two sample-time invariants make this exact in most cases: +# * multivariate treatment is never binary, and +# * adaptive_coding is forced FALSE unless the treatment is binary, so +# adaptive_coding == TRUE implies binary_treatment == TRUE. +# A univariate, default-coded treatment is genuinely ambiguous from JSON alone +# (it would require the original Z_train), so we conservatively return FALSE. +# Reads directly from the JSON object so it does not depend on the order in +# which the calling load path populates model_params. +.inferBinaryTreatmentFromJson <- function(json_obj, has_field_fn) { + multivariate <- if (has_field_fn("multivariate_treatment")) { + json_obj$get_boolean("multivariate_treatment") + } else { + FALSE + } + if (isTRUE(multivariate)) { + return(FALSE) + } + adaptive <- if (has_field_fn("adaptive_coding")) { + json_obj$get_boolean("adaptive_coding") + } else { + FALSE + } + isTRUE(adaptive) } -#' @title Convert JSON List to BCF Model -#' @param json_object_list List of objects of type `CppJson` containing Json representation of a BCF model +# In-place v0 -> v1 migration for a BCF model envelope: positional forest keys -> +# named keys (forest_0 -> prognostic_forest, forest_1 -> treatment_forest, and when +# present forest_2 -> variance_forest). The mu/tau forests are always present. +.migrateBcfJsonV0ToV1 <- function(json_object, loaded_version) { + json_object$add_string("platform", inferPlatformV0(json_object, "R")) + json_object$rename_field( + "forest_0", + "prognostic_forest", + subfolder_name = "forests" + ) + json_object$rename_field( + "forest_1", + "treatment_forest", + subfolder_name = "forests" + ) + if (json_object$get_boolean_or_default("include_variance_forest", FALSE)) { + json_object$rename_field( + "forest_2", + "variance_forest", + subfolder_name = "forests" + ) + } + # R's legacy preprocessor key -> unified v1 key (no-op for Python v0 JSON, + # which already uses `covariate_preprocessor`). + json_object$rename_field("preprocessor_metadata", "covariate_preprocessor") + # Legacy parameter-trace field names -> canonical v1 names. All no-ops when the + # field is absent (e.g. non-adaptive-coding models have no b0/b1 samples). + json_object$rename_field("initial_sigma2", "sigma2_init") + json_object$rename_field( + "b_0_samples", + "b0_samples", + subfolder_name = "parameters" + ) + json_object$rename_field( + "b_1_samples", + "b1_samples", + subfolder_name = "parameters" + ) + # Relocate R's top-level rfx unique group ids into the random_effects subfolder + # (no-op for Python v0 JSON, which never wrote this field). + if (json_object$contains("rfx_unique_group_ids")) { + .rfx_uids <- json_object$get_string_vector("rfx_unique_group_ids") + json_object$add_string_vector( + "rfx_unique_group_ids", + .rfx_uids, + subfolder_name = "random_effects" + ) + json_object$erase_field("rfx_unique_group_ids") + } +} + +#' @description Create a BCFSamples object from JSON +#' @noRd +createBCFSamplesFromJson <- function(json) { + bcf_samples <- BCFSamples$new() + bcf_samples$from_json(json) + bcf_samples +} + +#' @title Convert JSON to BCF Model +#' @param json_object Object of type `CppJson` containing Json representation of a BCF model #' @export #' @rdname BCFSerialization -createBCFModelFromCombinedJson <- function(json_object_list) { +createBCFModelFromJson <- function(json_object) { # Initialize the BCF model output <- list() - # For scalar / preprocessing details which aren't sample-dependent, - # defer to the first json - json_object_default <- json_object_list[[1]] - # Version inference and presence-check helpers - .ver <- inferStochtreeJsonVersion(json_object_default) + .ver <- inferStochtreeJsonVersion(json_object) + resolveSchemaVersion(json_object, migrate = .migrateBcfJsonV0ToV1) + cross_platform <- enforceCrossPlatformGate(json_object, "R") has_field <- function(name) { - json_contains_field_cpp(json_object_default$json_ptr, name) + json_contains_field_cpp(json_object$json_ptr, name) } has_subfolder_field <- function(subfolder, name) { - json_contains_field_subfolder_cpp( - json_object_default$json_ptr, - subfolder, - name - ) - } - - # Unpack the forests - output[["forests_mu"]] <- loadForestContainerCombinedJson( - json_object_list, - "forest_0" - ) - output[["forests_tau"]] <- loadForestContainerCombinedJson( - json_object_list, - "forest_1" - ) - include_variance_forest <- json_object_default$get_boolean( - "include_variance_forest" - ) - if (include_variance_forest) { - output[["forests_variance"]] <- loadForestContainerCombinedJson( - json_object_list, - "forest_2" - ) - } - - # Unpack metadata - train_set_metadata = list() - train_set_metadata[["num_numeric_vars"]] <- json_object_default$get_scalar( - "num_numeric_vars" - ) - train_set_metadata[[ - "num_ordered_cat_vars" - ]] <- json_object_default$get_scalar("num_ordered_cat_vars") - train_set_metadata[[ - "num_unordered_cat_vars" - ]] <- json_object_default$get_scalar("num_unordered_cat_vars") - if (train_set_metadata[["num_numeric_vars"]] > 0) { - train_set_metadata[[ - "numeric_vars" - ]] <- json_object_default$get_string_vector("numeric_vars") - } - if (train_set_metadata[["num_ordered_cat_vars"]] > 0) { - train_set_metadata[[ - "ordered_cat_vars" - ]] <- json_object_default$get_string_vector("ordered_cat_vars") - train_set_metadata[[ - "ordered_unique_levels" - ]] <- json_object_default$get_string_list( - "ordered_unique_levels", - train_set_metadata[["ordered_cat_vars"]] - ) - } - if (train_set_metadata[["num_unordered_cat_vars"]] > 0) { - train_set_metadata[[ - "unordered_cat_vars" - ]] <- json_object_default$get_string_vector("unordered_cat_vars") - train_set_metadata[[ - "unordered_unique_levels" - ]] <- json_object_default$get_string_list( - "unordered_unique_levels", - train_set_metadata[["unordered_cat_vars"]] - ) + json_contains_field_subfolder_cpp(json_object$json_ptr, subfolder, name) } - output[["train_set_metadata"]] <- train_set_metadata # Unpack model params - model_params = list() - model_params[["outcome_scale"]] <- json_object_default$get_scalar( - "outcome_scale" - ) - model_params[["outcome_mean"]] <- json_object_default$get_scalar( - "outcome_mean" - ) - model_params[["standardize"]] <- json_object_default$get_boolean( - "standardize" + model_params <- list() + include_variance_forest <- json_object$get_boolean( + "include_variance_forest" ) - if (has_field("sigma2_init")) { - model_params[["initial_sigma2"]] <- json_object_default$get_scalar( - "sigma2_init" - ) - } else { - model_params[["initial_sigma2"]] <- json_object_default$get_scalar( - "initial_sigma2" - ) - warning(sprintf( - "JSON field 'initial_sigma2' is deprecated; please re-save the model to use 'sigma2_init' (inferred version: %s).", - .ver - )) - } - model_params[["sample_sigma2_global"]] <- json_object_default$get_boolean( + model_params[["outcome_scale"]] <- json_object$get_scalar("outcome_scale") + model_params[["outcome_mean"]] <- json_object$get_scalar("outcome_mean") + model_params[["standardize"]] <- json_object$get_boolean("standardize") + # Legacy `initial_sigma2` -> `sigma2_init` is handled by the v0 -> v1 migration + # (.migrateBcfJsonV0ToV1), so by this point the canonical key is always present. + model_params[["initial_sigma2"]] <- json_object$get_scalar("sigma2_init") + model_params[["sample_sigma2_global"]] <- json_object$get_boolean( "sample_sigma2_global" ) - model_params[["sample_sigma2_leaf_mu"]] <- json_object_default$get_boolean( + model_params[["sample_sigma2_leaf_mu"]] <- json_object$get_boolean( "sample_sigma2_leaf_mu" ) - model_params[["sample_sigma2_leaf_tau"]] <- json_object_default$get_boolean( + model_params[["sample_sigma2_leaf_tau"]] <- json_object$get_boolean( "sample_sigma2_leaf_tau" ) model_params[["include_variance_forest"]] <- include_variance_forest - model_params[["propensity_covariate"]] <- json_object_default$get_string( + model_params[["propensity_covariate"]] <- json_object$get_string( "propensity_covariate" ) - model_params[["has_rfx"]] <- json_object_default$get_boolean("has_rfx") + model_params[["has_rfx"]] <- json_object$get_boolean("has_rfx") if (has_field("has_rfx_basis")) { - model_params[["has_rfx_basis"]] <- json_object_default$get_boolean( - "has_rfx_basis" - ) - model_params[["num_rfx_basis"]] <- json_object_default$get_scalar( - "num_rfx_basis" - ) + model_params[["has_rfx_basis"]] <- json_object$get_boolean("has_rfx_basis") + model_params[["num_rfx_basis"]] <- json_object$get_scalar("num_rfx_basis") } else { model_params[["has_rfx_basis"]] <- FALSE model_params[["num_rfx_basis"]] <- 1 @@ -5474,37 +3428,16 @@ createBCFModelFromCombinedJson <- function(json_object_list) { .ver )) } - model_params[["num_covariates"]] <- json_object_default$get_scalar( - "num_covariates" - ) - if (has_field("num_chains")) { - model_params[["num_chains"]] <- json_object_default$get_scalar("num_chains") - } else { - model_params[["num_chains"]] <- 1 - warning(sprintf( - "Field 'num_chains' not found in BCF JSON (inferred version: %s). Defaulting to 1.", - .ver - )) - } - if (has_field("keep_every")) { - model_params[["keep_every"]] <- json_object_default$get_scalar("keep_every") - } else { - model_params[["keep_every"]] <- 1 - warning(sprintf( - "Field 'keep_every' not found in BCF JSON (inferred version: %s). Defaulting to 1.", - .ver - )) - } - model_params[["adaptive_coding"]] <- json_object_default$get_boolean( + model_params[["adaptive_coding"]] <- json_object$get_boolean( "adaptive_coding" ) if (has_field("binary_treatment")) { - model_params[["binary_treatment"]] <- json_object_default$get_boolean( + model_params[["binary_treatment"]] <- json_object$get_boolean( "binary_treatment" ) } else { model_params[["binary_treatment"]] <- .inferBinaryTreatmentFromJson( - json_object_default, + json_object, has_field ) warning(sprintf( @@ -5514,14 +3447,12 @@ createBCFModelFromCombinedJson <- function(json_object_list) { )) } if (has_field("treatment_dim")) { - model_params[["treatment_dim"]] <- json_object_default$get_scalar( - "treatment_dim" - ) + model_params[["treatment_dim"]] <- json_object$get_scalar("treatment_dim") } else { model_params[["treatment_dim"]] <- 1 if ( has_field("multivariate_treatment") && - isTRUE(json_object_default$get_boolean("multivariate_treatment")) + isTRUE(json_object$get_boolean("multivariate_treatment")) ) { warning(sprintf( "Field 'treatment_dim' not found in BCF JSON (inferred version: %s) for a multivariate-treatment model. Defaulting to 1.", @@ -5530,9 +3461,7 @@ createBCFModelFromCombinedJson <- function(json_object_list) { } } if (has_field("sample_tau_0")) { - model_params[["sample_tau_0"]] <- json_object_default$get_boolean( - "sample_tau_0" - ) + model_params[["sample_tau_0"]] <- json_object$get_boolean("sample_tau_0") } else { model_params[["sample_tau_0"]] <- FALSE warning(sprintf( @@ -5541,7 +3470,7 @@ createBCFModelFromCombinedJson <- function(json_object_list) { )) } if (has_field("multivariate_treatment")) { - model_params[["multivariate_treatment"]] <- json_object_default$get_boolean( + model_params[["multivariate_treatment"]] <- json_object$get_boolean( "multivariate_treatment" ) } else { @@ -5552,9 +3481,9 @@ createBCFModelFromCombinedJson <- function(json_object_list) { )) } if (has_field("internal_propensity_model")) { - model_params[[ + model_params[["internal_propensity_model"]] <- json_object$get_boolean( "internal_propensity_model" - ]] <- json_object_default$get_boolean("internal_propensity_model") + ) } else { model_params[["internal_propensity_model"]] <- FALSE warning(sprintf( @@ -5562,8 +3491,31 @@ createBCFModelFromCombinedJson <- function(json_object_list) { .ver )) } + model_params[["num_gfr"]] <- json_object$get_scalar("num_gfr") + model_params[["num_burnin"]] <- json_object$get_scalar("num_burnin") + model_params[["num_mcmc"]] <- json_object$get_scalar("num_mcmc") + model_params[["num_samples"]] <- json_object$get_scalar("num_samples") + model_params[["num_covariates"]] <- json_object$get_scalar("num_covariates") + if (has_field("num_chains")) { + model_params[["num_chains"]] <- json_object$get_scalar("num_chains") + } else { + model_params[["num_chains"]] <- 1 + warning(sprintf( + "Field 'num_chains' not found in BCF JSON (inferred version: %s). Defaulting to 1.", + .ver + )) + } + if (has_field("keep_every")) { + model_params[["keep_every"]] <- json_object$get_scalar("keep_every") + } else { + model_params[["keep_every"]] <- 1 + warning(sprintf( + "Field 'keep_every' not found in BCF JSON (inferred version: %s). Defaulting to 1.", + .ver + )) + } if (has_field("probit_outcome_model")) { - model_params[["probit_outcome_model"]] <- json_object_default$get_boolean( + model_params[["probit_outcome_model"]] <- json_object$get_boolean( "probit_outcome_model" ) } else { @@ -5574,14 +3526,8 @@ createBCFModelFromCombinedJson <- function(json_object_list) { )) } if (has_subfolder_field("outcome_model", "outcome")) { - outcome_model_outcome <- json_object_default$get_string( - "outcome", - "outcome_model" - ) - outcome_model_link <- json_object_default$get_string( - "link", - "outcome_model" - ) + outcome_model_outcome <- json_object$get_string("outcome", "outcome_model") + outcome_model_link <- json_object$get_string("link", "outcome_model") } else { outcome_model_outcome <- "continuous" outcome_model_link <- "identity" @@ -5595,7 +3541,7 @@ createBCFModelFromCombinedJson <- function(json_object_list) { link = outcome_model_link ) if (has_field("rfx_model_spec")) { - model_params[["rfx_model_spec"]] <- json_object_default$get_string( + model_params[["rfx_model_spec"]] <- json_object$get_string( "rfx_model_spec" ) } else { @@ -5607,179 +3553,39 @@ createBCFModelFromCombinedJson <- function(json_object_list) { )) } } - - # Combine values that are sample-specific - for (i in 1:length(json_object_list)) { - json_object <- json_object_list[[i]] - if (i == 1) { - model_params[["num_gfr"]] <- json_object$get_scalar("num_gfr") - model_params[["num_burnin"]] <- json_object$get_scalar("num_burnin") - model_params[["num_mcmc"]] <- json_object$get_scalar("num_mcmc") - model_params[["num_samples"]] <- json_object$get_scalar( - "num_samples" - ) - } else { - prev_json <- json_object_list[[i - 1]] - model_params[["num_gfr"]] <- model_params[["num_gfr"]] + - json_object$get_scalar("num_gfr") - model_params[["num_burnin"]] <- model_params[["num_burnin"]] + - json_object$get_scalar("num_burnin") - model_params[["num_mcmc"]] <- model_params[["num_mcmc"]] + - json_object$get_scalar("num_mcmc") - model_params[["num_samples"]] <- model_params[["num_samples"]] + - json_object$get_scalar("num_samples") - } - } output[["model_params"]] <- model_params - # Unpack sampled parameters - if (model_params[["sample_sigma2_global"]]) { - for (i in 1:length(json_object_list)) { - json_object <- json_object_list[[i]] - if (i == 1) { - output[["sigma2_global_samples"]] <- json_object$get_vector( - "sigma2_global_samples", - "parameters" - ) - } else { - output[["sigma2_global_samples"]] <- c( - output[["sigma2_global_samples"]], - json_object$get_vector( - "sigma2_global_samples", - "parameters" - ) - ) - } - } - } - if (model_params[["sample_sigma2_leaf_mu"]]) { - for (i in 1:length(json_object_list)) { - json_object <- json_object_list[[i]] - if (i == 1) { - output[["sigma2_leaf_mu_samples"]] <- json_object$get_vector( - "sigma2_leaf_mu_samples", - "parameters" - ) - } else { - output[["sigma2_leaf_mu_samples"]] <- c( - output[["sigma2_leaf_mu_samples"]], - json_object$get_vector( - "sigma2_leaf_mu_samples", - "parameters" - ) - ) - } - } - } - if (model_params[["sample_sigma2_leaf_tau"]]) { - for (i in 1:length(json_object_list)) { - json_object <- json_object_list[[i]] - if (i == 1) { - output[["sigma2_leaf_tau_samples"]] <- json_object$get_vector( - "sigma2_leaf_tau_samples", - "parameters" - ) - } else { - output[["sigma2_leaf_tau_samples"]] <- c( - output[["sigma2_leaf_tau_samples"]], - json_object$get_vector( - "sigma2_leaf_tau_samples", - "parameters" - ) - ) - } - } - } - if (model_params[["sample_sigma2_leaf_tau"]]) { - for (i in 1:length(json_object_list)) { - json_object <- json_object_list[[i]] - if (i == 1) { - output[["sigma2_leaf_tau_samples"]] <- json_object$get_vector( - "sigma2_leaf_tau_samples", - "parameters" - ) - } else { - output[["sigma2_leaf_tau_samples"]] <- c( - output[["sigma2_leaf_tau_samples"]], - json_object$get_vector( - "sigma2_leaf_tau_samples", - "parameters" - ) - ) - } - } - } - .b_use_new_names <- has_subfolder_field("parameters", "b1_samples") - if (model_params[["adaptive_coding"]]) { - if (!.b_use_new_names) { - warning(sprintf( - "JSON fields 'b_1_samples'/'b_0_samples' are deprecated; please re-save the model to use 'b1_samples'/'b0_samples' (inferred version: %s).", - .ver - )) - } - for (i in 1:length(json_object_list)) { - json_object <- json_object_list[[i]] - if (i == 1) { - output[["b_1_samples"]] <- json_object$get_vector( - if (.b_use_new_names) "b1_samples" else "b_1_samples", - "parameters" - ) - output[["b_0_samples"]] <- json_object$get_vector( - if (.b_use_new_names) "b0_samples" else "b_0_samples", - "parameters" - ) - } else { - output[["b_1_samples"]] <- c( - output[["b_1_samples"]], - json_object$get_vector( - if (.b_use_new_names) "b1_samples" else "b_1_samples", - "parameters" - ) - ) - output[["b_0_samples"]] <- c( - output[["b_0_samples"]], - json_object$get_vector( - if (.b_use_new_names) "b0_samples" else "b_0_samples", - "parameters" - ) - ) - } - } - } - if (model_params[["sample_tau_0"]]) { - tau_0_dim <- as.integer(json_object_default$get_scalar("tau_0_dim")) - for (i in 1:length(json_object_list)) { - json_object <- json_object_list[[i]] - tau_0_mat_i <- matrix( - json_object$get_vector("tau_0_samples", "parameters"), - nrow = tau_0_dim - ) - if (i == 1) { - output[["tau_0_samples"]] <- tau_0_mat_i - } else { - output[["tau_0_samples"]] <- cbind( - output[["tau_0_samples"]], - tau_0_mat_i - ) - } - } - } + # Unpack samples + output[["samples"]] <- createBCFSamplesFromJson(json_object) - # Unpack random effects + # Unpack random effects group IDs if (model_params[["has_rfx"]]) { - output[[ - "rfx_unique_group_ids" - ]] <- json_object_default$get_string_vector("rfx_unique_group_ids") - output[["rfx_samples"]] <- loadRandomEffectSamplesCombinedJson( - json_object_list, - 0 + output[["rfx_unique_group_ids"]] <- resolveRfxUniqueGroupIds( + json_object, + output[["samples"]]$materialize_rfx() + ) + } + + # Unpack propensity model (if it exists) + if (model_params[["internal_propensity_model"]]) { + bart_propensity_string <- json_object$get_string( + "bart_propensity_model" + ) + output[["bart_propensity_model"]] <- createBARTModelFromJsonString( + bart_propensity_string ) } # Unpack covariate preprocessor - if (has_field("preprocessor_metadata")) { - preprocessor_metadata_string <- json_object_default$get_string( - "preprocessor_metadata" + if (cross_platform) { + # Identity metadata for the cross-platform all-numeric path (gate enforced); + # the foreign native preprocessor is not reconstructed. + output[["train_set_metadata"]] <- buildIdentityPreprocessorMetadata( + json_object + ) + } else if (has_field("covariate_preprocessor")) { + preprocessor_metadata_string <- json_object$get_string( + "covariate_preprocessor" ) output[["train_set_metadata"]] <- createPreprocessorFromJsonString( preprocessor_metadata_string @@ -5787,7 +3593,7 @@ createBCFModelFromCombinedJson <- function(json_object_list) { } else { output[["train_set_metadata"]] <- NULL warning(sprintf( - "Field 'preprocessor_metadata' not found in BCF JSON (inferred version: %s). Preprocessor is unavailable; prediction may fail.", + "Field 'covariate_preprocessor' not found in BCF JSON (inferred version: %s). Preprocessor is unavailable; prediction may fail.", .ver )) } @@ -5796,41 +3602,52 @@ createBCFModelFromCombinedJson <- function(json_object_list) { return(output) } -#' @title Convert JSON String List to BCF Model -#' @param json_string_list List of JSON strings which can be parsed to objects of type `CppJson` containing Json representation of a BCF model +#' @title Convert JSON File to BCF Model +#' @param json_filename String of filepath, must end in ".json" #' @export #' @rdname BCFSerialization -createBCFModelFromCombinedJsonString <- function(json_string_list) { +createBCFModelFromJsonFile <- function(json_filename) { + # Load a `CppJson` object from file + bcf_json <- createCppJsonFile(json_filename) + + # Create and return the BCF object + bcf_object <- createBCFModelFromJson(bcf_json) + + return(bcf_object) +} + +#' @title Convert JSON String to BCF Model +#' @param json_string JSON string dump +#' @export +#' @rdname BCFSerialization +createBCFModelFromJsonString <- function(json_string) { + # Load a `CppJson` object from string + bcf_json <- createCppJsonString(json_string) + + # Create and return the BCF object + bcf_object <- createBCFModelFromJson(bcf_json) + + return(bcf_object) +} + +#' @title Convert JSON List to BCF Model +#' @param json_object_list List of objects of type `CppJson` containing Json representation of a BCF model +#' @export +#' @rdname BCFSerialization +createBCFModelFromCombinedJson <- function(json_object_list) { # Initialize the BCF model output <- list() - # Convert JSON strings - json_object_list <- list() - for (i in 1:length(json_string_list)) { - json_string <- json_string_list[[i]] - json_object_list[[i]] <- createCppJsonString(json_string) - # Add runtime check for separately serialized propensity models - # We don't support merging BCF models with independent propensity models - # this way at the moment - if ( - json_contains_field_cpp( - json_object_list[[i]]$json_ptr, - "internal_propensity_model" - ) && - json_object_list[[i]]$get_boolean("internal_propensity_model") - ) { - stop( - "Combining separate BCF models with cached internal propensity models is currently unsupported. To make this work, please first train a propensity model and then pass the propensities as data to the separate BCF models before sampling." - ) - } - } - # For scalar / preprocessing details which aren't sample-dependent, # defer to the first json json_object_default <- json_object_list[[1]] # Version inference and presence-check helpers .ver <- inferStochtreeJsonVersion(json_object_default) + for (.jo in json_object_list) { + resolveSchemaVersion(.jo, migrate = .migrateBcfJsonV0ToV1) + } + cross_platform <- enforceCrossPlatformGate(json_object_default, "R") has_field <- function(name) { json_contains_field_cpp(json_object_default$json_ptr, name) } @@ -5842,67 +3659,11 @@ createBCFModelFromCombinedJsonString <- function(json_string_list) { ) } - # Unpack the forests - output[["forests_mu"]] <- loadForestContainerCombinedJson( - json_object_list, - "forest_0" - ) - output[["forests_tau"]] <- loadForestContainerCombinedJson( - json_object_list, - "forest_1" - ) + # Unpack model params + model_params <- list() include_variance_forest <- json_object_default$get_boolean( "include_variance_forest" ) - if (include_variance_forest) { - output[["forests_variance"]] <- loadForestContainerCombinedJson( - json_object_list, - "forest_2" - ) - } - - # Unpack metadata - train_set_metadata = list() - train_set_metadata[["num_numeric_vars"]] <- json_object_default$get_scalar( - "num_numeric_vars" - ) - train_set_metadata[[ - "num_ordered_cat_vars" - ]] <- json_object_default$get_scalar("num_ordered_cat_vars") - train_set_metadata[[ - "num_unordered_cat_vars" - ]] <- json_object_default$get_scalar("num_unordered_cat_vars") - if (train_set_metadata[["num_numeric_vars"]] > 0) { - train_set_metadata[[ - "numeric_vars" - ]] <- json_object_default$get_string_vector("numeric_vars") - } - if (train_set_metadata[["num_ordered_cat_vars"]] > 0) { - train_set_metadata[[ - "ordered_cat_vars" - ]] <- json_object_default$get_string_vector("ordered_cat_vars") - train_set_metadata[[ - "ordered_unique_levels" - ]] <- json_object_default$get_string_list( - "ordered_unique_levels", - train_set_metadata[["ordered_cat_vars"]] - ) - } - if (train_set_metadata[["num_unordered_cat_vars"]] > 0) { - train_set_metadata[[ - "unordered_cat_vars" - ]] <- json_object_default$get_string_vector("unordered_cat_vars") - train_set_metadata[[ - "unordered_unique_levels" - ]] <- json_object_default$get_string_list( - "unordered_unique_levels", - train_set_metadata[["unordered_cat_vars"]] - ) - } - output[["train_set_metadata"]] <- train_set_metadata - - # Unpack model params - model_params = list() model_params[["outcome_scale"]] <- json_object_default$get_scalar( "outcome_scale" ) @@ -5912,19 +3673,11 @@ createBCFModelFromCombinedJsonString <- function(json_string_list) { model_params[["standardize"]] <- json_object_default$get_boolean( "standardize" ) - if (has_field("sigma2_init")) { - model_params[["initial_sigma2"]] <- json_object_default$get_scalar( - "sigma2_init" - ) - } else { - model_params[["initial_sigma2"]] <- json_object_default$get_scalar( - "initial_sigma2" - ) - warning(sprintf( - "JSON field 'initial_sigma2' is deprecated; please re-save the model to use 'sigma2_init' (inferred version: %s).", - .ver - )) - } + # Legacy `initial_sigma2` -> `sigma2_init` is handled by the v0 -> v1 migration + # (.migrateBcfJsonV0ToV1), so by this point the canonical key is always present. + model_params[["initial_sigma2"]] <- json_object_default$get_scalar( + "sigma2_init" + ) model_params[["sample_sigma2_global"]] <- json_object_default$get_boolean( "sample_sigma2_global" ) @@ -5975,17 +3728,6 @@ createBCFModelFromCombinedJsonString <- function(json_string_list) { .ver )) } - if (has_field("multivariate_treatment")) { - model_params[["multivariate_treatment"]] <- json_object_default$get_boolean( - "multivariate_treatment" - ) - } else { - model_params[["multivariate_treatment"]] <- FALSE - warning(sprintf( - "Field 'multivariate_treatment' not found in BCF JSON (inferred version: %s). Defaulting to FALSE.", - .ver - )) - } model_params[["adaptive_coding"]] <- json_object_default$get_boolean( "adaptive_coding" ) @@ -6031,6 +3773,17 @@ createBCFModelFromCombinedJsonString <- function(json_string_list) { .ver )) } + if (has_field("multivariate_treatment")) { + model_params[["multivariate_treatment"]] <- json_object_default$get_boolean( + "multivariate_treatment" + ) + } else { + model_params[["multivariate_treatment"]] <- FALSE + warning(sprintf( + "Field 'multivariate_treatment' not found in BCF JSON (inferred version: %s). Defaulting to FALSE.", + .ver + )) + } if (has_field("internal_propensity_model")) { model_params[[ "internal_propensity_model" @@ -6112,154 +3865,36 @@ createBCFModelFromCombinedJsonString <- function(json_string_list) { } output[["model_params"]] <- model_params - # Unpack sampled parameters - if (model_params[["sample_sigma2_global"]]) { - for (i in 1:length(json_object_list)) { - json_object <- json_object_list[[i]] - if (i == 1) { - output[["sigma2_global_samples"]] <- json_object$get_vector( - "sigma2_global_samples", - "parameters" - ) - } else { - output[["sigma2_global_samples"]] <- c( - output[["sigma2_global_samples"]], - json_object$get_vector( - "sigma2_global_samples", - "parameters" - ) - ) - } - } - } - if (model_params[["sample_sigma2_leaf_mu"]]) { - for (i in 1:length(json_object_list)) { - json_object <- json_object_list[[i]] - if (i == 1) { - output[["sigma2_leaf_mu_samples"]] <- json_object$get_vector( - "sigma2_leaf_mu_samples", - "parameters" - ) - } else { - output[["sigma2_leaf_mu_samples"]] <- c( - output[["sigma2_leaf_mu_samples"]], - json_object$get_vector( - "sigma2_leaf_mu_samples", - "parameters" - ) - ) - } - } - } - if (model_params[["sample_sigma2_leaf_tau"]]) { - for (i in 1:length(json_object_list)) { - json_object <- json_object_list[[i]] - if (i == 1) { - output[["sigma2_leaf_tau_samples"]] <- json_object$get_vector( - "sigma2_leaf_tau_samples", - "parameters" - ) - } else { - output[["sigma2_leaf_tau_samples"]] <- c( - output[["sigma2_leaf_tau_samples"]], - json_object$get_vector( - "sigma2_leaf_tau_samples", - "parameters" - ) - ) - } - } - } - if (model_params[["sample_sigma2_leaf_tau"]]) { - for (i in 1:length(json_object_list)) { - json_object <- json_object_list[[i]] - if (i == 1) { - output[["sigma2_leaf_tau_samples"]] <- json_object$get_vector( - "sigma2_leaf_tau_samples", - "parameters" - ) - } else { - output[["sigma2_leaf_tau_samples"]] <- c( - output[["sigma2_leaf_tau_samples"]], - json_object$get_vector( - "sigma2_leaf_tau_samples", - "parameters" - ) - ) - } - } - } - .b_use_new_names <- has_subfolder_field("parameters", "b1_samples") - if (model_params[["adaptive_coding"]]) { - if (!.b_use_new_names) { - warning(sprintf( - "JSON fields 'b_1_samples'/'b_0_samples' are deprecated; please re-save the model to use 'b1_samples'/'b0_samples' (inferred version: %s).", - .ver - )) - } - for (i in 1:length(json_object_list)) { - json_object <- json_object_list[[i]] - if (i == 1) { - output[["b_1_samples"]] <- json_object$get_vector( - if (.b_use_new_names) "b1_samples" else "b_1_samples", - "parameters" - ) - output[["b_0_samples"]] <- json_object$get_vector( - if (.b_use_new_names) "b0_samples" else "b_0_samples", - "parameters" - ) - } else { - output[["b_1_samples"]] <- c( - output[["b_1_samples"]], - json_object$get_vector( - if (.b_use_new_names) "b1_samples" else "b_1_samples", - "parameters" - ) - ) - output[["b_0_samples"]] <- c( - output[["b_0_samples"]], - json_object$get_vector( - if (.b_use_new_names) "b0_samples" else "b_0_samples", - "parameters" - ) - ) - } - } - } - if (model_params[["sample_tau_0"]]) { - tau_0_dim <- as.integer(json_object_default$get_scalar("tau_0_dim")) - for (i in 1:length(json_object_list)) { - json_object <- json_object_list[[i]] - tau_0_mat_i <- matrix( - json_object$get_vector("tau_0_samples", "parameters"), - nrow = tau_0_dim - ) - if (i == 1) { - output[["tau_0_samples"]] <- tau_0_mat_i - } else { - output[["tau_0_samples"]] <- cbind( - output[["tau_0_samples"]], - tau_0_mat_i - ) - } + # Unpack samples + for (i in 1:length(json_object_list)) { + json_object <- json_object_list[[i]] + if (i == 1) { + combined_samples <- createBCFSamplesFromJson(json_object) + } else { + additional_samples <- createBCFSamplesFromJson(json_object) + combined_samples$merge(additional_samples) } } + output[["samples"]] <- combined_samples - # Unpack random effects + # Unpack random effects group IDs if (model_params[["has_rfx"]]) { - output[[ - "rfx_unique_group_ids" - ]] <- json_object_default$get_string_vector("rfx_unique_group_ids") - output[["rfx_samples"]] <- loadRandomEffectSamplesCombinedJson( - json_object_list, - 0 + output[["rfx_unique_group_ids"]] <- resolveRfxUniqueGroupIds( + json_object_default, + output[["samples"]]$materialize_rfx() ) } # Unpack covariate preprocessor - if (has_field("preprocessor_metadata")) { + if (cross_platform) { + # Identity metadata for the cross-platform all-numeric path (gate enforced); + # the foreign native preprocessor is not reconstructed. + output[["train_set_metadata"]] <- buildIdentityPreprocessorMetadata( + json_object_default + ) + } else if (has_field("covariate_preprocessor")) { preprocessor_metadata_string <- json_object_default$get_string( - "preprocessor_metadata" + "covariate_preprocessor" ) output[["train_set_metadata"]] <- createPreprocessorFromJsonString( preprocessor_metadata_string @@ -6267,7 +3902,7 @@ createBCFModelFromCombinedJsonString <- function(json_string_list) { } else { output[["train_set_metadata"]] <- NULL warning(sprintf( - "Field 'preprocessor_metadata' not found in BCF JSON (inferred version: %s). Preprocessor is unavailable; prediction may fail.", + "Field 'covariate_preprocessor' not found in BCF JSON (inferred version: %s). Preprocessor is unavailable; prediction may fail.", .ver )) } @@ -6275,3 +3910,33 @@ createBCFModelFromCombinedJsonString <- function(json_string_list) { class(output) <- "bcfmodel" return(output) } + +#' @title Convert JSON String List to BCF Model +#' @param json_string_list List of JSON strings which can be parsed to objects of type `CppJson` containing Json representation of a BCF model +#' @export +#' @rdname BCFSerialization +createBCFModelFromCombinedJsonString <- function(json_string_list) { + # Convert JSON strings + json_object_list <- list() + for (i in 1:length(json_string_list)) { + json_string <- json_string_list[[i]] + json_object_list[[i]] <- createCppJsonString(json_string) + # Add runtime check for separately serialized propensity models + # We don't support merging BCF models with independent propensity models + # this way at the moment + if ( + json_contains_field_cpp( + json_object_list[[i]]$json_ptr, + "internal_propensity_model" + ) && + json_object_list[[i]]$get_boolean("internal_propensity_model") + ) { + stop( + "Combining separate BCF models with cached internal propensity models is currently unsupported. To make this work, please first train a propensity model and then pass the propensities as data to the separate BCF models before sampling." + ) + } + } + + # Create BCF model from list of JSON objects + createBCFModelFromCombinedJson(json_object_list) +} diff --git a/R/cpp11.R b/R/cpp11.R index c24913f0..c2c603e2 100644 --- a/R/cpp11.R +++ b/R/cpp11.R @@ -1,5 +1,21 @@ # Generated by cpp11: do not edit by hand +bart_sample_cpp <- function(samples, X_train, y_train, X_test, n_train, n_test, p, basis_train, basis_test, basis_dim, obs_weights_train, obs_weights_test, rfx_group_ids_train, rfx_group_ids_test, rfx_basis_train, rfx_basis_test, rfx_num_groups, rfx_basis_dim, num_gfr, num_burnin, keep_every, num_mcmc, num_chains, config_input) { + .Call(`_stochtree_bart_sample_cpp`, samples, X_train, y_train, X_test, n_train, n_test, p, basis_train, basis_test, basis_dim, obs_weights_train, obs_weights_test, rfx_group_ids_train, rfx_group_ids_test, rfx_basis_train, rfx_basis_test, rfx_num_groups, rfx_basis_dim, num_gfr, num_burnin, keep_every, num_mcmc, num_chains, config_input) +} + +bart_predict_cpp <- function(bart_samples_ptr, bart_model_metadata, X, leaf_basis, n, p, num_basis, obs_weights, rfx_group_ids, rfx_basis, rfx_num_groups, rfx_basis_dim, posterior, scale, predict_y_hat, predict_mean_forest, predict_variance_forest, predict_random_effects) { + .Call(`_stochtree_bart_predict_cpp`, bart_samples_ptr, bart_model_metadata, X, leaf_basis, n, p, num_basis, obs_weights, rfx_group_ids, rfx_basis, rfx_num_groups, rfx_basis_dim, posterior, scale, predict_y_hat, predict_mean_forest, predict_variance_forest, predict_random_effects) +} + +bcf_sample_cpp <- function(samples, X_train, Z_train, y_train, X_test, Z_test, n_train, n_test, p, treatment_dim, obs_weights_train, obs_weights_test, rfx_group_ids_train, rfx_group_ids_test, rfx_basis_train, rfx_basis_test, rfx_num_groups, rfx_basis_dim, num_gfr, num_burnin, keep_every, num_mcmc, num_chains, adaptive_coding, config_input) { + .Call(`_stochtree_bcf_sample_cpp`, samples, X_train, Z_train, y_train, X_test, Z_test, n_train, n_test, p, treatment_dim, obs_weights_train, obs_weights_test, rfx_group_ids_train, rfx_group_ids_test, rfx_basis_train, rfx_basis_test, rfx_num_groups, rfx_basis_dim, num_gfr, num_burnin, keep_every, num_mcmc, num_chains, adaptive_coding, config_input) +} + +bcf_predict_cpp <- function(bcf_samples_ptr, bcf_model_metadata, X, Z, n, p, treatment_dim, obs_weights, rfx_group_ids, rfx_basis, rfx_num_groups, rfx_basis_dim, posterior, scale, predict_y_hat, predict_mu_x, predict_tau_x, predict_prognostic_function, predict_cate, predict_conditional_variance, predict_random_effects) { + .Call(`_stochtree_bcf_predict_cpp`, bcf_samples_ptr, bcf_model_metadata, X, Z, n, p, treatment_dim, obs_weights, rfx_group_ids, rfx_basis, rfx_num_groups, rfx_basis_dim, posterior, scale, predict_y_hat, predict_mu_x, predict_tau_x, predict_prognostic_function, predict_cate, predict_conditional_variance, predict_random_effects) +} + create_forest_dataset_cpp <- function() { .Call(`_stochtree_create_forest_dataset_cpp`) } @@ -296,6 +312,302 @@ root_reset_rfx_tracker_cpp <- function(tracker, dataset, residual, rfx_model) { invisible(.Call(`_stochtree_root_reset_rfx_tracker_cpp`, tracker, dataset, residual, rfx_model)) } +bart_samples_cpp <- function() { + .Call(`_stochtree_bart_samples_cpp`) +} + +bart_samples_from_json_cpp <- function(json) { + .Call(`_stochtree_bart_samples_from_json_cpp`, json) +} + +append_bart_samples_to_json_cpp <- function(samples, json) { + invisible(.Call(`_stochtree_append_bart_samples_to_json_cpp`, samples, json)) +} + +bart_samples_num_samples_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_num_samples_cpp`, samples) +} + +bart_samples_y_bar_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_y_bar_cpp`, samples) +} + +bart_samples_y_std_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_y_std_cpp`, samples) +} + +bart_samples_has_yhat_train_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_has_yhat_train_cpp`, samples) +} + +bart_samples_has_mean_forest_predictions_train_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_has_mean_forest_predictions_train_cpp`, samples) +} + +bart_samples_has_variance_forest_predictions_train_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_has_variance_forest_predictions_train_cpp`, samples) +} + +bart_samples_has_rfx_predictions_train_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_has_rfx_predictions_train_cpp`, samples) +} + +bart_samples_has_yhat_test_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_has_yhat_test_cpp`, samples) +} + +bart_samples_has_mean_forest_predictions_test_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_has_mean_forest_predictions_test_cpp`, samples) +} + +bart_samples_has_variance_forest_predictions_test_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_has_variance_forest_predictions_test_cpp`, samples) +} + +bart_samples_has_rfx_predictions_test_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_has_rfx_predictions_test_cpp`, samples) +} + +bart_samples_has_cloglog_cutpoint_samples_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_has_cloglog_cutpoint_samples_cpp`, samples) +} + +bart_samples_has_global_var_samples_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_has_global_var_samples_cpp`, samples) +} + +bart_samples_has_leaf_scale_samples_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_has_leaf_scale_samples_cpp`, samples) +} + +bart_samples_has_mean_forest_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_has_mean_forest_cpp`, samples) +} + +bart_samples_has_rfx_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_has_rfx_cpp`, samples) +} + +bart_samples_has_variance_forest_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_has_variance_forest_cpp`, samples) +} + +bart_samples_mean_forest_predictions_train_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_mean_forest_predictions_train_cpp`, samples) +} + +bart_samples_variance_forest_predictions_train_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_variance_forest_predictions_train_cpp`, samples) +} + +bart_samples_rfx_predictions_train_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_rfx_predictions_train_cpp`, samples) +} + +bart_samples_yhat_train_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_yhat_train_cpp`, samples) +} + +bart_samples_mean_forest_predictions_test_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_mean_forest_predictions_test_cpp`, samples) +} + +bart_samples_variance_forest_predictions_test_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_variance_forest_predictions_test_cpp`, samples) +} + +bart_samples_rfx_predictions_test_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_rfx_predictions_test_cpp`, samples) +} + +bart_samples_yhat_test_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_yhat_test_cpp`, samples) +} + +bart_samples_global_var_samples_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_global_var_samples_cpp`, samples) +} + +bart_samples_leaf_scale_samples_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_leaf_scale_samples_cpp`, samples) +} + +bart_samples_cloglog_cutpoint_samples_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_cloglog_cutpoint_samples_cpp`, samples) +} + +bart_samples_materialize_mean_forest_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_materialize_mean_forest_cpp`, samples) +} + +bart_samples_materialize_variance_forest_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_materialize_variance_forest_cpp`, samples) +} + +bart_samples_materialize_rfx_container_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_materialize_rfx_container_cpp`, samples) +} + +bart_samples_materialize_rfx_label_mapper_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_materialize_rfx_label_mapper_cpp`, samples) +} + +bart_samples_mean_forest_ptr_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_mean_forest_ptr_cpp`, samples) +} + +bart_samples_variance_forest_ptr_cpp <- function(samples) { + .Call(`_stochtree_bart_samples_variance_forest_ptr_cpp`, samples) +} + +bart_samples_merge_cpp <- function(samples, other) { + invisible(.Call(`_stochtree_bart_samples_merge_cpp`, samples, other)) +} + +bcf_samples_cpp <- function() { + .Call(`_stochtree_bcf_samples_cpp`) +} + +bcf_samples_from_json_cpp <- function(json) { + .Call(`_stochtree_bcf_samples_from_json_cpp`, json) +} + +append_bcf_samples_to_json_cpp <- function(samples, json) { + invisible(.Call(`_stochtree_append_bcf_samples_to_json_cpp`, samples, json)) +} + +bcf_samples_num_samples_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_num_samples_cpp`, samples) +} + +bcf_samples_treatment_dim_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_treatment_dim_cpp`, samples) +} + +bcf_samples_y_bar_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_y_bar_cpp`, samples) +} + +bcf_samples_y_std_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_y_std_cpp`, samples) +} + +bcf_samples_has_mu_forest_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_has_mu_forest_cpp`, samples) +} + +bcf_samples_has_tau_forest_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_has_tau_forest_cpp`, samples) +} + +bcf_samples_has_variance_forest_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_has_variance_forest_cpp`, samples) +} + +bcf_samples_has_rfx_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_has_rfx_cpp`, samples) +} + +bcf_samples_global_var_samples_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_global_var_samples_cpp`, samples) +} + +bcf_samples_leaf_scale_mu_samples_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_leaf_scale_mu_samples_cpp`, samples) +} + +bcf_samples_leaf_scale_tau_samples_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_leaf_scale_tau_samples_cpp`, samples) +} + +bcf_samples_tau_0_samples_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_tau_0_samples_cpp`, samples) +} + +bcf_samples_b0_samples_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_b0_samples_cpp`, samples) +} + +bcf_samples_b1_samples_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_b1_samples_cpp`, samples) +} + +bcf_samples_yhat_train_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_yhat_train_cpp`, samples) +} + +bcf_samples_mu_forest_predictions_train_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_mu_forest_predictions_train_cpp`, samples) +} + +bcf_samples_tau_forest_predictions_train_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_tau_forest_predictions_train_cpp`, samples) +} + +bcf_samples_variance_forest_predictions_train_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_variance_forest_predictions_train_cpp`, samples) +} + +bcf_samples_rfx_predictions_train_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_rfx_predictions_train_cpp`, samples) +} + +bcf_samples_yhat_test_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_yhat_test_cpp`, samples) +} + +bcf_samples_mu_forest_predictions_test_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_mu_forest_predictions_test_cpp`, samples) +} + +bcf_samples_tau_forest_predictions_test_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_tau_forest_predictions_test_cpp`, samples) +} + +bcf_samples_variance_forest_predictions_test_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_variance_forest_predictions_test_cpp`, samples) +} + +bcf_samples_rfx_predictions_test_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_rfx_predictions_test_cpp`, samples) +} + +bcf_samples_materialize_mu_forest_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_materialize_mu_forest_cpp`, samples) +} + +bcf_samples_materialize_tau_forest_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_materialize_tau_forest_cpp`, samples) +} + +bcf_samples_materialize_variance_forest_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_materialize_variance_forest_cpp`, samples) +} + +bcf_samples_materialize_rfx_container_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_materialize_rfx_container_cpp`, samples) +} + +bcf_samples_materialize_rfx_label_mapper_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_materialize_rfx_label_mapper_cpp`, samples) +} + +bcf_samples_mu_forest_ptr_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_mu_forest_ptr_cpp`, samples) +} + +bcf_samples_tau_forest_ptr_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_tau_forest_ptr_cpp`, samples) +} + +bcf_samples_variance_forest_ptr_cpp <- function(samples) { + .Call(`_stochtree_bcf_samples_variance_forest_ptr_cpp`, samples) +} + +bcf_samples_merge_cpp <- function(samples, other) { + invisible(.Call(`_stochtree_bcf_samples_merge_cpp`, samples, other)) +} + sum_cpp <- function(x) { .Call(`_stochtree_sum_cpp`, x) } @@ -812,6 +1124,22 @@ json_contains_field_cpp <- function(json_ptr, field_name) { .Call(`_stochtree_json_contains_field_cpp`, json_ptr, field_name) } +json_erase_field_cpp <- function(json_ptr, field_name) { + invisible(.Call(`_stochtree_json_erase_field_cpp`, json_ptr, field_name)) +} + +json_erase_field_subfolder_cpp <- function(json_ptr, subfolder_name, field_name) { + invisible(.Call(`_stochtree_json_erase_field_subfolder_cpp`, json_ptr, subfolder_name, field_name)) +} + +json_rename_field_cpp <- function(json_ptr, old_name, new_name) { + invisible(.Call(`_stochtree_json_rename_field_cpp`, json_ptr, old_name, new_name)) +} + +json_rename_field_subfolder_cpp <- function(json_ptr, subfolder_name, old_name, new_name) { + invisible(.Call(`_stochtree_json_rename_field_subfolder_cpp`, json_ptr, subfolder_name, old_name, new_name)) +} + json_extract_double_subfolder_cpp <- function(json_ptr, subfolder_name, field_name) { .Call(`_stochtree_json_extract_double_subfolder_cpp`, json_ptr, subfolder_name, field_name) } @@ -868,8 +1196,8 @@ json_extract_string_vector_cpp <- function(json_ptr, field_name) { .Call(`_stochtree_json_extract_string_vector_cpp`, json_ptr, field_name) } -json_add_forest_cpp <- function(json_ptr, forest_samples) { - .Call(`_stochtree_json_add_forest_cpp`, json_ptr, forest_samples) +json_add_forest_cpp <- function(json_ptr, forest_samples, forest_label) { + .Call(`_stochtree_json_add_forest_cpp`, json_ptr, forest_samples, forest_label) } json_increment_rfx_count_cpp <- function(json_ptr) { @@ -884,6 +1212,10 @@ json_add_rfx_label_mapper_cpp <- function(json_ptr, label_mapper) { .Call(`_stochtree_json_add_rfx_label_mapper_cpp`, json_ptr, label_mapper) } +rfx_label_mapper_unique_group_ids_cpp <- function(label_mapper) { + .Call(`_stochtree_rfx_label_mapper_unique_group_ids_cpp`, label_mapper) +} + json_add_rfx_groupids_cpp <- function(json_ptr, groupids) { .Call(`_stochtree_json_add_rfx_groupids_cpp`, json_ptr, groupids) } diff --git a/R/forest.R b/R/forest.R index 40028a55..92b7bc06 100644 --- a/R/forest.R +++ b/R/forest.R @@ -203,7 +203,7 @@ ForestSamples <- R6::R6Class( #' @description #' Create a new `ForestContainer` object from a json object #' @param json_object Object of class `CppJson` - #' @param json_forest_label Label referring to a particular forest (i.e. "forest_0") in the overall json hierarchy + #' @param json_forest_label Label referring to a particular forest (i.e. "mean_forest") in the overall json hierarchy #' @return A new `ForestContainer` object. load_from_json = function(json_object, json_forest_label) { self$forest_container_ptr <- forest_container_from_json_cpp( @@ -215,7 +215,7 @@ ForestSamples <- R6::R6Class( #' @description #' Append to a `ForestContainer` object from a json object #' @param json_object Object of class `CppJson` - #' @param json_forest_label Label referring to a particular forest (i.e. "forest_0") in the overall json hierarchy + #' @param json_forest_label Label referring to a particular forest (i.e. "mean_forest") in the overall json hierarchy #' @return None append_from_json = function(json_object, json_forest_label) { forest_container_append_from_json_cpp( @@ -228,7 +228,7 @@ ForestSamples <- R6::R6Class( #' @description #' Create a new `ForestContainer` object from a json object #' @param json_string JSON string which parses into object of class `CppJson` - #' @param json_forest_label Label referring to a particular forest (i.e. "forest_0") in the overall json hierarchy + #' @param json_forest_label Label referring to a particular forest (i.e. "mean_forest") in the overall json hierarchy #' @return A new `ForestContainer` object. load_from_json_string = function(json_string, json_forest_label) { self$forest_container_ptr <- forest_container_from_json_string_cpp( @@ -240,7 +240,7 @@ ForestSamples <- R6::R6Class( #' @description #' Append to a `ForestContainer` object from a json object #' @param json_string JSON string which parses into object of class `CppJson` - #' @param json_forest_label Label referring to a particular forest (i.e. "forest_0") in the overall json hierarchy + #' @param json_forest_label Label referring to a particular forest (i.e. "mean_forest") in the overall json hierarchy #' @return None append_from_json_string = function(json_string, json_forest_label) { forest_container_append_from_json_string_cpp( diff --git a/R/generics.R b/R/generics.R index 83bc9ade..886ccb6f 100644 --- a/R/generics.R +++ b/R/generics.R @@ -43,3 +43,26 @@ getRandomEffectSamples <- function(object, ...) { extractParameter <- function(object, term) { UseMethod("extractParameter") } + +#' @title Extract Forest Samples Generic Function +#' @description +#' Generic function for extracting forest samples from a model object (BCF, BART, etc...) +#' +#' @param object Fitted model object from which to extract forest samples +#' @param term Name of the forest to extract (e.g., `"mean"`, `"variance"`, etc.) +#' @return Forest sample array +#' @export +#' +#' @examples +#' n <- 100 +#' p <- 10 +#' X <- matrix(runif(n*p), ncol = p) +#' rfx_group_ids <- sample(1:2, size = n, replace = TRUE) +#' rfx_basis <- rep(1.0, n) +#' y <- (-5 + 10*(X[,1] > 0.5)) + (-2*(rfx_group_ids==1)+2*(rfx_group_ids==2)) + rnorm(n) +#' bart_model <- bart(X_train=X, y_train=y, rfx_group_ids_train=rfx_group_ids, +#' rfx_basis_train = rfx_basis, num_gfr=0, num_mcmc=10) +#' mean_forest <- extractForest(bart_model, "mean") +extractForest <- function(object, term) { + UseMethod("extractForest") +} diff --git a/R/kernel.R b/R/kernel.R index 82588b29..46503582 100644 --- a/R/kernel.R +++ b/R/kernel.R @@ -100,28 +100,28 @@ computeForestLeafIndices <- function( if (!model_object$model_params$include_mean_forest) { stop("Mean forest was not sampled in the bart model provided") } - forest_container <- model_object$mean_forests + forest_container <- model_object$samples$mean_forest_view() } else if (forest_type == "variance") { if (!model_object$model_params$include_variance_forest) { stop( "Variance forest was not sampled in the bart model provided" ) } - forest_container <- model_object$variance_forests + forest_container <- model_object$samples$variance_forest_view() } } else if (model_type == "bcf") { stopifnot(forest_type %in% c("prognostic", "treatment", "variance")) if (forest_type == "prognostic") { - forest_container <- model_object$forests_mu + forest_container <- model_object$samples$mu_forest_view() } else if (forest_type == "treatment") { - forest_container <- model_object$forests_tau + forest_container <- model_object$samples$tau_forest_view() } else if (forest_type == "variance") { if (!model_object$model_params$include_variance_forest) { stop( "Variance forest was not sampled in the bcf model provided" ) } - forest_container <- model_object$variance_forests + forest_container <- model_object$samples$variance_forest_view() } } else { forest_container <- model_object @@ -229,7 +229,7 @@ computeForestLeafVariances <- function( "Leaf scale parameter was not sampled for the mean forest in the bart model provided" ) } - leaf_scale_vector <- model_object$sigma2_leaf_samples + leaf_scale_vector <- model_object$samples$leaf_scale_samples() } else if (forest_type == "variance") { if (!model_object$model_params$include_variance_forest) { stop( @@ -248,14 +248,14 @@ computeForestLeafVariances <- function( "Leaf scale parameter was not sampled for the prognostic forest in the bcf model provided" ) } - leaf_scale_vector <- model_object$sigma2_leaf_mu_samples + leaf_scale_vector <- model_object$samples$leaf_scale_mu_samples() } else if (forest_type == "treatment") { if (!model_object$model_params$sample_sigma2_leaf_tau) { stop( "Leaf scale parameter was not sampled for the treatment effect forest in the bcf model provided" ) } - leaf_scale_vector <- model_object$sigma2_leaf_tau_samples + leaf_scale_vector <- model_object$samples$leaf_scale_tau_samples() } else if (forest_type == "variance") { if (!model_object$model_params$include_variance_forest) { stop( @@ -331,28 +331,28 @@ computeForestMaxLeafIndex <- function( if (!model_object$model_params$include_mean_forest) { stop("Mean forest was not sampled in the bart model provided") } - forest_container <- model_object$mean_forests + forest_container <- model_object$samples$mean_forest_view() } else if (forest_type == "variance") { if (!model_object$model_params$include_variance_forest) { stop( "Variance forest was not sampled in the bart model provided" ) } - forest_container <- model_object$variance_forests + forest_container <- model_object$samples$variance_forest_view() } } else if (model_type == "bcf") { stopifnot(forest_type %in% c("prognostic", "treatment", "variance")) if (forest_type == "prognostic") { - forest_container <- model_object$forests_mu + forest_container <- model_object$samples$mu_forest_view() } else if (forest_type == "treatment") { - forest_container <- model_object$forests_tau + forest_container <- model_object$samples$tau_forest_view() } else if (forest_type == "variance") { if (!model_object$model_params$include_variance_forest) { stop( "Variance forest was not sampled in the bcf model provided" ) } - forest_container <- model_object$variance_forests + forest_container <- model_object$samples$variance_forest_view() } } else { forest_container <- model_object diff --git a/R/posterior_transformation.R b/R/posterior_transformation.R index 8f4ab8c8..c440ea6d 100644 --- a/R/posterior_transformation.R +++ b/R/posterior_transformation.R @@ -658,7 +658,7 @@ sampleBCFPosteriorPredictive <- function( if (samples_global_variance) { ppd_variance <- matrix( rep( - model_object$sigma2_global_samples, + model_object$samples$global_var_samples(), each = num_observations ), nrow = num_observations @@ -834,7 +834,7 @@ sampleBARTPosteriorPredictive <- function( if (samples_global_variance) { ppd_variance <- matrix( rep( - model_object$sigma2_global_samples, + model_object$samples$global_var_samples(), each = num_observations ), nrow = num_observations @@ -973,9 +973,9 @@ posterior_predictive_heuristic_multiplier <- function( #' #' The treatment effect terms follow a three-level hierarchy: #' \itemize{ -#' \item `"tau"` returns `tau_0 + tau(X)`: the parametric treatment intercept (if sampled) plus the treatment forest. This matches `model$tau_hat_train` / `model$tau_hat_test`. +#' \item `"tau"` returns `tau_0 + tau(X)`: the parametric treatment intercept (if sampled) plus the treatment forest. This matches `extractParameter(model, "tau_hat_train")` / `extractParameter(model, "tau_hat_test")`. #' \item `"cate"` additionally folds in the random slope on treatment when random effects are fit with `rfx_model_spec = "intercept_plus_treatment"`; otherwise it is identical to `"tau"`. -#' \item The raw forest-only component (without `tau_0`) is not directly returned by this method; use `model$forests_tau` to access it. +#' \item The raw forest-only component (without `tau_0`) is not directly returned by this method; extract the treatment forest with `model$samples$materialize_tau_forest()` to access it. #' } #' #' Similarly for the prognostic term: `"mu"` returns the prognostic forest only, while `"prognostic_function"` additionally folds in the random intercept when `rfx_model_spec` is `"intercept_only"` or `"intercept_plus_treatment"`; otherwise the two are identical. diff --git a/R/samples.R b/R/samples.R new file mode 100644 index 00000000..7596b8b0 --- /dev/null +++ b/R/samples.R @@ -0,0 +1,473 @@ +# R6 wrapper around the C++ BARTSamples object + +#' @description +#' Container holding a sampled BART model's forests and parameter traces as a single C++ object. +#' @noRd +BARTSamples <- R6::R6Class( + classname = "BARTSamples", + cloneable = FALSE, + public = list( + #' @field samples_ptr External pointer to a C++ BARTSamples object + samples_ptr = NULL, + + #' @description + #' Initialize an empty BARTSamples object in C++ and wrap an external pointer to the object. + initialize = function() { + self$samples_ptr <- bart_samples_cpp() + }, + + #' @description + #' Initialize a BARTSamples object from JSON and wrap an external pointer to the object. + from_json = function(json) { + self$samples_ptr <- bart_samples_from_json_cpp(json$json_ptr) + }, + + #' @description + #' Convert a BARTSamples object to JSON and include it in a `CppJson` object wrapping a C++ JSON representation. + append_to_json = function(json) { + append_bart_samples_to_json_cpp(self$samples_ptr, json$json_ptr) + }, + + #' @description Number of retained posterior samples. + num_samples = function() bart_samples_num_samples_cpp(self$samples_ptr), + + #' @description Outcome mean used for standardization. + y_bar = function() bart_samples_y_bar_cpp(self$samples_ptr), + + #' @description Outcome standard deviation used for standardization. + y_std = function() bart_samples_y_std_cpp(self$samples_ptr), + + #' @description Whether global error scale samples are present. + has_global_var_samples = function() { + bart_samples_has_global_var_samples_cpp(self$samples_ptr) + }, + + #' @description Whether leaf scale samples are present. + has_leaf_scale_samples = function() { + bart_samples_has_leaf_scale_samples_cpp(self$samples_ptr) + }, + + #' @description Whether outcome predictions for the training set are present. + has_yhat_train = function() { + bart_samples_has_yhat_train_cpp(self$samples_ptr) + }, + + #' @description Whether mean forest predictions for the training set are present. + has_mean_forest_predictions_train = function() { + bart_samples_has_mean_forest_predictions_train_cpp(self$samples_ptr) + }, + + #' @description Whether variance forest predictions for the training set are present. + has_variance_forest_predictions_train = function() { + bart_samples_has_variance_forest_predictions_train_cpp(self$samples_ptr) + }, + + #' @description Whether outcome predictions for the test set are present. + has_yhat_test = function() { + bart_samples_has_yhat_test_cpp(self$samples_ptr) + }, + + #' @description Whether mean forest predictions for the test set are present. + has_mean_forest_predictions_test = function() { + bart_samples_has_mean_forest_predictions_test_cpp(self$samples_ptr) + }, + + #' @description Whether variance forest predictions for the test set are present. + has_variance_forest_predictions_test = function() { + bart_samples_has_variance_forest_predictions_test_cpp(self$samples_ptr) + }, + + #' @description Whether a mean forest is present. + has_mean_forest = function() { + bart_samples_has_mean_forest_cpp(self$samples_ptr) + }, + + #' @description Whether a variance forest is present. + has_variance_forest = function() { + bart_samples_has_variance_forest_cpp(self$samples_ptr) + }, + + #' @description Whether random effects are present. + has_rfx = function() { + bart_samples_has_rfx_cpp(self$samples_ptr) + }, + + #' @description Global error variance samples (length `num_samples`, or empty). + global_var_samples = function() { + bart_samples_global_var_samples_cpp(self$samples_ptr) + }, + + #' @description Leaf scale samples (length `num_samples`, or empty). + leaf_scale_samples = function() { + bart_samples_leaf_scale_samples_cpp(self$samples_ptr) + }, + + #' @description Whether cloglog cutpoint samples are present. + has_cloglog_cutpoint_samples = function() { + bart_samples_has_cloglog_cutpoint_samples_cpp(self$samples_ptr) + }, + + #' @description Cloglog ordinal cutpoint samples (`(num_classes - 1)` x `num_samples`), or NULL. + cloglog_cutpoint_samples = function() { + if (!self$has_cloglog_cutpoint_samples()) { + return(NULL) + } + bart_samples_cloglog_cutpoint_samples_cpp(self$samples_ptr) + }, + + #' @description Mean forest predictions for the training set (length `num_samples` * `num_train`, or empty). + y_hat_train = function() { + bart_samples_yhat_train_cpp(self$samples_ptr) + }, + + #' @description Mean forest predictions for the training set (length `num_samples` * `num_train`, or empty). + mean_forest_predictions_train = function() { + bart_samples_mean_forest_predictions_train_cpp(self$samples_ptr) + }, + + #' @description Variance forest predictions for the training set (length `num_samples` * `num_train`, or empty). + variance_forest_predictions_train = function() { + bart_samples_variance_forest_predictions_train_cpp(self$samples_ptr) + }, + + #' @description Random effects predictions for the training set (length `num_samples` * `num_train`, or empty). + rfx_predictions_train = function() { + bart_samples_rfx_predictions_train_cpp(self$samples_ptr) + }, + + #' @description Mean forest predictions for the test set (length `num_samples` * `num_test`, or empty). + y_hat_test = function() { + bart_samples_yhat_test_cpp(self$samples_ptr) + }, + + #' @description Mean forest predictions for the test set (length `num_samples` * `num_test`, or empty). + mean_forest_predictions_test = function() { + bart_samples_mean_forest_predictions_test_cpp(self$samples_ptr) + }, + + #' @description Variance forest predictions for the test set (length `num_samples` * `num_test`, or empty). + variance_forest_predictions_test = function() { + bart_samples_variance_forest_predictions_test_cpp(self$samples_ptr) + }, + + #' @description Random effects predictions for the test set (length `num_samples` * `num_test`, or empty). + rfx_predictions_test = function() { + bart_samples_rfx_predictions_test_cpp(self$samples_ptr) + }, + + #' @description Materialize a standalone deep copy of the mean forest as a `ForestSamples` + #' (or NULL if absent). + materialize_mean_forest = function() { + if (!self$has_mean_forest()) { + return(NULL) + } + fc <- ForestSamples$new(0, 1, FALSE, FALSE) + fc$forest_container_ptr <- bart_samples_materialize_mean_forest_cpp( + self$samples_ptr + ) + fc + }, + + #' @description Materialize a standalone deep copy of the variance forest as a `ForestSamples` + #' (or NULL if absent). + materialize_variance_forest = function() { + if (!self$has_variance_forest()) { + return(NULL) + } + fc <- ForestSamples$new(0, 1, FALSE, FALSE) + fc$forest_container_ptr <- bart_samples_materialize_variance_forest_cpp( + self$samples_ptr + ) + fc + }, + + #' @description Materialize a standalone deep copy of the random effects samples as a `RandomEffectSamples` + #' (or NULL if absent). + materialize_rfx = function() { + if (!self$has_rfx()) { + return(NULL) + } + fc <- RandomEffectSamples$new() + fc$rfx_container_ptr <- bart_samples_materialize_rfx_container_cpp( + self$samples_ptr + ) + fc$label_mapper_ptr <- bart_samples_materialize_rfx_label_mapper_cpp( + self$samples_ptr + ) + fc + }, + + #' @description Borrowed (non-owning) external pointer to the mean forest container, for + #' read-through prediction. Must not outlive this object. + mean_forest_ptr = function() { + bart_samples_mean_forest_ptr_cpp(self$samples_ptr) + }, + + #' @description Borrowed (non-owning) external pointer to the variance forest container. + variance_forest_ptr = function() { + bart_samples_variance_forest_ptr_cpp(self$samples_ptr) + }, + + #' @description Non-owning `ForestSamples` view over the mean forest (borrowed pointer, + #' no deep copy) for internal read-only consumers (serialization, kernels). NULL if absent. + mean_forest_view = function() { + if (!self$has_mean_forest()) { + return(NULL) + } + fc <- ForestSamples$new(0, 1, FALSE, FALSE) + fc$forest_container_ptr <- self$mean_forest_ptr() + fc + }, + + #' @description Non-owning `ForestSamples` view over the variance forest. NULL if absent. + variance_forest_view = function() { + if (!self$has_variance_forest()) { + return(NULL) + } + fc <- ForestSamples$new(0, 1, FALSE, FALSE) + fc$forest_container_ptr <- self$variance_forest_ptr() + fc + }, + + #' @description Append another chain's draws onto this one (multi-chain combine). + #' @param other Another `BARTSamples` with matching structure/standardization. + merge = function(other) { + bart_samples_merge_cpp(self$samples_ptr, other$samples_ptr) + } + ) +) + +#' @description +#' Container holding a sampled BCF model's forests and parameter traces as a single C++ object. +#' @noRd +BCFSamples <- R6::R6Class( + classname = "BCFSamples", + cloneable = FALSE, + public = list( + #' @field samples_ptr External pointer to a C++ BCFSamples object + samples_ptr = NULL, + + #' @description + #' Initialize an empty BCFSamples object in C++ and wrap an external pointer to the object. + initialize = function() { + self$samples_ptr <- bcf_samples_cpp() + }, + + #' @description + #' Initialize a BCFSamples object from JSON and wrap an external pointer to the object. + from_json = function(json) { + self$samples_ptr <- bcf_samples_from_json_cpp(json$json_ptr) + }, + + #' @description + #' Convert a BCFSamples object to JSON and return a `CppJson` object wrapping the C++ JSON representation. + append_to_json = function(json) { + append_bcf_samples_to_json_cpp(self$samples_ptr, json$json_ptr) + }, + + #' @description Number of retained posterior samples. + num_samples = function() bcf_samples_num_samples_cpp(self$samples_ptr), + + #' @description Treatment dimension. + treatment_dim = function() bcf_samples_treatment_dim_cpp(self$samples_ptr), + + #' @description Outcome mean used for standardization. + y_bar = function() bcf_samples_y_bar_cpp(self$samples_ptr), + + #' @description Outcome standard deviation used for standardization. + y_std = function() bcf_samples_y_std_cpp(self$samples_ptr), + + #' @description Whether a prognostic forest is present. + has_mu_forest = function() bcf_samples_has_mu_forest_cpp(self$samples_ptr), + + #' @description Whether a treatment forest is present. + has_tau_forest = function() { + bcf_samples_has_tau_forest_cpp(self$samples_ptr) + }, + + #' @description Whether a variance forest is present. + has_variance_forest = function() { + bcf_samples_has_variance_forest_cpp(self$samples_ptr) + }, + + #' @description Whether random effects are present. + has_rfx = function() { + bcf_samples_has_rfx_cpp(self$samples_ptr) + }, + + #' @description Global error variance samples. + global_var_samples = function() { + bcf_samples_global_var_samples_cpp(self$samples_ptr) + }, + + #' @description Prognostic leaf scale samples. + leaf_scale_mu_samples = function() { + bcf_samples_leaf_scale_mu_samples_cpp(self$samples_ptr) + }, + + #' @description Treatment leaf scale samples. + leaf_scale_tau_samples = function() { + bcf_samples_leaf_scale_tau_samples_cpp(self$samples_ptr) + }, + + #' @description Adaptive-coding b0 samples. + b0_samples = function() bcf_samples_b0_samples_cpp(self$samples_ptr), + + #' @description Adaptive-coding b1 samples. + b1_samples = function() bcf_samples_b1_samples_cpp(self$samples_ptr), + + #' @description Treatment intercept (tau_0) samples (flat). + tau_0_samples = function() bcf_samples_tau_0_samples_cpp(self$samples_ptr), + + #' @description Mean forest predictions for the training set (length `num_samples` * `num_train`, or empty). + y_hat_train = function() { + bcf_samples_yhat_train_cpp(self$samples_ptr) + }, + + #' @description Prognostic forest predictions for the training set (length `num_samples` * `num_train`, or empty). + mu_forest_predictions_train = function() { + bcf_samples_mu_forest_predictions_train_cpp(self$samples_ptr) + }, + + #' @description Treatment effect forest predictions for the training set (length `num_samples` * `num_treatment` * `num_train`, or `num_samples` * `num_train` if `num_treatment` <= 1, or empty). + tau_forest_predictions_train = function() { + bcf_samples_tau_forest_predictions_train_cpp(self$samples_ptr) + }, + + #' @description Variance forest predictions for the training set (length `num_samples` * `num_train`, or empty). + variance_forest_predictions_train = function() { + bcf_samples_variance_forest_predictions_train_cpp(self$samples_ptr) + }, + + #' @description Random effects predictions for the training set (length `num_samples` * `num_train`, or empty). + rfx_predictions_train = function() { + bcf_samples_rfx_predictions_train_cpp(self$samples_ptr) + }, + + #' @description Mean forest predictions for the test set (length `num_samples` * `num_test`, or empty). + y_hat_test = function() { + bcf_samples_yhat_test_cpp(self$samples_ptr) + }, + + #' @description Prognostic forest predictions for the test set (length `num_samples` * `num_test`, or empty). + mu_forest_predictions_test = function() { + bcf_samples_mu_forest_predictions_test_cpp(self$samples_ptr) + }, + + #' @description Treatment effect forest predictions for the test set (length `num_samples` * `num_treatment` * `num_test`, or `num_samples` * `num_test` if `num_treatment` <= 1, or empty). + tau_forest_predictions_test = function() { + bcf_samples_tau_forest_predictions_test_cpp(self$samples_ptr) + }, + + #' @description Variance forest predictions for the test set (length `num_samples` * `num_test`, or empty). + variance_forest_predictions_test = function() { + bcf_samples_variance_forest_predictions_test_cpp(self$samples_ptr) + }, + + #' @description Random effects predictions for the test set (length `num_samples` * `num_test`, or empty). + rfx_predictions_test = function() { + bcf_samples_rfx_predictions_test_cpp(self$samples_ptr) + }, + + #' @description Materialize a deep copy of the prognostic forest as a `ForestSamples`. + materialize_mu_forest = function() { + if (!self$has_mu_forest()) { + return(NULL) + } + fc <- ForestSamples$new(0, 1, FALSE, FALSE) + fc$forest_container_ptr <- bcf_samples_materialize_mu_forest_cpp( + self$samples_ptr + ) + fc + }, + + #' @description Materialize a deep copy of the treatment forest as a `ForestSamples`. + materialize_tau_forest = function() { + if (!self$has_tau_forest()) { + return(NULL) + } + fc <- ForestSamples$new(0, 1, FALSE, FALSE) + fc$forest_container_ptr <- bcf_samples_materialize_tau_forest_cpp( + self$samples_ptr + ) + fc + }, + + #' @description Materialize a deep copy of the variance forest as a `ForestSamples` (or NULL). + materialize_variance_forest = function() { + if (!self$has_variance_forest()) { + return(NULL) + } + fc <- ForestSamples$new(0, 1, FALSE, FALSE) + fc$forest_container_ptr <- bcf_samples_materialize_variance_forest_cpp( + self$samples_ptr + ) + fc + }, + + #' @description Materialize a standalone deep copy of the random effects samples as a `RandomEffectSamples` + #' (or NULL if absent). + materialize_rfx = function() { + if (!self$has_rfx()) { + return(NULL) + } + fc <- RandomEffectSamples$new() + fc$rfx_container_ptr <- bcf_samples_materialize_rfx_container_cpp( + self$samples_ptr + ) + fc$label_mapper_ptr <- bcf_samples_materialize_rfx_label_mapper_cpp( + self$samples_ptr + ) + fc + }, + + #' @description Borrowed (non-owning) external pointer to the prognostic forest container. + mu_forest_ptr = function() bcf_samples_mu_forest_ptr_cpp(self$samples_ptr), + + #' @description Borrowed (non-owning) external pointer to the treatment forest container. + tau_forest_ptr = function() { + bcf_samples_tau_forest_ptr_cpp(self$samples_ptr) + }, + + #' @description Borrowed (non-owning) external pointer to the variance forest container. + variance_forest_ptr = function() { + bcf_samples_variance_forest_ptr_cpp(self$samples_ptr) + }, + + #' @description Non-owning `ForestSamples` view over the prognostic forest. NULL if absent. + mu_forest_view = function() { + if (!self$has_mu_forest()) { + return(NULL) + } + fc <- ForestSamples$new(0, 1, FALSE, FALSE) + fc$forest_container_ptr <- self$mu_forest_ptr() + fc + }, + + #' @description Non-owning `ForestSamples` view over the treatment forest. NULL if absent. + tau_forest_view = function() { + if (!self$has_tau_forest()) { + return(NULL) + } + fc <- ForestSamples$new(0, 1, FALSE, FALSE) + fc$forest_container_ptr <- self$tau_forest_ptr() + fc + }, + + #' @description Non-owning `ForestSamples` view over the variance forest. NULL if absent. + variance_forest_view = function() { + if (!self$has_variance_forest()) { + return(NULL) + } + fc <- ForestSamples$new(0, 1, FALSE, FALSE) + fc$forest_container_ptr <- self$variance_forest_ptr() + fc + }, + + #' @description Append another chain's draws onto this one (multi-chain combine). + #' @param other Another `BCFSamples` with matching structure/standardization. + merge = function(other) { + bcf_samples_merge_cpp(self$samples_ptr, other$samples_ptr) + } + ) +) diff --git a/R/serialization.R b/R/serialization.R index 05232081..91a99a23 100644 --- a/R/serialization.R +++ b/R/serialization.R @@ -1,3 +1,180 @@ +# ----------------------------------------------------------------------------- +# Serialized model envelope schema version (RFC 0005). +# +# Integer identifying the structure of the serialized BART/BCF JSON envelope. +# Bumped ONLY on a breaking change (rename / remove / re-type a field, change a +# field's meaning, or change a structural convention). Additive, safely-defaulted +# fields do NOT bump it -- they are handled by "augmentation" on read (see the +# defaults registry below and the `get_*_or_default` methods on `CppJson`). +# +# Kept in sync with the Python `SCHEMA_VERSION` (stochtree/serialization.py). The +# two are independent constants by design (each language owns its serde); their +# agreement is enforced by the cross-platform golden fixtures, not by sharing a value. +STOCHTREE_SCHEMA_VERSION <- 1L + +# ----------------------------------------------------------------------------- +# Augmentation defaults registry (schema_version = 1) +# +# Every OPTIONAL envelope field that may be ABSENT from a model written by an +# earlier release at the SAME schema_version must be read with a default that +# reproduces that model's pre-field behavior. This list is the single source of +# truth for "which fields need defaulting"; read those fields via the `CppJson` +# `get_scalar_or_default` / `get_boolean_or_default` / `get_string_or_default` +# helpers (never a bare required getter). +# +# field default (behavior when absent) +# "outcome" "continuous" +# "link" "identity" +# "multivariate_treatment" FALSE +# "internal_propensity_model" FALSE +# "has_rfx_basis" FALSE +# ... (add a row whenever you add an additive field) +# +# If a new field has NO behavior-preserving default, it is NOT additive: bump +# STOCHTREE_SCHEMA_VERSION and add a migrate_v{N}_to_v{N+1} step instead. +# ----------------------------------------------------------------------------- + +# Read the envelope schema_version and enforce the RFC 0005 reader rules. Returns +# the loaded version (0 = legacy / absent). Errors if the model was written by a +# newer stochtree than this installation supports. A value below the current version +# is the hook for the migration ladder (no rungs exist yet at version 1; legacy v0 +# models are handled by field-presence default-filling during parsing). +resolveSchemaVersion <- function(json_object, migrate = NULL) { + loaded <- json_object$get_integer_or_default("schema_version", 0L) + if (loaded > STOCHTREE_SCHEMA_VERSION) { + stop(sprintf( + paste0( + "This model was serialized with schema_version=%d, but this installation of ", + "stochtree supports up to schema_version=%d. Please upgrade stochtree to load it." + ), + loaded, + STOCHTREE_SCHEMA_VERSION + )) + } + if (loaded < STOCHTREE_SCHEMA_VERSION && !is.null(migrate)) { + # Dispatch owns the migration: upgrade the JSON in place to the current + # schema before the (v1-only) parser runs. + migrate(json_object, loaded) + } + loaded +} + +# Infer the writer platform of a legacy (v0) envelope from structural +# fingerprints. R wrote `preprocessor_metadata` and a top-level +# `rfx_unique_group_ids`; Python wrote `covariate_preprocessor`. Called before +# the v0 -> v1 renames, so the legacy keys are still present. Falls back to +# `default` (the loading platform) when no fingerprint is decisive. +inferPlatformV0 <- function(json_object, default) { + if (json_object$contains("covariate_preprocessor")) { + return("python") + } + if ( + json_object$contains("preprocessor_metadata") || + json_object$contains("rfx_unique_group_ids") + ) { + return("R") + } + default +} + +# Return TRUE iff this is a cross-platform load (writer `platform` != loader). +# For a cross-platform load, verify the model is portable (all-numeric covariates) +# and rfx-compatible (integer group ids); otherwise stop() with a clear, actionable +# error. Same-platform loads return FALSE and ignore the flags. The portability +# flag is peeked from the covariate_preprocessor JSON via the C++ parser -- the +# foreign native preprocessor is never reconstructed. An absent portability flag +# is treated as non-portable (e.g. legacy v0 models), so such a cross-platform load +# is refused rather than mis-handled. +enforceCrossPlatformGate <- function(json_object, loader_platform) { + writer_platform <- json_object$get_string_or_default("platform", loader_platform) + if (identical(writer_platform, loader_platform)) { + return(FALSE) + } + portable <- FALSE + non_portable_columns <- "" + if (json_object$contains("covariate_preprocessor")) { + prep <- createCppJsonString(json_object$get_string("covariate_preprocessor")) + portable <- prep$get_boolean_or_default("cross_platform_portable", FALSE) + non_portable_columns <- prep$get_string_or_default("non_portable_columns", "") + } + compatible <- json_object$get_boolean_or_default( + "cross_platform_compatible", + TRUE, + subfolder_name = "random_effects" + ) + if (!portable) { + stop(sprintf( + paste0( + "Cannot load a model written on platform '%s' from '%s': it was fit with ", + "non-numeric covariates (columns: %s), which are not cross-platform portable. ", + "Load it on '%s', or refit / re-save with all-numeric covariates." + ), + writer_platform, + loader_platform, + if (nzchar(non_portable_columns)) non_portable_columns else "unknown", + writer_platform + )) + } + if (!compatible) { + stop(sprintf( + paste0( + "Cannot load a model written on platform '%s' from '%s': it has string-labeled ", + "random effects, which are not cross-platform compatible (only integer-valued ", + "group ids are). Load it on '%s'." + ), + writer_platform, + loader_platform, + writer_platform + )) + } + TRUE +} + +# Minimal identity preprocessor metadata for a cross-platform all-numeric load. +# The foreign native preprocessor can't be reconstructed, but the model is +# all-numeric (gate enforced), so prediction on a numeric matrix passes through; +# the feature count is read from the foreign preprocessor for the dimension check. +buildIdentityPreprocessorMetadata <- function(json_object) { + n_features <- 0L + if (json_object$contains("covariate_preprocessor")) { + prep <- createCppJsonString(json_object$get_string("covariate_preprocessor")) + n_features <- prep$get_integer_or_default( + "num_original_features", + prep$get_integer_or_default("num_numeric_vars", 0L) + ) + } + # A complete all-numeric metadata: predict only needs `num_numeric_vars`, but + # the extra fields let a cross-loaded model be re-saved (savePreprocessorToJson + # requires feature_types / original_var_indices / numeric_vars). Original column + # names are unknown cross-platform, so use positional placeholders. + list( + num_numeric_vars = n_features, + num_ordered_cat_vars = 0, + num_unordered_cat_vars = 0, + feature_types = rep(0, n_features), + original_var_indices = seq_len(n_features), + numeric_vars = as.character(seq_len(n_features)) + ) +} + +# Resolve rfx group-id levels: read `rfx_unique_group_ids` when present (R-written +# models), or reconstruct sorted integer levels from the rfx label mapper when +# absent (e.g. a cross-platform Python model, which never writes that field). +resolveRfxUniqueGroupIds <- function(json_object, rfx_samples) { + if ( + json_object$contains("rfx_unique_group_ids", subfolder_name = "random_effects") + ) { + json_object$get_string_vector( + "rfx_unique_group_ids", + subfolder_name = "random_effects" + ) + } else { + as.character( + rfx_label_mapper_unique_group_ids_cpp(rfx_samples$label_mapper_ptr) + ) + } +} + #' Forest Container Serialization Routines #' @name ForestSamplesSerialization #' @description @@ -6,7 +183,7 @@ #' loading a single [ForestSamples] container from a broader BART / BCF model (which may include multiple forests and other parametric terms). #' #' `loadForestContainerJson` converts a [CppJson] object representing a BART or BCF model into a [ForestSamples] container -#' by extracting the JSON indexed by a forest label (i.e. `"forest_0"`) and deserializing it into a [ForestSamples] object. +#' by extracting the JSON indexed by a forest label (i.e. `"mean_forest"`) and deserializing it into a [ForestSamples] object. #' #' Both `loadForestContainerJson` and `loadForestContainerCombinedJson` operate similarly, but on a list of [CppJson] or JSON string #' representations of BART / BCF models with the same structure. @@ -24,9 +201,9 @@ #' bart_json_string <- saveBARTModelToJsonString(bart_model) #' bart_json_list <- list(bart_json) #' bart_json_string_list <- list(bart_json_string) -#' mean_forest <- loadForestContainerJson(bart_json, "forest_0") -#' mean_forest <- loadForestContainerCombinedJson(bart_json_list, "forest_0") -#' mean_forest <- loadForestContainerCombinedJsonString(bart_json_string_list, "forest_0") +#' mean_forest <- loadForestContainerJson(bart_json, "mean_forest") +#' mean_forest <- loadForestContainerCombinedJson(bart_json_list, "mean_forest") +#' mean_forest <- loadForestContainerCombinedJsonString(bart_json_string_list, "mean_forest") #' NULL #> NULL @@ -118,14 +295,18 @@ CppJson <- R6::R6Class( #' @description #' Convert a forest container to json and add to the current `CppJson` object #' @param forest_samples `ForestSamples` R class + #' @param forest_label Key under the `forests` subfolder to store this + #' container. Defaults to an auto-numbered `forest_` label when empty. #' @return None - add_forest = function(forest_samples) { + add_forest = function(forest_samples, forest_label = "") { forest_label <- json_add_forest_cpp( self$json_ptr, - forest_samples$forest_container_ptr + forest_samples$forest_container_ptr, + forest_label ) self$num_forests <- self$num_forests + 1 self$forest_labels <- c(self$forest_labels, forest_label) + invisible(forest_label) }, #' @description @@ -448,6 +629,111 @@ CppJson <- R6::R6Class( return(result) }, + #' @description + #' Whether the json object contains a field "field_name" (with optional subfolder "subfolder_name") + #' @param field_name The name of the field to check for + #' @param subfolder_name (Optional) Name of the subfolder / hierarchy + #' @return Logical, `TRUE` if the field is present + contains = function(field_name, subfolder_name = NULL) { + if (is.null(subfolder_name)) { + json_contains_field_cpp(self$json_ptr, field_name) + } else { + json_contains_field_subfolder_cpp(self$json_ptr, subfolder_name, field_name) + } + }, + + #' @description + #' Retrieve a scalar value, returning `default` if the field is absent. Used to + #' "augment" older JSON within a schema version (see RFC 0005). + #' @param field_name The name of the field + #' @param default Value returned if the field is absent + #' @param subfolder_name (Optional) Name of the subfolder / hierarchy + #' @return The stored value, or `default` if absent + get_scalar_or_default = function(field_name, default, subfolder_name = NULL) { + if (self$contains(field_name, subfolder_name)) { + self$get_scalar(field_name, subfolder_name) + } else { + default + } + }, + + #' @description + #' Retrieve an integer value, returning `default` if the field is absent. + #' @param field_name The name of the field + #' @param default Value returned if the field is absent + #' @param subfolder_name (Optional) Name of the subfolder / hierarchy + #' @return The stored value, or `default` if absent + get_integer_or_default = function(field_name, default, subfolder_name = NULL) { + if (self$contains(field_name, subfolder_name)) { + self$get_integer(field_name, subfolder_name) + } else { + default + } + }, + + #' @description + #' Retrieve a boolean value, returning `default` if the field is absent. + #' @param field_name The name of the field + #' @param default Value returned if the field is absent + #' @param subfolder_name (Optional) Name of the subfolder / hierarchy + #' @return The stored value, or `default` if absent + get_boolean_or_default = function(field_name, default, subfolder_name = NULL) { + if (self$contains(field_name, subfolder_name)) { + self$get_boolean(field_name, subfolder_name) + } else { + default + } + }, + + #' @description + #' Retrieve a string value, returning `default` if the field is absent. + #' @param field_name The name of the field + #' @param default Value returned if the field is absent + #' @param subfolder_name (Optional) Name of the subfolder / hierarchy + #' @return The stored value, or `default` if absent + get_string_or_default = function(field_name, default, subfolder_name = NULL) { + if (self$contains(field_name, subfolder_name)) { + self$get_string(field_name, subfolder_name) + } else { + default + } + }, + + #' @description + #' Rename a field "old_name" to "new_name" (with optional subfolder). No-op if + #' "old_name" is absent. Used by JSON schema migrations (RFC 0005). + #' @param old_name Current field name + #' @param new_name New field name + #' @param subfolder_name (Optional) Name of the subfolder / hierarchy + #' @return None + rename_field = function(old_name, new_name, subfolder_name = NULL) { + if (is.null(subfolder_name)) { + json_rename_field_cpp(self$json_ptr, old_name, new_name) + } else { + json_rename_field_subfolder_cpp( + self$json_ptr, + subfolder_name, + old_name, + new_name + ) + } + invisible(self) + }, + + #' @description + #' Erase a field "field_name" (with optional subfolder). No-op if absent. + #' @param field_name The name of the field to erase + #' @param subfolder_name (Optional) Name of the subfolder / hierarchy + #' @return None + erase_field = function(field_name, subfolder_name = NULL) { + if (is.null(subfolder_name)) { + json_erase_field_cpp(self$json_ptr, field_name) + } else { + json_erase_field_subfolder_cpp(self$json_ptr, subfolder_name, field_name) + } + invisible(self) + }, + #' @description #' Retrieve a vector from the json object under the name "field_name" (with optional subfolder "subfolder_name") #' @param field_name The name of the field to be accessed from json @@ -598,7 +884,7 @@ CppJson <- R6::R6Class( #' @title Load Forest Samples from JSON #' @rdname ForestSamplesSerialization #' @param json_object Object of class `CppJson` -#' @param json_forest_label Label referring to a particular forest (i.e. "forest_0") in the overall json hierarchy (must exist in every json object in a list if a list is provided) +#' @param json_forest_label Label referring to a particular forest (i.e. "mean_forest") in the overall json hierarchy (must exist in every json object in a list if a list is provided) #' @export loadForestContainerJson <- function(json_object, json_forest_label) { invisible(output <- ForestSamples$new(0, 1, T)) @@ -609,7 +895,7 @@ loadForestContainerJson <- function(json_object, json_forest_label) { #' @title Combine JSON Model Objects into ForestSamples #' @rdname ForestSamplesSerialization #' @param json_object_list List of objects of class `CppJson` -#' @param json_forest_label Label referring to a particular forest (i.e. "forest_0") in the overall json hierarchy (must exist in every json object in a list if a list is provided) +#' @param json_forest_label Label referring to a particular forest (i.e. "mean_forest") in the overall json hierarchy (must exist in every json object in a list if a list is provided) #' @export loadForestContainerCombinedJson <- function( json_object_list, @@ -630,7 +916,7 @@ loadForestContainerCombinedJson <- function( #' @title Combine JSON Strings into ForestSamples #' @rdname ForestSamplesSerialization #' @param json_string_list List of strings that parse into objects of type `CppJson` -#' @param json_forest_label Label referring to a particular forest (i.e. "forest_0") in the overall json hierarchy (must exist in every json object in a list if a list is provided) +#' @param json_forest_label Label referring to a particular forest (i.e. "mean_forest") in the overall json hierarchy (must exist in every json object in a list if a list is provided) #' @export loadForestContainerCombinedJsonString <- function( json_string_list, diff --git a/R/utils.R b/R/utils.R index 352a0cf3..93c88155 100644 --- a/R/utils.R +++ b/R/utils.R @@ -554,6 +554,25 @@ savePreprocessorToJson <- function(object) { } } + # Cross-platform portability (RFC 0005): portable iff every covariate is + # numeric (no categorical encoding). The cross-platform loader peeks at this + # flag via a plain JSON parse without reconstructing the native preprocessor. + portable <- (object$num_ordered_cat_vars == 0) && + (object$num_unordered_cat_vars == 0) + jsonobj$add_boolean("cross_platform_portable", portable) + if (!portable) { + non_portable_columns <- c( + if (object$num_ordered_cat_vars > 0) object$ordered_cat_vars else NULL, + if (object$num_unordered_cat_vars > 0) object$unordered_cat_vars else NULL + ) + # Human-readable string for the cross-platform refusal message (the gate + # only branches on the boolean above). + jsonobj$add_string( + "non_portable_columns", + paste(non_portable_columns, collapse = ", ") + ) + } + return(jsonobj) } @@ -1236,8 +1255,14 @@ inferStochtreeJsonVersion <- function(json_object) { return("<0.3.2") } - # rfx_model_spec and preprocessor_metadata were added in ~0.3.0 - if (!has_field("rfx_model_spec") || !has_field("preprocessor_metadata")) { + # rfx_model_spec and the covariate preprocessor were added in ~0.3.0. The + # preprocessor key is `preprocessor_metadata` in legacy R JSON and + # `covariate_preprocessor` in v1 / Python JSON, so accept either. + if ( + !has_field("rfx_model_spec") || + (!has_field("preprocessor_metadata") && + !has_field("covariate_preprocessor")) + ) { return("<0.3.0") } diff --git a/debug/api_debug.cpp b/debug/api_debug.cpp deleted file mode 100644 index 5426fe7e..00000000 --- a/debug/api_debug.cpp +++ /dev/null @@ -1,810 +0,0 @@ -/*! Copyright (c) 2024 stochtree authors*/ -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace StochTree{ - -void GenerateDGP1(std::vector& covariates, std::vector& basis, std::vector& outcome, std::vector& rfx_basis, std::vector& rfx_groups, std::vector& feature_types, std::mt19937& gen, int& n, int& x_cols, int& omega_cols, int& y_cols, int& rfx_basis_cols, int& num_rfx_groups, bool rfx_included, int random_seed = -1) { - // Data dimensions - n = 1000; - x_cols = 2; - omega_cols = 1; - y_cols = 1; - if (rfx_included) { - num_rfx_groups = 2; - rfx_basis_cols = 1; - } else { - num_rfx_groups = 0; - rfx_basis_cols = 0; - } - - // Resize data - covariates.resize(n * x_cols); - basis.resize(n * omega_cols); - rfx_basis.resize(n * rfx_basis_cols); - outcome.resize(n * y_cols); - rfx_groups.resize(n); - feature_types.resize(x_cols, FeatureType::kNumeric); - - // Random number generation - standard_normal normal_dist; - - // DGP parameters - std::vector betas{-10, -5, 5, 10}; - int num_partitions = betas.size(); - double f_x_omega; - double rfx; - double error; - - for (int i = 0; i < n; i++) { - for (int j = 0; j < x_cols; j++) { - covariates[i*x_cols + j] = standard_uniform_draw(gen); - } - - for (int j = 0; j < omega_cols; j++) { - basis[i*omega_cols + j] = standard_uniform_draw(gen); - } - - if (rfx_included) { - for (int j = 0; j < rfx_basis_cols; j++) { - rfx_basis[i * rfx_basis_cols + j] = 1; - } - - if (i % 2 == 0) { - rfx_groups[i] = 1; - } - else { - rfx_groups[i] = 2; - } - } - - for (int j = 0; j < y_cols; j++) { - if ((covariates[i * x_cols + 0] >= 0.0) && covariates[i * x_cols + 0] < 0.25) { - f_x_omega = betas[0] * basis[i * omega_cols + 0]; - } else if ((covariates[i * x_cols + 0] >= 0.25) && covariates[i * x_cols + 0] < 0.5) { - f_x_omega = betas[1] * basis[i * omega_cols + 0]; - } else if ((covariates[i * x_cols + 0] >= 0.5) && covariates[i * x_cols + 0] < 0.75) { - f_x_omega = betas[2] * basis[i * omega_cols + 0]; - } else { - f_x_omega = betas[3] * basis[i * omega_cols + 0]; - } - error = 0.1 * normal_dist(gen); - outcome[i * y_cols + j] = f_x_omega + error; - if (rfx_included) { - if (rfx_groups[i] == 1) { - rfx = 5.; - } - else { - rfx = -5.; - } - outcome[i * y_cols + j] += rfx; - } - } - } -} - -void int_to_binary_vector(int32_t input, std::vector& output, int32_t offset) { - bool terminated = false; - int numerator = input; - int vec_pos = 0; - if (numerator < 2) { - terminated = true; - output.at(offset + vec_pos) = numerator; - } - while (!terminated) { - std::div_t div_result = std::div(numerator, 2); - output.at(offset + vec_pos) = div_result.rem; - if (div_result.quot == 1) { - terminated = true; - output.at(offset + vec_pos + 1) = 1; - } else { - numerator = div_result.quot; - vec_pos += 1; - } - } -} - -void GenerateDGP2(std::vector& covariates, std::vector& basis, std::vector& outcome, std::vector& rfx_basis, std::vector& rfx_groups, std::vector& feature_types, std::mt19937& gen, int& n, int& x_cols, int& omega_cols, int& y_cols, int& rfx_basis_cols, int& num_rfx_groups, bool rfx_included, int random_seed = -1) { - // Data dimensions - int n1 = 50; - x_cols = 10; - int num_cells = std::pow(2, x_cols); - int p1 = 100; - if (p1 >= num_cells) Log::Fatal("p1 must be < 2^x_cols"); - n = n1*num_cells; - omega_cols = 0; - y_cols = 1; - if (rfx_included) { - num_rfx_groups = 2; - rfx_basis_cols = 1; - } - else { - num_rfx_groups = 0; - rfx_basis_cols = 0; - } - - // Resize data - covariates.resize(n * x_cols); - basis.resize(n * omega_cols); - rfx_basis.resize(n * rfx_basis_cols); - outcome.resize(n * y_cols); - rfx_groups.resize(n); - feature_types.resize(x_cols, FeatureType::kNumeric); - - // Random number generation - standard_normal normal_dist; - - // Generate a sequence of integers from 0 to num_cells - 1 - std::vector cell_run(num_cells); - std::iota(cell_run.begin(), cell_run.end(), 0); - - // Repeat this sequence n1 times as the "covariates" - std::vector cell_vector; - for (int i = 0; i < n1; i++) { - std::copy(cell_run.begin(), cell_run.end(), std::back_inserter(cell_vector)); - } - - // Convert cells to binary covariate columns (row-major) - std::vector covariates_binary(n * x_cols); - int32_t offset = 0; - for (size_t i = 0; i < n; i++) { - int_to_binary_vector(cell_vector.at(i), covariates_binary, offset); - offset += x_cols; - } - - // Add (folded) gaussian noise to the binary covariates - // std::vector covariates_numeric(n* x_cols); - std::vector noise1(n); - std::vector noise2(n); - int switch_flip; - for (size_t i = 0; i < n; i++) { - noise1.at(i) = std::abs(normal_dist(gen)); - noise2.at(i) = std::abs(normal_dist(gen)); - } - for (size_t i = 0; i < n; i++) { - for (int j = 0; j < x_cols; j++) { - switch_flip = covariates_binary.at(i * x_cols + j); - covariates.at(i * x_cols + j) = switch_flip * noise1.at(i) + (1 - switch_flip) * noise2.at(i); - } - } - - // DGP parameters - double intercept = 0.5; - std::vector cell_coefficients_sparse(p1-1); - for (int i = 0; i < cell_coefficients_sparse.size(); i++) { - cell_coefficients_sparse.at(i) = -10*std::abs(normal_dist(gen)); - } - std::vector cell_weights(num_cells, 1./num_cells); - std::vector cell_indices_sparse(p1 - 1); - walker_vose cell_selector(cell_weights.begin(), cell_weights.end()); - for (int i = 0; i < p1-1; i++) { - cell_indices_sparse.at(i) = cell_selector(gen); - } - std::vector cell_coefficients_full(num_cells, 0.); - for (int i = 0; i < p1 - 1; i++) { - cell_coefficients_full.at(cell_indices_sparse.at(i)) = cell_coefficients_sparse.at(i); - } - double f_x; - double rfx; - double error; - - // Outcome - for (int i = 0; i < n; i++) { - f_x = intercept + cell_coefficients_full.at(cell_vector.at(i)); - - if (rfx_included) { - for (int j = 0; j < rfx_basis_cols; j++) { - rfx_basis[i * rfx_basis_cols + j] = 1; - } - - if (i % 2 == 0) { - rfx_groups[i] = 1; - } - else { - rfx_groups[i] = 2; - } - } - - for (int j = 0; j < y_cols; j++) { - error = 0.1 * normal_dist(gen); - outcome[i * y_cols + j] = f_x + error; - if (rfx_included) { - if (rfx_groups[i] == 1) { - rfx = 5.; - } - else { - rfx = -5.; - } - outcome[i * y_cols + j] += rfx; - } - } - } -} - -void GenerateDGP3(std::vector& covariates, std::vector& basis, std::vector& outcome, std::vector& rfx_basis, std::vector& rfx_groups, std::vector& feature_types, std::mt19937& gen, int& n, int& x_cols, int& omega_cols, int& y_cols, int& rfx_basis_cols, int& num_rfx_groups, bool rfx_included, int random_seed = -1) { - // Data dimensions - n = 1000; - x_cols = 2; - omega_cols = 2; - y_cols = 1; - if (rfx_included) { - num_rfx_groups = 2; - rfx_basis_cols = 1; - } else { - num_rfx_groups = 0; - rfx_basis_cols = 0; - } - - // Resize data - covariates.resize(n * x_cols); - basis.resize(n * omega_cols); - rfx_basis.resize(n * rfx_basis_cols); - outcome.resize(n * y_cols); - rfx_groups.resize(n); - feature_types.resize(x_cols, FeatureType::kNumeric); - - // Random number generation - standard_normal normal_dist; - - // DGP parameters - std::vector betas{-10, -5, 5, 10}; - int num_partitions = betas.size(); - double f_x_omega; - double rfx; - double error; - - for (int i = 0; i < n; i++) { - for (int j = 0; j < x_cols; j++) { - covariates[i*x_cols + j] = standard_uniform_draw(gen); - } - - for (int j = 0; j < omega_cols; j++) { - basis[i*omega_cols + j] = standard_uniform_draw(gen); - } - - if (rfx_included) { - for (int j = 0; j < rfx_basis_cols; j++) { - rfx_basis[i * rfx_basis_cols + j] = 1; - } - - if (i % 2 == 0) { - rfx_groups[i] = 1; - } - else { - rfx_groups[i] = 2; - } - } - - for (int j = 0; j < y_cols; j++) { - if ((covariates[i * x_cols + 0] >= 0.0) && covariates[i * x_cols + 0] < 0.25) { - f_x_omega = betas[0] * basis[i * omega_cols + 0]; - } else if ((covariates[i * x_cols + 0] >= 0.25) && covariates[i * x_cols + 0] < 0.5) { - f_x_omega = betas[1] * basis[i * omega_cols + 0]; - } else if ((covariates[i * x_cols + 0] >= 0.5) && covariates[i * x_cols + 0] < 0.75) { - f_x_omega = betas[2] * basis[i * omega_cols + 0]; - } else { - f_x_omega = betas[3] * basis[i * omega_cols + 0]; - } - error = 0.1 * normal_dist(gen); - outcome[i * y_cols + j] = f_x_omega + error; - if (rfx_included) { - if (rfx_groups[i] == 1) { - rfx = 5.; - } - else { - rfx = -5.; - } - outcome[i * y_cols + j] += rfx; - } - } - } -} - -void GenerateDGP4(std::vector& covariates, std::vector& basis, std::vector& outcome, std::vector& rfx_basis, std::vector& rfx_groups, std::vector& feature_types, std::mt19937& gen, int& n, int& x_cols, int& omega_cols, int& y_cols, int& rfx_basis_cols, int& num_rfx_groups, bool rfx_included, int random_seed = -1) { - // Data dimensions - n = 400; - x_cols = 10; - omega_cols = 0; - y_cols = 1; - if (rfx_included) { - num_rfx_groups = 2; - rfx_basis_cols = 1; - } else { - num_rfx_groups = 0; - rfx_basis_cols = 0; - } - - // Resize data - covariates.resize(n * x_cols); - basis.resize(n * omega_cols); - rfx_basis.resize(n * rfx_basis_cols); - outcome.resize(n * y_cols); - rfx_groups.resize(n); - feature_types.resize(x_cols, FeatureType::kNumeric); - - // Random number generation - standard_normal normal_dist; - - // DGP parameters - std::vector betas{0.5, 1, 2, 3}; - int num_partitions = betas.size(); - double s_x; - double rfx; - double error; - - for (int i = 0; i < n; i++) { - for (int j = 0; j < x_cols; j++) { - covariates[i*x_cols + j] = standard_uniform_draw(gen); - } - - for (int j = 0; j < omega_cols; j++) { - basis[i*omega_cols + j] = standard_uniform_draw(gen); - } - - if (rfx_included) { - for (int j = 0; j < rfx_basis_cols; j++) { - rfx_basis[i * rfx_basis_cols + j] = 1; - } - - if (i % 2 == 0) { - rfx_groups[i] = 1; - } - else { - rfx_groups[i] = 2; - } - } - - for (int j = 0; j < y_cols; j++) { - if ((covariates[i * x_cols + 0] >= 0.0) && covariates[i * x_cols + 0] < 0.25) { - s_x = betas[0]; - } else if ((covariates[i * x_cols + 0] >= 0.25) && covariates[i * x_cols + 0] < 0.5) { - s_x = betas[1]; - } else if ((covariates[i * x_cols + 0] >= 0.5) && covariates[i * x_cols + 0] < 0.75) { - s_x = betas[2]; - } else { - s_x = betas[3]; - } - error = s_x * normal_dist(gen); - outcome[i * y_cols + j] = error; - if (rfx_included) { - if (rfx_groups[i] == 1) { - rfx = 5.; - } - else { - rfx = -5.; - } - outcome[i * y_cols + j] += rfx; - } - } - } -} - -void OutcomeOffsetScale(ColumnVector& residual, double& outcome_offset, double& outcome_scale) { - data_size_t n = residual.NumRows(); - double outcome_val = 0.0; - double outcome_sum = 0.0; - double outcome_sum_squares = 0.0; - double var_y = 0.0; - for (data_size_t i = 0; i < n; i++){ - outcome_val = residual.GetElement(i); - outcome_sum += outcome_val; - outcome_sum_squares += std::pow(outcome_val, 2.0); - } - var_y = outcome_sum_squares / static_cast(n) - std::pow(outcome_sum / static_cast(n), 2.0); - outcome_scale = std::sqrt(var_y); - outcome_offset = outcome_sum / static_cast(n); - double previous_residual; - for (data_size_t i = 0; i < n; i++){ - previous_residual = residual.GetElement(i); - residual.SetElement(i, (previous_residual - outcome_offset) / outcome_scale); - } -} - -void RunDebug(int dgp_num = 0, const ModelType model_type = kConstantLeafGaussian, - bool rfx_included = false, int num_gfr = 10, int num_mcmc = 100, int random_seed = -1, - std::string dataset_filename = "", int outcome_col = -1, std::string covariate_cols = "", - std::string basis_cols = "", int num_threads = -1) { - // Flag the data as row-major - bool row_major = true; - - // Determine whether we will generate data or read from file - bool data_from_file = false; - if (!dataset_filename.empty()) { - data_from_file = true; - } - - // Random number generation - std::mt19937 gen; - if (random_seed == -1) { - std::random_device rd; - std::mt19937 gen(rd()); - } - else { - std::mt19937 gen(random_seed); - } - - // Initialize dataset - ForestDataset dataset = ForestDataset(); - - // Initialize outcome - ColumnVector residual = ColumnVector(); - - // Empty data containers and dimensions (filled in by calling a specific DGP simulation function below) - int n; - int x_cols; - int omega_cols; - int y_cols; - int num_rfx_groups; - int rfx_basis_cols; - std::vector covariates_raw; - std::vector basis_raw; - std::vector rfx_basis_raw; - std::vector outcome_raw; - std::vector rfx_groups; - std::vector feature_types; - - // Check for DGP : ModelType compatibility - if ((model_type != kConstantLeafGaussian) && (dgp_num == 1)) { - Log::Fatal("dgp 2 is only compatible with a constant leaf model"); - } - - // Generate the data - int output_dimension; - bool is_leaf_constant; - if (!data_from_file) { - if (dgp_num == 0) { - GenerateDGP1(covariates_raw, basis_raw, outcome_raw, rfx_basis_raw, rfx_groups, feature_types, gen, n, x_cols, omega_cols, y_cols, rfx_basis_cols, num_rfx_groups, rfx_included, random_seed); - dataset.AddCovariates(covariates_raw.data(), n, x_cols, row_major); - dataset.AddBasis(basis_raw.data(), n, omega_cols, row_major); - output_dimension = 1; - is_leaf_constant = false; - } else if (dgp_num == 1) { - GenerateDGP2(covariates_raw, basis_raw, outcome_raw, rfx_basis_raw, rfx_groups, feature_types, gen, n, x_cols, omega_cols, y_cols, rfx_basis_cols, num_rfx_groups, rfx_included, random_seed); - dataset.AddCovariates(covariates_raw.data(), n, x_cols, row_major); - output_dimension = 1; - is_leaf_constant = true; - } else if (dgp_num == 2) { - GenerateDGP3(covariates_raw, basis_raw, outcome_raw, rfx_basis_raw, rfx_groups, feature_types, gen, n, x_cols, omega_cols, y_cols, rfx_basis_cols, num_rfx_groups, rfx_included, random_seed); - dataset.AddCovariates(covariates_raw.data(), n, x_cols, row_major); - dataset.AddBasis(basis_raw.data(), n, omega_cols, row_major); - output_dimension = omega_cols; - is_leaf_constant = false; - } else if (dgp_num == 3) { - GenerateDGP4(covariates_raw, basis_raw, outcome_raw, rfx_basis_raw, rfx_groups, feature_types, gen, n, x_cols, omega_cols, y_cols, rfx_basis_cols, num_rfx_groups, rfx_included, random_seed); - dataset.AddCovariates(covariates_raw.data(), n, x_cols, row_major); - output_dimension = 1; - is_leaf_constant = true; - } else { - Log::Fatal("Invalid dgp_num"); - } - // Construct residual - residual = ColumnVector(outcome_raw.data(), n); - } else { - // Override RFX - rfx_included = false; - // Construct residual - residual = ColumnVector(dataset_filename, outcome_col); - y_cols = 0; - // Add covariates - dataset.AddCovariatesFromCSV(dataset_filename, covariate_cols); - n = dataset.NumObservations(); - x_cols = dataset.NumCovariates(); - feature_types.resize(x_cols, FeatureType::kNumeric); - if (!basis_cols.empty()) { - dataset.AddBasisFromCSV(dataset_filename, basis_cols); - output_dimension = dataset.NumBasis(); - is_leaf_constant = false; - omega_cols = dataset.NumBasis(); - } else { - output_dimension = 1; - is_leaf_constant = true; - omega_cols = 0; - } - } - - // Runtime check --- cannot have case / variance weights and be modeling heteroskedastic variance - if ((dgp_num == 3) && (dataset.HasVarWeights())) { - StochTree::Log::Fatal("Cannot provide variance / case weights when modeling heteroskedasticity with a forest"); - } - - // Center and scale the data - double outcome_offset; - double outcome_scale; - OutcomeOffsetScale(residual, outcome_offset, outcome_scale); - - // Prepare random effects sampling (if desired) - RandomEffectsDataset rfx_dataset; - std::vector rfx_init(n, 0); - RandomEffectsTracker rfx_tracker = RandomEffectsTracker(rfx_init); - MultivariateRegressionRandomEffectsModel rfx_model = MultivariateRegressionRandomEffectsModel(1, 1); - RandomEffectsContainer rfx_container; - LabelMapper label_mapper; - if (rfx_included) { - // Construct a random effects dataset - rfx_dataset = RandomEffectsDataset(); - rfx_dataset.AddBasis(rfx_basis_raw.data(), n, rfx_basis_cols, true); - rfx_dataset.AddGroupLabels(rfx_groups); - - // Construct random effects tracker / model / container - RandomEffectsTracker rfx_tracker = RandomEffectsTracker(rfx_groups); - MultivariateRegressionRandomEffectsModel rfx_model = MultivariateRegressionRandomEffectsModel(rfx_basis_cols, num_rfx_groups); - RandomEffectsContainer rfx_container = RandomEffectsContainer(rfx_basis_cols, num_rfx_groups); - LabelMapper label_mapper = LabelMapper(rfx_tracker.GetLabelMap()); - - // Set random effects model parameters - Eigen::VectorXd working_param_init(rfx_basis_cols); - Eigen::MatrixXd group_param_init(rfx_basis_cols, num_rfx_groups); - Eigen::MatrixXd working_param_cov_init(rfx_basis_cols, rfx_basis_cols); - Eigen::MatrixXd group_param_cov_init(rfx_basis_cols, rfx_basis_cols); - double variance_prior_shape = 1.; - double variance_prior_scale = 1.; - working_param_init << 1.; - group_param_init << 1., 1.; - working_param_cov_init << 1; - group_param_cov_init << 1; - rfx_model.SetWorkingParameter(working_param_init); - rfx_model.SetGroupParameters(group_param_init); - rfx_model.SetWorkingParameterCovariance(working_param_cov_init); - rfx_model.SetGroupParameterCovariance(group_param_cov_init); - rfx_model.SetVariancePriorShape(variance_prior_shape); - rfx_model.SetVariancePriorScale(variance_prior_scale); - } - - // Initialize an ensemble - int num_trees = 50; - bool forest_exponentiated; - if (model_type == kLogLinearVariance) { - forest_exponentiated = true; - } else { - forest_exponentiated = false; - } - // "Active" tree ensemble - TreeEnsemble active_forest = TreeEnsemble(num_trees, output_dimension, is_leaf_constant, forest_exponentiated); - // Stored forest samples - ForestContainer forest_samples = ForestContainer(num_trees, output_dimension, is_leaf_constant, forest_exponentiated); - - // Initialize a leaf model - double leaf_prior_mean = 0.; - double leaf_prior_scale = 1./num_trees; - - // Initialize forest sampling machinery - double alpha = 0.95; - double beta = 2.; - int min_samples_leaf = 1; - int max_depth = 10; - int cutpoint_grid_size = 100; - double a_rfx = 1.; - double b_rfx = 1.; - double a_leaf = 2.; - double b_leaf = 0.5; - double a_global = 0; - double b_global = 0; - double a_0 = 1.5; - double a_forest = num_trees / (a_0 * a_0) + 0.5; - double b_forest = num_trees / (a_0 * a_0); - - // Set leaf model parameters - double leaf_scale; - double leaf_scale_init = 1.; - Eigen::MatrixXd leaf_scale_matrix(omega_cols, omega_cols); - Eigen::MatrixXd leaf_scale_matrix_init(omega_cols, omega_cols); - if (omega_cols > 0) { - leaf_scale_matrix_init = Eigen::MatrixXd::Identity(omega_cols, omega_cols); - // leaf_scale_matrix_init << 1.0, 0.0, 0.0, 1.0; - leaf_scale_matrix = leaf_scale_matrix_init / num_trees; - } - - // Set global variance - double global_variance; - double global_variance_init = 1.0; - - // Set variable weights - double const_var_wt = static_cast(1. / x_cols); - std::vector variable_weights(x_cols, const_var_wt); - - // Initialize tracker and tree prior - ForestTracker tracker = ForestTracker(dataset.GetCovariates(), feature_types, num_trees, n); - TreePrior tree_prior = TreePrior(alpha, beta, min_samples_leaf, max_depth); - - // Initialize variance models - GlobalHomoskedasticVarianceModel global_var_model = GlobalHomoskedasticVarianceModel(); - LeafNodeHomoskedasticVarianceModel leaf_var_model = LeafNodeHomoskedasticVarianceModel(); - - // Initialize storage for samples of variance - std::vector global_variance_samples{}; - std::vector leaf_variance_samples{}; - - // Initialize leaf model - double init_val; - double init_val_glob; - std::vector init_vec; - if (model_type == kConstantLeafGaussian) { - init_val_glob = ComputeMeanOutcome(residual); - init_val = init_val_glob / static_cast(num_trees); - active_forest.SetLeafValue(init_val); - UpdateResidualEntireForest(tracker, dataset, residual, &active_forest, false, std::minus()); - tracker.UpdatePredictions(&active_forest, dataset); - } else if (model_type == kUnivariateRegressionLeafGaussian) { - init_val_glob = ComputeMeanOutcome(residual); - init_val = init_val_glob / static_cast(num_trees); - active_forest.SetLeafValue(init_val); - UpdateResidualEntireForest(tracker, dataset, residual, &active_forest, true, std::minus()); - tracker.UpdatePredictions(&active_forest, dataset); - } else if (model_type == kMultivariateRegressionLeafGaussian) { - init_val_glob = ComputeMeanOutcome(residual); - init_val = init_val_glob / static_cast(num_trees); - init_vec = std::vector(omega_cols, init_val); - active_forest.SetLeafVector(init_vec); - UpdateResidualEntireForest(tracker, dataset, residual, &active_forest, true, std::minus()); - tracker.UpdatePredictions(&active_forest, dataset); - } else if (model_type == kLogLinearVariance) { - init_val_glob = ComputeVarianceOutcome(residual) * 0.4; - init_val = std::log(init_val_glob) / static_cast(num_trees); - active_forest.SetLeafValue(init_val); - tracker.UpdatePredictions(&active_forest, dataset); - std::vector initial_preds(n, init_val_glob); - dataset.AddVarianceWeights(initial_preds.data(), n); - } - - // Prepare the samplers - LeafModelVariant leaf_model = leafModelFactory(model_type, leaf_scale, leaf_scale_matrix, a_forest, b_forest); - int num_features_subsample = x_cols; - - // Initialize vector of sweep update indices - std::vector sweep_indices(num_trees); - std::iota(sweep_indices.begin(), sweep_indices.end(), 0); - - // Run the GFR sampler - if (num_gfr > 0) { - for (int i = 0; i < num_gfr; i++) { - if (i == 0) { - global_variance = global_variance_init; - leaf_scale = leaf_scale_init; - } - else { - global_variance = global_variance_samples[i - 1]; - leaf_scale = leaf_variance_samples[i - 1]; - } - - // Sample tree ensemble - if (model_type == ModelType::kConstantLeafGaussian) { - GFRSampleOneIter(active_forest, tracker, forest_samples, std::get(leaf_model), dataset, residual, tree_prior, gen, variable_weights, sweep_indices, global_variance, feature_types, cutpoint_grid_size, true, true, true, num_features_subsample, num_threads); - } else if (model_type == ModelType::kUnivariateRegressionLeafGaussian) { - GFRSampleOneIter(active_forest, tracker, forest_samples, std::get(leaf_model), dataset, residual, tree_prior, gen, variable_weights, sweep_indices, global_variance, feature_types, cutpoint_grid_size, true, true, true, num_features_subsample, num_threads); - } else if (model_type == ModelType::kMultivariateRegressionLeafGaussian) { - GFRSampleOneIter(active_forest, tracker, forest_samples, std::get(leaf_model), dataset, residual, tree_prior, gen, variable_weights, sweep_indices, global_variance, feature_types, cutpoint_grid_size, true, true, true, num_features_subsample, num_threads, omega_cols); - } else if (model_type == ModelType::kLogLinearVariance) { - GFRSampleOneIter(active_forest, tracker, forest_samples, std::get(leaf_model), dataset, residual, tree_prior, gen, variable_weights, sweep_indices, global_variance, feature_types, cutpoint_grid_size, true, true, false, num_features_subsample, num_threads); - } - - if (rfx_included) { - // Sample random effects - rfx_model.SampleRandomEffects(rfx_dataset, residual, rfx_tracker, global_variance, gen); - rfx_container.AddSample(rfx_model); - } - - // Sample leaf node variance - leaf_variance_samples.push_back(leaf_var_model.SampleVarianceParameter(&active_forest, a_leaf, b_leaf, gen)); - - // Sample global variance - global_variance_samples.push_back(global_var_model.SampleVarianceParameter(residual.GetData(), a_global, b_global, gen)); - } - } - - // Run the MCMC sampler - if (num_mcmc > 0) { - for (int i = num_gfr; i < num_gfr + num_mcmc; i++) { - if (i == 0) { - global_variance = global_variance_init; - leaf_scale = leaf_scale_init; - } - else { - global_variance = global_variance_samples[i - 1]; - leaf_scale = leaf_variance_samples[i - 1]; - } - - // Sample tree ensemble - if (model_type == ModelType::kConstantLeafGaussian) { - MCMCSampleOneIter(active_forest, tracker, forest_samples, std::get(leaf_model), dataset, residual, tree_prior, gen, variable_weights, sweep_indices, global_variance, true, true, true, num_threads); - } else if (model_type == ModelType::kUnivariateRegressionLeafGaussian) { - MCMCSampleOneIter(active_forest, tracker, forest_samples, std::get(leaf_model), dataset, residual, tree_prior, gen, variable_weights, sweep_indices, global_variance, true, true, true, num_threads); - } else if (model_type == ModelType::kMultivariateRegressionLeafGaussian) { - MCMCSampleOneIter(active_forest, tracker, forest_samples, std::get(leaf_model), dataset, residual, tree_prior, gen, variable_weights, sweep_indices, global_variance, true, true, true, num_threads, omega_cols); - } else if (model_type == ModelType::kLogLinearVariance) { - MCMCSampleOneIter(active_forest, tracker, forest_samples, std::get(leaf_model), dataset, residual, tree_prior, gen, variable_weights, sweep_indices, global_variance, true, true, false, num_threads); - } - - if (rfx_included) { - // Sample random effects - rfx_model.SampleRandomEffects(rfx_dataset, residual, rfx_tracker, global_variance, gen); - rfx_container.AddSample(rfx_model); - } - - // Sample leaf node variance - leaf_variance_samples.push_back(leaf_var_model.SampleVarianceParameter(&active_forest, a_leaf, b_leaf, gen)); - - // Sample global variance - global_variance_samples.push_back(global_var_model.SampleVarianceParameter(residual.GetData(), a_global, b_global, gen)); - } - } - - // Predict from the tree ensemble - int num_samples = num_gfr + num_mcmc; - std::vector pred_orig = forest_samples.Predict(dataset); - - if (rfx_included) { - // Predict from the random effects dataset - std::vector rfx_predictions(n * num_samples); - rfx_container.Predict(rfx_dataset, label_mapper, rfx_predictions); - } - - // Write model to a file - std::string filename = "model.json"; - forest_samples.SaveToJsonFile(filename); - - // Read and parse json from file - ForestContainer forest_samples_parsed = ForestContainer(num_trees, output_dimension, is_leaf_constant); - forest_samples_parsed.LoadFromJsonFile(filename); - - // Make sure we can predict from both the original (above) and parsed forest containers - std::vector pred_parsed = forest_samples_parsed.Predict(dataset); -} - -} // namespace StochTree - -int main(int argc, char* argv[]) { - // Unpack command line arguments - int dgp_num = std::stoi(argv[1]); - if ((dgp_num != 0) && (dgp_num != 1) && (dgp_num != 2) && (dgp_num != 3)) { - StochTree::Log::Fatal("The first command line argument must be 0, 1, 2, or 3"); - } - int model_type_int = static_cast(std::stoi(argv[2])); - if ((model_type_int != 0) && (model_type_int != 1) && (model_type_int != 2) && (model_type_int != 3)) { - StochTree::Log::Fatal("The second command line argument must be 0, 1, 2, or 3"); - } - StochTree::ModelType model_type = static_cast(model_type_int); - int rfx_int = std::stoi(argv[3]); - if ((rfx_int != 0) && (rfx_int != 1)) { - StochTree::Log::Fatal("The third command line argument must be 0 or 1"); - } - bool rfx_included = static_cast(rfx_int); - int num_gfr = std::stoi(argv[4]); - if (num_gfr < 0) { - StochTree::Log::Fatal("The fourth command line argument must be >= 0"); - } - int num_mcmc = std::stoi(argv[5]); - if (num_mcmc < 0) { - StochTree::Log::Fatal("The fifth command line argument must be >= 0"); - } - int random_seed = std::stoi(argv[6]); - if (random_seed < -1) { - StochTree::Log::Fatal("The sixth command line argument must be >= -0"); - } - std::string dataset_filename = argv[7]; - int outcome_col = std::stoi(argv[8]); - std::string covariate_cols = argv[9]; - std::string basis_cols = argv[10]; - int num_threads = std::stoi(argv[11]); - - // Run the debug program - StochTree::RunDebug(dgp_num, model_type, rfx_included, num_gfr, num_mcmc, random_seed, - dataset_filename, outcome_col, covariate_cols, basis_cols, num_threads); -} diff --git a/debug/bart_debug.cpp b/debug/bart_debug.cpp new file mode 100644 index 00000000..90d5d309 --- /dev/null +++ b/debug/bart_debug.cpp @@ -0,0 +1,333 @@ +/* + * BART debug program. The first CLI argument selects the scenario (default: 0). + * + * Usage: bart_debug [--scenario N] [--n N] [--n_test N] [--p N] [--num_trees N] + * [--num_gfr N] [--num_mcmc N] [--seed N] + * + * 0 Homoskedastic constant-leaf BART + * DGP: y = sin(2*pi*x1) + 0.5*x2 - 1.5*x3 + eps, eps ~ N(0,1) + * + * 1 Homoskedastic constant-leaf probit BART + * DGP: Z = sin(2*pi*x1) + 0.5*x2 - 1.5*x3 + eps, eps ~ N(0,1) + * y = 1{Z > 0} + * + * 2 Homoskedastic multivariate-leaf BART (2-column basis) + * DGP: tau1(x) = piecewise on x1 in {-2,-1,1,2} (quartile bins) + * tau2(x) = piecewise on x1 in {1,2,-1,-2} (quartile bins) + * z1, z2 ~ Uniform(0,1) + * y = tau1(x)*z1 + tau2(x)*z2 + eps, eps ~ N(0,1) + * + * Add scenarios here as the BARTSampler API develops (heteroskedastic, + * random effects, etc.). + */ + +#include +#include + +#include +#include +#include +#include +#include +#include "stochtree/meta.h" + +static constexpr double kPi = 3.14159265358979323846; + +// ---- Data ------------------------------------------------------------ + +struct RegressionDataset { + std::vector X; + std::vector y; +}; + +struct ProbitDataset { + std::vector X; + std::vector y; + std::vector Z; +}; + +struct MultivariateRegressionDataset { + std::vector X; + std::vector y; + std::vector basis; // col-major n x 2: first n = z1, next n = z2 + std::vector f_true; +}; + +// DGP: y = sin(2*pi*x1) + 0.5*x2 - 1.5*x3 + N(0,1) +static RegressionDataset generate_regression_data(int n, int p, std::mt19937& rng) { + std::uniform_real_distribution unif(0.0, 1.0); + std::normal_distribution normal(0.0, 1.0); + RegressionDataset d; + d.X.resize(n * p); + d.y.resize(n); + for (int i = 0; i < n; i++) + for (int j = 0; j < p; j++) + d.X[j * n + i] = unif(rng); + for (int i = 0; i < n; i++) + d.y[i] = std::sin(2.0 * kPi * d.X[i]) + 0.5 * d.X[1 * n + i] - 1.5 * d.X[2 * n + i] + normal(rng); + return d; +} + +// DGP: +// Z = sin(2*pi*x1) + 0.5*x2 - 1.5*x3 + N(0,1) +// y = 1{Z > 0} +static ProbitDataset generate_probit_data(int n, int p, std::mt19937& rng) { + std::uniform_real_distribution unif(0.0, 1.0); + std::normal_distribution normal(0.0, 1.0); + ProbitDataset d; + d.X.resize(n * p); + d.y.resize(n); + d.Z.resize(n); + for (int i = 0; i < n; i++) + for (int j = 0; j < p; j++) + d.X[j * n + i] = unif(rng); + for (int i = 0; i < n; i++) { + d.Z[i] = std::sin(2.0 * kPi * d.X[i]) + 0.5 * d.X[1 * n + i] - 1.5 * d.X[2 * n + i] + normal(rng); + d.y[i] = (d.Z[i] > 0) ? 1.0 : 0.0; + } + return d; +} + +// DGP: tau1(x) and tau2(x) are piecewise-constant on x1 quartile bins; +// z1, z2 ~ Uniform(0,1); y = tau1*z1 + tau2*z2 + N(0,1) +static MultivariateRegressionDataset generate_multivariate_regression_data(int n, int p, std::mt19937& rng) { + std::uniform_real_distribution unif(0.0, 1.0); + std::normal_distribution normal(0.0, 1.0); + MultivariateRegressionDataset d; + d.X.resize(n * p); + d.y.resize(n); + d.basis.resize(n * 2); + d.f_true.resize(n); + for (int i = 0; i < n; i++) + for (int j = 0; j < p; j++) + d.X[j * n + i] = unif(rng); + for (int i = 0; i < n; i++) { + double z1 = unif(rng); + double z2 = unif(rng); + d.basis[i] = z1; + d.basis[n + i] = z2; + double x0 = d.X[i]; + double tau1 = (x0 < 0.25) ? -2.0 : (x0 < 0.50) ? -1.0 : (x0 < 0.75) ? 1.0 : 2.0; + double tau2 = (x0 < 0.25) ? 1.0 : (x0 < 0.50) ? 2.0 : (x0 < 0.75) ? -1.0 : -2.0; + d.f_true[i] = tau1 * z1 + tau2 * z2; + d.y[i] = d.f_true[i] + normal(rng); + } + return d; +} + +// ---- Reporter -------------------------------------------------------- +// +// Reads directly from BARTSamples (already un-standardized by BARTSamplerFit). +// test_ref is the comparison target: +// identity link — original outcome scale (y) +// probit link — binary outcome (0/1); reports Brier score and accuracy + +static void report_bart(const StochTree::BARTSamples& samples, + const std::vector& test_ref, + StochTree::LinkFunction link, + const char* scenario_name) { + const int num_samples = samples.num_samples; + const int n_test = samples.num_test; + std::cout << "\n" << scenario_name << ":\n"; + if (link == StochTree::LinkFunction::Identity) { + double rmse_sum = 0.0; + for (int i = 0; i < n_test; i++) { + double y_hat = 0.0; + for (int j = 0; j < num_samples; j++) + y_hat += samples.mean_forest_predictions_test[static_cast(j * n_test + i)] / num_samples; + double err = (y_hat * samples.y_std + samples.y_bar) - test_ref[i]; + rmse_sum += err * err; + } + std::cout << " RMSE (test): " << std::sqrt(rmse_sum / n_test) << "\n"; + if (!samples.global_error_variance_samples.empty()) { + std::cout << " sigma (last): " << std::sqrt(samples.global_error_variance_samples.back()) * samples.y_std << "\n" + << " sigma (truth): 1.0\n"; + } + } else { + // mean_forest_predictions_test is on the linear (latent) scale; apply norm_cdf for probit probability + double brier_sum = 0.0; + int correct = 0; + for (int i = 0; i < n_test; i++) { + double latent = 0.0; + for (int j = 0; j < num_samples; j++) + latent += samples.mean_forest_predictions_test[static_cast(j * n_test + i)] / num_samples; + double p = StochTree::norm_cdf(latent * samples.y_std + samples.y_bar); + double diff = p - test_ref[i]; + brier_sum += diff * diff; + correct += ((p >= 0.5) == (test_ref[i] >= 0.5)) ? 1 : 0; + } + std::cout << " Brier (test): " << brier_sum / n_test << "\n" + << " Acc (test): " << static_cast(correct) / n_test << "\n"; + } +} + +// ---- Scenario 0: homoskedastic constant-leaf BART -------------------- + +static void run_scenario_0(int n, int n_test, int p, int num_trees, int num_gfr, int num_mcmc, int seed) { + std::mt19937 rng(seed < 0 ? std::random_device{}() : static_cast(seed)); + RegressionDataset train = generate_regression_data(n, p, rng); + RegressionDataset test = generate_regression_data(n_test, p, rng); + + StochTree::BARTData data; + data.X_train = train.X.data(); + data.y_train = train.y.data(); + data.n_train = n; + data.p = p; + data.X_test = test.X.data(); + data.n_test = n_test; + + StochTree::BARTConfig config; + config.num_trees_mean = num_trees; + config.random_seed = seed; + config.mean_leaf_model_type = StochTree::MeanLeafModelType::GaussianConstant; + config.link_function = StochTree::LinkFunction::Identity; + config.standardize_outcome = true; + config.sample_sigma2_global = true; + config.var_weights_mean = std::vector(p, 1.0 / p); + config.feature_types = std::vector(p, StochTree::FeatureType::kNumeric); + config.sweep_update_indices_mean = std::vector(num_trees, 0); + std::iota(config.sweep_update_indices_mean.begin(), config.sweep_update_indices_mean.end(), 0); + + StochTree::BARTSamples samples; + StochTree::BARTSampler sampler(samples, config, data); + sampler.run_gfr(samples, num_gfr, true); + sampler.run_mcmc(samples, 0, 1, num_mcmc); + sampler.postprocess_samples(samples); + report_bart(samples, test.y, StochTree::LinkFunction::Identity, "Scenario 0 (Homoskedastic BART)"); +} + +// ---- Scenario 1: constant-leaf probit BART --------------------------- + +static void run_scenario_1(int n, int n_test, int p, int num_trees, int num_gfr, int num_mcmc, int seed) { + std::mt19937 rng(seed < 0 ? std::random_device{}() : static_cast(seed)); + ProbitDataset train = generate_probit_data(n, p, rng); + ProbitDataset test = generate_probit_data(n_test, p, rng); + + StochTree::BARTData data; + data.X_train = train.X.data(); + data.y_train = train.y.data(); + data.n_train = n; + data.p = p; + data.X_test = test.X.data(); + data.n_test = n_test; + + StochTree::BARTConfig config; + config.num_trees_mean = num_trees; + config.random_seed = seed; + config.mean_leaf_model_type = StochTree::MeanLeafModelType::GaussianConstant; + config.link_function = StochTree::LinkFunction::Probit; + config.sample_sigma2_global = false; + config.var_weights_mean = std::vector(p, 1.0 / p); + config.feature_types = std::vector(p, StochTree::FeatureType::kNumeric); + config.sweep_update_indices_mean = std::vector(num_trees, 0); + std::iota(config.sweep_update_indices_mean.begin(), config.sweep_update_indices_mean.end(), 0); + + StochTree::BARTSamples samples; + StochTree::BARTSampler sampler(samples, config, data); + sampler.run_gfr(samples, num_gfr, true); + sampler.run_mcmc(samples, 0, 1, num_mcmc); + report_bart(samples, test.y, StochTree::LinkFunction::Probit, "Scenario 1 (Probit BART)"); +} + +// ---- Scenario 2: multivariate-leaf BART (2-column basis) ------------- + +static void run_scenario_2(int n, int n_test, int p, int num_trees, int num_gfr, int num_mcmc, int seed) { + std::mt19937 rng(seed < 0 ? std::random_device{}() : static_cast(seed)); + MultivariateRegressionDataset train = generate_multivariate_regression_data(n, p, rng); + MultivariateRegressionDataset test = generate_multivariate_regression_data(n_test, p, rng); + const int basis_dim = 2; + + StochTree::BARTData data; + data.X_train = train.X.data(); + data.y_train = train.y.data(); + data.n_train = n; + data.p = p; + data.basis_train = train.basis.data(); + data.basis_test = test.basis.data(); + data.basis_dim = basis_dim; + data.X_test = test.X.data(); + data.n_test = n_test; + + StochTree::BARTConfig config; + config.num_trees_mean = num_trees; + config.random_seed = seed; + config.mean_leaf_model_type = StochTree::MeanLeafModelType::GaussianMultivariateRegression; + config.leaf_dim_mean = basis_dim; + config.leaf_constant_mean = false; + config.link_function = StochTree::LinkFunction::Identity; + config.standardize_outcome = true; + config.sample_sigma2_global = true; + config.var_weights_mean = std::vector(p, 1.0 / p); + config.feature_types = std::vector(p, StochTree::FeatureType::kNumeric); + config.sweep_update_indices_mean = std::vector(num_trees, 0); + std::iota(config.sweep_update_indices_mean.begin(), config.sweep_update_indices_mean.end(), 0); + config.sigma2_leaf_mean_matrix = {0.5, 0.0, 0.0, 0.5}; // 0.5 * I_{2x2}, col-major + + StochTree::BARTSamples samples; + StochTree::BARTSampler sampler(samples, config, data); + sampler.run_gfr(samples, num_gfr, true); + sampler.run_mcmc(samples, 0, 1, num_mcmc); + sampler.postprocess_samples(samples); + report_bart(samples, test.f_true, StochTree::LinkFunction::Identity, + "Scenario 2 (Multivariate-leaf BART)"); +} + +// ---- Main ----------------------------------------------------------- + +int main(int argc, char** argv) { + int scenario = 0; + int n = 500; + int n_test = 100; + int p = 5; + int num_trees = 200; + int num_gfr = 10; + int num_mcmc = 100; + int seed = 1234; + + for (int i = 1; i < argc; ++i) { + std::string arg = argv[i]; + if ((arg == "--scenario" || arg == "--n" || arg == "--n_test" || arg == "--p" || + arg == "--num_trees" || arg == "--num_gfr" || arg == "--num_mcmc" || arg == "--seed") && + i + 1 < argc) { + int val = std::stoi(argv[++i]); + if (arg == "--scenario") + scenario = val; + else if (arg == "--n") + n = val; + else if (arg == "--n_test") + n_test = val; + else if (arg == "--p") + p = val; + else if (arg == "--num_trees") + num_trees = val; + else if (arg == "--num_gfr") + num_gfr = val; + else if (arg == "--num_mcmc") + num_mcmc = val; + else if (arg == "--seed") + seed = val; + } else { + std::cerr << "Unknown or incomplete argument: " << arg << "\n" + << "Usage: bart_debug [--scenario N] [--n N] [--n_test N] [--p N]" + " [--num_trees N] [--num_gfr N] [--num_mcmc N] [--seed N]\n"; + return 1; + } + } + + switch (scenario) { + case 0: + run_scenario_0(n, n_test, p, num_trees, num_gfr, num_mcmc, seed); + break; + case 1: + run_scenario_1(n, n_test, p, num_trees, num_gfr, num_mcmc, seed); + break; + case 2: + run_scenario_2(n, n_test, p, num_trees, num_gfr, num_mcmc, seed); + break; + default: + std::cerr << "Unknown scenario " << scenario + << ". Available: 0 (Homoskedastic BART), 1 (Probit BART), 2 (Multivariate-leaf BART)\n"; + return 1; + } + return 0; +} diff --git a/debug/bcf_debug.cpp b/debug/bcf_debug.cpp new file mode 100644 index 00000000..c5b18c8b --- /dev/null +++ b/debug/bcf_debug.cpp @@ -0,0 +1,464 @@ +/* + * BCF debug program. The first CLI argument selects the scenario (default: 0). + * + * Usage: bcf_debug [--scenario N] [--n N] [--n_test N] [--p N] + * [--num_trees_mu N] [--num_trees_tau N] + * [--num_gfr N] [--num_mcmc N] [--seed N] + * + * 0 Two-forest BCF: constant-leaf mu, univariate-leaf tau (Z as basis) + * DGP: mu(x) = 2*sin(pi*x1) + 0.5*x2 + * tau(x) = 1 + x3 + * z ~ Bernoulli(0.5) + * y = mu(x) + tau(x)*z + N(0, 0.5^2) + * + * 1 Two-forest BCF with probit link: constant-leaf mu, univariate-leaf tau + * DGP: mu(x) = 2*sin(pi*x1) + 0.5*x2 + * tau(x) = 1 + x3 + * z ~ Bernoulli(0.5) + * W = mu(x) + tau(x)*z + N(0, 1) + * y = 1{W > 0} + * + * 2 Two-forest BCF: constant-leaf mu, multivariate-leaf tau (2 treatments) + * DGP: mu(x) = 2*sin(pi*x1) + 0.5*x2 + * tau1(x) = 1 + x3, tau2(x) = 0.5 - x4 + * z1, z2 ~ Bernoulli(0.5) + * y = mu(x) + tau1(x)*z1 + tau2(x)*z2 + N(0, 0.5^2) + * + * Add scenarios here as the BCFSampler API develops (propensity covariate, + * adaptive coding, random effects, etc.). + */ + +#include +#include + +#include +#include +#include +#include +#include +#include "stochtree/meta.h" + +static constexpr double kPi = 3.14159265358979323846; + +// ---- Data ------------------------------------------------------------ + +struct SimpleBCFDataset { + std::vector X; + std::vector y; + std::vector z; + std::vector mu_true; + std::vector tau_true; +}; + +struct ProbitBCFDataset { + std::vector X; + std::vector y; + std::vector latent_outcome; + std::vector z; + std::vector mu_true; + std::vector tau_true; +}; + +struct MultivariateBCFDataset { + std::vector X; + std::vector y; + std::vector z; // col-major n x 2: first n = z1, next n = z2 + std::vector mu_true; + std::vector tau1_true; + std::vector tau2_true; +}; + +// DGP: mu(x) = 2*sin(pi*x1) + 0.5*x2, tau(x) = 1 + x3 +// z ~ Bernoulli(0.5), y = mu + tau*z + N(0, 0.25) +static SimpleBCFDataset generate_simple_bcf_data(int n, int p, std::mt19937& rng) { + std::uniform_real_distribution unif(0.0, 1.0); + std::normal_distribution normal(0.0, 1.0); + std::bernoulli_distribution bern(0.5); + SimpleBCFDataset d; + d.X.resize(n * p); + d.y.resize(n); + d.z.resize(n); + d.mu_true.resize(n); + d.tau_true.resize(n); + for (int i = 0; i < n; i++) + for (int j = 0; j < p; j++) + d.X[j * n + i] = unif(rng); + for (int i = 0; i < n; i++) { + d.z[i] = bern(rng) ? 1.0 : 0.0; + d.mu_true[i] = 2.0 * std::sin(kPi * d.X[i]) + 0.5 * d.X[1 * n + i]; + d.tau_true[i] = 1.0 + d.X[2 * n + i]; + d.y[i] = d.mu_true[i] + d.tau_true[i] * d.z[i] + 0.5 * normal(rng); + } + return d; +} + +// DGP: same mu/tau; W = mu + tau*z + N(0,1); y = 1{W > 0} +static ProbitBCFDataset generate_probit_bcf_data(int n, int p, std::mt19937& rng) { + std::uniform_real_distribution unif(0.0, 1.0); + std::normal_distribution normal(0.0, 1.0); + std::bernoulli_distribution bern(0.5); + ProbitBCFDataset d; + d.X.resize(n * p); + d.y.resize(n); + d.z.resize(n); + d.mu_true.resize(n); + d.tau_true.resize(n); + d.latent_outcome.resize(n); + for (int i = 0; i < n; i++) + for (int j = 0; j < p; j++) + d.X[j * n + i] = unif(rng); + for (int i = 0; i < n; i++) { + d.z[i] = bern(rng) ? 1.0 : 0.0; + d.mu_true[i] = 2.0 * std::sin(kPi * d.X[i]) + 0.5 * d.X[1 * n + i]; + d.tau_true[i] = 1.0 + d.X[2 * n + i]; + d.latent_outcome[i] = d.mu_true[i] + d.tau_true[i] * d.z[i] + normal(rng); + d.y[i] = (d.latent_outcome[i] > 0.0) ? 1.0 : 0.0; + } + return d; +} + +// DGP: mu(x) = 2*sin(pi*x1) + 0.5*x2, tau1(x) = 1 + x3, tau2(x) = 0.5 - x4 +// z1, z2 ~ Bernoulli(0.5), y = mu + tau1*z1 + tau2*z2 + N(0, 0.25) +static MultivariateBCFDataset generate_multivariate_bcf_data(int n, int p, std::mt19937& rng) { + std::uniform_real_distribution unif(0.0, 1.0); + std::normal_distribution normal(0.0, 1.0); + std::bernoulli_distribution bern(0.5); + MultivariateBCFDataset d; + d.X.resize(n * p); + d.y.resize(n); + d.z.resize(n * 2); + d.mu_true.resize(n); + d.tau1_true.resize(n); + d.tau2_true.resize(n); + for (int i = 0; i < n; i++) + for (int j = 0; j < p; j++) + d.X[j * n + i] = unif(rng); + for (int i = 0; i < n; i++) { + double z1 = bern(rng) ? 1.0 : 0.0; + double z2 = bern(rng) ? 1.0 : 0.0; + d.z[i] = z1; + d.z[n + i] = z2; + d.mu_true[i] = 2.0 * std::sin(kPi * d.X[i]) + 0.5 * d.X[1 * n + i]; + d.tau1_true[i] = 1.0 + d.X[2 * n + i]; + d.tau2_true[i] = 0.5 - d.X[3 * n + i]; + d.y[i] = d.mu_true[i] + d.tau1_true[i] * z1 + d.tau2_true[i] * z2 + 0.5 * normal(rng); + } + return d; +} + +// ---- Reporter -------------------------------------------------------- +// +// Reads directly from BCFSamples (already un-standardized by BCFSamplerFit). +// mu_ref — true prognostic function (original outcome scale) +// tau_ref — true CATE (treatment effect scale, no y_bar offset) +// y_ref — binary outcome (0/1) for probit; continuous outcome for identity +// link — link function (Identity or Probit) +// +// mu/tau are always evaluated with RMSE (all test units). For the outcome: +// identity — RMSE from y_hat_test (original scale) +// probit — Brier score and accuracy from y_hat_test (already a probability) + +static void report_bcf(const StochTree::BCFSamples& samples, + const std::vector& mu_ref, + const std::vector& tau_ref, + const std::vector& y_ref, + StochTree::LinkFunction link, + const char* scenario_name) { + const int num_samples = samples.num_samples; + const int n_test = samples.num_test; + double mu_rmse_sum = 0.0, tau_rmse_sum = 0.0; + for (int i = 0; i < n_test; i++) { + double mu_hat = 0.0, tau_hat = 0.0; + for (int j = 0; j < num_samples; j++) { + const auto k = static_cast(j * n_test + i); + mu_hat += samples.mu_forest_predictions_test[k] / num_samples; + tau_hat += samples.tau_forest_predictions_test[k] / num_samples; + } + double mu_pred = mu_hat * samples.y_std + samples.y_bar; + double tau_pred = tau_hat * samples.y_std; + mu_rmse_sum += (mu_pred - mu_ref[i]) * (mu_pred - mu_ref[i]); + tau_rmse_sum += (tau_pred - tau_ref[i]) * (tau_pred - tau_ref[i]); + } + std::cout << "\n" + << scenario_name << ":\n" + << " mu RMSE (test): " << std::sqrt(mu_rmse_sum / n_test) << "\n" + << " tau RMSE (test): " << std::sqrt(tau_rmse_sum / n_test) << "\n"; + + if (link == StochTree::LinkFunction::Identity) { + double y_rmse_sum = 0.0; + for (int i = 0; i < n_test; i++) { + double y_hat = 0.0; + for (int j = 0; j < num_samples; j++) + y_hat += samples.y_hat_test[static_cast(j * n_test + i)] / num_samples; + y_rmse_sum += (y_hat - y_ref[i]) * (y_hat - y_ref[i]); + } + std::cout << " y RMSE (test): " << std::sqrt(y_rmse_sum / n_test) << "\n"; + if (!samples.global_error_variance_samples.empty()) { + std::cout << " sigma (last): " + << std::sqrt(samples.global_error_variance_samples.back()) * samples.y_std << "\n"; + } + } else { + double brier_sum = 0.0; + int correct = 0; + for (int i = 0; i < n_test; i++) { + double latent = 0.0; + for (int j = 0; j < num_samples; j++) + latent += samples.y_hat_test[static_cast(j * n_test + i)] / num_samples; + double p = StochTree::norm_cdf(latent); + double diff = p - y_ref[i]; + brier_sum += diff * diff; + correct += ((p >= 0.5) == (y_ref[i] >= 0.5)) ? 1 : 0; + } + std::cout << " Brier (test): " << brier_sum / n_test << "\n" + << " Acc (test): " << static_cast(correct) / n_test << "\n"; + } +} + +static void report_bcf_multivariate(const StochTree::BCFSamples& samples, + const std::vector& mu_ref, + const std::vector& tau1_ref, + const std::vector& tau2_ref, + const std::vector& y_ref, + const char* scenario_name) { + const int num_samples = samples.num_samples; + const int n_test = samples.num_test; + const int treatment_dim = 2; + double mu_rmse_sum = 0.0, tau1_rmse_sum = 0.0, tau2_rmse_sum = 0.0, y_rmse_sum = 0.0; + for (int i = 0; i < n_test; i++) { + double mu_hat = 0.0, tau1_hat = 0.0, tau2_hat = 0.0; + for (int j = 0; j < num_samples; j++) { + mu_hat += samples.mu_forest_predictions_test[static_cast(j * n_test + i)] / num_samples; + // tau layout: j * n_test * treatment_dim + n_test * treatment_idx + i + tau1_hat += samples.tau_forest_predictions_test[static_cast(j * n_test * treatment_dim + i)] / num_samples; + tau2_hat += samples.tau_forest_predictions_test[static_cast(j * n_test * treatment_dim + n_test + i)] / num_samples; + } + double mu_pred = mu_hat * samples.y_std + samples.y_bar; + double tau1_pred = tau1_hat * samples.y_std; + double tau2_pred = tau2_hat * samples.y_std; + mu_rmse_sum += (mu_pred - mu_ref[i]) * (mu_pred - mu_ref[i]); + tau1_rmse_sum += (tau1_pred - tau1_ref[i]) * (tau1_pred - tau1_ref[i]); + tau2_rmse_sum += (tau2_pred - tau2_ref[i]) * (tau2_pred - tau2_ref[i]); + } + for (int i = 0; i < n_test; i++) { + double y_hat = 0.0; + for (int j = 0; j < num_samples; j++) + y_hat += samples.y_hat_test[static_cast(j * n_test + i)] / num_samples; + y_rmse_sum += (y_hat - y_ref[i]) * (y_hat - y_ref[i]); + } + std::cout << "\n" + << scenario_name << ":\n" + << " mu RMSE (test): " << std::sqrt(mu_rmse_sum / n_test) << "\n" + << " tau1 RMSE (test): " << std::sqrt(tau1_rmse_sum / n_test) << "\n" + << " tau2 RMSE (test): " << std::sqrt(tau2_rmse_sum / n_test) << "\n" + << " y RMSE (test): " << std::sqrt(y_rmse_sum / n_test) << "\n"; + if (!samples.global_error_variance_samples.empty()) { + std::cout << " sigma (last): " + << std::sqrt(samples.global_error_variance_samples.back()) * samples.y_std << "\n" + << " sigma (truth): 0.5\n"; + } +} + +// ---- Scenario 0: constant-leaf mu + univariate-leaf tau (identity link) --- + +static void run_scenario_0(int n, int n_test, int p, + int num_trees_mu, int num_trees_tau, + int num_gfr, int num_mcmc, int seed) { + std::mt19937 rng(seed < 0 ? std::random_device{}() : static_cast(seed)); + SimpleBCFDataset train = generate_simple_bcf_data(n, p, rng); + SimpleBCFDataset test = generate_simple_bcf_data(n_test, p, rng); + + StochTree::BCFData data; + data.X_train = train.X.data(); + data.y_train = train.y.data(); + data.treatment_train = train.z.data(); + data.n_train = n; + data.p = p; + data.treatment_dim = 1; + data.X_test = test.X.data(); + data.treatment_test = test.z.data(); + data.n_test = n_test; + + StochTree::BCFConfig config; + config.num_trees_mu = num_trees_mu; + config.num_trees_tau = num_trees_tau; + config.random_seed = seed; + config.tau_leaf_model_type = StochTree::MeanLeafModelType::GaussianUnivariateRegression; + config.link_function = StochTree::LinkFunction::Identity; + config.standardize_outcome = true; + config.sample_sigma2_global = true; + config.var_weights_mu = std::vector(p, 1.0 / p); + config.var_weights_tau = std::vector(p, 1.0 / p); + config.feature_types = std::vector(p, StochTree::FeatureType::kNumeric); + config.sweep_update_indices_mu = std::vector(num_trees_mu, 0); + config.sweep_update_indices_tau = std::vector(num_trees_tau, 0); + std::iota(config.sweep_update_indices_mu.begin(), config.sweep_update_indices_mu.end(), 0); + std::iota(config.sweep_update_indices_tau.begin(), config.sweep_update_indices_tau.end(), 0); + + StochTree::BCFSamples samples; + StochTree::BCFSampler sampler(samples, config, data); + sampler.run_gfr(samples, num_gfr, /*keep_gfr=*/true); + sampler.run_mcmc(samples, /*num_burnin=*/0, /*keep_every=*/1, /*num_mcmc=*/num_mcmc); + sampler.postprocess_samples(samples); + report_bcf(samples, test.mu_true, test.tau_true, test.y, + StochTree::LinkFunction::Identity, + "Scenario 0 (BCF: constant mu + univariate tau, identity link)"); + std::cout << " sigma (truth): 0.5\n"; +} + +// ---- Scenario 1: probit BCF (constant-leaf mu + univariate-leaf tau) ---- + +static void run_scenario_1(int n, int n_test, int p, + int num_trees_mu, int num_trees_tau, + int num_gfr, int num_mcmc, int seed) { + std::mt19937 rng(seed < 0 ? std::random_device{}() : static_cast(seed)); + ProbitBCFDataset train = generate_probit_bcf_data(n, p, rng); + ProbitBCFDataset test = generate_probit_bcf_data(n_test, p, rng); + + StochTree::BCFData data; + data.X_train = train.X.data(); + data.y_train = train.y.data(); + data.treatment_train = train.z.data(); + data.n_train = n; + data.p = p; + data.treatment_dim = 1; + data.X_test = test.X.data(); + data.treatment_test = test.z.data(); + data.n_test = n_test; + + StochTree::BCFConfig config; + config.num_trees_mu = num_trees_mu; + config.num_trees_tau = num_trees_tau; + config.random_seed = seed; + config.tau_leaf_model_type = StochTree::MeanLeafModelType::GaussianUnivariateRegression; + config.link_function = StochTree::LinkFunction::Probit; + config.sample_sigma2_global = false; + config.var_weights_mu = std::vector(p, 1.0 / p); + config.var_weights_tau = std::vector(p, 1.0 / p); + config.feature_types = std::vector(p, StochTree::FeatureType::kNumeric); + config.sweep_update_indices_mu = std::vector(num_trees_mu, 0); + config.sweep_update_indices_tau = std::vector(num_trees_tau, 0); + std::iota(config.sweep_update_indices_mu.begin(), config.sweep_update_indices_mu.end(), 0); + std::iota(config.sweep_update_indices_tau.begin(), config.sweep_update_indices_tau.end(), 0); + + StochTree::BCFSamples samples; + StochTree::BCFSampler sampler(samples, config, data); + sampler.run_gfr(samples, num_gfr, /*keep_gfr=*/true); + sampler.run_mcmc(samples, /*num_burnin=*/0, /*keep_every=*/1, /*num_mcmc=*/num_mcmc); + sampler.postprocess_samples(samples); + report_bcf(samples, test.mu_true, test.tau_true, test.y, + StochTree::LinkFunction::Probit, + "Scenario 1 (BCF: constant mu + univariate tau, probit link)"); +} + +// ---- Scenario 2: constant-leaf mu + multivariate-leaf tau (identity link) --- + +static void run_scenario_2(int n, int n_test, int p, + int num_trees_mu, int num_trees_tau, + int num_gfr, int num_mcmc, int seed) { + std::mt19937 rng(seed < 0 ? std::random_device{}() : static_cast(seed)); + MultivariateBCFDataset train = generate_multivariate_bcf_data(n, p, rng); + MultivariateBCFDataset test = generate_multivariate_bcf_data(n_test, p, rng); + const int treatment_dim = 2; + + StochTree::BCFData data; + data.X_train = train.X.data(); + data.y_train = train.y.data(); + data.treatment_train = train.z.data(); + data.n_train = n; + data.p = p; + data.treatment_dim = treatment_dim; + data.X_test = test.X.data(); + data.treatment_test = test.z.data(); + data.n_test = n_test; + + StochTree::BCFConfig config; + config.num_trees_mu = num_trees_mu; + config.num_trees_tau = num_trees_tau; + config.random_seed = seed; + config.tau_leaf_model_type = StochTree::MeanLeafModelType::GaussianMultivariateRegression; + config.leaf_dim_tau = treatment_dim; + config.leaf_constant_tau = false; + config.link_function = StochTree::LinkFunction::Identity; + config.standardize_outcome = true; + config.sample_sigma2_global = true; + config.var_weights_mu = std::vector(p, 1.0 / p); + config.var_weights_tau = std::vector(p, 1.0 / p); + config.feature_types = std::vector(p, StochTree::FeatureType::kNumeric); + config.sweep_update_indices_mu = std::vector(num_trees_mu, 0); + config.sweep_update_indices_tau = std::vector(num_trees_tau, 0); + std::iota(config.sweep_update_indices_mu.begin(), config.sweep_update_indices_mu.end(), 0); + std::iota(config.sweep_update_indices_tau.begin(), config.sweep_update_indices_tau.end(), 0); + config.sigma2_leaf_tau_matrix = {0.5, 0.0, 0.0, 0.5}; // 0.5 * I_{2x2}, col-major + + StochTree::BCFSamples samples; + StochTree::BCFSampler sampler(samples, config, data); + sampler.run_gfr(samples, num_gfr, /*keep_gfr=*/true); + sampler.run_mcmc(samples, /*num_burnin=*/0, /*keep_every=*/1, /*num_mcmc=*/num_mcmc); + sampler.postprocess_samples(samples); + report_bcf_multivariate(samples, test.mu_true, test.tau1_true, test.tau2_true, test.y, + "Scenario 2 (BCF: constant mu + multivariate tau, identity link)"); +} + +// ---- Main ----------------------------------------------------------- + +int main(int argc, char** argv) { + int scenario = 0; + int n = 500; + int n_test = 100; + int p = 5; + int num_trees_mu = 200; + int num_trees_tau = 50; + int num_gfr = 20; + int num_mcmc = 100; + int seed = 1234; + + for (int i = 1; i < argc; ++i) { + std::string arg = argv[i]; + if ((arg == "--scenario" || arg == "--n" || arg == "--n_test" || arg == "--p" || + arg == "--num_trees_mu" || arg == "--num_trees_tau" || arg == "--num_gfr" || + arg == "--num_mcmc" || arg == "--seed") && + i + 1 < argc) { + int val = std::stoi(argv[++i]); + if (arg == "--scenario") + scenario = val; + else if (arg == "--n") + n = val; + else if (arg == "--n_test") + n_test = val; + else if (arg == "--p") + p = val; + else if (arg == "--num_trees_mu") + num_trees_mu = val; + else if (arg == "--num_trees_tau") + num_trees_tau = val; + else if (arg == "--num_gfr") + num_gfr = val; + else if (arg == "--num_mcmc") + num_mcmc = val; + else if (arg == "--seed") + seed = val; + } else { + std::cerr << "Unknown or incomplete argument: " << arg << "\n" + << "Usage: bcf_debug [--scenario N] [--n N] [--n_test N] [--p N]" + " [--num_trees_mu N] [--num_trees_tau N] [--num_gfr N] [--num_mcmc N] [--seed N]\n"; + return 1; + } + } + + switch (scenario) { + case 0: + run_scenario_0(n, n_test, p, num_trees_mu, num_trees_tau, num_gfr, num_mcmc, seed); + break; + case 1: + run_scenario_1(n, n_test, p, num_trees_mu, num_trees_tau, num_gfr, num_mcmc, seed); + break; + case 2: + run_scenario_2(n, n_test, p, num_trees_mu, num_trees_tau, num_gfr, num_mcmc, seed); + break; + default: + std::cerr << "Unknown scenario " << scenario + << ". Available: 0 (BCF: identity), 1 (BCF: probit), 2 (BCF: multivariate tau)\n"; + return 1; + } + return 0; +} diff --git a/debug/benchmark_cpp_vs_py_sampler.py b/debug/benchmark_cpp_vs_py_sampler.py new file mode 100644 index 00000000..616fb036 --- /dev/null +++ b/debug/benchmark_cpp_vs_py_sampler.py @@ -0,0 +1,141 @@ +"""Benchmark: C++ sampler loop vs. Python sampler loop. + +Compares runtime and test-set RMSE across run_cpp=True / False in BARTModel.sample(). + +Usage: + source venv/bin/activate + python debug/benchmark_cpp_vs_py_sampler.py +""" + +import argparse +import time +import numpy as np +from stochtree import BARTModel + +parser = argparse.ArgumentParser() +parser.add_argument("--num-chains", type=int, default=1) +args = parser.parse_args() +num_chains = args.num_chains + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +rng = np.random.default_rng(1234) + +n = 10000 +p = 10 +X = rng.uniform(size=(n, p)) +f_X = ( + np.where((X[:, 0] >= 0.00) & (X[:, 0] < 0.25), -7.5, 0) + + np.where((X[:, 0] >= 0.25) & (X[:, 0] < 0.50), -2.5, 0) + + np.where((X[:, 0] >= 0.50) & (X[:, 0] < 0.75), 2.5, 0) + + np.where((X[:, 0] >= 0.75) & (X[:, 0] < 1.00), 7.5, 0) +) +noise_sd = 1.0 +y = f_X + rng.normal(scale=noise_sd, size=n) + +test_frac = 0.2 +n_test = round(test_frac * n) +n_train = n - n_test +test_inds = rng.choice(n, size=n_test, replace=False) +train_inds = np.setdiff1d(np.arange(n), test_inds) + +X_train, X_test = X[train_inds], X[test_inds] +y_train, y_test = y[train_inds], y[test_inds] +f_test = f_X[test_inds] + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr = 10 +num_mcmc = 100 +num_trees = 200 +n_reps = 3 + +print( + f"n_train={n_train} n_test={n_test} p={p} " + f"num_trees={num_trees} num_gfr={num_gfr} num_mcmc={num_mcmc} " + f"num_chains={num_chains} reps={n_reps}\n" +) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + RMSE +# --------------------------------------------------------------------------- +def run_once(run_cpp, num_gfr, num_mcmc, seed): + m = BARTModel() + t0 = time.perf_counter() + m.sample( + X_train=X_train, + y_train=y_train, + num_gfr=num_gfr, + num_burnin=0, + num_mcmc=num_mcmc, + general_params={"random_seed": seed, "num_chains": num_chains}, + mean_forest_params={"num_trees": num_trees}, + run_cpp=run_cpp, + ) + elapsed_sample = time.perf_counter() - t0 + + t1 = time.perf_counter() + preds = m.predict(X=X_test, run_cpp=run_cpp) + elapsed_predict = time.perf_counter() - t1 + + yhat = preds["y_hat"].mean(axis=1) + rmse = float(np.sqrt(np.mean((yhat - y_test) ** 2))) + rmse_f = float(np.sqrt(np.mean((yhat - f_test) ** 2))) + return {"elapsed_sample": elapsed_sample, "elapsed_predict": elapsed_predict, "rmse": rmse, "rmse_f": rmse_f} + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds = [1000 + i for i in range(1, n_reps + 1)] + +results_cpp = [] +results_py = [] + +print("Running C++ sampler (run_cpp=True)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_cpp.append(run_once(run_cpp=True, num_gfr=num_gfr, num_mcmc=num_mcmc, seed=seed)) + +print("\nRunning Python sampler (run_cpp=False)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_py.append(run_once(run_cpp=False, num_gfr=num_gfr, num_mcmc=num_mcmc, seed=seed)) + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +def summarise(results): + keys = ["elapsed_sample", "elapsed_predict", "rmse", "rmse_f"] + out = {k: float(np.mean([r[k] for r in results])) for k in keys} + out["elapsed"] = out["elapsed_sample"] + out["elapsed_predict"] + out["elapsed_sd"] = float(np.std( + [r["elapsed_sample"] + r["elapsed_predict"] for r in results], ddof=1 + )) + return out + +s_cpp = summarise(results_cpp) +s_py = summarise(results_py) +rows = [("cpp (run_cpp=True)", s_cpp), ("py (run_cpp=False)", s_py)] + +print("\n--- Results ---") +print( + f"{'Sampler':<22} {'Total (s)':>10} {'Samp (s)':>10} {'Pred (s)':>10} {'SD':>8} " + f"{'RMSE (obs)':>12} {'RMSE (f)':>12}" +) +print("-" * 94) +for label, s in rows: + print( + f"{label:<22} {s['elapsed']:>10.3f} {s['elapsed_sample']:>10.3f} " + f"{s['elapsed_predict']:>10.3f} {s['elapsed_sd']:>8.3f} " + f"{s['rmse']:>12.4f} {s['rmse_f']:>12.4f}" + ) + +speedup = s_py["elapsed"] / s_cpp["elapsed"] +print(f"\nSpeedup (py / cpp): {speedup:.2f}x") +print( + f"RMSE delta (cpp - py): " + f"obs={s_cpp['rmse'] - s_py['rmse']:.4f} " + f"f={s_cpp['rmse_f'] - s_py['rmse_f']:.4f}" +) diff --git a/debug/benchmark_cpp_vs_py_sampler_bcf.py b/debug/benchmark_cpp_vs_py_sampler_bcf.py new file mode 100644 index 00000000..114342c1 --- /dev/null +++ b/debug/benchmark_cpp_vs_py_sampler_bcf.py @@ -0,0 +1,177 @@ +"""Benchmark: C++ sampler loop vs. Python sampler loop -- BCF. + +Compares runtime and test-set RMSE across run_cpp=True / False in BCFModel.sample(). +Simplest BCF case: univariate binary treatment, no RFX, no adaptive coding, +no treatment intercept. + +DGP: mu(X) is a step function on X[:,0]. tau(X) = 2 + 4*X[:,2] (linear CATE). + pi(X) = 0.2 + 0.6*X[:,3] (mild confounding). Z ~ Bernoulli(pi(X)). + y = mu(X) + tau(X)*Z + noise. + +Usage: + source venv/bin/activate + python debug/benchmark_cpp_vs_py_sampler_bcf.py +""" + +import argparse +import time +import numpy as np +from stochtree import BCFModel + +parser = argparse.ArgumentParser() +parser.add_argument("--num-chains", type=int, default=1) +args = parser.parse_args() +num_chains = args.num_chains + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +rng = np.random.default_rng(1234) + +n = 2000 +p = 10 +X = rng.uniform(size=(n, p)) + +mu_X = ( + np.where((X[:, 0] >= 0.00) & (X[:, 0] < 0.25), -7.5, 0) + + np.where((X[:, 0] >= 0.25) & (X[:, 0] < 0.50), -2.5, 0) + + np.where((X[:, 0] >= 0.50) & (X[:, 0] < 0.75), 2.5, 0) + + np.where((X[:, 0] >= 0.75) & (X[:, 0] < 1.00), 7.5, 0) +) +tau_X = 2 + 4 * X[:, 2] +pi_X = 0.2 + 0.6 * X[:, 3] +Z = rng.binomial(1, pi_X).astype(float) + +noise_sd = 1.0 +y = mu_X + tau_X * Z + rng.normal(scale=noise_sd, size=n) + +test_frac = 0.2 +n_test = round(test_frac * n) +n_train = n - n_test +test_inds = rng.choice(n, size=n_test, replace=False) +train_inds = np.setdiff1d(np.arange(n), test_inds) + +X_train, X_test = X[train_inds], X[test_inds] +Z_train, Z_test = Z[train_inds], Z[test_inds] +pi_train, pi_test = pi_X[train_inds], pi_X[test_inds] +y_train, y_test = y[train_inds], y[test_inds] +mu_test = mu_X[test_inds] +tau_test = tau_X[test_inds] +f_test = mu_test + tau_test * Z_test + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr = 10 +num_burnin = 0 +num_mcmc = 100 +num_trees_mu = 200 +num_trees_tau = 50 +n_reps = 3 + +print( + f"n_train={n_train} n_test={n_test} p={p}\n" + f"mu_trees={num_trees_mu} tau_trees={num_trees_tau} " + f"num_gfr={num_gfr} num_burnin={num_burnin} num_mcmc={num_mcmc} " + f"num_chains={num_chains} reps={n_reps}\n" +) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + accuracy metrics +# --------------------------------------------------------------------------- +def run_once(run_cpp: bool, seed: int) -> dict: + m = BCFModel() + t0 = time.perf_counter() + m.sample( + X_train=X_train, + Z_train=Z_train, + y_train=y_train, + propensity_train=pi_train, + num_gfr=num_gfr, + num_burnin=num_burnin, + num_mcmc=num_mcmc, + general_params={ + "adaptive_coding": False, + "random_seed": seed, + "num_chains": num_chains, + "propensity_covariate": "prognostic", + }, + prognostic_forest_params={"num_trees": num_trees_mu}, + treatment_effect_forest_params={ + "num_trees": num_trees_tau, + "sample_intercept": False, + }, + run_cpp=run_cpp, + ) + elapsed_sample = time.perf_counter() - t0 + + t1 = time.perf_counter() + preds = m.predict(X=X_test, Z=Z_test, propensity=pi_test, run_cpp=run_cpp) + elapsed_predict = time.perf_counter() - t1 + + y_hat = preds["y_hat"].mean(axis=1) + tau_hat = preds["tau_hat"].mean(axis=1) + + return { + "elapsed_sample": elapsed_sample, + "elapsed_predict": elapsed_predict, + "rmse_y": float(np.sqrt(np.mean((y_hat - y_test) ** 2))), + "rmse_f": float(np.sqrt(np.mean((y_hat - f_test) ** 2))), + "rmse_tau": float(np.sqrt(np.mean((tau_hat - tau_test) ** 2))), + } + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds = [1000 + i for i in range(1, n_reps + 1)] + +results_cpp = [] +results_py = [] + +print("Running C++ sampler (run_cpp=True)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_cpp.append(run_once(run_cpp=True, seed=seed)) + +print("\nRunning Python sampler (run_cpp=False)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_py.append(run_once(run_cpp=False, seed=seed)) + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +def summarise(results: list) -> dict: + keys = ["elapsed_sample", "elapsed_predict", "rmse_y", "rmse_f", "rmse_tau"] + out = {k: float(np.mean([r[k] for r in results])) for k in keys} + out["elapsed"] = out["elapsed_sample"] + out["elapsed_predict"] + out["elapsed_sd"] = float(np.std( + [r["elapsed_sample"] + r["elapsed_predict"] for r in results], ddof=1 + )) + return out + +s_cpp = summarise(results_cpp) +s_py = summarise(results_py) +rows = [("cpp (run_cpp=True)", s_cpp), ("py (run_cpp=False)", s_py)] + +print("\n--- Results ---") +print( + f"{'Sampler':<22} {'Total (s)':>10} {'Samp (s)':>10} {'Pred (s)':>10} {'SD':>8} " + f"{'RMSE (obs)':>12} {'RMSE (f)':>12} {'RMSE (tau)':>12}" +) +print("-" * 104) +for label, s in rows: + print( + f"{label:<22} {s['elapsed']:>10.3f} {s['elapsed_sample']:>10.3f} " + f"{s['elapsed_predict']:>10.3f} {s['elapsed_sd']:>8.3f} " + f"{s['rmse_y']:>12.4f} {s['rmse_f']:>12.4f} {s['rmse_tau']:>12.4f}" + ) + +speedup = s_py["elapsed"] / s_cpp["elapsed"] +print(f"\nSpeedup (py / cpp): {speedup:.2f}x") +print( + f"RMSE delta (cpp - py): " + f"obs={s_cpp['rmse_y'] - s_py['rmse_y']:.4f} " + f"f={s_cpp['rmse_f'] - s_py['rmse_f']:.4f} " + f"tau={s_cpp['rmse_tau'] - s_py['rmse_tau']:.4f}" +) diff --git a/debug/benchmark_cpp_vs_py_sampler_bcf_adaptive_coding.py b/debug/benchmark_cpp_vs_py_sampler_bcf_adaptive_coding.py new file mode 100644 index 00000000..d384d685 --- /dev/null +++ b/debug/benchmark_cpp_vs_py_sampler_bcf_adaptive_coding.py @@ -0,0 +1,220 @@ +"""Benchmark: C++ sampler loop vs. Python sampler loop -- BCF with adaptive coding. + +Exercises SampleAdaptiveCodingParameters() for binary treatment. Verifies: + - The C++ path runs without error. + - b_0_samples and b_1_samples are populated with the right shape. + - mu_hat + Z * tau_hat == y_hat (internal decomposition check on train set). + - CATE RMSE (cpp) is close to CATE RMSE (python) -- large differences indicate + a residual accounting bug in SampleAdaptiveCodingParameters or the + mu/tau prediction split. + - Speedup is reported for reference, though the primary goal is correctness. + +DGP: + mu(X) = step function on X[:,0] + tau_forest(X) = 2 * X[:,2] (forest component, scale-free) + true_b_0 = -0.5 + true_b_1 = 1.5 + CATE(X) = (true_b_1 - true_b_0) * tau_forest(X) = 4 * X[:,2] + pi(X) = 0.2 + 0.6 * X[:,3] + Z ~ Bernoulli(pi(X)) + y = mu(X) + (true_b_0*(1-Z) + true_b_1*Z) * tau_forest(X) + noise + +Usage: + source venv/bin/activate + python debug/benchmark_cpp_vs_py_sampler_bcf_adaptive_coding.py +""" + +import argparse +import time +import numpy as np +from stochtree import BCFModel + +parser = argparse.ArgumentParser() +parser.add_argument("--num-chains", type=int, default=1) +args = parser.parse_args() +num_chains = args.num_chains + +# --------------------------------------------------------------------------- +# Shared RNG and dataset sizes +# --------------------------------------------------------------------------- +rng = np.random.default_rng(1234) + +n = 2000 +p = 10 +noise_sd = 1.0 +test_frac = 0.2 +n_test = round(test_frac * n) +n_train = n - n_test + +num_gfr = 10 +num_burnin = 0 +num_mcmc = 100 +num_trees_mu = 200 +num_trees_tau = 50 +n_reps = 3 + +# --------------------------------------------------------------------------- +# DGP: binary treatment with adaptive coding +# --------------------------------------------------------------------------- +X_all = rng.uniform(size=(n, p)) + +mu_X = ( + np.where((X_all[:, 0] >= 0.00) & (X_all[:, 0] < 0.25), -7.5, 0) + + np.where((X_all[:, 0] >= 0.25) & (X_all[:, 0] < 0.50), -2.5, 0) + + np.where((X_all[:, 0] >= 0.50) & (X_all[:, 0] < 0.75), 2.5, 0) + + np.where((X_all[:, 0] >= 0.75) & (X_all[:, 0] < 1.00), 7.5, 0) +) +TRUE_B_0 = -0.5 +TRUE_B_1 = 1.5 +tau_forest_X = 2.0 * X_all[:, 2] # forest component (no coding) +tau_X = (TRUE_B_1 - TRUE_B_0) * tau_forest_X # CATE = (b_1 - b_0) * tau_forest +pi_X = 0.2 + 0.6 * X_all[:, 3] +Z_all = rng.binomial(1, pi_X).astype(float) +coded_basis = TRUE_B_0 * (1 - Z_all) + TRUE_B_1 * Z_all +y_all = mu_X + coded_basis * tau_forest_X + rng.normal(scale=noise_sd, size=n) + +test_inds = rng.choice(n, size=n_test, replace=False) +train_inds = np.setdiff1d(np.arange(n), test_inds) + +X_train = X_all[train_inds]; X_test = X_all[test_inds] +Z_train = Z_all[train_inds]; Z_test = Z_all[test_inds] +pi_train = pi_X[train_inds]; pi_test = pi_X[test_inds] +y_train = y_all[train_inds]; y_test = y_all[test_inds] +mu_test = mu_X[test_inds] +tau_test = tau_X[test_inds] # true CATE = (b_1-b_0)*tau_forest(X_test) +# true fitted value: mu(X) + (b_0*(1-Z) + b_1*Z) * tau_forest(X) +coded_test = TRUE_B_0 * (1 - Z_test) + TRUE_B_1 * Z_test +f_test = mu_test + coded_test * tau_forest_X[test_inds] + +print( + f"n_train={n_train} n_test={n_test} p={p}\n" + f"mu_trees={num_trees_mu} tau_trees={num_trees_tau} " + f"num_gfr={num_gfr} num_burnin={num_burnin} num_mcmc={num_mcmc} " + f"num_chains={num_chains} reps={n_reps}\n" + f"true_b_0={TRUE_B_0} true_b_1={TRUE_B_1}\n" +) + +# --------------------------------------------------------------------------- +# Runner +# --------------------------------------------------------------------------- +def run_once(run_cpp: bool, seed: int) -> dict: + m = BCFModel() + t0 = time.perf_counter() + m.sample( + X_train=X_train, + Z_train=Z_train, + y_train=y_train, + propensity_train=pi_train, + num_gfr=num_gfr, + num_burnin=num_burnin, + num_mcmc=num_mcmc, + general_params={ + "adaptive_coding": True, + "random_seed": seed, + "num_chains": num_chains, + "propensity_covariate": "prognostic", + }, + prognostic_forest_params={"num_trees": num_trees_mu}, + treatment_effect_forest_params={"num_trees": num_trees_tau}, + run_cpp=run_cpp, + ) + elapsed_sample = time.perf_counter() - t0 + + # Internal consistency check on train set (still available on the model object) + max_decomp_err_train = float(np.max(np.abs( + m.y_hat_train - (m.mu_hat_train + Z_train[:, None] * m.tau_hat_train) + ))) + + t1 = time.perf_counter() + preds = m.predict(X=X_test, Z=Z_test, propensity=pi_test, run_cpp=run_cpp) + elapsed_predict = time.perf_counter() - t1 + + yhat = preds["y_hat"].mean(axis=1) + tauhat = preds["tau_hat"].mean(axis=1) + + max_decomp_err_test = float(np.max(np.abs( + preds["y_hat"] - (preds["mu_hat"] + Z_test[:, None] * preds["tau_hat"]) + ))) + + b0_mean = float(np.mean(m.b0_samples)) + b1_mean = float(np.mean(m.b1_samples)) + + return { + "elapsed_sample": elapsed_sample, + "elapsed_predict": elapsed_predict, + "b0_mean": b0_mean, + "b1_mean": b1_mean, + "b0_shape": m.b0_samples.shape, + "b1_shape": m.b1_samples.shape, + "max_decomp_err_train": max_decomp_err_train, + "max_decomp_err_test": max_decomp_err_test, + "rmse_y": float(np.sqrt(np.mean((yhat - y_test) ** 2))), + "rmse_f": float(np.sqrt(np.mean((yhat - f_test) ** 2))), + "rmse_tau": float(np.sqrt(np.mean((tauhat - tau_test) ** 2))), + } + +# --------------------------------------------------------------------------- +# Run +# --------------------------------------------------------------------------- +seeds = [1000 + i for i in range(1, n_reps + 1)] + +print("=" * 60) +print("BINARY TREATMENT (adaptive_coding=True)") +print(f"True b_0={TRUE_B_0} b_1={TRUE_B_1} CATE=(b1-b0)*tau_forest(X)") +print("=" * 60) + +results_cpp = [] +results_py = [] + +print("Running C++ sampler (run_cpp=True)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_cpp.append(run_once(run_cpp=True, seed=seed)) + +print("\nRunning Python sampler (run_cpp=False)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_py.append(run_once(run_cpp=False, seed=seed)) + +def summarise(results): + keys = ["elapsed_sample", "elapsed_predict", "b0_mean", "b1_mean", + "max_decomp_err_train", "max_decomp_err_test", + "rmse_y", "rmse_f", "rmse_tau"] + out = {k: float(np.mean([r[k] for r in results])) for k in keys} + out["elapsed"] = out["elapsed_sample"] + out["elapsed_predict"] + out["elapsed_sd"] = float(np.std( + [r["elapsed_sample"] + r["elapsed_predict"] for r in results], ddof=1 + )) + out["b0_shape"] = results[0]["b0_shape"] + out["b1_shape"] = results[0]["b1_shape"] + return out + +s_cpp = summarise(results_cpp) +s_py = summarise(results_py) + +print("\n--- Results ---") +print(f"b0_samples shape cpp={s_cpp['b0_shape']} py={s_py['b0_shape']}") +print(f"b1_samples shape cpp={s_cpp['b1_shape']} py={s_py['b1_shape']}") +print() +print(f"{'Sampler':<22} {'Total (s)':>9} {'Samp (s)':>9} {'Pred (s)':>9} {'SD':>6} " + f"{'b_0 mean':>8} {'b_1 mean':>8} " + f"{'max_decomp_tr':>13} {'max_decomp_te':>13} " + f"{'RMSE(y)':>8} {'RMSE(f)':>8} {'RMSE(tau)':>10}") +print("-" * 140) +for label, s in [("cpp (run_cpp=True)", s_cpp), ("py (run_cpp=False)", s_py)]: + print( + f"{label:<22} {s['elapsed']:>9.3f} {s['elapsed_sample']:>9.3f} " + f"{s['elapsed_predict']:>9.3f} {s['elapsed_sd']:>6.3f} " + f"{s['b0_mean']:>8.4f} {s['b1_mean']:>8.4f} " + f"{s['max_decomp_err_train']:>13.2e} {s['max_decomp_err_test']:>13.2e} " + f"{s['rmse_y']:>8.4f} {s['rmse_f']:>8.4f} {s['rmse_tau']:>10.4f}" + ) +print(f"True b_0={TRUE_B_0:.4f} b_1={TRUE_B_1:.4f}") +speedup = s_py["elapsed"] / s_cpp["elapsed"] +print(f"Speedup (py / cpp): {speedup:.2f}x") +print( + f"RMSE delta (cpp - py): " + f"y={s_cpp['rmse_y'] - s_py['rmse_y']:.4f} " + f"f={s_cpp['rmse_f'] - s_py['rmse_f']:.4f} " + f"tau={s_cpp['rmse_tau'] - s_py['rmse_tau']:.4f}" +) diff --git a/debug/benchmark_cpp_vs_py_sampler_bcf_multivariate.py b/debug/benchmark_cpp_vs_py_sampler_bcf_multivariate.py new file mode 100644 index 00000000..46ab9c39 --- /dev/null +++ b/debug/benchmark_cpp_vs_py_sampler_bcf_multivariate.py @@ -0,0 +1,183 @@ +"""Benchmark: C++ sampler loop vs. Python sampler loop -- multivariate treatment BCF. + +Compares runtime and accuracy across run_cpp=True / False in BCFModel.sample() +with a 2-column binary treatment (multivariate BCF). + +DGP: propensity pi(X) = [0.25 + 0.5*X[:,0], 0.75 - 0.5*X[:,1]] (2-column). + mu(X) = 5*pi[:,0] + 2*pi[:,1] + 2*X[:,2]. + tau(X) = [X[:,1], X[:,2]] (2-column CATE). + Z ~ Bernoulli(pi(X)) column-wise (binary, shape n x 2). + y = mu(X) + sum(Z * tau(X), axis=1) + noise. +Adaptive coding is disabled; sigma2_leaf is not sampled for the tau forest. + +Usage: + source venv/bin/activate + python debug/benchmark_cpp_vs_py_sampler_bcf_multivariate.py +""" + +import argparse +import time +import numpy as np +from stochtree import BCFModel + +parser = argparse.ArgumentParser() +parser.add_argument("--num-chains", type=int, default=1) +args = parser.parse_args() +num_chains = args.num_chains + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +rng = np.random.default_rng(1234) + +n = 2000 +p = 5 +snr = 2.0 + +X = rng.uniform(size=(n, p)) +pi_x = np.c_[0.25 + 0.5 * X[:, 0], 0.75 - 0.5 * X[:, 1]] +mu_x = pi_x[:, 0] * 5 + pi_x[:, 1] * 2 + 2 * X[:, 2] +tau_x = np.c_[X[:, 1], X[:, 2]] +Z = (rng.uniform(size=(n, 2)) < pi_x).astype(float) +E_XZ = mu_x + (Z * tau_x).sum(axis=1) +y = E_XZ + rng.normal(size=n) * (E_XZ.std() / snr) + +test_frac = 0.2 +n_test = round(test_frac * n) +n_train = n - n_test +test_inds = rng.choice(n, size=n_test, replace=False) +train_inds = np.setdiff1d(np.arange(n), test_inds) + +X_train, X_test = X[train_inds], X[test_inds] +Z_train, Z_test = Z[train_inds], Z[test_inds] +pi_train, pi_test = pi_x[train_inds], pi_x[test_inds] +y_train, y_test = y[train_inds], y[test_inds] +mu_test = mu_x[test_inds] +tau_test = tau_x[test_inds] +f_test = mu_test + (Z_test * tau_test).sum(axis=1) + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr = 10 +num_burnin = 0 +num_mcmc = 100 +num_trees_mu = 200 +num_trees_tau = 50 +n_reps = 3 + +print( + f"n_train={n_train} n_test={n_test} p={p} treatment_dim=2\n" + f"mu_trees={num_trees_mu} tau_trees={num_trees_tau} " + f"num_gfr={num_gfr} num_burnin={num_burnin} num_mcmc={num_mcmc} " + f"num_chains={num_chains} reps={n_reps}\n" +) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + accuracy metrics +# --------------------------------------------------------------------------- +def run_once(run_cpp: bool, seed: int) -> dict: + m = BCFModel() + t0 = time.perf_counter() + m.sample( + X_train=X_train, + Z_train=Z_train, + y_train=y_train, + propensity_train=pi_train, + num_gfr=num_gfr, + num_burnin=num_burnin, + num_mcmc=num_mcmc, + general_params={ + "adaptive_coding": False, + "random_seed": seed, + "num_chains": num_chains, + }, + prognostic_forest_params={"num_trees": num_trees_mu}, + treatment_effect_forest_params={ + "num_trees": num_trees_tau, + "sample_sigma2_leaf": False, + "sample_intercept": False, + }, + run_cpp=run_cpp, + ) + elapsed_sample = time.perf_counter() - t0 + + t1 = time.perf_counter() + preds = m.predict(X=X_test, Z=Z_test, propensity=pi_test, run_cpp=run_cpp) + elapsed_predict = time.perf_counter() - t1 + + mu_hat = preds["mu_hat"].mean(axis=1) + # tau_hat is (n, treatment_dim, num_samples) for multivariate treatment + tau_hat1 = preds["tau_hat"][:, 0, :].mean(axis=1) + tau_hat2 = preds["tau_hat"][:, 1, :].mean(axis=1) + y_hat = preds["y_hat"].mean(axis=1) + + return { + "elapsed_sample": elapsed_sample, + "elapsed_predict": elapsed_predict, + "rmse_y": float(np.sqrt(np.mean((y_hat - y_test) ** 2))), + "rmse_f": float(np.sqrt(np.mean((y_hat - f_test) ** 2))), + "rmse_mu": float(np.sqrt(np.mean((mu_hat - mu_test) ** 2))), + "rmse_tau1": float(np.sqrt(np.mean((tau_hat1 - tau_test[:, 0]) ** 2))), + "rmse_tau2": float(np.sqrt(np.mean((tau_hat2 - tau_test[:, 1]) ** 2))), + } + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds = [1000 + i for i in range(1, n_reps + 1)] + +results_cpp = [] +results_py = [] + +print("Running C++ sampler (run_cpp=True)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_cpp.append(run_once(run_cpp=True, seed=seed)) + +print("\nRunning Python sampler (run_cpp=False)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_py.append(run_once(run_cpp=False, seed=seed)) + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +def summarise(results: list) -> dict: + keys = ["elapsed_sample", "elapsed_predict", "rmse_y", "rmse_f", "rmse_mu", "rmse_tau1", "rmse_tau2"] + out = {k: float(np.mean([r[k] for r in results])) for k in keys} + out["elapsed"] = out["elapsed_sample"] + out["elapsed_predict"] + out["elapsed_sd"] = float(np.std( + [r["elapsed_sample"] + r["elapsed_predict"] for r in results], ddof=1 + )) + return out + +s_cpp = summarise(results_cpp) +s_py = summarise(results_py) +rows = [("cpp (run_cpp=True)", s_cpp), ("py (run_cpp=False)", s_py)] + +print("\n--- Results ---") +print( + f"{'Sampler':<22} {'Total (s)':>9} {'Samp (s)':>9} {'Pred (s)':>9} {'SD':>8} " + f"{'RMSE(y)':>9} {'RMSE(f)':>9} {'RMSE(mu)':>9} " + f"{'RMSE(tau1)':>10} {'RMSE(tau2)':>10}" +) +print("-" * 115) +for label, s in rows: + print( + f"{label:<22} {s['elapsed']:>9.3f} {s['elapsed_sample']:>9.3f} " + f"{s['elapsed_predict']:>9.3f} {s['elapsed_sd']:>8.3f} " + f"{s['rmse_y']:>9.4f} {s['rmse_f']:>9.4f} {s['rmse_mu']:>9.4f} " + f"{s['rmse_tau1']:>10.4f} {s['rmse_tau2']:>10.4f}" + ) + +speedup = s_py["elapsed"] / s_cpp["elapsed"] +print(f"\nSpeedup (py / cpp): {speedup:.2f}x") +print( + f"RMSE delta (cpp - py): " + f"y={s_cpp['rmse_y'] - s_py['rmse_y']:.4f} " + f"f={s_cpp['rmse_f'] - s_py['rmse_f']:.4f} " + f"mu={s_cpp['rmse_mu'] - s_py['rmse_mu']:.4f} " + f"tau1={s_cpp['rmse_tau1'] - s_py['rmse_tau1']:.4f} " + f"tau2={s_cpp['rmse_tau2'] - s_py['rmse_tau2']:.4f}" +) diff --git a/debug/benchmark_cpp_vs_py_sampler_bcf_probit.py b/debug/benchmark_cpp_vs_py_sampler_bcf_probit.py new file mode 100644 index 00000000..1de1267a --- /dev/null +++ b/debug/benchmark_cpp_vs_py_sampler_bcf_probit.py @@ -0,0 +1,195 @@ +"""Benchmark: C++ sampler loop vs. Python sampler loop -- probit BCF. + +Compares runtime, Brier score on outcome, and latent-scale tau RMSE across +run_cpp=True / False in BCFModel.sample(). + +DGP (latent-index model): + w = mu(X) + tau(X)*Z + eps, eps ~ N(0, 1) + y = 1(w > 0) + mu(X) = 1 + 2*X[:,0] + X[:,1] (confounded with propensity) + tau(X) = 0.5 + X[:,2] (latent-scale CATE) + pi(X) = 0.4 + 0.2*X[:,0] (mild confounding) + Z ~ Bernoulli(pi(X)) + +Metrics: + Brier score: mean((mean_s Phi(mu_hat[i,s] + tau_hat[i,s]*Z[i]) - y[i])^2) + RMSE(tau): sqrt(mean((mean_s tau_hat_test[i,s] - tau_test[i])^2)) + +Usage: + source venv/bin/activate + python debug/benchmark_cpp_vs_py_sampler_bcf_probit.py +""" + +import argparse +import time +import numpy as np +from scipy.stats import norm +from stochtree import BCFModel, OutcomeModel + +parser = argparse.ArgumentParser() +parser.add_argument("--num-chains", type=int, default=1) +args = parser.parse_args() +num_chains = args.num_chains + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +rng = np.random.default_rng(1234) + +n = 2000 +p = 5 +X = rng.uniform(size=(n, p)) + +mu_X = 1 + 2 * X[:, 0] + X[:, 1] +tau_X = 0.5 + X[:, 2] +pi_X = 0.4 + 0.2 * X[:, 0] +Z = rng.binomial(1, pi_X).astype(float) + +w = mu_X + tau_X * Z + rng.standard_normal(n) +y = (w > 0).astype(float) + +test_frac = 0.2 +n_test = round(test_frac * n) +n_train = n - n_test +test_inds = rng.choice(n, size=n_test, replace=False) +train_inds = np.setdiff1d(np.arange(n), test_inds) + +X_train, X_test = X[train_inds], X[test_inds] +Z_train, Z_test = Z[train_inds], Z[test_inds] +pi_train, pi_test = pi_X[train_inds], pi_X[test_inds] +y_train, y_test = y[train_inds], y[test_inds] +tau_test = tau_X[test_inds] + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr = 10 +num_burnin = 0 +num_mcmc = 100 +num_trees_mu = 200 +num_trees_tau = 50 +n_reps = 3 + +print( + f"n_train={n_train} n_test={n_test} p={p}\n" + f"mu_trees={num_trees_mu} tau_trees={num_trees_tau} " + f"num_gfr={num_gfr} num_burnin={num_burnin} num_mcmc={num_mcmc} " + f"num_chains={num_chains} reps={n_reps}\n" +) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + accuracy metrics +# --------------------------------------------------------------------------- +def run_once(run_cpp: bool, seed: int) -> dict: + m = BCFModel() + t0 = time.perf_counter() + m.sample( + X_train=X_train, + Z_train=Z_train, + y_train=y_train, + propensity_train=pi_train, + num_gfr=num_gfr, + num_burnin=num_burnin, + num_mcmc=num_mcmc, + general_params={ + "random_seed": seed, + "num_chains": num_chains, + "outcome_model": OutcomeModel(outcome="binary", link="probit"), + "sample_sigma2_global": False, + }, + prognostic_forest_params={ + "num_trees": num_trees_mu, + "sample_sigma2_leaf": False, + }, + treatment_effect_forest_params={ + "num_trees": num_trees_tau, + "sample_sigma2_leaf": False, + "sample_intercept": False, + }, + run_cpp=run_cpp, + ) + elapsed_sample = time.perf_counter() - t0 + + # Request latent-scale mu and tau (scale="linear", no probit transform applied) + t1 = time.perf_counter() + preds = m.predict( + X=X_test, Z=Z_test, propensity=pi_test, + terms=["mu", "tau"], scale="linear", run_cpp=run_cpp + ) + elapsed_predict = time.perf_counter() - t1 + + # mu_hat, tau_hat: (n_test, num_samples) — latent scale + mu_hat = preds["mu_hat"] # (n_test, num_samples) + tau_hat = preds["tau_hat"] # (n_test, num_samples) + + # P(Y=1 | X, Z, sample s) = Phi(mu_hat[i,s] + tau_hat[i,s] * Z_test[i]) + linear_pred = mu_hat + tau_hat * Z_test[:, np.newaxis] + p_hat_samples = norm.cdf(linear_pred) # (n_test, num_samples) + p_hat_mean = p_hat_samples.mean(axis=1) # (n_test,) + + tau_hat_mean = tau_hat.mean(axis=1) # (n_test,) + + brier = float(np.mean((p_hat_mean - y_test) ** 2)) + rmse_tau = float(np.sqrt(np.mean((tau_hat_mean - tau_test) ** 2))) + + return { + "elapsed_sample": elapsed_sample, + "elapsed_predict": elapsed_predict, + "brier": brier, + "rmse_tau": rmse_tau, + } + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds = [1000 + i for i in range(1, n_reps + 1)] + +results_cpp = [] +results_py = [] + +print("Running C++ sampler (run_cpp=True)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_cpp.append(run_once(run_cpp=True, seed=seed)) + +print("\nRunning Python sampler (run_cpp=False)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_py.append(run_once(run_cpp=False, seed=seed)) + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +def summarise(results: list) -> dict: + keys = ["elapsed_sample", "elapsed_predict", "brier", "rmse_tau"] + out = {k: float(np.mean([r[k] for r in results])) for k in keys} + out["elapsed"] = out["elapsed_sample"] + out["elapsed_predict"] + out["elapsed_sd"] = float(np.std( + [r["elapsed_sample"] + r["elapsed_predict"] for r in results], ddof=1 + )) + return out + +s_cpp = summarise(results_cpp) +s_py = summarise(results_py) +rows = [("cpp (run_cpp=True)", s_cpp), ("py (run_cpp=False)", s_py)] + +print("\n--- Results ---") +print( + f"{'Sampler':<22} {'Total (s)':>10} {'Samp (s)':>10} {'Pred (s)':>10} {'SD':>10} " + f"{'Brier':>10} {'RMSE (tau)':>12}" +) +print("-" * 92) +for label, s in rows: + print( + f"{label:<22} {s['elapsed']:>10.3f} {s['elapsed_sample']:>10.3f} " + f"{s['elapsed_predict']:>10.3f} {s['elapsed_sd']:>10.3f} " + f"{s['brier']:>10.4f} {s['rmse_tau']:>12.4f}" + ) + +speedup = s_py["elapsed"] / s_cpp["elapsed"] +print(f"\nSpeedup (py / cpp): {speedup:.2f}x") +print( + f"Delta (cpp - py): " + f"brier={s_cpp['brier'] - s_py['brier']:.4f} " + f"rmse_tau={s_cpp['rmse_tau'] - s_py['rmse_tau']:.4f}" +) diff --git a/debug/benchmark_cpp_vs_py_sampler_bcf_tau_intercept.py b/debug/benchmark_cpp_vs_py_sampler_bcf_tau_intercept.py new file mode 100644 index 00000000..10583902 --- /dev/null +++ b/debug/benchmark_cpp_vs_py_sampler_bcf_tau_intercept.py @@ -0,0 +1,341 @@ +"""Benchmark: C++ sampler loop vs. Python sampler loop -- BCF with treatment intercept. + +Exercises SampleParametricTreatmentEffect() for both univariate and multivariate +treatment by toggling sample_intercept=True. Verifies: + - The C++ path runs without error. + - tau_0_samples is populated and has the right shape. + - CATE RMSE (cpp) is close to CATE RMSE (python) -- large differences indicate + a residual accounting bug in the new intercept step. + - Speedup is reported for reference, though the primary goal here is correctness. + +DGP (univariate): + mu(X) = step function on X[:,0] + tau_0 = 1.5 (true global intercept) + tau(X) = tau_0 + 2*X[:,2] (full CATE) + pi(X) = 0.2 + 0.6*X[:,3] + Z ~ Bernoulli(pi(X)) + y = mu(X) + (tau_0 + tau(X)) * Z + noise + +DGP (multivariate, treatment_dim=2): + Same mu(X); tau_0 = [0.5, 1.0]; tau_k(X) = tau_0[k] + X[:,k+1] + +Usage: + source venv/bin/activate + python debug/benchmark_cpp_vs_py_sampler_bcf_tau_intercept.py +""" + +import argparse +import time +import numpy as np +from stochtree import BCFModel + +parser = argparse.ArgumentParser() +parser.add_argument("--num-chains", type=int, default=1) +args = parser.parse_args() +num_chains = args.num_chains + +# --------------------------------------------------------------------------- +# Shared RNG and dataset sizes +# --------------------------------------------------------------------------- +rng = np.random.default_rng(1234) + +n = 2000 +p = 10 +noise_sd = 1.0 +test_frac = 0.2 +n_test = round(test_frac * n) +n_train = n - n_test + +num_gfr = 10 +num_burnin = 0 +num_mcmc = 100 +num_trees_mu = 200 +num_trees_tau = 50 +n_reps = 3 + +# --------------------------------------------------------------------------- +# DGP: univariate binary treatment with global tau_0 +# --------------------------------------------------------------------------- +X_all = rng.uniform(size=(n, p)) + +mu_X = ( + np.where((X_all[:, 0] >= 0.00) & (X_all[:, 0] < 0.25), -7.5, 0) + + np.where((X_all[:, 0] >= 0.25) & (X_all[:, 0] < 0.50), -2.5, 0) + + np.where((X_all[:, 0] >= 0.50) & (X_all[:, 0] < 0.75), 2.5, 0) + + np.where((X_all[:, 0] >= 0.75) & (X_all[:, 0] < 1.00), 7.5, 0) +) +TRUE_TAU0_UNIVARIATE = 1.5 +tau_forest_X = 2.0 * X_all[:, 2] # forest component only +tau_X = TRUE_TAU0_UNIVARIATE + tau_forest_X # full CATE = tau_0 + tau(X) +pi_X = 0.2 + 0.6 * X_all[:, 3] +Z_all = rng.binomial(1, pi_X).astype(float) +y_all = mu_X + tau_X * Z_all + rng.normal(scale=noise_sd, size=n) + +test_inds = rng.choice(n, size=n_test, replace=False) +train_inds = np.setdiff1d(np.arange(n), test_inds) + +X_train_u = X_all[train_inds]; X_test_u = X_all[test_inds] +Z_train_u = Z_all[train_inds]; Z_test_u = Z_all[test_inds] +pi_train_u = pi_X[train_inds]; pi_test_u = pi_X[test_inds] +y_train_u = y_all[train_inds]; y_test_u = y_all[test_inds] +mu_test_u = mu_X[test_inds] +tau_test_u = tau_X[test_inds] +f_test_u = mu_test_u + tau_test_u * Z_test_u + +# --------------------------------------------------------------------------- +# DGP: multivariate (2-column) treatment with per-arm global tau_0 +# --------------------------------------------------------------------------- +X_all_mv = rng.uniform(size=(n, p)) +TRUE_TAU0_MV = np.array([0.5, 1.0]) +pi_mv = np.c_[0.25 + 0.5 * X_all_mv[:, 0], 0.75 - 0.5 * X_all_mv[:, 1]] +mu_mv = pi_mv[:, 0] * 5 + pi_mv[:, 1] * 2 + 2 * X_all_mv[:, 2] +tau_forest_mv = np.c_[X_all_mv[:, 1], X_all_mv[:, 2]] # forest component only +tau_mv = TRUE_TAU0_MV + tau_forest_mv # full CATE +Z_mv = (rng.uniform(size=(n, 2)) < pi_mv).astype(float) +y_mv = mu_mv + (Z_mv * tau_mv).sum(axis=1) + rng.normal(size=n) * noise_sd + +test_inds_mv = rng.choice(n, size=n_test, replace=False) +train_inds_mv = np.setdiff1d(np.arange(n), test_inds_mv) + +X_train_mv = X_all_mv[train_inds_mv]; X_test_mv = X_all_mv[test_inds_mv] +Z_train_mv = Z_mv[train_inds_mv]; Z_test_mv = Z_mv[test_inds_mv] +pi_train_mv = pi_mv[train_inds_mv]; pi_test_mv = pi_mv[test_inds_mv] +y_train_mv = y_mv[train_inds_mv]; y_test_mv = y_mv[test_inds_mv] +mu_test_mv = mu_mv[test_inds_mv] +tau_test_mv = tau_mv[test_inds_mv] +f_test_mv = mu_test_mv + (Z_test_mv * tau_test_mv).sum(axis=1) + +print( + f"n_train={n_train} n_test={n_test} p={p}\n" + f"mu_trees={num_trees_mu} tau_trees={num_trees_tau} " + f"num_gfr={num_gfr} num_burnin={num_burnin} num_mcmc={num_mcmc} " + f"num_chains={num_chains} reps={n_reps}\n" +) + +# --------------------------------------------------------------------------- +# Runner: univariate treatment with sample_intercept=True +# --------------------------------------------------------------------------- +def run_once_univariate(run_cpp: bool, seed: int) -> dict: + m = BCFModel() + t0 = time.perf_counter() + m.sample( + X_train=X_train_u, + Z_train=Z_train_u, + y_train=y_train_u, + propensity_train=pi_train_u, + num_gfr=num_gfr, + num_burnin=num_burnin, + num_mcmc=num_mcmc, + general_params={ + "adaptive_coding": False, + "random_seed": seed, + "num_chains": num_chains, + "propensity_covariate": "prognostic", + }, + prognostic_forest_params={"num_trees": num_trees_mu}, + treatment_effect_forest_params={ + "num_trees": num_trees_tau, + "sample_intercept": True, + }, + run_cpp=run_cpp, + ) + elapsed_sample = time.perf_counter() - t0 + + # tau_0_samples comes from the model object (set at sample time) + tau_0_samples = getattr(m, "tau_0_samples", None) + tau_0_mean = float(np.mean(m.tau_0_samples)) if tau_0_samples is not None else float("nan") + + t1 = time.perf_counter() + preds = m.predict(X=X_test_u, Z=Z_test_u, propensity=pi_test_u, run_cpp=run_cpp) + elapsed_predict = time.perf_counter() - t1 + + y_hat = preds["y_hat"].mean(axis=1) + tau_hat = preds["tau_hat"].mean(axis=1) + + return { + "elapsed_sample": elapsed_sample, + "elapsed_predict": elapsed_predict, + "tau_0_mean": tau_0_mean, + "tau_0_shape": tuple(m.tau_0_samples.shape) if tau_0_samples is not None else None, + "rmse_y": float(np.sqrt(np.mean((y_hat - y_test_u) ** 2))), + "rmse_f": float(np.sqrt(np.mean((y_hat - f_test_u) ** 2))), + "rmse_tau": float(np.sqrt(np.mean((tau_hat - tau_test_u) ** 2))), + } + +# --------------------------------------------------------------------------- +# Runner: multivariate (2-column) treatment with sample_intercept=True +# --------------------------------------------------------------------------- +def run_once_multivariate(run_cpp: bool, seed: int) -> dict: + m = BCFModel() + t0 = time.perf_counter() + m.sample( + X_train=X_train_mv, + Z_train=Z_train_mv, + y_train=y_train_mv, + propensity_train=pi_train_mv, + num_gfr=num_gfr, + num_burnin=num_burnin, + num_mcmc=num_mcmc, + general_params={ + "adaptive_coding": False, + "random_seed": seed, + "num_chains": num_chains, + }, + prognostic_forest_params={"num_trees": num_trees_mu}, + treatment_effect_forest_params={ + "num_trees": num_trees_tau, + "sample_sigma2_leaf": False, + "sample_intercept": True, + }, + run_cpp=run_cpp, + ) + elapsed_sample = time.perf_counter() - t0 + + # tau_0_samples: (treatment_dim, num_mcmc) for multivariate + tau_0_samples = getattr(m, "tau_0_samples", None) + tau_0_mean = ( + m.tau_0_samples.mean(axis=1).tolist() if tau_0_samples is not None else [float("nan")] * 2 + ) + + t1 = time.perf_counter() + preds = m.predict(X=X_test_mv, Z=Z_test_mv, propensity=pi_test_mv, run_cpp=run_cpp) + elapsed_predict = time.perf_counter() - t1 + + # tau_hat is (n, treatment_dim, num_samples) for multivariate treatment + tau_hat1 = preds["tau_hat"][:, 0, :].mean(axis=1) + tau_hat2 = preds["tau_hat"][:, 1, :].mean(axis=1) + y_hat = preds["y_hat"].mean(axis=1) + + return { + "elapsed_sample": elapsed_sample, + "elapsed_predict": elapsed_predict, + "tau_0_mean": tau_0_mean, + "tau_0_shape": tuple(m.tau_0_samples.shape) if tau_0_samples is not None else None, + "rmse_y": float(np.sqrt(np.mean((y_hat - y_test_mv) ** 2))), + "rmse_f": float(np.sqrt(np.mean((y_hat - f_test_mv) ** 2))), + "rmse_tau1": float(np.sqrt(np.mean((tau_hat1 - tau_test_mv[:, 0]) ** 2))), + "rmse_tau2": float(np.sqrt(np.mean((tau_hat2 - tau_test_mv[:, 1]) ** 2))), + } + +# --------------------------------------------------------------------------- +# Run: univariate +# --------------------------------------------------------------------------- +seeds = [1000 + i for i in range(1, n_reps + 1)] + +print("=" * 60) +print("UNIVARIATE TREATMENT (sample_intercept=True)") +print(f"True tau_0 = {TRUE_TAU0_UNIVARIATE}") +print("=" * 60) + +results_cpp_u = [] +results_py_u = [] + +print("Running C++ sampler (run_cpp=True)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_cpp_u.append(run_once_univariate(run_cpp=True, seed=seed)) + +print("\nRunning Python sampler (run_cpp=False)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_py_u.append(run_once_univariate(run_cpp=False, seed=seed)) + +def summarise_u(results): + keys = ["elapsed_sample", "elapsed_predict", "rmse_y", "rmse_f", "rmse_tau"] + out = {k: float(np.mean([r[k] for r in results])) for k in keys} + out["elapsed"] = out["elapsed_sample"] + out["elapsed_predict"] + out["elapsed_sd"] = float(np.std( + [r["elapsed_sample"] + r["elapsed_predict"] for r in results], ddof=1 + )) + out["tau_0_mean"] = float(np.mean([r["tau_0_mean"] for r in results])) + out["tau_0_shape"] = results[0]["tau_0_shape"] + return out + +s_cpp_u = summarise_u(results_cpp_u) +s_py_u = summarise_u(results_py_u) + +print("\n--- Univariate Results ---") +print(f"tau_0_samples shape cpp={s_cpp_u['tau_0_shape']} py={s_py_u['tau_0_shape']}") +print( + f"{'Sampler':<22} {'Total (s)':>9} {'Samp (s)':>9} {'Pred (s)':>9} {'SD':>8} " + f"{'tau_0 mean':>10} {'RMSE(y)':>9} {'RMSE(f)':>9} {'RMSE(tau)':>10}" +) +print("-" * 108) +for label, s in [("cpp (run_cpp=True)", s_cpp_u), ("py (run_cpp=False)", s_py_u)]: + print( + f"{label:<22} {s['elapsed']:>9.3f} {s['elapsed_sample']:>9.3f} " + f"{s['elapsed_predict']:>9.3f} {s['elapsed_sd']:>8.3f} " + f"{s['tau_0_mean']:>10.4f} {s['rmse_y']:>9.4f} {s['rmse_f']:>9.4f} {s['rmse_tau']:>10.4f}" + ) +print(f"True tau_0: {TRUE_TAU0_UNIVARIATE:.4f}") +speedup_u = s_py_u["elapsed"] / s_cpp_u["elapsed"] +print(f"Speedup (py / cpp): {speedup_u:.2f}x") +print( + f"RMSE delta (cpp - py): " + f"y={s_cpp_u['rmse_y'] - s_py_u['rmse_y']:.4f} " + f"f={s_cpp_u['rmse_f'] - s_py_u['rmse_f']:.4f} " + f"tau={s_cpp_u['rmse_tau'] - s_py_u['rmse_tau']:.4f}" +) + +# --------------------------------------------------------------------------- +# Run: multivariate +# --------------------------------------------------------------------------- +print() +print("=" * 60) +print("MULTIVARIATE TREATMENT (treatment_dim=2, sample_intercept=True)") +print(f"True tau_0 = {TRUE_TAU0_MV.tolist()}") +print("=" * 60) + +results_cpp_mv = [] +results_py_mv = [] + +print("Running C++ sampler (run_cpp=True)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_cpp_mv.append(run_once_multivariate(run_cpp=True, seed=seed)) + +print("\nRunning Python sampler (run_cpp=False)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_py_mv.append(run_once_multivariate(run_cpp=False, seed=seed)) + +def summarise_mv(results): + keys = ["elapsed_sample", "elapsed_predict", "rmse_y", "rmse_f", "rmse_tau1", "rmse_tau2"] + out = {k: float(np.mean([r[k] for r in results])) for k in keys} + out["elapsed"] = out["elapsed_sample"] + out["elapsed_predict"] + out["elapsed_sd"] = float(np.std( + [r["elapsed_sample"] + r["elapsed_predict"] for r in results], ddof=1 + )) + out["tau_0_mean_0"] = float(np.mean([r["tau_0_mean"][0] for r in results])) + out["tau_0_mean_1"] = float(np.mean([r["tau_0_mean"][1] for r in results])) + out["tau_0_shape"] = results[0]["tau_0_shape"] + return out + +s_cpp_mv = summarise_mv(results_cpp_mv) +s_py_mv = summarise_mv(results_py_mv) + +print("\n--- Multivariate Results ---") +print(f"tau_0_samples shape cpp={s_cpp_mv['tau_0_shape']} py={s_py_mv['tau_0_shape']}") +print( + f"{'Sampler':<22} {'Total (s)':>9} {'Samp (s)':>9} {'Pred (s)':>9} {'SD':>8} " + f"{'tau_0[0]':>9} {'tau_0[1]':>9} " + f"{'RMSE(y)':>8} {'RMSE(f)':>8} {'RMSE(tau1)':>10} {'RMSE(tau2)':>10}" +) +print("-" * 125) +for label, s in [("cpp (run_cpp=True)", s_cpp_mv), ("py (run_cpp=False)", s_py_mv)]: + print( + f"{label:<22} {s['elapsed']:>9.3f} {s['elapsed_sample']:>9.3f} " + f"{s['elapsed_predict']:>9.3f} {s['elapsed_sd']:>8.3f} " + f"{s['tau_0_mean_0']:>9.4f} {s['tau_0_mean_1']:>9.4f} " + f"{s['rmse_y']:>8.4f} {s['rmse_f']:>8.4f} {s['rmse_tau1']:>10.4f} {s['rmse_tau2']:>10.4f}" + ) +print(f"True tau_0: [{TRUE_TAU0_MV[0]:.4f}, {TRUE_TAU0_MV[1]:.4f}]") +speedup_mv = s_py_mv["elapsed"] / s_cpp_mv["elapsed"] +print(f"Speedup (py / cpp): {speedup_mv:.2f}x") +print( + f"RMSE delta (cpp - py): " + f"y={s_cpp_mv['rmse_y'] - s_py_mv['rmse_y']:.4f} " + f"f={s_cpp_mv['rmse_f'] - s_py_mv['rmse_f']:.4f} " + f"tau1={s_cpp_mv['rmse_tau1'] - s_py_mv['rmse_tau1']:.4f} " + f"tau2={s_cpp_mv['rmse_tau2'] - s_py_mv['rmse_tau2']:.4f}" +) diff --git a/debug/benchmark_cpp_vs_py_sampler_cloglog.py b/debug/benchmark_cpp_vs_py_sampler_cloglog.py new file mode 100644 index 00000000..0820751a --- /dev/null +++ b/debug/benchmark_cpp_vs_py_sampler_cloglog.py @@ -0,0 +1,157 @@ +"""Benchmark: C++ sampler loop vs. Python sampler loop – cloglog BART. + +Compares runtime, Brier score, and RMSE-to-truth (vs. true P(Y=1|X)) across +run_cpp=True / False in BARTModel.sample(). + +DGP uses the cloglog link: P(Y=1|X) = 1 - exp(-exp(f(X))). +f(X) is a sum of smooth sinusoidal terms across two covariates, keeping +probabilities in the moderate range [0.25, 0.75] for stable mixing. +GFR is disabled (num_gfr=0) since it interacts poorly with cloglog binary. + +Usage: + conda activate stochtree-book # or: source venv/bin/activate + python debug/benchmark_cpp_vs_py_sampler_cloglog.py +""" + +import argparse +import time +import numpy as np +from stochtree import BARTModel, OutcomeModel + +parser = argparse.ArgumentParser() +parser.add_argument("--num-chains", type=int, default=1) +args = parser.parse_args() +num_chains = args.num_chains + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +rng = np.random.default_rng(1234) + +n = 2000 +p = 10 +X = rng.uniform(size=(n, p)) + +# Latent mean on the cloglog (log-log) scale. +# f_X is centred near -0.5 so that P(Y=1|X) = 1 - exp(-exp(f_X)) stays +# in [~0.25, ~0.75], avoiding extreme probabilities that inflate GFR depth. +f_X = ( + 0.6 * np.sin(2 * np.pi * X[:, 0]) + + 0.4 * np.cos(2 * np.pi * X[:, 1]) + - 0.5 +) +p_X = 1.0 - np.exp(-np.exp(f_X)) # true P(Y = 1 | X) +y = rng.binomial(1, p_X).astype(float) # observed binary outcome + +test_frac = 0.2 +n_test = round(test_frac * n) +n_train = n - n_test +test_inds = rng.choice(n, size=n_test, replace=False) +train_inds = np.setdiff1d(np.arange(n), test_inds) + +X_train, X_test = X[train_inds], X[test_inds] +y_train, y_test = y[train_inds], y[test_inds] +p_test = p_X[test_inds] + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr = 0 +num_burnin = 100 +num_mcmc = 100 +num_trees = 200 +n_reps = 3 + +print( + f"n_train={n_train} n_test={n_test} p={p} " + f"num_trees={num_trees} num_gfr={num_gfr} num_burnin={num_burnin} " + f"num_mcmc={num_mcmc} num_chains={num_chains} reps={n_reps}\n" +) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + metrics +# --------------------------------------------------------------------------- +def run_once(run_cpp: bool, seed: int) -> dict: + m = BARTModel() + t0 = time.perf_counter() + m.sample( + X_train=X_train, + y_train=y_train, + num_gfr=num_gfr, + num_burnin=num_burnin, + num_mcmc=num_mcmc, + mean_forest_params={"num_trees": num_trees, "sample_sigma2_leaf": False}, + general_params={ + "random_seed": seed, + "outcome_model": OutcomeModel(outcome="binary", link="cloglog"), + "sample_sigma2_global": False, + "num_chains": num_chains, + }, + run_cpp=run_cpp, + ) + elapsed_sample = time.perf_counter() - t0 + + t1 = time.perf_counter() + preds = m.predict(X=X_test, scale="probability", run_cpp=run_cpp) + elapsed_predict = time.perf_counter() - t1 + + p_hat = preds["y_hat"].mean(axis=1) # (n_test,) + + brier = float(np.mean((p_hat - y_test) ** 2)) + rmse_p = float(np.sqrt(np.mean((p_hat - p_test) ** 2))) + + return {"elapsed_sample": elapsed_sample, "elapsed_predict": elapsed_predict, "brier": brier, "rmse_p": rmse_p} + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds = [100 + i for i in range(1, n_reps + 1)] + +results_cpp = [] +results_py = [] + +print("Running C++ sampler (run_cpp=True)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_cpp.append(run_once(run_cpp=True, seed=seed)) + +print("\nRunning Python sampler (run_cpp=False)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_py.append(run_once(run_cpp=False, seed=seed)) + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +def summarise(results: list) -> dict: + keys = ["elapsed_sample", "elapsed_predict", "brier", "rmse_p"] + out = {k: float(np.mean([r[k] for r in results])) for k in keys} + out["elapsed"] = out["elapsed_sample"] + out["elapsed_predict"] + out["elapsed_sd"] = float(np.std( + [r["elapsed_sample"] + r["elapsed_predict"] for r in results], ddof=1 + )) + return out + +s_cpp = summarise(results_cpp) +s_py = summarise(results_py) +rows = [("cpp (run_cpp=True)", s_cpp), ("py (run_cpp=False)", s_py)] + +print("\n--- Results ---") +print( + f"{'Sampler':<22} {'Total (s)':>10} {'Samp (s)':>10} {'Pred (s)':>10} {'SD':>8} " + f"{'Brier':>12} {'RMSE (vs truth)':>15}" +) +print("-" * 97) +for label, s in rows: + print( + f"{label:<22} {s['elapsed']:>10.3f} {s['elapsed_sample']:>10.3f} " + f"{s['elapsed_predict']:>10.3f} {s['elapsed_sd']:>8.3f} " + f"{s['brier']:>12.4f} {s['rmse_p']:>15.4f}" + ) + +speedup = s_py["elapsed"] / s_cpp["elapsed"] +print(f"\nSpeedup (py / cpp): {speedup:.2f}x") +print( + f"Brier delta (cpp - py): {s_cpp['brier'] - s_py['brier']:.4f}\n" + f"RMSE-p delta (cpp - py): {s_cpp['rmse_p'] - s_py['rmse_p']:.4f}" +) diff --git a/debug/benchmark_cpp_vs_py_sampler_heteroskedastic.py b/debug/benchmark_cpp_vs_py_sampler_heteroskedastic.py new file mode 100644 index 00000000..a3c42690 --- /dev/null +++ b/debug/benchmark_cpp_vs_py_sampler_heteroskedastic.py @@ -0,0 +1,172 @@ +"""Benchmark: C++ sampler loop vs. Python sampler loop – heteroskedastic BART. + +Compares runtime, mean-forest RMSE (vs. true f(X)) and RMSE of the estimated +conditional standard deviation (vs. the true s(X)) across run_cpp=True / +False in BARTModel.sample() with both a mean forest and a variance forest +(num_trees_variance > 0). + +DGP: f(X) is a step function of X[:,0]; s(X) varies by quadrant of X[:,0] +and linearly with X[:,2], matching the heteroskedastic_bart.R debug script. + +Note: A variance-only model (num_trees_mean=0, num_trees_variance>0) is now +supported in the C++ path. The mean-forest RMSE is reported as NaN in that +case since there is no mean forest to evaluate. + +Usage: + source venv/bin/activate + python debug/benchmark_cpp_vs_py_sampler_heteroskedastic.py +""" + +import argparse +import time +import numpy as np +from stochtree import BARTModel + +parser = argparse.ArgumentParser() +parser.add_argument("--num-chains", type=int, default=1) +args = parser.parse_args() +num_chains = args.num_chains + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +rng = np.random.default_rng(1234) + +n = 2000 +p = 10 +X = rng.uniform(size=(n, p)) + +# True conditional mean and conditional std dev +f_X = ( + np.where((X[:, 0] >= 0.00) & (X[:, 0] < 0.25), -3.0, 0) + + np.where((X[:, 0] >= 0.25) & (X[:, 0] < 0.50), -1.0, 0) + + np.where((X[:, 0] >= 0.50) & (X[:, 0] < 0.75), 1.0, 0) + + np.where((X[:, 0] >= 0.75) & (X[:, 0] < 1.00), 3.0, 0) +) +s_X = ( + np.where((X[:, 0] >= 0.00) & (X[:, 0] < 0.25), 0.5 * X[:, 2], 0) + + np.where((X[:, 0] >= 0.25) & (X[:, 0] < 0.50), 1.0 * X[:, 2], 0) + + np.where((X[:, 0] >= 0.50) & (X[:, 0] < 0.75), 2.0 * X[:, 2], 0) + + np.where((X[:, 0] >= 0.75) & (X[:, 0] < 1.00), 3.0 * X[:, 2], 0) +) +y = f_X + rng.standard_normal(n) * s_X + +test_frac = 0.2 +n_test = round(test_frac * n) +n_train = n - n_test +test_inds = rng.choice(n, size=n_test, replace=False) +train_inds = np.setdiff1d(np.arange(n), test_inds) + +X_train, X_test = X[train_inds], X[test_inds] +y_train, y_test = y[train_inds], y[test_inds] +f_test = f_X[test_inds] +s_test = s_X[test_inds] + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr = 10 +num_burnin = 0 +num_mcmc = 100 +num_trees_mean = 200 +num_trees_variance = 50 +n_reps = 3 + +print( + f"n_train={n_train} n_test={n_test} p={p} " + f"num_trees_mean={num_trees_mean} num_trees_variance={num_trees_variance} " + f"num_gfr={num_gfr} num_burnin={num_burnin} num_mcmc={num_mcmc} " + f"num_chains={num_chains} reps={n_reps}\n" +) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + metrics +# --------------------------------------------------------------------------- +def run_once(run_cpp: bool, seed: int) -> dict: + m = BARTModel() + t0 = time.perf_counter() + m.sample( + X_train=X_train, + y_train=y_train, + num_gfr=num_gfr, + num_burnin=num_burnin, + num_mcmc=num_mcmc, + general_params={ + "random_seed": seed, + "sample_sigma2_global": False, + "num_chains": num_chains, + }, + mean_forest_params={"num_trees": num_trees_mean}, + variance_forest_params={"num_trees": num_trees_variance}, + run_cpp=run_cpp, + ) + elapsed_sample = time.perf_counter() - t0 + + t1 = time.perf_counter() + preds = m.predict(X=X_test, run_cpp=run_cpp) + elapsed_predict = time.perf_counter() - t1 + + # mean-forest RMSE vs. true f(X) – only defined when a mean forest was fitted + if num_trees_mean > 0: + f_hat = preds["y_hat"].mean(axis=1) + rmse_f = float(np.sqrt(np.mean((f_hat - f_test) ** 2))) + else: + rmse_f = float("nan") + # variance_forest_predictions has shape (n_test, num_mcmc); posterior mean of cond. std dev + s_hat = np.sqrt(preds["variance_forest_predictions"]).mean(axis=1) + rmse_s = float(np.sqrt(np.mean((s_hat - s_test) ** 2))) + return {"elapsed_sample": elapsed_sample, "elapsed_predict": elapsed_predict, "rmse_f": rmse_f, "rmse_s": rmse_s} + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds = [1000 + i for i in range(1, n_reps + 1)] + +results_cpp = [] +results_py = [] + +print("Running C++ sampler (run_cpp=True)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_cpp.append(run_once(run_cpp=True, seed=seed)) + +print("\nRunning Python sampler (run_cpp=False)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_py.append(run_once(run_cpp=False, seed=seed)) + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +def summarise(results: list) -> dict: + keys = ["elapsed_sample", "elapsed_predict", "rmse_f", "rmse_s"] + out = {k: float(np.nanmean([r[k] for r in results])) for k in keys} + out["elapsed"] = out["elapsed_sample"] + out["elapsed_predict"] + out["elapsed_sd"] = float(np.std( + [r["elapsed_sample"] + r["elapsed_predict"] for r in results], ddof=1 + )) + return out + +s_cpp = summarise(results_cpp) +s_py = summarise(results_py) +rows = [("cpp (run_cpp=True)", s_cpp), ("py (run_cpp=False)", s_py)] + +print("\n--- Results ---") +print( + f"{'Sampler':<22} {'Total (s)':>10} {'Samp (s)':>10} {'Pred (s)':>10} {'SD':>8} " + f"{'RMSE f(X)':>12} {'RMSE s(X)':>12}" +) +print("-" * 96) +for label, s in rows: + print( + f"{label:<22} {s['elapsed']:>10.3f} {s['elapsed_sample']:>10.3f} " + f"{s['elapsed_predict']:>10.3f} {s['elapsed_sd']:>8.3f} " + f"{s['rmse_f']:>12.4f} {s['rmse_s']:>12.4f}" + ) + +speedup = s_py["elapsed"] / s_cpp["elapsed"] +print(f"\nSpeedup (py / cpp): {speedup:.2f}x") +print( + f"RMSE f(X) delta (cpp - py): {s_cpp['rmse_f'] - s_py['rmse_f']:.4f}\n" + f"RMSE s(X) delta (cpp - py): {s_cpp['rmse_s'] - s_py['rmse_s']:.4f}" +) diff --git a/debug/benchmark_cpp_vs_py_sampler_leaf_regression.py b/debug/benchmark_cpp_vs_py_sampler_leaf_regression.py new file mode 100644 index 00000000..23c0bd65 --- /dev/null +++ b/debug/benchmark_cpp_vs_py_sampler_leaf_regression.py @@ -0,0 +1,157 @@ +"""Benchmark: C++ sampler loop vs. Python sampler loop -- univariate leaf regression. + +Compares runtime and test-set RMSE across run_cpp=True / False in BARTModel.sample() +with a univariate leaf regression basis (leaf_basis_train with one column). + +DGP: f(X, Z) = tau(X) * Z, where tau(X) is a step function of X[:,0] and +Z is drawn uniform [0, 1]. A constant noise term is added. The leaf basis +passed to the sampler is just Z (shape n x 1). + +Usage: + source venv/bin/activate + python debug/benchmark_cpp_vs_py_sampler_leaf_regression.py +""" + +import argparse +import time +import numpy as np +from stochtree import BARTModel + +parser = argparse.ArgumentParser() +parser.add_argument("--num-chains", type=int, default=1) +args = parser.parse_args() +num_chains = args.num_chains + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +rng = np.random.default_rng(1234) + +n = 2000 +p = 10 +X = rng.uniform(size=(n, p)) +Z = rng.uniform(size=n) # scalar moderating variable / leaf basis + +# Heterogeneous slope on Z, partitioned by X[:,0] +tau_X = ( + np.where((X[:, 0] >= 0.00) & (X[:, 0] < 0.25), -2.0, 0) + + np.where((X[:, 0] >= 0.25) & (X[:, 0] < 0.50), -1.0, 0) + + np.where((X[:, 0] >= 0.50) & (X[:, 0] < 0.75), 1.0, 0) + + np.where((X[:, 0] >= 0.75) & (X[:, 0] < 1.00), 2.0, 0) +) +f_XZ = tau_X * Z +noise_sd = 1.0 +y = f_XZ + rng.normal(scale=noise_sd, size=n) + +test_frac = 0.2 +n_test = round(test_frac * n) +n_train = n - n_test +test_inds = rng.choice(n, size=n_test, replace=False) +train_inds = np.setdiff1d(np.arange(n), test_inds) + +X_train, X_test = X[train_inds], X[test_inds] +Z_train, Z_test = Z[train_inds], Z[test_inds] +y_train, y_test = y[train_inds], y[test_inds] +f_test = f_XZ[test_inds] + +# Leaf basis matrices (n x 1) +basis_train = Z_train.reshape(-1, 1) +basis_test = Z_test.reshape(-1, 1) + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr = 10 +num_burnin = 0 +num_mcmc = 100 +num_trees = 200 +n_reps = 3 + +print( + f"n_train={n_train} n_test={n_test} p={p} " + f"num_trees={num_trees} num_gfr={num_gfr} num_burnin={num_burnin} " + f"num_mcmc={num_mcmc} num_chains={num_chains} reps={n_reps}\n" +) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + RMSE +# --------------------------------------------------------------------------- +def run_once(run_cpp: bool, seed: int) -> dict: + m = BARTModel() + t0 = time.perf_counter() + m.sample( + X_train=X_train, + y_train=y_train, + leaf_basis_train=basis_train, + num_gfr=num_gfr, + num_burnin=num_burnin, + num_mcmc=num_mcmc, + general_params={"random_seed": seed, "num_chains": num_chains}, + mean_forest_params={"num_trees": num_trees}, + run_cpp=run_cpp, + ) + elapsed_sample = time.perf_counter() - t0 + + t1 = time.perf_counter() + preds = m.predict(X=X_test, leaf_basis=basis_test, run_cpp=run_cpp) + elapsed_predict = time.perf_counter() - t1 + + yhat = preds["y_hat"].mean(axis=1) + rmse = float(np.sqrt(np.mean((yhat - y_test) ** 2))) + rmse_f = float(np.sqrt(np.mean((yhat - f_test) ** 2))) + return {"elapsed_sample": elapsed_sample, "elapsed_predict": elapsed_predict, "rmse": rmse, "rmse_f": rmse_f} + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds = [1000 + i for i in range(1, n_reps + 1)] + +results_cpp = [] +results_py = [] + +print("Running C++ sampler (run_cpp=True)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_cpp.append(run_once(run_cpp=True, seed=seed)) + +print("\nRunning Python sampler (run_cpp=False)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_py.append(run_once(run_cpp=False, seed=seed)) + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +def summarise(results: list) -> dict: + keys = ["elapsed_sample", "elapsed_predict", "rmse", "rmse_f"] + out = {k: float(np.mean([r[k] for r in results])) for k in keys} + out["elapsed"] = out["elapsed_sample"] + out["elapsed_predict"] + out["elapsed_sd"] = float(np.std( + [r["elapsed_sample"] + r["elapsed_predict"] for r in results], ddof=1 + )) + return out + +s_cpp = summarise(results_cpp) +s_py = summarise(results_py) +rows = [("cpp (run_cpp=True)", s_cpp), ("py (run_cpp=False)", s_py)] + +print("\n--- Results ---") +print( + f"{'Sampler':<22} {'Total (s)':>10} {'Samp (s)':>10} {'Pred (s)':>10} {'SD':>8} " + f"{'RMSE (obs)':>12} {'RMSE f(X,Z)':>13}" +) +print("-" * 96) +for label, s in rows: + print( + f"{label:<22} {s['elapsed']:>10.3f} {s['elapsed_sample']:>10.3f} " + f"{s['elapsed_predict']:>10.3f} {s['elapsed_sd']:>8.3f} " + f"{s['rmse']:>12.4f} {s['rmse_f']:>13.4f}" + ) + +speedup = s_py["elapsed"] / s_cpp["elapsed"] +print(f"\nSpeedup (py / cpp): {speedup:.2f}x") +print( + f"RMSE delta (cpp - py): " + f"obs={s_cpp['rmse'] - s_py['rmse']:.4f} " + f"f={s_cpp['rmse_f'] - s_py['rmse_f']:.4f}" +) diff --git a/debug/benchmark_cpp_vs_py_sampler_multivariate_leaf_regression.py b/debug/benchmark_cpp_vs_py_sampler_multivariate_leaf_regression.py new file mode 100644 index 00000000..651ec333 --- /dev/null +++ b/debug/benchmark_cpp_vs_py_sampler_multivariate_leaf_regression.py @@ -0,0 +1,172 @@ +"""Benchmark: C++ sampler loop vs. Python sampler loop -- multivariate leaf regression. + +Compares runtime and test-set RMSE across run_cpp=True / False in BARTModel.sample() +with a 2-column leaf basis (multivariate leaf regression). + +DGP: f(X, Z) = tau_1(X)*Z_1 + tau_2(X)*Z_2, where tau_1/tau_2 are step functions +of X[:,0] and Z_1, Z_2 are drawn uniform [0, 1]. A constant noise term is added. +The leaf basis passed to the sampler is [Z_1, Z_2] (shape n x 2). + +Usage: + source venv/bin/activate + python debug/benchmark_cpp_vs_py_sampler_multivariate_leaf_regression.py +""" + +import argparse +import time +import numpy as np +from stochtree import BARTModel + +parser = argparse.ArgumentParser() +parser.add_argument("--num-chains", type=int, default=1) +args = parser.parse_args() +num_chains = args.num_chains + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +rng = np.random.default_rng(1234) + +n = 2000 +p = 10 +X = rng.uniform(size=(n, p)) +Z1 = rng.uniform(size=n) +Z2 = rng.uniform(size=n) + +# Heterogeneous slopes on Z1 and Z2, partitioned by X[:,0] +tau1_X = ( + np.where((X[:, 0] >= 0.00) & (X[:, 0] < 0.25), -2.0, 0) + + np.where((X[:, 0] >= 0.25) & (X[:, 0] < 0.50), -1.0, 0) + + np.where((X[:, 0] >= 0.50) & (X[:, 0] < 0.75), 1.0, 0) + + np.where((X[:, 0] >= 0.75) & (X[:, 0] < 1.00), 2.0, 0) +) +tau2_X = ( + np.where((X[:, 0] >= 0.00) & (X[:, 0] < 0.25), 1.0, 0) + + np.where((X[:, 0] >= 0.25) & (X[:, 0] < 0.50), 2.0, 0) + + np.where((X[:, 0] >= 0.50) & (X[:, 0] < 0.75), -1.0, 0) + + np.where((X[:, 0] >= 0.75) & (X[:, 0] < 1.00), -2.0, 0) +) +f_XZ = tau1_X * Z1 + tau2_X * Z2 +noise_sd = 1.0 +y = f_XZ + rng.normal(scale=noise_sd, size=n) + +test_frac = 0.2 +n_test = round(test_frac * n) +n_train = n - n_test +test_inds = rng.choice(n, size=n_test, replace=False) +train_inds = np.setdiff1d(np.arange(n), test_inds) + +X_train, X_test = X[train_inds], X[test_inds] +Z1_train, Z1_test = Z1[train_inds], Z1[test_inds] +Z2_train, Z2_test = Z2[train_inds], Z2[test_inds] +y_train, y_test = y[train_inds], y[test_inds] +f_test = f_XZ[test_inds] + +# Leaf basis matrices (n x 2) +basis_train = np.column_stack([Z1_train, Z2_train]) +basis_test = np.column_stack([Z1_test, Z2_test]) + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr = 10 +num_burnin = 0 +num_mcmc = 100 +num_trees = 200 +n_reps = 3 + +# Optional: pass a 2x2 prior covariance for the leaf coefficients +sigma2_leaf_init = np.array([[0.5, 0.0], [0.0, 0.5]]) + +print( + f"n_train={n_train} n_test={n_test} p={p} basis_dim=2\n" + f"num_trees={num_trees} num_gfr={num_gfr} num_burnin={num_burnin} " + f"num_mcmc={num_mcmc} num_chains={num_chains} reps={n_reps}\n" +) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + RMSE +# --------------------------------------------------------------------------- +def run_once(run_cpp: bool, seed: int) -> dict: + m = BARTModel() + t0 = time.perf_counter() + m.sample( + X_train=X_train, + y_train=y_train, + leaf_basis_train=basis_train, + num_gfr=num_gfr, + num_burnin=num_burnin, + num_mcmc=num_mcmc, + general_params={"random_seed": seed, "num_chains": num_chains}, + mean_forest_params={ + "num_trees": num_trees, + "sigma2_leaf_init": sigma2_leaf_init, + "sample_sigma2_leaf": False, + }, + run_cpp=run_cpp, + ) + elapsed_sample = time.perf_counter() - t0 + + t1 = time.perf_counter() + preds = m.predict(X=X_test, leaf_basis=basis_test, run_cpp=run_cpp) + elapsed_predict = time.perf_counter() - t1 + + yhat = preds["y_hat"].mean(axis=1) + rmse = float(np.sqrt(np.mean((yhat - y_test) ** 2))) + rmse_f = float(np.sqrt(np.mean((yhat - f_test) ** 2))) + return {"elapsed_sample": elapsed_sample, "elapsed_predict": elapsed_predict, "rmse": rmse, "rmse_f": rmse_f} + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds = [1000 + i for i in range(1, n_reps + 1)] + +results_cpp = [] +results_py = [] + +print("Running C++ sampler (run_cpp=True)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_cpp.append(run_once(run_cpp=True, seed=seed)) + +print("\nRunning Python sampler (run_cpp=False)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_py.append(run_once(run_cpp=False, seed=seed)) + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +def summarise(results: list) -> dict: + keys = ["elapsed_sample", "elapsed_predict", "rmse", "rmse_f"] + out = {k: float(np.mean([r[k] for r in results])) for k in keys} + out["elapsed"] = out["elapsed_sample"] + out["elapsed_predict"] + out["elapsed_sd"] = float(np.std( + [r["elapsed_sample"] + r["elapsed_predict"] for r in results], ddof=1 + )) + return out + +s_cpp = summarise(results_cpp) +s_py = summarise(results_py) +rows = [("cpp (run_cpp=True)", s_cpp), ("py (run_cpp=False)", s_py)] + +print("\n--- Results ---") +print( + f"{'Sampler':<22} {'Total (s)':>10} {'Samp (s)':>10} {'Pred (s)':>10} {'SD':>8} " + f"{'RMSE (obs)':>12} {'RMSE f(X,Z)':>13}" +) +print("-" * 96) +for label, s in rows: + print( + f"{label:<22} {s['elapsed']:>10.3f} {s['elapsed_sample']:>10.3f} " + f"{s['elapsed_predict']:>10.3f} {s['elapsed_sd']:>8.3f} " + f"{s['rmse']:>12.4f} {s['rmse_f']:>13.4f}" + ) + +speedup = s_py["elapsed"] / s_cpp["elapsed"] +print(f"\nSpeedup (py / cpp): {speedup:.2f}x") +print( + f"RMSE delta (cpp - py): " + f"obs={s_cpp['rmse'] - s_py['rmse']:.4f} " + f"f={s_cpp['rmse_f'] - s_py['rmse_f']:.4f}" +) diff --git a/debug/benchmark_cpp_vs_py_sampler_ordinal_cloglog.py b/debug/benchmark_cpp_vs_py_sampler_ordinal_cloglog.py new file mode 100644 index 00000000..c3b86196 --- /dev/null +++ b/debug/benchmark_cpp_vs_py_sampler_ordinal_cloglog.py @@ -0,0 +1,178 @@ +"""Benchmark: C++ sampler loop vs. Python sampler loop – ordinal cloglog BART. + +Compares runtime, mean Brier score, and mean RMSE-to-truth (vs. true class +probabilities) across run_cpp=True / False in BARTModel.sample(). + +DGP uses 4 ordinal categories with a cloglog link. +f(X) is a smooth sinusoidal function of two covariates so that all four +categories are well-populated. Cutpoints are spaced to yield roughly equal +marginal class frequencies. GFR is disabled (num_gfr=0). + +Usage: + conda activate stochtree-book # or: source venv/bin/activate + python debug/benchmark_cpp_vs_py_sampler_ordinal_cloglog.py +""" + +import argparse +import time +import numpy as np +from stochtree import BARTModel, OutcomeModel + +parser = argparse.ArgumentParser() +parser.add_argument("--num-chains", type=int, default=1) +args = parser.parse_args() +num_chains = args.num_chains + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +rng = np.random.default_rng(1234) + +n = 2000 +p = 10 +X = rng.uniform(size=(n, p)) + +# Latent mean on the cloglog scale (smooth, two covariates) +f_X = 0.6 * np.sin(2 * np.pi * X[:, 0]) + 0.4 * np.cos(2 * np.pi * X[:, 1]) + +# Fixed log-scale cutpoints spaced to give roughly equal marginal class freqs. +# With f_X in roughly [-1, 1], gamma = [0, log(2), log(4)] puts the four +# cumulative boundaries at moderate probability levels. +K = 4 +gamma_true = np.array([0.0, np.log(2), np.log(4)]) + +# True cumulative probabilities: P(Y <= k | X) = 1 - exp(-exp(f_X - gamma_k)) +# Shape: (n, K-1) +cum_prob = 1.0 - np.exp(-np.exp(f_X[:, None] - gamma_true[None, :])) + +# True class probabilities: P(Y = k | X), shape (n, K) +p_X = np.column_stack([ + cum_prob[:, 0], + cum_prob[:, 1] - cum_prob[:, 0], + cum_prob[:, 2] - cum_prob[:, 1], + 1.0 - cum_prob[:, 2], +]) + +# Draw ordinal outcomes (0-indexed: 0, 1, 2, 3) +u = rng.uniform(size=n) +y = ( + (u > cum_prob[:, 0]).astype(int) + + (u > cum_prob[:, 1]).astype(int) + + (u > cum_prob[:, 2]).astype(int) +).astype(float) + +test_frac = 0.2 +n_test = round(test_frac * n) +n_train = n - n_test +test_inds = rng.choice(n, size=n_test, replace=False) +train_inds = np.setdiff1d(np.arange(n), test_inds) + +X_train, X_test = X[train_inds], X[test_inds] +y_train, y_test = y[train_inds], y[test_inds] +p_test = p_X[test_inds] # (n_test, K) true class probabilities + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr = 0 +num_burnin = 100 +num_mcmc = 100 +num_trees = 200 +n_reps = 3 + +print( + f"K={K} n_train={n_train} n_test={n_test} p={p} " + f"num_trees={num_trees} num_gfr={num_gfr} num_burnin={num_burnin} " + f"num_mcmc={num_mcmc} num_chains={num_chains} reps={n_reps}\n" +) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + metrics +# --------------------------------------------------------------------------- +def run_once(run_cpp: bool, seed: int) -> dict: + m = BARTModel() + t0 = time.perf_counter() + m.sample( + X_train=X_train, + y_train=y_train, + num_gfr=num_gfr, + num_burnin=num_burnin, + num_mcmc=num_mcmc, + mean_forest_params={"num_trees": num_trees, "sample_sigma2_leaf": False}, + general_params={ + "random_seed": seed, + "outcome_model": OutcomeModel(outcome="ordinal", link="cloglog"), + "sample_sigma2_global": False, + "num_chains": num_chains, + }, + run_cpp=run_cpp, + ) + elapsed_sample = time.perf_counter() - t0 + + t1 = time.perf_counter() + # For ordinal probability scale, "y_hat" has shape (n_test, K, num_mcmc) + preds = m.predict(X=X_test, scale="probability", run_cpp=run_cpp) + elapsed_predict = time.perf_counter() - t1 + + p_hat = preds["y_hat"].mean(axis=2) # (n_test, K) posterior mean + + # Mean Brier score across all cells + brier = float(np.mean((p_hat - p_test) ** 2)) + # Per-class RMSE vs. true probs, averaged over classes + rmse_p = float(np.mean(np.sqrt(np.mean((p_hat - p_test) ** 2, axis=0)))) + + return {"elapsed_sample": elapsed_sample, "elapsed_predict": elapsed_predict, "brier": brier, "rmse_p": rmse_p} + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds = [1000 + i for i in range(1, n_reps + 1)] + +results_cpp = [] +results_py = [] + +print("Running C++ sampler (run_cpp=True)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_cpp.append(run_once(run_cpp=True, seed=seed)) + +print("\nRunning Python sampler (run_cpp=False)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_py.append(run_once(run_cpp=False, seed=seed)) + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +def summarise(results: list) -> dict: + keys = ["elapsed_sample", "elapsed_predict", "brier", "rmse_p"] + out = {k: float(np.mean([r[k] for r in results])) for k in keys} + out["elapsed"] = out["elapsed_sample"] + out["elapsed_predict"] + out["elapsed_sd"] = float(np.std( + [r["elapsed_sample"] + r["elapsed_predict"] for r in results], ddof=1 + )) + return out + +s_cpp = summarise(results_cpp) +s_py = summarise(results_py) +rows = [("cpp (run_cpp=True)", s_cpp), ("py (run_cpp=False)", s_py)] + +print("\n--- Results ---") +print( + f"{'Sampler':<22} {'Total (s)':>10} {'Samp (s)':>10} {'Pred (s)':>10} {'SD':>8} " + f"{'Brier':>12} {'RMSE (vs truth)':>15}" +) +print("-" * 97) +for label, s in rows: + print( + f"{label:<22} {s['elapsed']:>10.3f} {s['elapsed_sample']:>10.3f} " + f"{s['elapsed_predict']:>10.3f} {s['elapsed_sd']:>8.3f} " + f"{s['brier']:>12.4f} {s['rmse_p']:>15.4f}" + ) + +speedup = s_py["elapsed"] / s_cpp["elapsed"] +print(f"\nSpeedup (py / cpp): {speedup:.2f}x") +print( + f"Brier delta (cpp - py): {s_cpp['brier'] - s_py['brier']:.4f}\n" + f"RMSE-p delta (cpp - py): {s_cpp['rmse_p'] - s_py['rmse_p']:.4f}" +) diff --git a/debug/benchmark_cpp_vs_py_sampler_probit.py b/debug/benchmark_cpp_vs_py_sampler_probit.py new file mode 100644 index 00000000..7d20eba1 --- /dev/null +++ b/debug/benchmark_cpp_vs_py_sampler_probit.py @@ -0,0 +1,152 @@ +"""Benchmark: C++ sampler loop vs. Python sampler loop – probit BART. + +Compares runtime, Brier score, and RMSE-to-truth (vs. pnorm(f_X)) across +run_cpp=True / False in BARTModel.sample(). + +Usage: + source venv/bin/activate # or: conda activate stochtree-book + python debug/benchmark_cpp_vs_py_sampler_probit.py +""" + +import argparse +import time +import numpy as np +from scipy.stats import norm +from stochtree import BARTModel, OutcomeModel + +parser = argparse.ArgumentParser() +parser.add_argument("--num-chains", type=int, default=1) +args = parser.parse_args() +num_chains = args.num_chains + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +rng = np.random.default_rng(1234) + +n = 2000 +p = 10 +X = rng.uniform(size=(n, p)) + +# Latent mean on the probit (standard-normal) scale – same step function as +# the continuous benchmark, keeping values well within identifiable range. +f_X = ( + np.where((X[:, 0] >= 0.00) & (X[:, 0] < 0.25), -7.5, 0) + + np.where((X[:, 0] >= 0.25) & (X[:, 0] < 0.50), -2.5, 0) + + np.where((X[:, 0] >= 0.50) & (X[:, 0] < 0.75), 2.5, 0) + + np.where((X[:, 0] >= 0.75) & (X[:, 0] < 1.00), 7.5, 0) +) +p_X = norm.cdf(f_X) # true P(Y = 1 | X) +y = rng.binomial(1, p_X).astype(float) # observed binary outcome + +test_frac = 0.2 +n_test = round(test_frac * n) +n_train = n - n_test +test_inds = rng.choice(n, size=n_test, replace=False) +train_inds = np.setdiff1d(np.arange(n), test_inds) + +X_train, X_test = X[train_inds], X[test_inds] +y_train, y_test = y[train_inds], y[test_inds] +p_test = p_X[test_inds] + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr = 10 +num_mcmc = 100 +num_trees = 200 +n_reps = 3 + +print( + f"n_train={n_train} n_test={n_test} p={p} " + f"num_trees={num_trees} num_gfr={num_gfr} num_mcmc={num_mcmc} " + f"num_chains={num_chains} reps={n_reps}\n" +) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + metrics +# --------------------------------------------------------------------------- +def run_once(run_cpp: bool, num_gfr: int, num_mcmc: int, seed: int) -> dict: + m = BARTModel() + t0 = time.perf_counter() + m.sample( + X_train=X_train, + y_train=y_train, + num_gfr=num_gfr, + num_burnin=0, + num_mcmc=num_mcmc, + mean_forest_params={"num_trees": num_trees}, + general_params={ + "random_seed": seed, + "outcome_model": OutcomeModel(outcome="binary", link="probit"), + "sample_sigma2_global": False, + "num_chains": num_chains, + }, + run_cpp=run_cpp, + ) + elapsed_sample = time.perf_counter() - t0 + + t1 = time.perf_counter() + preds = m.predict(X=X_test, scale="probability", run_cpp=run_cpp) + elapsed_predict = time.perf_counter() - t1 + + p_hat = preds["y_hat"].mean(axis=1) # (n_test,) + + brier = float(np.mean((p_hat - y_test) ** 2)) + rmse_p = float(np.sqrt(np.mean((p_hat - p_test) ** 2))) + + return {"elapsed_sample": elapsed_sample, "elapsed_predict": elapsed_predict, "brier": brier, "rmse_p": rmse_p} + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds = [1000 + i for i in range(1, n_reps + 1)] + +results_cpp = [] +results_py = [] + +print("Running C++ sampler (run_cpp=True)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_cpp.append(run_once(run_cpp=True, num_gfr=num_gfr, num_mcmc=num_mcmc, seed=seed)) + +print("\nRunning Python sampler (run_cpp=False)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_py.append(run_once(run_cpp=False, num_gfr=num_gfr, num_mcmc=num_mcmc, seed=seed)) + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +def summarise(results: list) -> dict: + keys = ["elapsed_sample", "elapsed_predict", "brier", "rmse_p"] + out = {k: float(np.mean([r[k] for r in results])) for k in keys} + out["elapsed"] = out["elapsed_sample"] + out["elapsed_predict"] + out["elapsed_sd"] = float(np.std( + [r["elapsed_sample"] + r["elapsed_predict"] for r in results], ddof=1 + )) + return out + +s_cpp = summarise(results_cpp) +s_py = summarise(results_py) +rows = [("cpp (run_cpp=True)", s_cpp), ("py (run_cpp=False)", s_py)] + +print("\n--- Results ---") +print( + f"{'Sampler':<22} {'Total (s)':>10} {'Samp (s)':>10} {'Pred (s)':>10} {'SD':>8} " + f"{'Brier':>12} {'RMSE (vs pnorm)':>16}" +) +print("-" * 98) +for label, s in rows: + print( + f"{label:<22} {s['elapsed']:>10.3f} {s['elapsed_sample']:>10.3f} " + f"{s['elapsed_predict']:>10.3f} {s['elapsed_sd']:>8.3f} " + f"{s['brier']:>12.4f} {s['rmse_p']:>16.4f}" + ) + +speedup = s_py["elapsed"] / s_cpp["elapsed"] +print(f"\nSpeedup (py / cpp): {speedup:.2f}x") +print( + f"Brier delta (cpp - py): {s_cpp['brier'] - s_py['brier']:.4f}\n" + f"RMSE-p delta (cpp - py): {s_cpp['rmse_p'] - s_py['rmse_p']:.4f}" +) diff --git a/debug/benchmark_cpp_vs_py_sampler_rfx.py b/debug/benchmark_cpp_vs_py_sampler_rfx.py new file mode 100644 index 00000000..32fe0fea --- /dev/null +++ b/debug/benchmark_cpp_vs_py_sampler_rfx.py @@ -0,0 +1,162 @@ +"""Benchmark: C++ sampler loop vs. Python sampler loop – BART with random effects. + +Compares runtime and RMSE (vs. test outcomes and vs. true mean) across +run_cpp=True / False in BARTModel.sample(). + +DGP: continuous outcome with an additive intercept-only random effect. + y_i = f(X_i) + alpha_{g_i} + eps_i, eps ~ N(0, 0.5^2) + f(X) is a piecewise-constant step function on X[:,0]. + Group intercepts alpha_g ~ N(0, 1), 10 groups. + +Usage: + source venv/bin/activate + python debug/benchmark_cpp_vs_py_sampler_rfx.py +""" + +import argparse +import time +import numpy as np +from stochtree import BARTModel + +parser = argparse.ArgumentParser() +parser.add_argument("--num-chains", type=int, default=1) +args = parser.parse_args() +num_chains = args.num_chains + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +rng = np.random.default_rng(1234) + +n = 2000 +p = 10 +num_groups = 10 + +X = rng.uniform(size=(n, p)) +f_X = np.where(X[:, 0] < 0.25, -7.5, + np.where(X[:, 0] < 0.5, -2.5, + np.where(X[:, 0] < 0.75, 2.5, 7.5))) + +group_ids = rng.integers(0, num_groups, size=n).astype(np.int32) +rfx_coefs = rng.normal(0, 1, size=num_groups) +rfx_term = rfx_coefs[group_ids] + +mu_true = f_X + rfx_term +y = mu_true + rng.normal(0, 0.5, size=n) + +test_frac = 0.2 +n_test = round(test_frac * n) +test_inds = rng.choice(n, size=n_test, replace=False) +train_inds = np.setdiff1d(np.arange(n), test_inds) + +X_train, X_test = X[train_inds], X[test_inds] +y_train, y_test = y[train_inds], y[test_inds] +group_ids_train = group_ids[train_inds] +group_ids_test = group_ids[test_inds] +rfx_basis_train = np.ones((len(train_inds), 1)) +rfx_basis_test = np.ones((n_test, 1)) +mu_test = mu_true[test_inds] + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr = 10 +num_burnin = 100 +num_mcmc = 100 +num_trees = 200 +n_reps = 3 + +print( + f"n_train={len(train_inds)} n_test={n_test} p={p} num_groups={num_groups} " + f"num_trees={num_trees} num_gfr={num_gfr} num_burnin={num_burnin} " + f"num_mcmc={num_mcmc} num_chains={num_chains} reps={n_reps}\n" +) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + metrics +# --------------------------------------------------------------------------- +def run_once(run_cpp: bool, seed: int) -> dict: + m = BARTModel() + t0 = time.perf_counter() + m.sample( + X_train=X_train, + y_train=y_train, + num_gfr=num_gfr, + num_burnin=num_burnin, + num_mcmc=num_mcmc, + rfx_group_ids_train=group_ids_train, + rfx_basis_train=rfx_basis_train, + mean_forest_params={"num_trees": num_trees}, + general_params={"random_seed": seed, "num_chains": num_chains}, + run_cpp=run_cpp, + ) + elapsed_sample = time.perf_counter() - t0 + + t1 = time.perf_counter() + preds = m.predict( + X=X_test, + rfx_group_ids=group_ids_test, + rfx_basis=rfx_basis_test, + run_cpp=run_cpp, + ) + elapsed_predict = time.perf_counter() - t1 + + y_hat = preds["y_hat"].mean(axis=1) + rmse_y = float(np.sqrt(np.mean((y_hat - y_test) ** 2))) + rmse_mu = float(np.sqrt(np.mean((y_hat - mu_test) ** 2))) + + return {"elapsed_sample": elapsed_sample, "elapsed_predict": elapsed_predict, "rmse_y": rmse_y, "rmse_mu": rmse_mu} + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds = [100 + i for i in range(1, n_reps + 1)] + +results_cpp = [] +results_py = [] + +print("Running C++ sampler (run_cpp=True)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_cpp.append(run_once(run_cpp=True, seed=seed)) + +print("\nRunning Python sampler (run_cpp=False)...") +for i, seed in enumerate(seeds, 1): + print(f" rep {i}/{n_reps}") + results_py.append(run_once(run_cpp=False, seed=seed)) + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +def summarise(results: list) -> dict: + keys = ["elapsed_sample", "elapsed_predict", "rmse_y", "rmse_mu"] + out = {k: float(np.mean([r[k] for r in results])) for k in keys} + out["elapsed"] = out["elapsed_sample"] + out["elapsed_predict"] + out["elapsed_sd"] = float(np.std( + [r["elapsed_sample"] + r["elapsed_predict"] for r in results], ddof=1 + )) + return out + +s_cpp = summarise(results_cpp) +s_py = summarise(results_py) +rows = [("cpp (run_cpp=True)", s_cpp), ("py (run_cpp=False)", s_py)] + +print("\n--- Results ---") +print( + f"{'Sampler':<22} {'Total (s)':>10} {'Samp (s)':>10} {'Pred (s)':>10} {'SD':>8} " + f"{'RMSE (y)':>10} {'RMSE (mu)':>10}" +) +print("-" * 92) +for label, s in rows: + print( + f"{label:<22} {s['elapsed']:>10.3f} {s['elapsed_sample']:>10.3f} " + f"{s['elapsed_predict']:>10.3f} {s['elapsed_sd']:>8.3f} " + f"{s['rmse_y']:>10.4f} {s['rmse_mu']:>10.4f}" + ) + +speedup = s_py["elapsed"] / s_cpp["elapsed"] +print(f"\nSpeedup (py / cpp): {speedup:.2f}x") +print( + f"RMSE-y delta (cpp - py): {s_cpp['rmse_y'] - s_py['rmse_y']:.4f}\n" + f"RMSE-mu delta (cpp - py): {s_cpp['rmse_mu'] - s_py['rmse_mu']:.4f}" +) diff --git a/debug/benchmark_cpp_vs_r_sampler.R b/debug/benchmark_cpp_vs_r_sampler.R new file mode 100644 index 00000000..fd21ba6c --- /dev/null +++ b/debug/benchmark_cpp_vs_r_sampler.R @@ -0,0 +1,188 @@ +## Benchmark: C++ sampler loop vs. R sampler loop +## Compares runtime and test-set RMSE across run_cpp = TRUE / FALSE in bart(). +## +## Usage: Rscript debug/benchmark_cpp_vs_r_sampler.R +## or source() from an interactive session after devtools::load_all('.') +library(stochtree) + +args <- commandArgs(trailingOnly = TRUE) +num_chains <- 1L +idx <- grep("^--num-chains=", args) +if (length(idx)) { + num_chains <- as.integer(sub("^--num-chains=", "", args[idx[1]])) +} + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +set.seed(1234) + +n <- 10000 +p <- 10 +X <- matrix(runif(n * p), ncol = p) +f_X <- (((0.00 <= X[, 1]) & (X[, 1] < 0.25)) * + (-7.5) + + ((0.25 <= X[, 1]) & (X[, 1] < 0.50)) * (-2.5) + + ((0.50 <= X[, 1]) & (X[, 1] < 0.75)) * (2.5) + + ((0.75 <= X[, 1]) & (X[, 1] < 1.00)) * (7.5)) +noise_sd <- 1 +y <- f_X + rnorm(n, 0, noise_sd) + +test_frac <- 0.2 +n_test <- round(test_frac * n) +n_train <- n - n_test +test_inds <- sort(sample(seq_len(n), n_test, replace = FALSE)) +train_inds <- setdiff(seq_len(n), test_inds) + +X_train <- X[train_inds, ] +X_test <- X[test_inds, ] +y_train <- y[train_inds] +y_test <- y[test_inds] +f_test <- f_X[test_inds] + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr <- 10 +num_mcmc <- 100 +num_trees <- 200 +n_reps <- 3 # repeated runs for stable timing + +cat(sprintf( + "n_train=%d n_test=%d p=%d num_trees=%d num_gfr=%d num_mcmc=%d num_chains=%d reps=%d\n\n", + n_train, + n_test, + p, + num_trees, + num_gfr, + num_mcmc, + num_chains, + n_reps +)) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + RMSE +# --------------------------------------------------------------------------- +run_once <- function(run_cpp, num_gfr, num_mcmc, seed = -1) { + t0 <- proc.time() + m <- bart( + X_train = X_train, + y_train = y_train, + num_gfr = num_gfr, + num_burnin = 0, + num_mcmc = num_mcmc, + mean_forest_params = list(num_trees = num_trees), + general_params = list(random_seed = seed, num_chains = num_chains), + run_cpp = run_cpp + ) + elapsed_sample <- (proc.time() - t0)[["elapsed"]] + + t1 <- proc.time() + preds <- predict( + m, + X = X_test, + run_cpp = run_cpp + ) + elapsed_predict <- (proc.time() - t1)[["elapsed"]] + + yhat <- rowMeans(preds$y_hat) + rmse <- sqrt(mean((yhat - y_test)^2)) + rmse_f <- sqrt(mean((yhat - f_test)^2)) + + list( + elapsed = elapsed_sample + elapsed_predict, + elapsed_sample = elapsed_sample, + elapsed_predict = elapsed_predict, + rmse = rmse, + rmse_f = rmse_f + ) +} + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds <- 1000 + seq_len(n_reps) + +results_cpp <- vector("list", n_reps) +results_r <- vector("list", n_reps) + +cat("Running C++ sampler (run_cpp = TRUE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_cpp[[i]] <- run_once( + run_cpp = TRUE, + num_gfr = num_gfr, + num_mcmc = num_mcmc, + seed = seeds[i] + ) +} + +cat("\nRunning R sampler (run_cpp = FALSE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_r[[i]] <- run_once( + run_cpp = FALSE, + num_gfr = num_gfr, + num_mcmc = num_mcmc, + seed = seeds[i] + ) +} + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +summarise <- function(results, label) { + elapsed <- sapply(results, `[[`, "elapsed") + elapsed_sample <- sapply(results, `[[`, "elapsed_sample") + elapsed_predict <- sapply(results, `[[`, "elapsed_predict") + rmse <- sapply(results, `[[`, "rmse") + rmse_f <- sapply(results, `[[`, "rmse_f") + data.frame( + sampler = label, + elapsed_mean = mean(elapsed), + elapsed_sd = sd(elapsed), + elapsed_sample_mean = mean(elapsed_sample), + elapsed_predict_mean = mean(elapsed_predict), + rmse_mean = mean(rmse), + rmse_f_mean = mean(rmse_f), + row.names = NULL + ) +} + +res <- rbind( + summarise(results_cpp, "cpp (run_cpp=TRUE)"), + summarise(results_r, "R (run_cpp=FALSE)") +) + +cat("\n--- Results ---\n") +cat(sprintf( + "%-22s %10s %10s %11s %10s %12s %12s\n", + "Sampler", + "Total (s)", + "Sample (s)", + "Predict (s)", + " SD", + "RMSE (obs)", + "RMSE (f)" +)) +cat(strrep("-", 96), "\n") +for (i in seq_len(nrow(res))) { + cat(sprintf( + "%-22s %10.3f %10.3f %11.3f %10.3f %12.4f %12.4f\n", + res$sampler[i], + res$elapsed_mean[i], + res$elapsed_sample_mean[i], + res$elapsed_predict_mean[i], + res$elapsed_sd[i], + res$rmse_mean[i], + res$rmse_f_mean[i] + )) +} + +speedup <- res$elapsed_mean[2] / res$elapsed_mean[1] +cat(sprintf("\nSpeedup (R / C++): %.2fx\n", speedup)) +cat(sprintf( + "RMSE delta (cpp - R): obs=%.4f f=%.4f\n", + res$rmse_mean[1] - res$rmse_mean[2], + res$rmse_f_mean[1] - res$rmse_f_mean[2] +)) diff --git a/debug/benchmark_cpp_vs_r_sampler_bcf.R b/debug/benchmark_cpp_vs_r_sampler_bcf.R new file mode 100644 index 00000000..fb9cd3e4 --- /dev/null +++ b/debug/benchmark_cpp_vs_r_sampler_bcf.R @@ -0,0 +1,229 @@ +## Benchmark: C++ sampler loop vs. R sampler loop (BCF) +## Simplest BCF case: univariate binary treatment, no RFX, no adaptive coding, +## no treatment intercept. +## +## Usage: Rscript debug/benchmark_cpp_vs_r_sampler_bcf.R +## or source() from an interactive session after devtools::load_all('.') +library(stochtree) + +args <- commandArgs(trailingOnly = TRUE) +num_chains <- 1L +idx <- grep("^--num-chains=", args) +if (length(idx)) { + num_chains <- as.integer(sub("^--num-chains=", "", args[idx[1]])) +} + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +set.seed(1234) + +n <- 2000 +p <- 10 +X <- matrix(runif(n * p), ncol = p) + +# Prognostic function: step function on X[,1] +mu_X <- (((0.00 <= X[, 1]) & (X[, 1] < 0.25)) * + (-7.5) + + ((0.25 <= X[, 1]) & (X[, 1] < 0.50)) * (-2.5) + + ((0.50 <= X[, 1]) & (X[, 1] < 0.75)) * (2.5) + + ((0.75 <= X[, 1]) & (X[, 1] < 1.00)) * (7.5)) + +# Treatment effect: linear in X[,2] +tau_X <- 2 + 4 * X[, 2] + +# Propensity score: mild confounding via X[,3] +pi_X <- 0.2 + 0.6 * X[, 3] +Z <- rbinom(n, 1, pi_X) + +noise_sd <- 1 +y <- mu_X + tau_X * Z + rnorm(n, 0, noise_sd) + +test_frac <- 0.2 +n_test <- round(test_frac * n) +n_train <- n - n_test +test_inds <- sort(sample(seq_len(n), n_test, replace = FALSE)) +train_inds <- setdiff(seq_len(n), test_inds) + +X_train <- X[train_inds, ] +X_test <- X[test_inds, ] +Z_train <- Z[train_inds] +Z_test <- Z[test_inds] +y_train <- y[train_inds] +y_test <- y[test_inds] +pi_train <- pi_X[train_inds] +pi_test <- pi_X[test_inds] +mu_test <- mu_X[test_inds] +tau_test <- tau_X[test_inds] +f_test <- mu_test + tau_test * Z_test # E[y|X,Z] + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr <- 10 +num_mcmc <- 100 +num_trees_mu <- 200 +num_trees_tau <- 50 +n_reps <- 3 + +cat(sprintf( + "n_train=%d n_test=%d p=%d mu_trees=%d tau_trees=%d num_gfr=%d num_mcmc=%d num_chains=%d reps=%d\n\n", + n_train, + n_test, + p, + num_trees_mu, + num_trees_tau, + num_gfr, + num_mcmc, + num_chains, + n_reps +)) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + RMSE +# --------------------------------------------------------------------------- +run_once <- function(run_cpp, num_gfr, num_mcmc, seed = -1) { + t0 <- proc.time() + m <- bcf( + X_train = X_train, + Z_train = Z_train, + y_train = y_train, + propensity_train = pi_train, + num_gfr = num_gfr, + num_burnin = 0, + num_mcmc = num_mcmc, + prognostic_forest_params = list(num_trees = num_trees_mu), + treatment_effect_forest_params = list( + num_trees = num_trees_tau, + sample_intercept = FALSE + ), + general_params = list( + random_seed = seed, + num_chains = num_chains, + adaptive_coding = FALSE, + propensity_covariate = "prognostic" + ), + run_cpp = run_cpp + ) + elapsed_sample <- (proc.time() - t0)[["elapsed"]] + + t1 <- proc.time() + preds <- predict( + m, + X = X_test, + Z = Z_test, + propensity = pi_test, + run_cpp = run_cpp + ) + elapsed_predict <- (proc.time() - t1)[["elapsed"]] + + yhat <- rowMeans(preds$y_hat) + tauhat <- rowMeans(preds$tau_hat) + + rmse_y <- sqrt(mean((yhat - y_test)^2)) + rmse_f <- sqrt(mean((yhat - f_test)^2)) + rmse_tau <- sqrt(mean((tauhat - tau_test)^2)) + + list( + elapsed = elapsed_sample + elapsed_predict, + elapsed_sample = elapsed_sample, + elapsed_predict = elapsed_predict, + rmse_y = rmse_y, + rmse_f = rmse_f, + rmse_tau = rmse_tau + ) +} + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds <- 1000 + seq_len(n_reps) + +results_cpp <- vector("list", n_reps) +results_r <- vector("list", n_reps) + +cat("Running C++ sampler (run_cpp = TRUE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_cpp[[i]] <- run_once( + run_cpp = TRUE, + num_gfr = num_gfr, + num_mcmc = num_mcmc, + seed = seeds[i] + ) +} + +cat("\nRunning R sampler (run_cpp = FALSE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_r[[i]] <- run_once( + run_cpp = FALSE, + num_gfr = num_gfr, + num_mcmc = num_mcmc, + seed = seeds[i] + ) +} + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +summarise <- function(results, label) { + elapsed <- sapply(results, `[[`, "elapsed") + elapsed_sample <- sapply(results, `[[`, "elapsed_sample") + elapsed_predict <- sapply(results, `[[`, "elapsed_predict") + rmse_y <- sapply(results, `[[`, "rmse_y") + rmse_f <- sapply(results, `[[`, "rmse_f") + rmse_tau <- sapply(results, `[[`, "rmse_tau") + data.frame( + sampler = label, + elapsed_mean = mean(elapsed), + elapsed_sd = sd(elapsed), + elapsed_sample_mean = mean(elapsed_sample), + elapsed_predict_mean = mean(elapsed_predict), + rmse_y_mean = mean(rmse_y), + rmse_f_mean = mean(rmse_f), + rmse_tau_mean = mean(rmse_tau), + row.names = NULL + ) +} + +res <- rbind( + summarise(results_cpp, "cpp (run_cpp=TRUE)"), + summarise(results_r, "R (run_cpp=FALSE)") +) + +cat("\n--- Results ---\n") +cat(sprintf( + "%-22s %10s %10s %10s %10s %12s %12s %12s\n", + "Sampler", + "Total (s)", + "Sample (s)", + "Predict (s)", + " SD", + "RMSE (obs)", + "RMSE (f)", + "RMSE (tau)" +)) +cat(strrep("-", 108), "\n") +for (i in seq_len(nrow(res))) { + cat(sprintf( + "%-22s %10.3f %10.3f %11.3f %10.3f %12.4f %12.4f %12.4f\n", + res$sampler[i], + res$elapsed_mean[i], + res$elapsed_sample_mean[i], + res$elapsed_predict_mean[i], + res$elapsed_sd[i], + res$rmse_y_mean[i], + res$rmse_f_mean[i], + res$rmse_tau_mean[i] + )) +} + +speedup <- res$elapsed_mean[2] / res$elapsed_mean[1] +cat(sprintf("\nSpeedup (R / C++): %.2fx\n", speedup)) +cat(sprintf( + "RMSE delta (cpp - R): obs=%.4f f=%.4f tau=%.4f\n", + res$rmse_y_mean[1] - res$rmse_y_mean[2], + res$rmse_f_mean[1] - res$rmse_f_mean[2], + res$rmse_tau_mean[1] - res$rmse_tau_mean[2] +)) diff --git a/debug/benchmark_cpp_vs_r_sampler_bcf_adaptive_coding.R b/debug/benchmark_cpp_vs_r_sampler_bcf_adaptive_coding.R new file mode 100644 index 00000000..6097e86b --- /dev/null +++ b/debug/benchmark_cpp_vs_r_sampler_bcf_adaptive_coding.R @@ -0,0 +1,268 @@ +## Benchmark: C++ sampler loop vs. R sampler loop -- BCF with adaptive coding. +## +## Exercises SampleAdaptiveCodingParameters() for binary treatment. Verifies: +## - The C++ path runs without error. +## - b_0_samples and b_1_samples are populated with the right shape. +## - mu_hat + Z * tau_hat == y_hat (internal decomposition check). +## - CATE RMSE (cpp) is close to CATE RMSE (R) -- large differences indicate +## a residual accounting bug in SampleAdaptiveCodingParameters or the +## mu/tau prediction split. +## - Speedup is reported for reference, though the primary goal is correctness. +## +## DGP: +## mu(X) = step function on X[,1] +## tau_forest(X) = 2 * X[,3] (forest component, scale-free) +## true_b_0 = -0.5 +## true_b_1 = 1.5 +## CATE(X) = (true_b_1 - true_b_0) * tau_forest(X) = 4 * X[,3] +## pi(X) = 0.2 + 0.6 * X[,4] +## Z ~ Bernoulli(pi(X)) +## y = mu(X) + (true_b_0*(1-Z) + true_b_1*Z) * tau_forest(X) + noise +## +## Usage: Rscript debug/benchmark_cpp_vs_r_sampler_bcf_adaptive_coding.R +## or source() from an interactive session after devtools::load_all('.') +library(stochtree) + +args <- commandArgs(trailingOnly = TRUE) +num_chains <- 1L +idx <- grep("^--num-chains=", args) +if (length(idx)) { + num_chains <- as.integer(sub("^--num-chains=", "", args[idx[1]])) +} + +# --------------------------------------------------------------------------- +# Shared settings +# --------------------------------------------------------------------------- +set.seed(1234) + +n <- 2000 +p <- 10 +noise_sd <- 1.0 +test_frac <- 0.2 +n_test <- round(test_frac * n) +n_train <- n - n_test + +num_gfr <- 10 +num_burnin <- 0 +num_mcmc <- 100 +num_trees_mu <- 200 +num_trees_tau <- 50 +n_reps <- 3 + +# --------------------------------------------------------------------------- +# DGP: binary treatment with adaptive coding +# --------------------------------------------------------------------------- +X_all <- matrix(runif(n * p), ncol = p) + +mu_X <- (((0.00 <= X_all[, 1]) & (X_all[, 1] < 0.25)) * + (-7.5) + + ((0.25 <= X_all[, 1]) & (X_all[, 1] < 0.50)) * (-2.5) + + ((0.50 <= X_all[, 1]) & (X_all[, 1] < 0.75)) * (2.5) + + ((0.75 <= X_all[, 1]) & (X_all[, 1] < 1.00)) * (7.5)) +TRUE_B_0 <- -0.5 +TRUE_B_1 <- 1.5 +tau_forest_X <- 2.0 * X_all[, 3] # forest component (no coding) +tau_X <- (TRUE_B_1 - TRUE_B_0) * tau_forest_X # CATE = (b_1 - b_0) * tau_forest +pi_X <- 0.2 + 0.6 * X_all[, 4] +Z_all <- rbinom(n, 1, pi_X) +coded_basis <- TRUE_B_0 * (1 - Z_all) + TRUE_B_1 * Z_all +y_all <- mu_X + coded_basis * tau_forest_X + rnorm(n, 0, noise_sd) + +test_inds <- sort(sample(seq_len(n), n_test, replace = FALSE)) +train_inds <- setdiff(seq_len(n), test_inds) + +X_train <- X_all[train_inds, ] +X_test <- X_all[test_inds, ] +Z_train <- Z_all[train_inds] +Z_test <- Z_all[test_inds] +pi_train <- pi_X[train_inds] +pi_test <- pi_X[test_inds] +y_train <- y_all[train_inds] +y_test <- y_all[test_inds] +mu_test <- mu_X[test_inds] +tau_test <- tau_X[test_inds] # true CATE +coded_test <- TRUE_B_0 * (1 - Z_test) + TRUE_B_1 * Z_test +f_test <- mu_test + coded_test * tau_forest_X[test_inds] + +cat(sprintf( + "n_train=%d n_test=%d p=%d mu_trees=%d tau_trees=%d num_gfr=%d num_mcmc=%d num_chains=%d reps=%d\n", + n_train, + n_test, + p, + num_trees_mu, + num_trees_tau, + num_gfr, + num_mcmc, + num_chains, + n_reps +)) +cat(sprintf("true_b_0=%.1f true_b_1=%.1f\n\n", TRUE_B_0, TRUE_B_1)) + +# --------------------------------------------------------------------------- +# Runner +# --------------------------------------------------------------------------- +run_once <- function(run_cpp, seed = -1) { + t0 <- proc.time() + m <- bcf( + X_train = X_train, + Z_train = Z_train, + y_train = y_train, + propensity_train = pi_train, + num_gfr = num_gfr, + num_burnin = num_burnin, + num_mcmc = num_mcmc, + prognostic_forest_params = list(num_trees = num_trees_mu), + treatment_effect_forest_params = list(num_trees = num_trees_tau), + general_params = list( + adaptive_coding = TRUE, + random_seed = seed, + num_chains = num_chains, + propensity_covariate = "prognostic" + ), + run_cpp = run_cpp + ) + elapsed_sample <- (proc.time() - t0)[["elapsed"]] + + t1 <- proc.time() + preds <- predict( + m, + X = X_test, + Z = Z_test, + propensity = pi_test, + run_cpp = run_cpp + ) + elapsed_predict <- (proc.time() - t1)[["elapsed"]] + + # Internal consistency: y_hat == mu_hat + Z * tau_hat + max_decomp_err_train <- max(abs( + m$y_hat_train - (m$mu_hat_train + Z_train * m$tau_hat_train) + )) + max_decomp_err_test <- max(abs( + preds$y_hat - (preds$mu_hat + Z_test * preds$tau_hat) + )) + + yhat <- rowMeans(preds$y_hat) + tauhat <- rowMeans(preds$tau_hat) + + list( + elapsed = elapsed_sample + elapsed_predict, + elapsed_sample = elapsed_sample, + elapsed_predict = elapsed_predict, + b0_mean = mean(m$b_0_samples), + b1_mean = mean(m$b_1_samples), + b0_length = length(m$b_0_samples), + b1_length = length(m$b_1_samples), + max_decomp_err_train = max_decomp_err_train, + max_decomp_err_test = max_decomp_err_test, + rmse_y = sqrt(mean((yhat - y_test)^2)), + rmse_f = sqrt(mean((yhat - f_test)^2)), + rmse_tau = sqrt(mean((tauhat - tau_test)^2)) + ) +} + +# --------------------------------------------------------------------------- +# Run +# --------------------------------------------------------------------------- +seeds <- 1000 + seq_len(n_reps) + +cat(strrep("=", 60), "\n") +cat("BINARY TREATMENT (adaptive_coding=TRUE)\n") +cat(sprintf( + "True b_0=%.1f b_1=%.1f CATE=(b1-b0)*tau_forest(X)\n", + TRUE_B_0, + TRUE_B_1 +)) +cat(strrep("=", 60), "\n") + +results_cpp <- vector("list", n_reps) +results_r <- vector("list", n_reps) + +cat("Running C++ sampler (run_cpp = TRUE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_cpp[[i]] <- run_once(run_cpp = TRUE, seed = seeds[i]) +} + +cat("\nRunning R sampler (run_cpp = FALSE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_r[[i]] <- run_once(run_cpp = FALSE, seed = seeds[i]) +} + +summarise <- function(results, label) { + data.frame( + sampler = label, + elapsed_mean = mean(sapply(results, `[[`, "elapsed")), + elapsed_sd = sd(sapply(results, `[[`, "elapsed")), + elapsed_sample_mean = mean(sapply(results, `[[`, "elapsed_sample")), + elapsed_predict_mean = mean(sapply(results, `[[`, "elapsed_predict")), + b0_mean = mean(sapply(results, `[[`, "b0_mean")), + b1_mean = mean(sapply(results, `[[`, "b1_mean")), + max_decomp_err_train = mean(sapply(results, `[[`, "max_decomp_err_train")), + max_decomp_err_test = mean(sapply(results, `[[`, "max_decomp_err_test")), + rmse_y = mean(sapply(results, `[[`, "rmse_y")), + rmse_f = mean(sapply(results, `[[`, "rmse_f")), + rmse_tau = mean(sapply(results, `[[`, "rmse_tau")), + row.names = NULL + ) +} + +res <- rbind( + summarise(results_cpp, "cpp (run_cpp=TRUE)"), + summarise(results_r, "R (run_cpp=FALSE)") +) + +cat("\n--- Results ---\n") +cat(sprintf( + "b_0_samples length cpp=%d R=%d\n", + results_cpp[[1]]$b0_length, + results_r[[1]]$b0_length +)) +cat(sprintf( + "b_1_samples length cpp=%d R=%d\n", + results_cpp[[1]]$b1_length, + results_r[[1]]$b1_length +)) +cat("\n") +cat(sprintf( + "%-22s %8s %8s %8s %6s %8s %8s %13s %13s %8s %8s %10s\n", + "Sampler", + "Total (s)", + "Samp (s)", + "Pred (s)", + "SD", + "b_0 mean", + "b_1 mean", + "max_decomp_tr", + "max_decomp_te", + "RMSE(y)", + "RMSE(f)", + "RMSE(tau)" +)) +cat(strrep("-", 140), "\n") +for (i in seq_len(nrow(res))) { + cat(sprintf( + "%-22s %8.3f %8.3f %8.3f %6.3f %8.4f %8.4f %13.2e %13.2e %8.4f %8.4f %10.4f\n", + res$sampler[i], + res$elapsed_mean[i], + res$elapsed_sample_mean[i], + res$elapsed_predict_mean[i], + res$elapsed_sd[i], + res$b0_mean[i], + res$b1_mean[i], + res$max_decomp_err_train[i], + res$max_decomp_err_test[i], + res$rmse_y[i], + res$rmse_f[i], + res$rmse_tau[i] + )) +} +cat(sprintf("True b_0=%.4f b_1=%.4f\n", TRUE_B_0, TRUE_B_1)) +speedup <- res$elapsed_mean[2] / res$elapsed_mean[1] +cat(sprintf("Speedup (R / C++): %.2fx\n", speedup)) +cat(sprintf( + "RMSE delta (cpp - R): y=%.4f f=%.4f tau=%.4f\n", + res$rmse_y[1] - res$rmse_y[2], + res$rmse_f[1] - res$rmse_f[2], + res$rmse_tau[1] - res$rmse_tau[2] +)) diff --git a/debug/benchmark_cpp_vs_r_sampler_bcf_multivariate.R b/debug/benchmark_cpp_vs_r_sampler_bcf_multivariate.R new file mode 100644 index 00000000..0bd4b8ab --- /dev/null +++ b/debug/benchmark_cpp_vs_r_sampler_bcf_multivariate.R @@ -0,0 +1,229 @@ +## Benchmark: C++ sampler loop vs. R sampler loop -- multivariate treatment BCF. +## +## Compares runtime and accuracy across run_cpp = TRUE / FALSE in bcf() +## with a 2-column binary treatment (multivariate BCF). +## +## DGP: propensity pi(X) = cbind(0.25 + 0.5*X[,1], 0.75 - 0.5*X[,2]) (2-column). +## mu(X) = 5*pi[,1] + 2*pi[,2] + 2*X[,3]. +## tau(X) = cbind(X[,2], X[,3]) (2-column CATE). +## Z ~ Bernoulli(pi(X)) column-wise (binary, shape n x 2). +## y = mu(X) + rowSums(Z * tau(X)) + noise. +## Adaptive coding is disabled; sigma2_leaf is not sampled for the tau forest. +## +## Usage: +## Rscript debug/benchmark_cpp_vs_r_sampler_bcf_multivariate.R +## or source() from an interactive session after devtools::load_all('.') + +library(stochtree) + +args <- commandArgs(trailingOnly = TRUE) +num_chains <- 1L +idx <- grep("^--num-chains=", args) +if (length(idx)) { + num_chains <- as.integer(sub("^--num-chains=", "", args[idx[1]])) +} + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +set.seed(1234) + +n <- 2000 +p <- 5 +snr <- 2.0 + +X <- matrix(runif(n * p), ncol = p) +pi_x <- cbind(0.25 + 0.5 * X[, 1], 0.75 - 0.5 * X[, 2]) +mu_x <- pi_x[, 1] * 5 + pi_x[, 2] * 2 + 2 * X[, 3] +tau_x <- cbind(X[, 2], X[, 3]) +Z <- matrix( + as.numeric(matrix(runif(n * 2), ncol = 2) < pi_x), + ncol = 2 +) +E_XZ <- mu_x + rowSums(Z * tau_x) +y <- E_XZ + rnorm(n, sd = sd(E_XZ) / snr) + +test_frac <- 0.2 +n_test <- round(test_frac * n) +n_train <- n - n_test +test_inds <- sort(sample(seq_len(n), n_test, replace = FALSE)) +train_inds <- setdiff(seq_len(n), test_inds) + +X_train <- X[train_inds, ] +X_test <- X[test_inds, ] +Z_train <- Z[train_inds, ] +Z_test <- Z[test_inds, ] +pi_train <- pi_x[train_inds, ] +pi_test <- pi_x[test_inds, ] +y_train <- y[train_inds] +y_test <- y[test_inds] +mu_test <- mu_x[test_inds] +tau_test <- tau_x[test_inds, ] +f_test <- mu_test + rowSums(Z_test * tau_test) + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr <- 10 +num_burnin <- 0 +num_mcmc <- 100 +num_trees_mu <- 200 +num_trees_tau <- 50 +n_reps <- 3 + +cat(sprintf( + "n_train=%d n_test=%d p=%d treatment_dim=2\nmu_trees=%d tau_trees=%d num_gfr=%d num_burnin=%d num_mcmc=%d num_chains=%d reps=%d\n\n", + n_train, + n_test, + p, + num_trees_mu, + num_trees_tau, + num_gfr, + num_burnin, + num_mcmc, + num_chains, + n_reps +)) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + accuracy metrics +# --------------------------------------------------------------------------- +run_once <- function(run_cpp, seed = -1) { + t0 <- proc.time() + m <- bcf( + X_train = X_train, + Z_train = Z_train, + y_train = y_train, + propensity_train = pi_train, + num_gfr = num_gfr, + num_burnin = num_burnin, + num_mcmc = num_mcmc, + prognostic_forest_params = list(num_trees = num_trees_mu), + treatment_effect_forest_params = list( + num_trees = num_trees_tau, + sample_sigma2_leaf = FALSE, + sample_intercept = FALSE + ), + general_params = list( + adaptive_coding = FALSE, + random_seed = seed, + num_chains = num_chains + ), + run_cpp = run_cpp + ) + elapsed_sample <- (proc.time() - t0)[["elapsed"]] + + t1 <- proc.time() + preds <- predict( + m, + X = X_test, + Z = Z_test, + propensity = pi_test, + run_cpp = run_cpp + ) + elapsed_predict <- (proc.time() - t1)[["elapsed"]] + + y_hat <- rowMeans(preds$y_hat) + mu_hat <- rowMeans(preds$mu_hat) + # tau_hat: array(n_test, treatment_dim, num_samples) + tau_hat1 <- apply(preds$tau_hat[, 1, ], 1, mean) + tau_hat2 <- apply(preds$tau_hat[, 2, ], 1, mean) + + list( + elapsed = elapsed_sample + elapsed_predict, + elapsed_sample = elapsed_sample, + elapsed_predict = elapsed_predict, + rmse_y = sqrt(mean((y_hat - y_test)^2)), + rmse_f = sqrt(mean((y_hat - f_test)^2)), + rmse_mu = sqrt(mean((mu_hat - mu_test)^2)), + rmse_tau1 = sqrt(mean((tau_hat1 - tau_test[, 1])^2)), + rmse_tau2 = sqrt(mean((tau_hat2 - tau_test[, 2])^2)) + ) +} + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds <- 1000 + seq_len(n_reps) + +results_cpp <- vector("list", n_reps) +results_r <- vector("list", n_reps) + +cat("Running C++ sampler (run_cpp = TRUE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_cpp[[i]] <- run_once(run_cpp = TRUE, seed = seeds[i]) +} + +cat("\nRunning R sampler (run_cpp = FALSE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_r[[i]] <- run_once(run_cpp = FALSE, seed = seeds[i]) +} + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +summarise <- function(results, label) { + get <- function(key) sapply(results, `[[`, key) + data.frame( + sampler = label, + elapsed_mean = mean(get("elapsed")), + elapsed_sd = sd(get("elapsed")), + elapsed_sample_mean = mean(get("elapsed_sample")), + elapsed_predict_mean = mean(get("elapsed_predict")), + rmse_y_mean = mean(get("rmse_y")), + rmse_f_mean = mean(get("rmse_f")), + rmse_mu_mean = mean(get("rmse_mu")), + rmse_tau1_mean = mean(get("rmse_tau1")), + rmse_tau2_mean = mean(get("rmse_tau2")), + row.names = NULL + ) +} + +res <- rbind( + summarise(results_cpp, "cpp (run_cpp=TRUE)"), + summarise(results_r, "R (run_cpp=FALSE)") +) + +cat("\n--- Results ---\n") +cat(sprintf( + "%-22s %8s %8s %8s %8s %9s %9s %9s %10s %10s\n", + "Sampler", + "Total (s)", + "Samp (s)", + "Pred (s)", + "SD", + "RMSE(y)", + "RMSE(f)", + "RMSE(mu)", + "RMSE(tau1)", + "RMSE(tau2)" +)) +cat(strrep("-", 115), "\n") +for (i in seq_len(nrow(res))) { + cat(sprintf( + "%-22s %8.3f %8.3f %8.3f %8.3f %9.4f %9.4f %9.4f %10.4f %10.4f\n", + res$sampler[i], + res$elapsed_mean[i], + res$elapsed_sample_mean[i], + res$elapsed_predict_mean[i], + res$elapsed_sd[i], + res$rmse_y_mean[i], + res$rmse_f_mean[i], + res$rmse_mu_mean[i], + res$rmse_tau1_mean[i], + res$rmse_tau2_mean[i] + )) +} + +speedup <- res$elapsed_mean[2] / res$elapsed_mean[1] +cat(sprintf("\nSpeedup (R / C++): %.2fx\n", speedup)) +cat(sprintf( + "RMSE delta (cpp - R): y=%.4f f=%.4f mu=%.4f tau1=%.4f tau2=%.4f\n", + res$rmse_y_mean[1] - res$rmse_y_mean[2], + res$rmse_f_mean[1] - res$rmse_f_mean[2], + res$rmse_mu_mean[1] - res$rmse_mu_mean[2], + res$rmse_tau1_mean[1] - res$rmse_tau1_mean[2], + res$rmse_tau2_mean[1] - res$rmse_tau2_mean[2] +)) diff --git a/debug/benchmark_cpp_vs_r_sampler_bcf_probit.R b/debug/benchmark_cpp_vs_r_sampler_bcf_probit.R new file mode 100644 index 00000000..15ea7d07 --- /dev/null +++ b/debug/benchmark_cpp_vs_r_sampler_bcf_probit.R @@ -0,0 +1,229 @@ +## Benchmark: C++ sampler loop vs. R sampler loop -- probit BCF. +## +## Compares runtime, Brier score on outcome, and latent-scale tau RMSE across +## run_cpp = TRUE / FALSE in bcf(). +## +## DGP (latent-index model): +## w = mu(X) + tau(X)*Z + eps, eps ~ N(0, 1) +## y = 1(w > 0) +## mu(X) = 1 + 2*X[,1] + X[,2] (confounded with propensity) +## tau(X) = 0.5 + X[,3] (latent-scale CATE) +## pi(X) = 0.4 + 0.2*X[,1] (mild confounding) +## Z ~ Bernoulli(pi(X)) +## +## Metrics: +## Brier score: mean((mean_s Phi(mu_hat[i,s] + tau_hat[i,s]*Z[i]) - y[i])^2) +## RMSE(tau): sqrt(mean((mean_s tau_hat_test[i,s] - tau_test[i])^2)) +## +## Usage: +## Rscript debug/benchmark_cpp_vs_r_sampler_bcf_probit.R +## or source() from an interactive session after devtools::load_all('.') + +library(stochtree) + +args <- commandArgs(trailingOnly = TRUE) +num_chains <- 1L +idx <- grep("^--num-chains=", args) +if (length(idx)) { + num_chains <- as.integer(sub("^--num-chains=", "", args[idx[1]])) +} + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +set.seed(1234) + +n <- 2000 +p <- 5 + +X <- matrix(runif(n * p), ncol = p) +mu_X <- 1 + 2 * X[, 1] + X[, 2] +tau_X <- 0.5 + X[, 3] +pi_X <- 0.4 + 0.2 * X[, 1] +Z <- rbinom(n, 1, pi_X) + +w <- mu_X + tau_X * Z + rnorm(n) +y <- as.numeric(w > 0) + +test_frac <- 0.2 +n_test <- round(test_frac * n) +n_train <- n - n_test +test_inds <- sort(sample(seq_len(n), n_test, replace = FALSE)) +train_inds <- setdiff(seq_len(n), test_inds) + +X_train <- X[train_inds, ] +X_test <- X[test_inds, ] +Z_train <- Z[train_inds] +Z_test <- Z[test_inds] +pi_train <- pi_X[train_inds] +pi_test <- pi_X[test_inds] +y_train <- y[train_inds] +y_test <- y[test_inds] +tau_test <- tau_X[test_inds] + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr <- 10 +num_burnin <- 0 +num_mcmc <- 100 +num_trees_mu <- 200 +num_trees_tau <- 50 +n_reps <- 3 + +cat(sprintf( + "n_train=%d n_test=%d p=%d\nmu_trees=%d tau_trees=%d num_gfr=%d num_burnin=%d num_mcmc=%d num_chains=%d reps=%d\n\n", + n_train, + n_test, + p, + num_trees_mu, + num_trees_tau, + num_gfr, + num_burnin, + num_mcmc, + num_chains, + n_reps +)) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + accuracy metrics +# --------------------------------------------------------------------------- +run_once <- function(run_cpp, seed = -1) { + t0 <- proc.time() + m <- bcf( + X_train = X_train, + Z_train = Z_train, + y_train = y_train, + propensity_train = pi_train, + num_gfr = num_gfr, + num_burnin = num_burnin, + num_mcmc = num_mcmc, + prognostic_forest_params = list( + num_trees = num_trees_mu, + sample_sigma2_leaf = FALSE + ), + treatment_effect_forest_params = list( + num_trees = num_trees_tau, + sample_sigma2_leaf = FALSE, + sample_intercept = FALSE + ), + general_params = list( + random_seed = seed, + num_chains = num_chains, + outcome_model = OutcomeModel(outcome = "binary", link = "probit"), + sample_sigma2_global = FALSE + ), + run_cpp = run_cpp + ) + elapsed_sample <- (proc.time() - t0)[["elapsed"]] + + # Request latent-scale mu and tau (scale = "linear", no probit transform applied) + t1 <- proc.time() + preds <- predict( + m, + X = X_test, + Z = Z_test, + propensity = pi_test, + terms = c("mu", "tau"), + scale = "linear", + run_cpp = run_cpp + ) + elapsed_predict <- (proc.time() - t1)[["elapsed"]] + + # mu_hat, tau_hat: (n_test, num_samples) — latent scale + mu_hat <- preds$mu_hat # n_test x num_samples + tau_hat <- preds$tau_hat # n_test x num_samples + + # P(Y=1 | X, Z, sample s) = Phi(mu_hat[i,s] + tau_hat[i,s] * Z_test[i]) + linear_pred <- mu_hat + tau_hat * Z_test # broadcasts Z_test over columns + p_hat_samples <- pnorm(linear_pred) # n_test x num_samples + p_hat_mean <- rowMeans(p_hat_samples) # n_test + + tau_hat_mean <- rowMeans(tau_hat) # n_test + + brier <- mean((p_hat_mean - y_test)^2) + rmse_tau <- sqrt(mean((tau_hat_mean - tau_test)^2)) + + list( + elapsed = elapsed_sample + elapsed_predict, + elapsed_sample = elapsed_sample, + elapsed_predict = elapsed_predict, + brier = brier, + rmse_tau = rmse_tau + ) +} + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds <- 1000 + seq_len(n_reps) + +results_cpp <- vector("list", n_reps) +results_r <- vector("list", n_reps) + +cat("Running C++ sampler (run_cpp = TRUE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_cpp[[i]] <- run_once(run_cpp = TRUE, seed = seeds[i]) +} + +cat("\nRunning R sampler (run_cpp = FALSE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_r[[i]] <- run_once(run_cpp = FALSE, seed = seeds[i]) +} + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +summarise <- function(results, label) { + get <- function(key) sapply(results, `[[`, key) + data.frame( + sampler = label, + elapsed_mean = mean(get("elapsed")), + elapsed_sd = sd(get("elapsed")), + elapsed_sample_mean = mean(get("elapsed_sample")), + elapsed_predict_mean = mean(get("elapsed_predict")), + brier_mean = mean(get("brier")), + rmse_tau_mean = mean(get("rmse_tau")), + row.names = NULL + ) +} + +res <- rbind( + summarise(results_cpp, "cpp (run_cpp=TRUE)"), + summarise(results_r, "R (run_cpp=FALSE)") +) + +cat("\n--- Results ---\n") +cat(sprintf( + "%-22s %10s %10s %10s %10s %10s %12s\n", + "Sampler", + "Total (s)", + "Samp (s)", + "Pred (s)", + "SD", + "Brier", + "RMSE (tau)" +)) +cat(strrep("-", 92), "\n") +for (i in seq_len(nrow(res))) { + cat(sprintf( + "%-22s %10.3f %10.3f %10.3f %10.3f %10.4f %12.4f\n", + res$sampler[i], + res$elapsed_mean[i], + res$elapsed_sample_mean[i], + res$elapsed_predict_mean[i], + res$elapsed_sd[i], + res$brier_mean[i], + res$rmse_tau_mean[i] + )) +} + +speedup <- res$elapsed_mean[2] / res$elapsed_mean[1] +cat(sprintf("\nSpeedup (R / C++): %.2fx\n", speedup)) +cat(sprintf( + "Delta (cpp - R): brier=%.4f rmse_tau=%.4f\n", + res$brier_mean[1] - res$brier_mean[2], + res$rmse_tau_mean[1] - res$rmse_tau_mean[2] +)) diff --git a/debug/benchmark_cpp_vs_r_sampler_bcf_tau_intercept.R b/debug/benchmark_cpp_vs_r_sampler_bcf_tau_intercept.R new file mode 100644 index 00000000..8ef57cbf --- /dev/null +++ b/debug/benchmark_cpp_vs_r_sampler_bcf_tau_intercept.R @@ -0,0 +1,441 @@ +## Benchmark: C++ sampler loop vs. R sampler loop -- BCF with treatment intercept. +## +## Exercises SampleParametricTreatmentEffect() for both univariate and multivariate +## treatment by toggling sample_intercept=TRUE. Verifies: +## - The C++ path runs without error. +## - tau_0_samples is populated and has the right shape. +## - CATE RMSE (cpp) is close to CATE RMSE (R) -- large differences indicate +## a residual accounting bug in the new intercept step. +## - Speedup is reported for reference, though the primary goal is correctness. +## +## DGP (univariate): +## mu(X) = step function on X[,1] +## tau_0 = 1.5 (true global intercept) +## tau(X) = tau_0 + 2*X[,3] (full CATE = tau_0 + forest component) +## pi(X) = 0.2 + 0.6*X[,4] +## Z ~ Bernoulli(pi(X)) +## y = mu(X) + (tau_0 + tau_forest(X)) * Z + noise +## +## DGP (multivariate, treatment_dim=2): +## Same mu(X); tau_0 = c(0.5, 1.0); tau_k(X) = tau_0[k] + X[,k+1] +## +## Usage: Rscript debug/benchmark_cpp_vs_r_sampler_bcf_tau_intercept.R +## or source() from an interactive session after devtools::load_all('.') +library(stochtree) + +args <- commandArgs(trailingOnly = TRUE) +num_chains <- 1L +idx <- grep("^--num-chains=", args) +if (length(idx)) { + num_chains <- as.integer(sub("^--num-chains=", "", args[idx[1]])) +} + +# --------------------------------------------------------------------------- +# Shared settings +# --------------------------------------------------------------------------- +set.seed(1234) + +n <- 2000 +p <- 10 +noise_sd <- 1.0 +test_frac <- 0.2 +n_test <- round(test_frac * n) +n_train <- n - n_test + +num_gfr <- 10 +num_burnin <- 0 +num_mcmc <- 100 +num_trees_mu <- 200 +num_trees_tau <- 50 +n_reps <- 3 + +# --------------------------------------------------------------------------- +# DGP: univariate binary treatment with global tau_0 +# --------------------------------------------------------------------------- +X_all <- matrix(runif(n * p), ncol = p) + +mu_X <- (((0.00 <= X_all[, 1]) & (X_all[, 1] < 0.25)) * + (-7.5) + + ((0.25 <= X_all[, 1]) & (X_all[, 1] < 0.50)) * (-2.5) + + ((0.50 <= X_all[, 1]) & (X_all[, 1] < 0.75)) * (2.5) + + ((0.75 <= X_all[, 1]) & (X_all[, 1] < 1.00)) * (7.5)) +TRUE_TAU0_UNIVARIATE <- 1.5 +tau_forest_X <- 2.0 * X_all[, 3] # forest component only +tau_X <- TRUE_TAU0_UNIVARIATE + tau_forest_X # full CATE +pi_X <- 0.2 + 0.6 * X_all[, 4] +Z_all <- rbinom(n, 1, pi_X) +y_all <- mu_X + tau_X * Z_all + rnorm(n, 0, noise_sd) + +test_inds <- sort(sample(seq_len(n), n_test, replace = FALSE)) +train_inds <- setdiff(seq_len(n), test_inds) + +X_train_u <- X_all[train_inds, ] +X_test_u <- X_all[test_inds, ] +Z_train_u <- Z_all[train_inds] +Z_test_u <- Z_all[test_inds] +pi_train_u <- pi_X[train_inds] +pi_test_u <- pi_X[test_inds] +y_train_u <- y_all[train_inds] +y_test_u <- y_all[test_inds] +mu_test_u <- mu_X[test_inds] +tau_test_u <- tau_X[test_inds] +f_test_u <- mu_test_u + tau_test_u * Z_test_u + +# --------------------------------------------------------------------------- +# DGP: multivariate (2-column) treatment with per-arm global tau_0 +# --------------------------------------------------------------------------- +X_all_mv <- matrix(runif(n * p), ncol = p) +TRUE_TAU0_MV <- c(0.5, 1.0) + +pi_mv <- cbind(0.25 + 0.5 * X_all_mv[, 1], 0.75 - 0.5 * X_all_mv[, 2]) +mu_mv <- pi_mv[, 1] * 5 + pi_mv[, 2] * 2 + 2 * X_all_mv[, 3] +tau_forest_mv <- cbind(X_all_mv[, 2], X_all_mv[, 3]) # forest component only +tau_mv <- sweep(tau_forest_mv, 2, TRUE_TAU0_MV, "+") # full CATE +Z_mv <- (matrix(runif(n * 2), ncol = 2) < pi_mv) * 1.0 +y_mv <- mu_mv + rowSums(Z_mv * tau_mv) + rnorm(n, 0, noise_sd) + +test_inds_mv <- sort(sample(seq_len(n), n_test, replace = FALSE)) +train_inds_mv <- setdiff(seq_len(n), test_inds_mv) + +X_train_mv <- X_all_mv[train_inds_mv, ] +X_test_mv <- X_all_mv[test_inds_mv, ] +Z_train_mv <- Z_mv[train_inds_mv, ] +Z_test_mv <- Z_mv[test_inds_mv, ] +pi_train_mv <- pi_mv[train_inds_mv, ] +pi_test_mv <- pi_mv[test_inds_mv, ] +y_train_mv <- y_mv[train_inds_mv] +y_test_mv <- y_mv[test_inds_mv] +mu_test_mv <- mu_mv[test_inds_mv] +tau_test_mv <- tau_mv[test_inds_mv, ] +f_test_mv <- mu_test_mv + rowSums(Z_test_mv * tau_test_mv) + +cat(sprintf( + "n_train=%d n_test=%d p=%d mu_trees=%d tau_trees=%d num_gfr=%d num_mcmc=%d num_chains=%d reps=%d\n\n", + n_train, + n_test, + p, + num_trees_mu, + num_trees_tau, + num_gfr, + num_mcmc, + num_chains, + n_reps +)) + +# --------------------------------------------------------------------------- +# Runner: univariate treatment with sample_intercept=TRUE +# --------------------------------------------------------------------------- +run_once_univariate <- function(run_cpp, seed = -1) { + t0 <- proc.time() + m <- bcf( + X_train = X_train_u, + Z_train = Z_train_u, + y_train = y_train_u, + propensity_train = pi_train_u, + num_gfr = num_gfr, + num_burnin = num_burnin, + num_mcmc = num_mcmc, + prognostic_forest_params = list(num_trees = num_trees_mu), + treatment_effect_forest_params = list( + num_trees = num_trees_tau, + sample_intercept = TRUE + ), + general_params = list( + random_seed = seed, + num_chains = num_chains, + adaptive_coding = FALSE, + propensity_covariate = "prognostic" + ), + run_cpp = run_cpp + ) + elapsed_sample <- (proc.time() - t0)[["elapsed"]] + + t1 <- proc.time() + preds <- predict( + m, + X = X_test_u, + Z = Z_test_u, + propensity = pi_test_u, + run_cpp = run_cpp + ) + elapsed_predict <- (proc.time() - t1)[["elapsed"]] + + tau_0_shape <- if (!is.null(m$tau_0_samples)) dim(m$tau_0_samples) else NULL + tau_0_mean <- if (!is.null(m$tau_0_samples)) { + mean(m$tau_0_samples) + } else { + NA_real_ + } + + yhat <- rowMeans(preds$y_hat) + tauhat <- rowMeans(preds$tau_hat) + + list( + elapsed = elapsed_sample + elapsed_predict, + elapsed_sample = elapsed_sample, + elapsed_predict = elapsed_predict, + tau_0_mean = tau_0_mean, + tau_0_shape = tau_0_shape, + rmse_y = sqrt(mean((yhat - y_test_u)^2)), + rmse_f = sqrt(mean((yhat - f_test_u)^2)), + rmse_tau = sqrt(mean((tauhat - tau_test_u)^2)) + ) +} + +# --------------------------------------------------------------------------- +# Runner: multivariate (2-column) treatment with sample_intercept=TRUE +# --------------------------------------------------------------------------- +run_once_multivariate <- function(run_cpp, seed = -1) { + t0 <- proc.time() + m <- bcf( + X_train = X_train_mv, + Z_train = Z_train_mv, + y_train = y_train_mv, + propensity_train = pi_train_mv, + num_gfr = num_gfr, + num_burnin = num_burnin, + num_mcmc = num_mcmc, + prognostic_forest_params = list(num_trees = num_trees_mu), + treatment_effect_forest_params = list( + num_trees = num_trees_tau, + sample_sigma2_leaf = FALSE, + sample_intercept = TRUE + ), + general_params = list( + random_seed = seed, + num_chains = num_chains, + adaptive_coding = FALSE + ), + run_cpp = run_cpp + ) + elapsed_sample <- (proc.time() - t0)[["elapsed"]] + + t1 <- proc.time() + preds <- predict( + m, + X = X_test_mv, + Z = Z_test_mv, + propensity = pi_test_mv, + run_cpp = run_cpp + ) + elapsed_predict <- (proc.time() - t1)[["elapsed"]] + + tau_0_shape <- if (!is.null(m$tau_0_samples)) dim(m$tau_0_samples) else NULL + tau_0_mean_0 <- if (!is.null(m$tau_0_samples)) { + mean(m$tau_0_samples[1, ]) + } else { + NA_real_ + } + tau_0_mean_1 <- if (!is.null(m$tau_0_samples)) { + mean(m$tau_0_samples[2, ]) + } else { + NA_real_ + } + + yhat <- rowMeans(preds$y_hat) + tauhat1 <- rowMeans(preds$tau_hat[, 1, ]) + tauhat2 <- rowMeans(preds$tau_hat[, 2, ]) + + list( + elapsed = elapsed_sample + elapsed_predict, + elapsed_sample = elapsed_sample, + elapsed_predict = elapsed_predict, + tau_0_mean_0 = tau_0_mean_0, + tau_0_mean_1 = tau_0_mean_1, + tau_0_shape = tau_0_shape, + rmse_y = sqrt(mean((yhat - y_test_mv)^2)), + rmse_f = sqrt(mean((yhat - f_test_mv)^2)), + rmse_tau1 = sqrt(mean((tauhat1 - tau_test_mv[, 1])^2)), + rmse_tau2 = sqrt(mean((tauhat2 - tau_test_mv[, 2])^2)) + ) +} + +# --------------------------------------------------------------------------- +# Run: univariate +# --------------------------------------------------------------------------- +seeds <- 1000 + seq_len(n_reps) + +cat(strrep("=", 60), "\n") +cat("UNIVARIATE TREATMENT (sample_intercept=TRUE)\n") +cat(sprintf("True tau_0 = %.4f\n", TRUE_TAU0_UNIVARIATE)) +cat(strrep("=", 60), "\n") + +results_cpp_u <- vector("list", n_reps) +results_r_u <- vector("list", n_reps) + +cat("Running C++ sampler (run_cpp = TRUE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_cpp_u[[i]] <- run_once_univariate(run_cpp = TRUE, seed = seeds[i]) +} + +cat("\nRunning R sampler (run_cpp = FALSE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_r_u[[i]] <- run_once_univariate(run_cpp = FALSE, seed = seeds[i]) +} + +summarise_u <- function(results, label) { + data.frame( + sampler = label, + elapsed_mean = mean(sapply(results, `[[`, "elapsed")), + elapsed_sd = sd(sapply(results, `[[`, "elapsed")), + tau_0_mean = mean(sapply(results, `[[`, "tau_0_mean")), + rmse_y = mean(sapply(results, `[[`, "rmse_y")), + rmse_f = mean(sapply(results, `[[`, "rmse_f")), + rmse_tau = mean(sapply(results, `[[`, "rmse_tau")), + row.names = NULL + ) +} + +res_u <- rbind( + summarise_u(results_cpp_u, "cpp (run_cpp=TRUE)"), + summarise_u(results_r_u, "R (run_cpp=FALSE)") +) + +cat("\n--- Univariate Results ---\n") +shape_cpp <- results_cpp_u[[1]]$tau_0_shape +shape_r <- results_r_u[[1]]$tau_0_shape +cat(sprintf( + "tau_0_samples shape cpp=%s R=%s\n", + if (is.null(shape_cpp)) { + "NULL" + } else { + paste0("[", paste(shape_cpp, collapse = ","), "]") + }, + if (is.null(shape_r)) { + "NULL" + } else { + paste0("[", paste(shape_r, collapse = ","), "]") + } +)) +cat(sprintf( + "%-22s %8s %8s %10s %9s %9s %10s\n", + "Sampler", + "Time (s)", + "SD", + "tau_0 mean", + "RMSE(y)", + "RMSE(f)", + "RMSE(tau)" +)) +cat(strrep("-", 90), "\n") +for (i in seq_len(nrow(res_u))) { + cat(sprintf( + "%-22s %8.3f %8.3f %10.4f %9.4f %9.4f %10.4f\n", + res_u$sampler[i], + res_u$elapsed_mean[i], + res_u$elapsed_sd[i], + res_u$tau_0_mean[i], + res_u$rmse_y[i], + res_u$rmse_f[i], + res_u$rmse_tau[i] + )) +} +cat(sprintf("True tau_0: %.4f\n", TRUE_TAU0_UNIVARIATE)) +speedup_u <- res_u$elapsed_mean[2] / res_u$elapsed_mean[1] +cat(sprintf("Speedup (R / C++): %.2fx\n", speedup_u)) +cat(sprintf( + "RMSE delta (cpp - R): y=%.4f f=%.4f tau=%.4f\n", + res_u$rmse_y[1] - res_u$rmse_y[2], + res_u$rmse_f[1] - res_u$rmse_f[2], + res_u$rmse_tau[1] - res_u$rmse_tau[2] +)) + +# --------------------------------------------------------------------------- +# Run: multivariate +# --------------------------------------------------------------------------- +cat("\n") +cat(strrep("=", 60), "\n") +cat("MULTIVARIATE TREATMENT (treatment_dim=2, sample_intercept=TRUE)\n") +cat(sprintf("True tau_0 = [%.4f, %.4f]\n", TRUE_TAU0_MV[1], TRUE_TAU0_MV[2])) +cat(strrep("=", 60), "\n") + +results_cpp_mv <- vector("list", n_reps) +results_r_mv <- vector("list", n_reps) + +cat("Running C++ sampler (run_cpp = TRUE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_cpp_mv[[i]] <- run_once_multivariate(run_cpp = TRUE, seed = seeds[i]) +} + +cat("\nRunning R sampler (run_cpp = FALSE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_r_mv[[i]] <- run_once_multivariate(run_cpp = FALSE, seed = seeds[i]) +} + +summarise_mv <- function(results, label) { + data.frame( + sampler = label, + elapsed_mean = mean(sapply(results, `[[`, "elapsed")), + elapsed_sd = sd(sapply(results, `[[`, "elapsed")), + tau_0_mean_0 = mean(sapply(results, `[[`, "tau_0_mean_0")), + tau_0_mean_1 = mean(sapply(results, `[[`, "tau_0_mean_1")), + rmse_y = mean(sapply(results, `[[`, "rmse_y")), + rmse_f = mean(sapply(results, `[[`, "rmse_f")), + rmse_tau1 = mean(sapply(results, `[[`, "rmse_tau1")), + rmse_tau2 = mean(sapply(results, `[[`, "rmse_tau2")), + row.names = NULL + ) +} + +res_mv <- rbind( + summarise_mv(results_cpp_mv, "cpp (run_cpp=TRUE)"), + summarise_mv(results_r_mv, "R (run_cpp=FALSE)") +) + +cat("\n--- Multivariate Results ---\n") +shape_cpp_mv <- results_cpp_mv[[1]]$tau_0_shape +shape_r_mv <- results_r_mv[[1]]$tau_0_shape +cat(sprintf( + "tau_0_samples shape cpp=%s R=%s\n", + if (is.null(shape_cpp_mv)) { + "NULL" + } else { + paste0("[", paste(shape_cpp_mv, collapse = ","), "]") + }, + if (is.null(shape_r_mv)) { + "NULL" + } else { + paste0("[", paste(shape_r_mv, collapse = ","), "]") + } +)) +cat(sprintf( + "%-22s %8s %8s %9s %9s %8s %8s %10s %10s\n", + "Sampler", + "Time (s)", + "SD", + "tau_0[1]", + "tau_0[2]", + "RMSE(y)", + "RMSE(f)", + "RMSE(tau1)", + "RMSE(tau2)" +)) +cat(strrep("-", 105), "\n") +for (i in seq_len(nrow(res_mv))) { + cat(sprintf( + "%-22s %8.3f %8.3f %9.4f %9.4f %8.4f %8.4f %10.4f %10.4f\n", + res_mv$sampler[i], + res_mv$elapsed_mean[i], + res_mv$elapsed_sd[i], + res_mv$tau_0_mean_0[i], + res_mv$tau_0_mean_1[i], + res_mv$rmse_y[i], + res_mv$rmse_f[i], + res_mv$rmse_tau1[i], + res_mv$rmse_tau2[i] + )) +} +cat(sprintf("True tau_0: [%.4f, %.4f]\n", TRUE_TAU0_MV[1], TRUE_TAU0_MV[2])) +speedup_mv <- res_mv$elapsed_mean[2] / res_mv$elapsed_mean[1] +cat(sprintf("Speedup (R / C++): %.2fx\n", speedup_mv)) +cat(sprintf( + "RMSE delta (cpp - R): y=%.4f f=%.4f tau1=%.4f tau2=%.4f\n", + res_mv$rmse_y[1] - res_mv$rmse_y[2], + res_mv$rmse_f[1] - res_mv$rmse_f[2], + res_mv$rmse_tau1[1] - res_mv$rmse_tau1[2], + res_mv$rmse_tau2[1] - res_mv$rmse_tau2[2] +)) diff --git a/debug/benchmark_cpp_vs_r_sampler_cloglog.R b/debug/benchmark_cpp_vs_r_sampler_cloglog.R new file mode 100644 index 00000000..46e1a416 --- /dev/null +++ b/debug/benchmark_cpp_vs_r_sampler_cloglog.R @@ -0,0 +1,191 @@ +## Benchmark: C++ sampler loop vs. R sampler loop – cloglog BART +## Compares runtime, Brier score, and RMSE-to-truth (vs. true P(Y=1|X)) across +## run_cpp = TRUE / FALSE in bart(). +## +## DGP uses the cloglog link: P(Y=1|X) = 1 - exp(-exp(f(X))). +## f(X) is a smooth sinusoidal function of two covariates, keeping probabilities +## in [~0.25, ~0.75] for stable mixing. GFR is disabled (num_gfr = 0). +## +## Usage: Rscript debug/benchmark_cpp_vs_r_sampler_cloglog.R +## or source() from an interactive session after devtools::load_all('.') +library(stochtree) + +args <- commandArgs(trailingOnly = TRUE) +num_chains <- 1L +idx <- grep("^--num-chains=", args) +if (length(idx)) { + num_chains <- as.integer(sub("^--num-chains=", "", args[idx[1]])) +} + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +set.seed(1234) + +n <- 2000 +p <- 10 +X <- matrix(runif(n * p), ncol = p) + +# Latent mean on the cloglog (log-log) scale. +# f_X is centred near -0.5 so that P(Y=1|X) = 1 - exp(-exp(f_X)) stays +# in [~0.25, ~0.75], avoiding extreme probabilities that inflate tree depth. +f_X <- 0.6 * sin(2 * pi * X[, 1]) + 0.4 * cos(2 * pi * X[, 2]) - 0.5 +p_X <- 1 - exp(-exp(f_X)) # true P(Y = 1 | X) +y <- rbinom(n, 1L, p_X) # observed binary outcome + +test_frac <- 0.2 +n_test <- round(test_frac * n) +n_train <- n - n_test +test_inds <- sort(sample(seq_len(n), n_test, replace = FALSE)) +train_inds <- setdiff(seq_len(n), test_inds) + +X_train <- X[train_inds, ] +X_test <- X[test_inds, ] +y_train <- y[train_inds] +y_test <- y[test_inds] +p_test <- p_X[test_inds] + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr <- 0 +num_burnin <- 100 +num_mcmc <- 100 +num_trees <- 200 +n_reps <- 3 + +cat(sprintf( + "n_train=%d n_test=%d p=%d num_trees=%d num_gfr=%d num_burnin=%d num_mcmc=%d num_chains=%d reps=%d\n\n", + n_train, + n_test, + p, + num_trees, + num_gfr, + num_burnin, + num_mcmc, + num_chains, + n_reps +)) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + metrics +# --------------------------------------------------------------------------- +run_once <- function(run_cpp, seed = -1) { + t0 <- proc.time() + m <- bart( + X_train = X_train, + y_train = y_train, + num_gfr = num_gfr, + num_burnin = num_burnin, + num_mcmc = num_mcmc, + mean_forest_params = list( + num_trees = num_trees, + sample_sigma2_leaf = FALSE + ), + general_params = list( + random_seed = seed, + outcome_model = OutcomeModel(outcome = "binary", link = "cloglog"), + sample_sigma2_global = FALSE, + num_chains = num_chains + ), + run_cpp = run_cpp + ) + elapsed_sample <- (proc.time() - t0)[["elapsed"]] + + t1 <- proc.time() + preds <- predict( + m, + X = X_test, + type = "posterior", + terms = "y_hat", + scale = "probability", + run_cpp = run_cpp + ) + elapsed_predict <- (proc.time() - t1)[["elapsed"]] + + p_hat_mat <- if (is.list(preds)) preds$y_hat else preds + if (is.null(dim(p_hat_mat))) { + p_hat_mat <- matrix(p_hat_mat, ncol = 1) + } + p_hat <- rowMeans(p_hat_mat) + + brier <- mean((p_hat - y_test)^2) + rmse_p <- sqrt(mean((p_hat - p_test)^2)) + + list( + elapsed = elapsed_sample + elapsed_predict, + elapsed_sample = elapsed_sample, + elapsed_predict = elapsed_predict, + brier = brier, + rmse_p = rmse_p + ) +} + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds <- 1000 + seq_len(n_reps) + +results_cpp <- vector("list", n_reps) +results_r <- vector("list", n_reps) + +cat("Running C++ sampler (run_cpp = TRUE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_cpp[[i]] <- run_once(run_cpp = TRUE, seed = seeds[i]) +} + +cat("\nRunning R sampler (run_cpp = FALSE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_r[[i]] <- run_once(run_cpp = FALSE, seed = seeds[i]) +} + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +summarise <- function(results, label) { + elapsed <- sapply(results, `[[`, "elapsed") + elapsed_sample <- sapply(results, `[[`, "elapsed_sample") + elapsed_predict <- sapply(results, `[[`, "elapsed_predict") + brier <- sapply(results, `[[`, "brier") + rmse_p <- sapply(results, `[[`, "rmse_p") + data.frame( + sampler = label, + elapsed_mean = mean(elapsed), + elapsed_sd = sd(elapsed), + elapsed_sample_mean = mean(elapsed_sample), + elapsed_predict_mean = mean(elapsed_predict), + brier_mean = mean(brier), + rmse_p_mean = mean(rmse_p), + row.names = NULL + ) +} + +res <- rbind( + summarise(results_cpp, "cpp (run_cpp=TRUE)"), + summarise(results_r, "R (run_cpp=FALSE)") +) + +cat("\n--- Results ---\n") +cat(sprintf( + "%-22s %10s %10s %11s %10s %12s %15s\n", + "Sampler", "Total (s)", "Sample (s)", "Predict (s)", "SD", "Brier", "RMSE (vs truth)" +)) +cat(strrep("-", 97), "\n") +for (i in seq_len(nrow(res))) { + cat(sprintf( + "%-22s %10.3f %10.3f %11.3f %10.3f %12.4f %15.4f\n", + res$sampler[i], res$elapsed_mean[i], res$elapsed_sample_mean[i], + res$elapsed_predict_mean[i], res$elapsed_sd[i], + res$brier_mean[i], res$rmse_p_mean[i] + )) +} + +speedup <- res$elapsed_mean[2] / res$elapsed_mean[1] +cat(sprintf("\nSpeedup (R / C++): %.2fx\n", speedup)) +cat(sprintf( + "Brier delta (cpp - R): %.4f\nRMSE-p delta (cpp - R): %.4f\n", + res$brier_mean[1] - res$brier_mean[2], + res$rmse_p_mean[1] - res$rmse_p_mean[2] +)) diff --git a/debug/benchmark_cpp_vs_r_sampler_heteroskedastic.R b/debug/benchmark_cpp_vs_r_sampler_heteroskedastic.R new file mode 100644 index 00000000..a6574d31 --- /dev/null +++ b/debug/benchmark_cpp_vs_r_sampler_heteroskedastic.R @@ -0,0 +1,207 @@ +## Benchmark: C++ sampler loop vs. R sampler loop -- heteroskedastic BART. +## +## Compares runtime, mean-forest RMSE (vs. true f(X)) and RMSE of the estimated +## conditional standard deviation (vs. the true s(X)) across run_cpp = TRUE / +## FALSE in bart() with both a mean forest and a variance forest +## (num_trees_variance > 0). +## +## DGP: f(X) is a step function of X[,1]; s(X) varies by quadrant of X[,1] +## and linearly with X[,3], matching the heteroskedastic Python benchmark. +## +## A variance-only model (num_trees_mean = 0, num_trees_variance > 0) is +## supported in the C++ path. The mean-forest RMSE is reported as NA in +## that case since there is no mean forest to evaluate. +## +## Usage: +## Rscript debug/benchmark_cpp_vs_r_sampler_heteroskedastic.R +## or source() from an interactive session after devtools::load_all('.') + +library(stochtree) + +args <- commandArgs(trailingOnly = TRUE) +num_chains <- 1L +idx <- grep("^--num-chains=", args) +if (length(idx)) { + num_chains <- as.integer(sub("^--num-chains=", "", args[idx[1]])) +} + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +set.seed(1234) + +n <- 2000 +p <- 10 +X <- matrix(runif(n * p), ncol = p) + +# True conditional mean +f_X <- (((0.00 <= X[, 1]) & (X[, 1] < 0.25)) * + (-3.0) + + ((0.25 <= X[, 1]) & (X[, 1] < 0.50)) * (-1.0) + + ((0.50 <= X[, 1]) & (X[, 1] < 0.75)) * (1.0) + + ((0.75 <= X[, 1]) & (X[, 1] < 1.00)) * (3.0)) + +# True conditional standard deviation +s_X <- (((0.00 <= X[, 1]) & (X[, 1] < 0.25)) * + (0.5 * X[, 3]) + + ((0.25 <= X[, 1]) & (X[, 1] < 0.50)) * (1.0 * X[, 3]) + + ((0.50 <= X[, 1]) & (X[, 1] < 0.75)) * (2.0 * X[, 3]) + + ((0.75 <= X[, 1]) & (X[, 1] < 1.00)) * (3.0 * X[, 3])) + +y <- f_X + rnorm(n, 0, 1) * s_X + +test_frac <- 0.2 +n_test <- round(test_frac * n) +n_train <- n - n_test +test_inds <- sort(sample(seq_len(n), n_test, replace = FALSE)) +train_inds <- setdiff(seq_len(n), test_inds) + +X_train <- X[train_inds, ] +X_test <- X[test_inds, ] +y_train <- y[train_inds] +y_test <- y[test_inds] +f_test <- f_X[test_inds] +s_test <- s_X[test_inds] + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr <- 10 +num_burnin <- 0 +num_mcmc <- 100 +num_trees_mean <- 200 +num_trees_variance <- 50 +n_reps <- 3 + +cat(sprintf( + "n_train=%d n_test=%d p=%d num_trees_mean=%d num_trees_variance=%d num_gfr=%d num_burnin=%d num_mcmc=%d num_chains=%d reps=%d\n\n", + n_train, + n_test, + p, + num_trees_mean, + num_trees_variance, + num_gfr, + num_burnin, + num_mcmc, + num_chains, + n_reps +)) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + metrics +# --------------------------------------------------------------------------- +run_once <- function(run_cpp, seed) { + t0 <- proc.time() + m <- bart( + X_train = X_train, + y_train = y_train, + num_gfr = num_gfr, + num_burnin = num_burnin, + num_mcmc = num_mcmc, + general_params = list( + random_seed = seed, + sample_sigma2_global = FALSE, + num_chains = num_chains + ), + mean_forest_params = list(num_trees = num_trees_mean), + variance_forest_params = list(num_trees = num_trees_variance), + run_cpp = run_cpp + ) + elapsed_sample <- (proc.time() - t0)[["elapsed"]] + + t1 <- proc.time() + preds <- predict( + m, + X = X_test, + terms = c("y_hat", "variance_forest"), + run_cpp = run_cpp + ) + elapsed_predict <- (proc.time() - t1)[["elapsed"]] + + if (num_trees_mean > 0) { + f_hat <- rowMeans(preds$y_hat) + rmse_f <- sqrt(mean((f_hat - f_test)^2)) + } else { + rmse_f <- NA_real_ + } + + s_hat <- rowMeans(sqrt(preds$variance_forest_predictions)) + rmse_s <- sqrt(mean((s_hat - s_test)^2)) + + list( + elapsed = elapsed_sample + elapsed_predict, + elapsed_sample = elapsed_sample, + elapsed_predict = elapsed_predict, + rmse_f = rmse_f, + rmse_s = rmse_s + ) +} + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds <- 1000 + seq_len(n_reps) + +results_cpp <- vector("list", n_reps) +results_r <- vector("list", n_reps) + +cat("Running C++ sampler (run_cpp = TRUE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_cpp[[i]] <- run_once(run_cpp = TRUE, seed = seeds[i]) +} + +cat("\nRunning R sampler (run_cpp = FALSE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_r[[i]] <- run_once(run_cpp = FALSE, seed = seeds[i]) +} + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +summarise <- function(results, label) { + elapsed <- sapply(results, `[[`, "elapsed") + elapsed_sample <- sapply(results, `[[`, "elapsed_sample") + elapsed_predict <- sapply(results, `[[`, "elapsed_predict") + rmse_f <- sapply(results, `[[`, "rmse_f") + rmse_s <- sapply(results, `[[`, "rmse_s") + data.frame( + sampler = label, + elapsed_mean = mean(elapsed), + elapsed_sd = sd(elapsed), + elapsed_sample_mean = mean(elapsed_sample), + elapsed_predict_mean = mean(elapsed_predict), + rmse_f_mean = mean(rmse_f, na.rm = TRUE), + rmse_s_mean = mean(rmse_s), + row.names = NULL + ) +} + +res <- rbind( + summarise(results_cpp, "cpp (run_cpp=TRUE)"), + summarise(results_r, "R (run_cpp=FALSE)") +) + +cat("\n--- Results ---\n") +cat(sprintf( + "%-22s %10s %10s %11s %10s %12s %12s\n", + "Sampler", "Total (s)", "Sample (s)", "Predict (s)", "SD", "RMSE f(X)", "RMSE s(X)" +)) +cat(strrep("-", 94), "\n") +for (i in seq_len(nrow(res))) { + cat(sprintf( + "%-22s %10.3f %10.3f %11.3f %10.3f %12s %12.4f\n", + res$sampler[i], res$elapsed_mean[i], res$elapsed_sample_mean[i], + res$elapsed_predict_mean[i], res$elapsed_sd[i], + if (is.nan(res$rmse_f_mean[i])) "nan" else sprintf("%.4f", res$rmse_f_mean[i]), + res$rmse_s_mean[i] + )) +} + +speedup <- res$elapsed_mean[2] / res$elapsed_mean[1] +cat(sprintf("\nSpeedup (R / C++): %.2fx\n", speedup)) +cat(sprintf( + "RMSE s(X) delta (cpp - R): %.4f\n", + res$rmse_s_mean[1] - res$rmse_s_mean[2] +)) diff --git a/debug/benchmark_cpp_vs_r_sampler_leaf_regression.R b/debug/benchmark_cpp_vs_r_sampler_leaf_regression.R new file mode 100644 index 00000000..6998f67b --- /dev/null +++ b/debug/benchmark_cpp_vs_r_sampler_leaf_regression.R @@ -0,0 +1,190 @@ +## Benchmark: C++ sampler loop vs. R sampler loop -- univariate leaf regression. +## +## Compares runtime and test-set RMSE across run_cpp = TRUE / FALSE in bart() +## with a univariate leaf regression basis (leaf_basis_train with one column). +## +## DGP: f(X, Z) = tau(X) * Z, where tau(X) is a step function of X[,1] and +## Z is drawn uniform [0, 1]. A constant noise term is added. The leaf basis +## passed to the sampler is just Z (shape n x 1). +## +## Usage: +## Rscript debug/benchmark_cpp_vs_r_sampler_leaf_regression.R +## or source() from an interactive session after devtools::load_all('.') + +library(stochtree) + +args <- commandArgs(trailingOnly = TRUE) +num_chains <- 1L +idx <- grep("^--num-chains=", args) +if (length(idx)) { + num_chains <- as.integer(sub("^--num-chains=", "", args[idx[1]])) +} + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +set.seed(1234) + +n <- 2000 +p <- 10 +X <- matrix(runif(n * p), ncol = p) +Z <- runif(n) # scalar moderating variable / leaf basis + +# Heterogeneous slope on Z, partitioned by X[,1] +tau_X <- (((0.00 <= X[, 1]) & (X[, 1] < 0.25)) * + (-2.0) + + ((0.25 <= X[, 1]) & (X[, 1] < 0.50)) * (-1.0) + + ((0.50 <= X[, 1]) & (X[, 1] < 0.75)) * (1.0) + + ((0.75 <= X[, 1]) & (X[, 1] < 1.00)) * (2.0)) +f_XZ <- tau_X * Z +noise_sd <- 1.0 +y <- f_XZ + rnorm(n, 0, noise_sd) + +test_frac <- 0.2 +n_test <- round(test_frac * n) +n_train <- n - n_test +test_inds <- sort(sample(seq_len(n), n_test, replace = FALSE)) +train_inds <- setdiff(seq_len(n), test_inds) + +X_train <- X[train_inds, ] +X_test <- X[test_inds, ] +Z_train <- Z[train_inds] +Z_test <- Z[test_inds] +y_train <- y[train_inds] +y_test <- y[test_inds] +f_test <- f_XZ[test_inds] + +# Leaf basis matrices (n x 1) +basis_train <- matrix(Z_train, ncol = 1) +basis_test <- matrix(Z_test, ncol = 1) + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr <- 10 +num_burnin <- 0 +num_mcmc <- 100 +num_trees <- 200 +n_reps <- 3 + +cat(sprintf( + "n_train=%d n_test=%d p=%d num_trees=%d num_gfr=%d num_burnin=%d num_mcmc=%d num_chains=%d reps=%d\n\n", + n_train, + n_test, + p, + num_trees, + num_gfr, + num_burnin, + num_mcmc, + num_chains, + n_reps +)) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + RMSE +# --------------------------------------------------------------------------- +run_once <- function(run_cpp, seed) { + t0 <- proc.time() + m <- bart( + X_train = X_train, + y_train = y_train, + leaf_basis_train = basis_train, + num_gfr = num_gfr, + num_burnin = num_burnin, + num_mcmc = num_mcmc, + general_params = list(random_seed = seed, num_chains = num_chains), + mean_forest_params = list(num_trees = num_trees), + run_cpp = run_cpp + ) + elapsed_sample <- (proc.time() - t0)[["elapsed"]] + + t1 <- proc.time() + preds <- predict( + m, + X = X_test, + leaf_basis = basis_test, + run_cpp = run_cpp + ) + elapsed_predict <- (proc.time() - t1)[["elapsed"]] + + yhat <- rowMeans(preds$y_hat) + rmse <- sqrt(mean((yhat - y_test)^2)) + rmse_f <- sqrt(mean((yhat - f_test)^2)) + + list( + elapsed = elapsed_sample + elapsed_predict, + elapsed_sample = elapsed_sample, + elapsed_predict = elapsed_predict, + rmse = rmse, + rmse_f = rmse_f + ) +} + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds <- 1000 + seq_len(n_reps) + +results_cpp <- vector("list", n_reps) +results_r <- vector("list", n_reps) + +cat("Running C++ sampler (run_cpp = TRUE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_cpp[[i]] <- run_once(run_cpp = TRUE, seed = seeds[i]) +} + +cat("\nRunning R sampler (run_cpp = FALSE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_r[[i]] <- run_once(run_cpp = FALSE, seed = seeds[i]) +} + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +summarise <- function(results, label) { + elapsed <- sapply(results, `[[`, "elapsed") + elapsed_sample <- sapply(results, `[[`, "elapsed_sample") + elapsed_predict <- sapply(results, `[[`, "elapsed_predict") + rmse <- sapply(results, `[[`, "rmse") + rmse_f <- sapply(results, `[[`, "rmse_f") + data.frame( + sampler = label, + elapsed_mean = mean(elapsed), + elapsed_sd = sd(elapsed), + elapsed_sample_mean = mean(elapsed_sample), + elapsed_predict_mean = mean(elapsed_predict), + rmse_mean = mean(rmse), + rmse_f_mean = mean(rmse_f), + row.names = NULL + ) +} + +res <- rbind( + summarise(results_cpp, "cpp (run_cpp=TRUE)"), + summarise(results_r, "R (run_cpp=FALSE)") +) + +cat("\n--- Results ---\n") +cat(sprintf( + "%-22s %10s %10s %11s %10s %12s %13s\n", + "Sampler", "Total (s)", "Sample (s)", "Predict (s)", "SD", "RMSE (obs)", "RMSE f(X,Z)" +)) +cat(strrep("-", 94), "\n") +for (i in seq_len(nrow(res))) { + cat(sprintf( + "%-22s %10.3f %10.3f %11.3f %10.3f %12.4f %13.4f\n", + res$sampler[i], res$elapsed_mean[i], res$elapsed_sample_mean[i], + res$elapsed_predict_mean[i], res$elapsed_sd[i], + res$rmse_mean[i], res$rmse_f_mean[i] + )) +} + +speedup <- res$elapsed_mean[2] / res$elapsed_mean[1] +cat(sprintf("\nSpeedup (R / C++): %.2fx\n", speedup)) +cat(sprintf( + "RMSE delta (cpp - R): obs=%.4f f=%.4f\n", + res$rmse_mean[1] - res$rmse_mean[2], + res$rmse_f_mean[1] - res$rmse_f_mean[2] +)) diff --git a/debug/benchmark_cpp_vs_r_sampler_multivariate_leaf_regression.R b/debug/benchmark_cpp_vs_r_sampler_multivariate_leaf_regression.R new file mode 100644 index 00000000..5fd8a3b8 --- /dev/null +++ b/debug/benchmark_cpp_vs_r_sampler_multivariate_leaf_regression.R @@ -0,0 +1,205 @@ +## Benchmark: C++ sampler loop vs. R sampler loop -- multivariate leaf regression. +## +## Compares runtime and test-set RMSE across run_cpp = TRUE / FALSE in bart() +## with a 2-column leaf basis (multivariate leaf regression). +## +## DGP: f(X, Z) = tau_1(X)*Z_1 + tau_2(X)*Z_2, where tau_1/tau_2 are step functions +## of X[,1] and Z_1, Z_2 are drawn uniform [0, 1]. A constant noise term is added. +## The leaf basis passed to the sampler is cbind(Z_1, Z_2) (n x 2). +## +## Usage: +## Rscript debug/benchmark_cpp_vs_r_sampler_multivariate_leaf_regression.R +## or source() from an interactive session after devtools::load_all('.') + +library(stochtree) + +args <- commandArgs(trailingOnly = TRUE) +num_chains <- 1L +idx <- grep("^--num-chains=", args) +if (length(idx)) { + num_chains <- as.integer(sub("^--num-chains=", "", args[idx[1]])) +} + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +set.seed(1234) + +n <- 2000 +p <- 10 +X <- matrix(runif(n * p), ncol = p) +Z1 <- runif(n) +Z2 <- runif(n) + +# Heterogeneous slopes on Z1 and Z2, partitioned by X[,1] +tau1_X <- (((0.00 <= X[, 1]) & (X[, 1] < 0.25)) * + (-2.0) + + ((0.25 <= X[, 1]) & (X[, 1] < 0.50)) * (-1.0) + + ((0.50 <= X[, 1]) & (X[, 1] < 0.75)) * (1.0) + + ((0.75 <= X[, 1]) & (X[, 1] < 1.00)) * (2.0)) +tau2_X <- (((0.00 <= X[, 1]) & (X[, 1] < 0.25)) * + (1.0) + + ((0.25 <= X[, 1]) & (X[, 1] < 0.50)) * (2.0) + + ((0.50 <= X[, 1]) & (X[, 1] < 0.75)) * (-1.0) + + ((0.75 <= X[, 1]) & (X[, 1] < 1.00)) * (-2.0)) +f_XZ <- tau1_X * Z1 + tau2_X * Z2 +noise_sd <- 1.0 +y <- f_XZ + rnorm(n, 0, noise_sd) + +test_frac <- 0.2 +n_test <- round(test_frac * n) +n_train <- n - n_test +test_inds <- sort(sample(seq_len(n), n_test, replace = FALSE)) +train_inds <- setdiff(seq_len(n), test_inds) + +X_train <- X[train_inds, ] +X_test <- X[test_inds, ] +Z1_train <- Z1[train_inds] +Z1_test <- Z1[test_inds] +Z2_train <- Z2[train_inds] +Z2_test <- Z2[test_inds] +y_train <- y[train_inds] +y_test <- y[test_inds] +f_test <- f_XZ[test_inds] + +# Leaf basis matrices (n x 2) +basis_train <- cbind(Z1_train, Z2_train) +basis_test <- cbind(Z1_test, Z2_test) + +# Optional: 2x2 prior covariance for the leaf coefficients +sigma2_leaf_init <- matrix(c(0.5, 0.0, 0.0, 0.5), nrow = 2, ncol = 2) + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr <- 10 +num_burnin <- 0 +num_mcmc <- 100 +num_trees <- 200 +n_reps <- 3 + +cat(sprintf( + "n_train=%d n_test=%d p=%d basis_dim=2\nnum_trees=%d num_gfr=%d num_burnin=%d num_mcmc=%d num_chains=%d reps=%d\n\n", + n_train, + n_test, + p, + num_trees, + num_gfr, + num_burnin, + num_mcmc, + num_chains, + n_reps +)) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + RMSE +# --------------------------------------------------------------------------- +run_once <- function(run_cpp, seed) { + t0 <- proc.time() + m <- bart( + X_train = X_train, + y_train = y_train, + leaf_basis_train = basis_train, + num_gfr = num_gfr, + num_burnin = num_burnin, + num_mcmc = num_mcmc, + general_params = list(random_seed = seed, num_chains = num_chains), + mean_forest_params = list( + num_trees = num_trees, + sigma2_leaf_init = sigma2_leaf_init, + sample_sigma2_leaf = FALSE + ), + run_cpp = run_cpp + ) + elapsed_sample <- (proc.time() - t0)[["elapsed"]] + + t1 <- proc.time() + preds <- predict( + m, + X = X_test, + leaf_basis = basis_test, + run_cpp = run_cpp + ) + elapsed_predict <- (proc.time() - t1)[["elapsed"]] + + yhat <- rowMeans(preds$y_hat) + rmse <- sqrt(mean((yhat - y_test)^2)) + rmse_f <- sqrt(mean((yhat - f_test)^2)) + + list( + elapsed = elapsed_sample + elapsed_predict, + elapsed_sample = elapsed_sample, + elapsed_predict = elapsed_predict, + rmse = rmse, + rmse_f = rmse_f + ) +} + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds <- 1000 + seq_len(n_reps) + +results_cpp <- vector("list", n_reps) +results_r <- vector("list", n_reps) + +cat("Running C++ sampler (run_cpp = TRUE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_cpp[[i]] <- run_once(run_cpp = TRUE, seed = seeds[i]) +} + +cat("\nRunning R sampler (run_cpp = FALSE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_r[[i]] <- run_once(run_cpp = FALSE, seed = seeds[i]) +} + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +summarise <- function(results, label) { + elapsed <- sapply(results, `[[`, "elapsed") + elapsed_sample <- sapply(results, `[[`, "elapsed_sample") + elapsed_predict <- sapply(results, `[[`, "elapsed_predict") + rmse <- sapply(results, `[[`, "rmse") + rmse_f <- sapply(results, `[[`, "rmse_f") + data.frame( + sampler = label, + elapsed_mean = mean(elapsed), + elapsed_sd = sd(elapsed), + elapsed_sample_mean = mean(elapsed_sample), + elapsed_predict_mean = mean(elapsed_predict), + rmse_mean = mean(rmse), + rmse_f_mean = mean(rmse_f), + row.names = NULL + ) +} + +res <- rbind( + summarise(results_cpp, "cpp (run_cpp=TRUE)"), + summarise(results_r, "R (run_cpp=FALSE)") +) + +cat("\n--- Results ---\n") +cat(sprintf( + "%-22s %10s %10s %11s %10s %12s %13s\n", + "Sampler", "Total (s)", "Sample (s)", "Predict (s)", "SD", "RMSE (obs)", "RMSE f(X,Z)" +)) +cat(strrep("-", 94), "\n") +for (i in seq_len(nrow(res))) { + cat(sprintf( + "%-22s %10.3f %10.3f %11.3f %10.3f %12.4f %13.4f\n", + res$sampler[i], res$elapsed_mean[i], res$elapsed_sample_mean[i], + res$elapsed_predict_mean[i], res$elapsed_sd[i], + res$rmse_mean[i], res$rmse_f_mean[i] + )) +} + +speedup <- res$elapsed_mean[2] / res$elapsed_mean[1] +cat(sprintf("\nSpeedup (R / C++): %.2fx\n", speedup)) +cat(sprintf( + "RMSE delta (cpp - R): obs=%.4f f=%.4f\n", + res$rmse_mean[1] - res$rmse_mean[2], + res$rmse_f_mean[1] - res$rmse_f_mean[2] +)) diff --git a/debug/benchmark_cpp_vs_r_sampler_ordinal_cloglog.R b/debug/benchmark_cpp_vs_r_sampler_ordinal_cloglog.R new file mode 100644 index 00000000..c55760de --- /dev/null +++ b/debug/benchmark_cpp_vs_r_sampler_ordinal_cloglog.R @@ -0,0 +1,210 @@ +## Benchmark: C++ sampler loop vs. R sampler loop – ordinal cloglog BART +## Compares runtime, mean Brier score, and mean RMSE-to-truth (vs. true class +## probabilities) across run_cpp = TRUE / FALSE in bart(). +## +## DGP uses 4 ordinal categories with a cloglog link. +## f(X) is a smooth sinusoidal function of two covariates so that all four +## categories are well-populated. Cutpoints are spaced to yield roughly equal +## marginal class frequencies. GFR is disabled (num_gfr = 0). +## +## Usage: Rscript debug/benchmark_cpp_vs_r_sampler_ordinal_cloglog.R +## or source() from an interactive session after devtools::load_all('.') +library(stochtree) + +args <- commandArgs(trailingOnly = TRUE) +num_chains <- 1L +idx <- grep("^--num-chains=", args) +if (length(idx)) { + num_chains <- as.integer(sub("^--num-chains=", "", args[idx[1]])) +} + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +set.seed(1234) + +n <- 2000 +p <- 10 +X <- matrix(runif(n * p), ncol = p) + +# Latent mean on the cloglog scale (smooth, two covariates) +f_X <- 0.6 * sin(2 * pi * X[, 1]) + 0.4 * cos(2 * pi * X[, 2]) + +# Fixed log-scale cutpoints spaced to give roughly equal marginal class freqs. +# With f_X in roughly [-1, 1], setting gamma = c(0, log(2), log(4)) puts the +# four cumulative boundaries at moderate probability levels. +K <- 4 +gamma_true <- c(0, log(2), log(4)) + +# True cumulative probabilities: P(Y <= k | X) = 1 - exp(-exp(f_X - gamma_k)) +# True class probabilities: P(Y = k | X) = P(Y <= k) - P(Y <= k-1) +cum_prob <- outer(f_X, gamma_true, function(f, g) 1 - exp(-exp(f - g))) +p_X <- cbind( + cum_prob[, 1], + cum_prob[, 2] - cum_prob[, 1], + cum_prob[, 3] - cum_prob[, 2], + 1 - cum_prob[, 3] +) # n x K matrix of true class probs + +# Draw ordinal outcomes (1-indexed: 1, 2, 3, 4) +u <- runif(n) +y <- as.integer(u > cum_prob[, 1]) + + as.integer(u > cum_prob[, 2]) + + as.integer(u > cum_prob[, 3]) + + 1L + +test_frac <- 0.2 +n_test <- round(test_frac * n) +n_train <- n - n_test +test_inds <- sort(sample(seq_len(n), n_test, replace = FALSE)) +train_inds <- setdiff(seq_len(n), test_inds) + +X_train <- X[train_inds, ] +X_test <- X[test_inds, ] +y_train <- y[train_inds] +y_test <- y[test_inds] +p_test <- p_X[test_inds, ] # n_test x K matrix of true class probabilities + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr <- 0 +num_burnin <- 100 +num_mcmc <- 100 +num_trees <- 200 +n_reps <- 3 + +cat(sprintf( + "K=%d n_train=%d n_test=%d p=%d num_trees=%d num_gfr=%d num_burnin=%d num_mcmc=%d num_chains=%d reps=%d\n\n", + K, + n_train, + n_test, + p, + num_trees, + num_gfr, + num_burnin, + num_mcmc, + num_chains, + n_reps +)) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + metrics +# --------------------------------------------------------------------------- +run_once <- function(run_cpp, seed = -1) { + t0 <- proc.time() + m <- bart( + X_train = X_train, + y_train = y_train, + num_gfr = num_gfr, + num_burnin = num_burnin, + num_mcmc = num_mcmc, + mean_forest_params = list( + num_trees = num_trees, + sample_sigma2_leaf = FALSE + ), + general_params = list( + random_seed = seed, + outcome_model = OutcomeModel(outcome = "ordinal", link = "cloglog"), + sample_sigma2_global = FALSE, + num_chains = num_chains + ), + run_cpp = run_cpp + ) + elapsed_sample <- (proc.time() - t0)[["elapsed"]] + + t1 <- proc.time() + preds <- predict( + m, + X = X_test, + type = "posterior", + terms = "y_hat", + scale = "probability", + run_cpp = run_cpp + ) + elapsed_predict <- (proc.time() - t1)[["elapsed"]] + + # predict() returns list with $y_hat (C++) or raw array (R) for single-term + p_hat_arr <- if (is.list(preds)) preds$y_hat else preds + p_hat <- apply(p_hat_arr, c(1, 2), mean) # n_test x K posterior mean + + brier <- mean((p_hat - p_test)^2) + rmse_p <- mean(sqrt(colMeans((p_hat - p_test)^2))) + + list( + elapsed = elapsed_sample + elapsed_predict, + elapsed_sample = elapsed_sample, + elapsed_predict = elapsed_predict, + brier = brier, + rmse_p = rmse_p + ) +} + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds <- 1000 + seq_len(n_reps) + +results_cpp <- vector("list", n_reps) +results_r <- vector("list", n_reps) + +cat("Running C++ sampler (run_cpp = TRUE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_cpp[[i]] <- run_once(run_cpp = TRUE, seed = seeds[i]) +} + +cat("\nRunning R sampler (run_cpp = FALSE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_r[[i]] <- run_once(run_cpp = FALSE, seed = seeds[i]) +} + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +summarise <- function(results, label) { + elapsed <- sapply(results, `[[`, "elapsed") + elapsed_sample <- sapply(results, `[[`, "elapsed_sample") + elapsed_predict <- sapply(results, `[[`, "elapsed_predict") + brier <- sapply(results, `[[`, "brier") + rmse_p <- sapply(results, `[[`, "rmse_p") + data.frame( + sampler = label, + elapsed_mean = mean(elapsed), + elapsed_sd = sd(elapsed), + elapsed_sample_mean = mean(elapsed_sample), + elapsed_predict_mean = mean(elapsed_predict), + brier_mean = mean(brier), + rmse_p_mean = mean(rmse_p), + row.names = NULL + ) +} + +res <- rbind( + summarise(results_cpp, "cpp (run_cpp=TRUE)"), + summarise(results_r, "R (run_cpp=FALSE)") +) + +cat("\n--- Results ---\n") +cat(sprintf( + "%-22s %10s %10s %11s %10s %12s %15s\n", + "Sampler", "Total (s)", "Sample (s)", "Predict (s)", "SD", "Brier", "RMSE (vs truth)" +)) +cat(strrep("-", 97), "\n") +for (i in seq_len(nrow(res))) { + cat(sprintf( + "%-22s %10.3f %10.3f %11.3f %10.3f %12.4f %15.4f\n", + res$sampler[i], res$elapsed_mean[i], res$elapsed_sample_mean[i], + res$elapsed_predict_mean[i], res$elapsed_sd[i], + res$brier_mean[i], res$rmse_p_mean[i] + )) +} + +speedup <- res$elapsed_mean[2] / res$elapsed_mean[1] +cat(sprintf("\nSpeedup (R / C++): %.2fx\n", speedup)) +cat(sprintf( + "Brier delta (cpp - R): %.4f\nRMSE-p delta (cpp - R): %.4f\n", + res$brier_mean[1] - res$brier_mean[2], + res$rmse_p_mean[1] - res$rmse_p_mean[2] +)) diff --git a/debug/benchmark_cpp_vs_r_sampler_probit.R b/debug/benchmark_cpp_vs_r_sampler_probit.R new file mode 100644 index 00000000..022f0450 --- /dev/null +++ b/debug/benchmark_cpp_vs_r_sampler_probit.R @@ -0,0 +1,196 @@ +## Benchmark: C++ sampler loop vs. R sampler loop – probit BART +## Compares runtime, Brier score, and RMSE-to-truth (vs. pnorm(f_X)) across +## run_cpp = TRUE / FALSE in bart(). +## +## Usage: Rscript debug/benchmark_cpp_vs_r_sampler_probit.R +## or source() from an interactive session after devtools::load_all('.') +library(stochtree) + +args <- commandArgs(trailingOnly = TRUE) +num_chains <- 1L +idx <- grep("^--num-chains=", args) +if (length(idx)) { + num_chains <- as.integer(sub("^--num-chains=", "", args[idx[1]])) +} + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +set.seed(1234) + +n <- 2000 +p <- 10 +X <- matrix(runif(n * p), ncol = p) + +# Latent mean on the probit (standard-normal) scale – same step function as +# the continuous benchmark, keeping values well within identifiable range. +f_X <- (((0.00 <= X[, 1]) & (X[, 1] < 0.25)) * + (-7.5) + + ((0.25 <= X[, 1]) & (X[, 1] < 0.50)) * (-2.5) + + ((0.50 <= X[, 1]) & (X[, 1] < 0.75)) * (2.5) + + ((0.75 <= X[, 1]) & (X[, 1] < 1.00)) * (7.5)) +p_X <- pnorm(f_X) # true P(Y = 1 | X) +y <- rbinom(n, 1L, p_X) # observed binary outcome + +test_frac <- 0.2 +n_test <- round(test_frac * n) +n_train <- n - n_test +test_inds <- sort(sample(seq_len(n), n_test, replace = FALSE)) +train_inds <- setdiff(seq_len(n), test_inds) + +X_train <- X[train_inds, ] +X_test <- X[test_inds, ] +y_train <- y[train_inds] +y_test <- y[test_inds] +p_test <- p_X[test_inds] + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr <- 10 +num_mcmc <- 100 +num_trees <- 200 +n_reps <- 3 + +cat(sprintf( + "n_train=%d n_test=%d p=%d num_trees=%d num_gfr=%d num_mcmc=%d num_chains=%d reps=%d\n\n", + n_train, + n_test, + p, + num_trees, + num_gfr, + num_mcmc, + num_chains, + n_reps +)) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + metrics +# --------------------------------------------------------------------------- +run_once <- function(run_cpp, num_gfr, num_mcmc, seed = -1) { + t0 <- proc.time() + m <- bart( + X_train = X_train, + y_train = y_train, + num_gfr = num_gfr, + num_burnin = 0, + num_mcmc = num_mcmc, + mean_forest_params = list(num_trees = num_trees), + general_params = list( + random_seed = seed, + outcome_model = OutcomeModel(outcome = "binary", link = "probit"), + sample_sigma2_global = FALSE, + num_chains = num_chains + ), + run_cpp = run_cpp + ) + elapsed_sample <- (proc.time() - t0)[["elapsed"]] + + t1 <- proc.time() + p_hat_mat <- predict( + m, + X = X_test, + type = "posterior", + terms = "y_hat", + scale = "probability", + run_cpp = run_cpp + ) + elapsed_predict <- (proc.time() - t1)[["elapsed"]] + + raw <- if (is.list(p_hat_mat)) p_hat_mat$y_hat else p_hat_mat + if (is.null(dim(raw))) { + p_hat <- raw + } else { + p_hat <- rowMeans(raw) + } + + brier <- mean((p_hat - y_test)^2) + rmse_p <- sqrt(mean((p_hat - p_test)^2)) + + list( + elapsed = elapsed_sample + elapsed_predict, + elapsed_sample = elapsed_sample, + elapsed_predict = elapsed_predict, + brier = brier, + rmse_p = rmse_p + ) +} + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds <- 1000 + seq_len(n_reps) + +results_cpp <- vector("list", n_reps) +results_r <- vector("list", n_reps) + +cat("Running C++ sampler (run_cpp = TRUE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_cpp[[i]] <- run_once( + run_cpp = TRUE, + num_gfr = num_gfr, + num_mcmc = num_mcmc, + seed = seeds[i] + ) +} + +cat("\nRunning R sampler (run_cpp = FALSE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_r[[i]] <- run_once( + run_cpp = FALSE, + num_gfr = num_gfr, + num_mcmc = num_mcmc, + seed = seeds[i] + ) +} + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +summarise <- function(results, label) { + elapsed <- sapply(results, `[[`, "elapsed") + elapsed_sample <- sapply(results, `[[`, "elapsed_sample") + elapsed_predict <- sapply(results, `[[`, "elapsed_predict") + brier <- sapply(results, `[[`, "brier") + rmse_p <- sapply(results, `[[`, "rmse_p") + data.frame( + sampler = label, + elapsed_mean = mean(elapsed), + elapsed_sd = sd(elapsed), + elapsed_sample_mean = mean(elapsed_sample), + elapsed_predict_mean = mean(elapsed_predict), + brier_mean = mean(brier), + rmse_p_mean = mean(rmse_p), + row.names = NULL + ) +} + +res <- rbind( + summarise(results_cpp, "cpp (run_cpp=TRUE)"), + summarise(results_r, "R (run_cpp=FALSE)") +) + +cat("\n--- Results ---\n") +cat(sprintf( + "%-22s %10s %10s %11s %10s %12s %16s\n", + "Sampler", "Total (s)", "Sample (s)", "Predict (s)", "SD", "Brier", "RMSE (vs pnorm)" +)) +cat(strrep("-", 98), "\n") +for (i in seq_len(nrow(res))) { + cat(sprintf( + "%-22s %10.3f %10.3f %11.3f %10.3f %12.4f %16.4f\n", + res$sampler[i], res$elapsed_mean[i], res$elapsed_sample_mean[i], + res$elapsed_predict_mean[i], res$elapsed_sd[i], + res$brier_mean[i], res$rmse_p_mean[i] + )) +} + +speedup <- res$elapsed_mean[2] / res$elapsed_mean[1] +cat(sprintf("\nSpeedup (R / C++): %.2fx\n", speedup)) +cat(sprintf( + "Brier delta (cpp - R): %.4f\nRMSE-p delta (cpp - R): %.4f\n", + res$brier_mean[1] - res$brier_mean[2], + res$rmse_p_mean[1] - res$rmse_p_mean[2] +)) diff --git a/debug/benchmark_cpp_vs_r_sampler_rfx.R b/debug/benchmark_cpp_vs_r_sampler_rfx.R new file mode 100644 index 00000000..14f9479b --- /dev/null +++ b/debug/benchmark_cpp_vs_r_sampler_rfx.R @@ -0,0 +1,180 @@ +## Benchmark: C++ sampler loop vs. R sampler loop – BART with random effects. +## Compares runtime and RMSE (vs. test outcomes and vs. true mean) across +## run_cpp = TRUE / FALSE in bart(). +## +## DGP: continuous outcome with an additive intercept-only random effect. +## y_i = f(X_i) + alpha_{g_i} + eps_i, eps ~ N(0, 0.5^2) +## f(X) is a piecewise-constant step function on X[,1]. +## Group intercepts alpha_g ~ N(0, 1), 10 groups. +## +## Usage: Rscript debug/benchmark_cpp_vs_r_sampler_rfx.R +## or source() from an interactive session after devtools::load_all('.') +library(stochtree) + +args <- commandArgs(trailingOnly = TRUE) +num_chains <- 1L +idx <- grep("^--num-chains=", args) +if (length(idx)) { + num_chains <- as.integer(sub("^--num-chains=", "", args[idx[1]])) +} + +# --------------------------------------------------------------------------- +# Data-generating process +# --------------------------------------------------------------------------- +set.seed(1234) + +n <- 2000 +p <- 10 +num_groups <- 10 + +X <- matrix(runif(n * p), ncol = p) +f_X <- ifelse(X[, 1] < 0.25, -7.5, + ifelse(X[, 1] < 0.5, -2.5, + ifelse(X[, 1] < 0.75, 2.5, 7.5))) + +group_ids <- sample(seq_len(num_groups), n, replace = TRUE) +rfx_coefs <- rnorm(num_groups) +rfx_term <- rfx_coefs[group_ids] + +mu_true <- f_X + rfx_term +y <- mu_true + rnorm(n, 0, 0.5) + +test_frac <- 0.2 +n_test <- round(test_frac * n) +test_inds <- sort(sample(seq_len(n), n_test, replace = FALSE)) +train_inds <- setdiff(seq_len(n), test_inds) + +X_train <- X[train_inds, ] +X_test <- X[test_inds, ] +y_train <- y[train_inds] +y_test <- y[test_inds] +group_ids_train <- group_ids[train_inds] +group_ids_test <- group_ids[test_inds] +rfx_basis_train <- matrix(1, nrow = length(train_inds), ncol = 1) +rfx_basis_test <- matrix(1, nrow = n_test, ncol = 1) +mu_test <- mu_true[test_inds] + +# --------------------------------------------------------------------------- +# Benchmark settings +# --------------------------------------------------------------------------- +num_gfr <- 10 +num_burnin <- 100 +num_mcmc <- 100 +num_trees <- 200 +n_reps <- 3 + +cat(sprintf( + "n_train=%d n_test=%d p=%d num_groups=%d num_trees=%d num_gfr=%d num_burnin=%d num_mcmc=%d num_chains=%d reps=%d\n\n", + length(train_inds), n_test, p, num_groups, num_trees, num_gfr, num_burnin, num_mcmc, num_chains, n_reps +)) + +# --------------------------------------------------------------------------- +# Helper: run one configuration and return timing + metrics +# --------------------------------------------------------------------------- +run_once <- function(run_cpp, seed = -1) { + t0 <- proc.time() + m <- bart( + X_train = X_train, + y_train = y_train, + num_gfr = num_gfr, + num_burnin = num_burnin, + num_mcmc = num_mcmc, + rfx_group_ids_train = group_ids_train, + rfx_basis_train = rfx_basis_train, + mean_forest_params = list(num_trees = num_trees), + general_params = list(random_seed = seed, num_chains = num_chains), + run_cpp = run_cpp + ) + elapsed_sample <- (proc.time() - t0)[["elapsed"]] + + t1 <- proc.time() + preds <- predict( + m, + X = X_test, + rfx_group_ids = group_ids_test, + rfx_basis = rfx_basis_test, + run_cpp = run_cpp + ) + elapsed_predict <- (proc.time() - t1)[["elapsed"]] + + y_hat <- rowMeans(preds$y_hat) + rmse_y <- sqrt(mean((y_hat - y_test)^2)) + rmse_mu <- sqrt(mean((y_hat - mu_test)^2)) + + list( + elapsed = elapsed_sample + elapsed_predict, + elapsed_sample = elapsed_sample, + elapsed_predict = elapsed_predict, + rmse_y = rmse_y, + rmse_mu = rmse_mu + ) +} + +# --------------------------------------------------------------------------- +# Run benchmarks +# --------------------------------------------------------------------------- +seeds <- 1000 + seq_len(n_reps) + +results_cpp <- vector("list", n_reps) +results_r <- vector("list", n_reps) + +cat("Running C++ sampler (run_cpp = TRUE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_cpp[[i]] <- run_once(run_cpp = TRUE, seed = seeds[i]) +} + +cat("\nRunning R sampler (run_cpp = FALSE)...\n") +for (i in seq_len(n_reps)) { + cat(sprintf(" rep %d/%d\n", i, n_reps)) + results_r[[i]] <- run_once(run_cpp = FALSE, seed = seeds[i]) +} + +# --------------------------------------------------------------------------- +# Summarise +# --------------------------------------------------------------------------- +summarise <- function(results, label) { + elapsed <- sapply(results, `[[`, "elapsed") + elapsed_sample <- sapply(results, `[[`, "elapsed_sample") + elapsed_predict <- sapply(results, `[[`, "elapsed_predict") + rmse_y <- sapply(results, `[[`, "rmse_y") + rmse_mu <- sapply(results, `[[`, "rmse_mu") + data.frame( + sampler = label, + elapsed_mean = mean(elapsed), + elapsed_sd = sd(elapsed), + elapsed_sample_mean = mean(elapsed_sample), + elapsed_predict_mean = mean(elapsed_predict), + rmse_y_mean = mean(rmse_y), + rmse_mu_mean = mean(rmse_mu), + row.names = NULL + ) +} + +res <- rbind( + summarise(results_cpp, "cpp (run_cpp=TRUE)"), + summarise(results_r, "R (run_cpp=FALSE)") +) + +cat("\n--- Results ---\n") +cat(sprintf( + "%-22s %10s %10s %11s %10s %10s %10s\n", + "Sampler", "Total (s)", "Sample (s)", "Predict (s)", "SD", "RMSE (y)", "RMSE (mu)" +)) +cat(strrep("-", 90), "\n") +for (i in seq_len(nrow(res))) { + cat(sprintf( + "%-22s %10.3f %10.3f %11.3f %10.3f %10.4f %10.4f\n", + res$sampler[i], res$elapsed_mean[i], res$elapsed_sample_mean[i], + res$elapsed_predict_mean[i], res$elapsed_sd[i], + res$rmse_y_mean[i], res$rmse_mu_mean[i] + )) +} + +speedup <- res$elapsed_mean[2] / res$elapsed_mean[1] +cat(sprintf("\nSpeedup (R / C++): %.2fx\n", speedup)) +cat(sprintf( + "RMSE-y delta (cpp - R): %.4f\nRMSE-mu delta (cpp - R): %.4f\n", + res$rmse_y_mean[1] - res$rmse_y_mean[2], + res$rmse_mu_mean[1] - res$rmse_mu_mean[2] +)) diff --git a/debug/data/heterosked_test.csv b/debug/data/heterosked_test.csv deleted file mode 100644 index 9c20150b..00000000 --- a/debug/data/heterosked_test.csv +++ /dev/null @@ -1,101 +0,0 @@ -"y","X1","X2","X3","X4","X5","X6","X7","X8","X9","X10","f_x","s_x" -0.16602939170279,0.211290733190253,0.3512832315173,0.484317823080346,0.0410519035067409,0.544329481897876,0.208616065792739,0.715793566778302,0.519131515640765,0.804753824602813,0.242159801069647,NA,0.5 --0.492151343206358,0.0310312763322145,0.467259066179395,0.217408320866525,0.585395927308127,0.310897014569491,0.555933631025255,0.999427518807352,0.609017762588337,0.833850894821808,0.831058437004685,NA,0.5 -1.20354672980893,0.324771000305191,0.243206361541525,0.00510925496928394,0.225724841468036,0.818582467501983,0.428320674458519,0.52159065939486,0.743639749241993,0.846550665097311,0.274055331945419,NA,1 --0.298743475101736,0.059549157274887,0.788293640129268,0.664353932486847,0.174846360227093,0.396595876198262,0.0166858048178256,0.671366059686989,0.164389409590513,0.748013024684042,0.213287916500121,NA,0.5 -1.03648471667601,0.252165184123442,0.103475574404001,0.159992411732674,0.836252088425681,0.74254785827361,0.88382129278034,0.559503264725208,0.0441013514064252,0.333072878420353,0.317699365783483,NA,1 -0.372412663856177,0.277434553951025,0.696199390338734,0.228776063071564,0.719009180320427,0.411765250377357,0.557771438034251,0.434479911345989,0.98785528820008,0.964887821115553,0.372693746583536,NA,1 -0.678007080201085,0.901923469034955,0.656224591424689,0.984398937784135,0.372089076321572,0.862292707199231,0.308036184636876,0.0231759585440159,0.113736179191619,0.934551725862548,0.281830958090723,NA,3 --0.552663625102306,0.384898619260639,0.653318807482719,0.57212492544204,0.617288585519418,0.212964130798355,0.596753151156008,0.907653672154993,0.364099610829726,0.8535895885434,0.874561849748716,NA,1 --0.774647129926176,0.542285286355764,0.891476916149259,0.86847648001276,0.3822966220323,0.468129127752036,0.658936342457309,0.679395364364609,0.313021682668477,0.295605404768139,0.946900382405147,NA,2 --1.44840009823072,0.786929592490196,0.990138707915321,0.00964757986366749,0.11595363356173,0.0707218167372048,0.0610561305657029,0.741362863453105,0.26028509857133,0.727263244334608,0.252514223800972,NA,3 --0.000520860259289017,0.373750334838405,0.77888114657253,0.958520723273978,0.851794941583648,0.405048871645704,0.789868315216154,0.650866032578051,0.848741357214749,0.00563913723453879,0.714518729131669,NA,1 -4.91065105854111,0.578438652446494,0.44657018291764,0.96312881493941,0.472523927222937,0.0506300190463662,0.498262552311644,0.523629083996639,0.353911799145862,0.00572359748184681,0.927425773814321,NA,2 -0.246599975172457,0.323262227233499,0.280061168596148,0.221339529380202,0.189832986332476,0.364617298590019,0.677225903840736,0.126530721783638,0.237819334724918,0.558171211974695,0.957330684643239,NA,1 -0.216574141744894,0.128740054322407,0.535212400136515,0.49389783013612,0.27021420141682,0.782948551233858,0.722482954151928,0.256591950077564,0.287929816404358,0.411982560995966,0.477760225534439,NA,0.5 -3.02695240909448,0.830361526459455,0.771890457021073,0.678807897027582,0.249619299778715,0.304002076154575,0.733004137175158,0.739543926436454,0.798186061438173,0.900906472001225,0.851402544882149,NA,3 --0.321787256332478,0.630839701509103,0.987936206627637,0.945720319403335,0.0623901186045259,0.258995248703286,0.429870260413736,0.10198374139145,0.536068871850148,0.0391868108417839,0.294132520444691,NA,2 -0.212045418087806,0.0110131918918341,0.85342196887359,0.588755084900185,0.461576499743387,0.770185991190374,0.871162265306339,0.0518779626581818,0.869023777078837,0.722436351701617,0.799294896423817,NA,0.5 --1.21242344678739,0.425265174359083,0.357040172442794,0.340407437179238,0.488938638009131,0.354673355352134,0.615913023473695,0.24169555818662,0.993121116189286,0.728206684347242,0.014574522851035,NA,1 -1.28352116284694,0.657284670509398,0.787876920541748,0.395542721729726,0.422870381968096,0.305611111223698,0.258184634381905,0.99757966119796,0.214673584559932,0.472818041918799,0.445311633404344,NA,2 --0.574621916007152,0.938520508352667,0.00738264597021043,0.339304566383362,0.935923039447516,0.826553477905691,0.751117973355576,0.412732847733423,0.446600724477321,0.908122762804851,0.410085891839117,NA,3 --1.79033131705676,0.980508934007958,0.991732611320913,0.0741828014142811,0.45592138543725,0.0406530441250652,0.153601433616132,0.680239439709112,0.495365400798619,0.136085328413174,0.0960345070343465,NA,3 -1.2730466507008,0.452979361405596,0.353644639486447,0.955131257884204,0.253656769637018,0.531206650892273,0.689106163335964,0.535432658391073,0.512008368736133,0.523187090177089,0.483651001704857,NA,1 -0.509242586001192,0.0175574035383761,0.875898850383237,0.553155447356403,0.245557092595845,0.353349522221833,0.957418726291507,0.397177098784596,0.216923247324303,0.675873225554824,0.0169080917257816,NA,0.5 --0.136460417178799,0.152752772672102,0.813950020587072,0.820145434001461,0.618891098769382,0.972209425177425,0.535218630451709,0.251730428542942,0.206172946607694,0.795590546447784,0.940857806708664,NA,0.5 -0.485453387886637,0.402201347984374,0.646702455822378,0.656405615853146,0.228795171016827,0.628267531748861,0.0990123685915023,0.514316890388727,0.170408786740154,0.966783886775374,0.67154831183143,NA,1 -0.358253559632074,0.103477319702506,0.780801502522081,0.291534556541592,0.102512540761381,0.0295835479628295,0.529529393883422,0.729500360088423,0.752762402174994,0.442406627582386,0.0135988409165293,NA,0.5 --5.08545632040851,0.502724350197241,0.428052270086482,0.527379664825276,0.26013516378589,0.645728018833324,0.877794990781695,0.123491464182734,0.907502186950296,0.479111275169998,0.0677777659147978,NA,2 -0.707859361399802,0.197307132417336,0.620997583027929,0.118650285294279,0.991431219037622,0.912463267799467,0.897971634287387,0.922440506285056,0.905837482772768,0.0830823809374124,0.424855173798278,NA,0.5 --0.288322601219351,0.0948237683624029,0.572173510212451,0.900335347745568,0.536399732576683,0.609837994910777,0.291524545988068,0.537141105625778,0.524107180535793,0.156775680137798,0.559588886797428,NA,0.5 -0.549673124655909,0.0728990600910038,0.447939761448652,0.282676422968507,0.89773990213871,0.0633204986806959,0.865510883042589,0.105334107298404,0.133278978988528,0.190158660756424,0.180829501943663,NA,0.5 -0.540571216994743,0.237544966395944,0.707299687666818,0.419169025262818,0.743102890206501,0.269150409614667,0.262001751922071,0.799080436583608,0.2612558118999,0.273263912182301,0.450734896818176,NA,0.5 --0.0452946833490895,0.226565732620656,0.955050441902131,0.150834861211479,0.852100720163435,0.115236598765478,0.736819905461743,0.699184895725921,0.0959788486361504,0.278612334979698,0.56139239971526,NA,0.5 -0.803100296173604,0.530489934142679,0.194167217472568,0.529533125227317,0.713073023594916,0.926254247082397,0.580682525411248,0.297447438817471,0.769652892835438,0.668208419112489,0.192935517756268,NA,2 -4.75098156092376,0.510896486230195,0.674823730485514,0.756135024828836,0.539776347577572,0.319426174042746,0.557334935991094,0.713212842354551,0.714800932211801,0.606322152074426,0.17348047113046,NA,2 -1.14191042699979,0.704268886707723,0.301489194622263,0.67195270746015,0.743559376103804,0.791731091681868,0.890636668074876,0.923103271983564,0.125805326737463,0.261845856206492,0.053273844299838,NA,2 --1.96289199262713,0.760442031314597,0.742438767105341,0.355483126826584,0.519423028221354,0.103875080589205,0.619771863799542,0.889494609320536,0.9085327538196,0.0071903329808265,0.535958426771685,NA,3 --1.03924818368133,0.355457389494404,0.383036406245083,0.142556034494191,0.735553824109957,0.801241792505607,0.312697742134333,0.730286944191903,0.337912452174351,0.0597852817736566,0.109000069089234,NA,1 -2.56288078808684,0.640810209326446,0.781058971770108,0.45493664429523,0.475713788997382,0.998631661524996,0.845030306605622,0.513773418031633,0.416949922917411,0.517727062571794,0.653628408210352,NA,2 -0.512191613003144,0.841465181438252,0.340397675754502,0.183199171675369,0.348972214618698,0.399342444725335,0.673026149626821,0.462349094683304,0.73421816341579,0.642916572280228,0.501591531094164,NA,3 --0.390539543957083,0.670964063378051,0.323868047911674,0.873477897839621,0.235217976849526,0.678807153599337,0.459241795353591,0.494968741200864,0.0186729081906378,0.205307064112276,0.93625286151655,NA,2 --0.387138910565391,0.505246647400782,0.206667850259691,0.212133591994643,0.223405942320824,0.534676593495533,0.622734377160668,0.088978860527277,0.732178951846436,0.555771772982553,0.74612275743857,NA,2 --0.589758275320359,0.917571729980409,0.606019907863811,0.361401373520494,0.349297178443521,0.815126905683428,0.432852742262185,0.217298388015479,0.659059998346493,0.282440681708977,0.74272277392447,NA,3 --0.535819616401154,0.195396314607933,0.922371664317325,0.635648620780557,0.0705672230105847,0.389690730255097,0.417000237386674,0.0651042077224702,0.747380811953917,0.306937695480883,0.186355540528893,NA,0.5 -0.100718265891839,0.12325372453779,0.838415022008121,0.486661186208948,0.464539987500757,0.0181915941648185,0.543029651045799,0.084053430473432,0.314679026603699,0.367974560242146,0.29371877363883,NA,0.5 -4.4639440117062,0.943375065457076,0.92623546323739,0.376004029763862,0.170030419481918,0.136762315873057,0.569733060197905,0.480285328347236,0.851396167185158,0.530285589862615,0.0130437007173896,NA,3 --1.7739095259091,0.750196340726689,0.386715405387804,0.119214843027294,0.157079570228234,0.972443665377796,0.0673526108730584,0.901384715922177,0.521999530028552,0.226870314916596,0.699825004441664,NA,3 --3.30597720310619,0.631283310241997,0.137164937565103,0.283018422778696,0.0198266028892249,0.599854060914367,0.308085889788345,0.868349778000265,0.512785997241735,0.855990704614669,0.914151601027697,NA,2 -0.573932596174342,0.0070427656173706,0.641306987032294,0.176330357091501,0.567426423309371,0.860410574125126,0.408202450722456,0.201708807609975,0.791534950956702,0.511292593320832,0.750208221375942,NA,0.5 --2.21001462414232,0.557995214127004,0.356393828522414,0.992580118589103,0.0819355831481516,0.553154400084168,0.32692315033637,0.239232912426814,0.190619267290458,0.00340557587333024,0.911083268001676,NA,2 --2.29486992099795,0.626155448844656,0.925242279656231,0.2268055845052,0.104973732959479,0.319284403230995,0.0643279056530446,0.407658671727404,0.308903530938551,0.547128311824054,0.589992973487824,NA,2 -1.20099834449345,0.465869394131005,0.820835534017533,0.980574035085738,0.109229437308386,0.939030001172796,0.209905731026083,0.141733445925638,0.5528343396727,0.369989494327456,0.328800423070788,NA,1 -2.04880751638206,0.705077335238457,0.499480672879145,0.795235471567139,0.749453478259966,0.910042788600549,0.965336692053825,0.16686162375845,0.790705517167225,0.60525494068861,0.386548299575225,NA,2 --0.510808721017406,0.0382022578269243,0.0574023951776326,0.298518662340939,0.847815460059792,0.419254497392103,0.545455661835149,0.804183976259083,0.566720734350383,0.938392051728442,0.968750852160156,NA,0.5 -0.678353506101917,0.00150989321991801,0.282608818961307,0.227338050492108,0.629864166723564,0.0941448376979679,0.269607466645539,0.0547637669369578,0.734315895475447,0.624765485292301,0.949363355059177,NA,0.5 --3.84612254244577,0.817981925327331,0.499992654426023,0.763754336163402,0.459059632383287,0.466768078505993,0.0136354095302522,0.119035738054663,0.619225281523541,0.695366266882047,0.287478444399312,NA,3 -1.18831966729165,0.689650556538254,0.954330617329106,0.0839970302768052,0.226895423140377,0.289217309327796,0.483793317107484,0.826940205646679,0.489855382591486,0.196166640147567,0.896134687354788,NA,2 --0.984385859730434,0.221693887142465,0.386208922602236,0.671697184676304,0.865381807554513,0.56949229654856,0.96111100865528,0.867747253272682,0.981907822191715,0.349981465376914,0.436939367791638,NA,0.5 --0.395896870249781,0.995180341880769,0.518618949921802,0.838818890042603,0.500900116749108,0.120013452135026,0.757572711678222,0.317005023127422,0.970106801018119,0.118986214278266,0.764332336606458,NA,3 -1.38921159150761,0.674878076650202,0.494008522713557,0.11772581236437,0.903489429038018,0.538562575588003,0.725335141411051,0.373649602057412,0.479591501876712,0.292627403745428,0.750551311299205,NA,2 -2.42945595659824,0.953342667547986,0.083741452312097,0.756313804769889,0.162194768199697,0.247160838916898,0.287883190670982,0.55613048421219,0.396303126122802,0.824216928100213,0.150460966397077,NA,3 --3.6886999739904,0.687809974187985,0.753712829900905,0.101322782225907,0.698061282979324,0.231127486098558,0.223462254740298,0.0152061164844781,0.74869236536324,0.895096591673791,0.904631495475769,NA,2 -0.536585118223136,0.688023552997038,0.401733021717519,0.964100113371387,0.407351006288081,0.605407851049677,0.754583293804899,0.348921799799427,0.613906478276476,0.991704752435908,0.878143492620438,NA,2 --0.945724263155543,0.912634621141478,0.0984702135901898,0.383788489736617,0.127180437557399,0.202219387050718,0.871567155467346,0.128467405447736,0.565469443099573,0.198427167953923,0.211701851338148,NA,3 -1.88190395261656,0.982580098556355,0.440120175015181,0.976914363214746,0.384266132023185,0.651897041592747,0.424716322217137,0.656596934422851,0.675765822641551,0.502742925658822,0.0215359264984727,NA,3 -1.75984850787277,0.818847576156259,0.283998677041382,0.752549906261265,0.436714354204014,0.99236780917272,0.125404683407396,0.771643821382895,0.72715258016251,0.0586418081074953,0.712552866665646,NA,3 -0.840920362678104,0.147506986744702,0.812251603929326,0.144564733840525,0.968569115269929,0.245198084274307,0.901848326670006,0.778830538270995,0.769952379865572,0.63684784504585,0.340954556828365,NA,0.5 --1.9371212817847,0.811514051631093,0.0480760163627565,0.165090465452522,0.285216748947278,0.520478271646425,0.14143471699208,0.530164128635079,0.171509642153978,0.663002614164725,0.605542375706136,NA,3 -2.66745669003816,0.618318668100983,0.538341464707628,0.279087009374052,0.891984875546768,0.127626759931445,0.712945127626881,0.0132724172435701,0.173096344573423,0.185944046359509,0.588332365499809,NA,2 -0.619195277945595,0.861147953663021,0.00769214262254536,0.243746276246384,0.0888408594764769,0.991699742851779,0.997524467529729,0.361437901854515,0.604364247061312,0.258257536683232,0.485892802011222,NA,3 --1.20727540221404,0.716689873952419,0.612233497900888,0.836100687040016,0.588582671945915,0.491994495503604,0.792359412880614,0.0365804079920053,0.800699540879577,0.275558543158695,0.215889111161232,NA,2 -4.72985785755328,0.932213481282815,0.372599190333858,0.526720034424216,0.653838029364124,0.497327620862052,0.48135881498456,0.0302206429187208,0.234266108367592,0.33457084489055,0.598570736125112,NA,3 -0.288958968840092,0.0119730825535953,0.300384242553264,0.850864094216377,0.893871575593948,0.680723139783368,0.990262570325285,0.865412194048986,0.379822072573006,0.718840815359727,0.890456868568435,NA,0.5 -0.259681068955783,0.237852266291156,0.0348299730103463,0.718882874818519,0.653284071711823,0.112274453276768,0.0580400477629155,0.833725305041298,0.536504363641143,0.720546880736947,0.159137749811634,NA,0.5 -0.347321451075773,0.188694522483274,0.531638014828786,0.572584339883178,0.238210330018774,0.231171588879079,0.413350867340341,0.0659690995234996,0.564861474558711,0.89177248836495,0.0945249791257083,NA,0.5 -3.89576276409494,0.833090277388692,0.203691257163882,0.146478878101334,0.819324584677815,0.196580097544938,0.892589659895748,0.226715816417709,0.877559039741755,0.85475739534013,0.0848844614811242,NA,3 -0.355581450144155,0.431055790279061,0.416096288943663,0.0859726935159415,0.325928226578981,0.58786572702229,0.350675149122253,0.401109970640391,0.756839712616056,0.0408970382995903,0.338316329289228,NA,1 -1.60495074110366,0.621179772540927,0.686551147606224,0.273364662658423,0.874324111267924,0.725664374418557,0.859822448343039,0.737721871584654,0.271499434486032,0.237752126529813,0.236145419767126,NA,2 --2.6328605154421,0.692925159586594,0.32423769752495,0.0426305204164237,0.968091825256124,0.1722840978764,0.862542052520439,0.704603992169723,0.0835464387200773,0.100951640866697,0.527263561496511,NA,2 -0.0799409275378214,0.496741143520921,0.358587516937405,0.925880230497569,0.544319620588794,0.682037363760173,0.974792276741937,0.415831092977896,0.969078256050125,0.604709276696667,0.89061022317037,NA,1 -2.72669228705894,0.478368178475648,0.310989870224148,0.928430530708283,0.247262259013951,0.795121156377718,0.861811829963699,0.400008339434862,0.144553888821974,0.947825188748538,0.19086144422181,NA,1 -0.0790671158723385,0.127485178643838,0.703777753748,0.503316113026813,0.764946941519156,0.433781018713489,0.384083716897294,0.968963841674849,0.380835567833856,0.876958164153621,0.818229214288294,NA,0.5 --0.459476378753931,0.18172625801526,0.720858050510287,0.848017353564501,0.720434938790277,0.728625974617898,0.44146662694402,0.951146713923663,0.455086235888302,0.740384874632582,0.091537274653092,NA,0.5 --0.00279715401469337,0.887286564800888,0.161548246862367,0.0689359360840172,0.555442054290324,0.846620650961995,0.00943395867943764,0.108124481979758,0.89916443801485,0.677391635021195,0.0622330722399056,NA,3 -0.249461137396448,0.0212012422271073,0.359011579537764,0.719863031757995,0.249324724078178,0.707923346664757,0.122814457165077,0.939011327456683,0.158998906379566,0.995269876671955,0.441887193359435,NA,0.5 -0.0133950863816936,0.0725361246149987,0.0616531122941524,0.287096783984452,0.382249645655975,0.883280041394755,0.615435254527256,0.487508539808914,0.597117075929418,0.629250884056091,0.159072700887918,NA,0.5 -0.418448230677183,0.0969394627027214,0.411688929889351,0.449063409119844,0.488310252549127,0.259443101240322,0.987585310125723,0.441142725525424,0.242221797816455,0.250788792734966,0.619807209586725,NA,0.5 -1.16506424422402,0.786211173050106,0.369333798764274,0.840181995881721,0.386459692381322,0.285024101380259,0.641522094374523,0.849066619062796,0.0109615188557655,0.2850232264027,0.217773179057986,NA,3 -0.989575894880979,0.539474389981478,0.796261472860351,0.47632975759916,0.676875605247915,0.890021840808913,0.499118740204722,0.972148579079658,0.982196035329252,0.0354776941239834,0.547611665446311,NA,2 --5.41773583711259,0.991268366575241,0.841971538960934,0.831874355440959,0.383187236264348,0.318007141817361,0.644914947915822,0.261283464496955,0.146094920346513,0.701359197963029,0.458953364752233,NA,3 -1.14091061117272,0.381695970892906,0.186867855954915,0.592309525003657,0.0125671157147735,0.62503357976675,0.712586513254791,0.063052834244445,0.450469782575965,0.306621422059834,0.900312934769318,NA,1 --0.0430420048530519,0.126209627138451,0.0910242712125182,0.877147340215743,0.65187450312078,0.148876765510067,0.271347487112507,0.483428512467071,0.995978597551584,0.910394222242758,0.631069793598726,NA,0.5 -2.53541834119888,0.527330596232787,0.949522181181237,0.366707332897931,0.0310849468223751,0.659584330860525,0.315074123907834,0.177570374915376,0.579150864621624,0.698416418628767,0.336142445914447,NA,2 --0.836400788818656,0.152910969685763,0.588620311347768,0.183476264355704,0.471249244175851,0.456099348841235,0.358562543056905,0.0993230841122568,0.865682073868811,0.688885818701237,0.783728409092873,NA,0.5 --0.36119486883805,0.255377688677981,0.472820749040693,0.497897547669709,0.599874146282673,0.595044590299949,0.915363078936934,0.366866510361433,0.132163925329223,0.286865012487397,0.395963127491996,NA,1 -3.04844600215742,0.784051381517202,0.334584510419518,0.889806121354923,0.633161137346178,0.861800517654046,0.5874773489777,0.395519501296803,0.458229560870677,0.794260403839871,0.584855122258887,NA,3 --1.56680469940733,0.424370453692973,0.798977045807987,0.652883867966011,0.178207511780784,0.704748782562092,0.0875129597261548,0.638247390044853,0.396533044287935,0.528938032919541,0.909660720732063,NA,1 --0.125657442776976,0.367308006854728,0.554850550601259,0.471428145188838,0.812708863057196,0.960221603047103,0.897328518796712,0.436890807002783,0.851095250342041,0.439047307474539,0.350051162531599,NA,1 --0.081334269056118,0.763460468733683,0.334580871742219,0.455008901190013,0.758000182220712,0.850450441241264,0.255375979002565,0.451483712298796,0.409017519326881,0.0273381865117699,0.59320038347505,NA,3 -0.492823092528766,0.0655732385348529,0.881846528034657,0.0530098504386842,0.35191644448787,0.175170796457678,0.296730656642467,0.907911094138399,0.603398642269894,0.255270632915199,0.642264471622184,NA,0.5 --0.000915148320753323,0.613401755923405,0.360998818883672,0.247979114530608,0.784389181062579,0.354104121448472,0.209656548919156,0.0889754833187908,0.849633313016966,0.397846080595627,0.135037919739261,NA,2 diff --git a/debug/data/heterosked_train.csv b/debug/data/heterosked_train.csv deleted file mode 100644 index 92d47304..00000000 --- a/debug/data/heterosked_train.csv +++ /dev/null @@ -1,401 +0,0 @@ -"y","X1","X2","X3","X4","X5","X6","X7","X8","X9","X10","f_x","s_x" --0.282241503700301,0.0321366493590176,0.0392713348846883,0.491026453208178,0.829869156936184,0.419834628002718,0.674110037973151,0.440613425569609,0.00938279670663178,0.688462630612776,0.246391290798783,0,0.5 -0.560156335491632,0.633398831356317,0.713975282153115,0.0746546245645732,0.183240191545337,0.366673885611817,0.195146100129932,0.543941563926637,0.770088015124202,0.815632452024147,0.204938031733036,NA,2 --2.55310551844943,0.699684502789751,0.518513803370297,0.278569060144946,0.518335236469284,0.557218188419938,0.487766429549083,0.912966759642586,0.306013585766777,0.737159843789414,0.415715463226661,NA,2 -1.2493544072818,0.376545963343233,0.365289261564612,0.878427121788263,0.236202105879784,0.150846817996353,0.263488188385963,0.295650234911591,0.752220941707492,0.503766167676076,0.0375082376413047,NA,1 --0.350237222512917,0.685897701652721,0.439451805083081,0.346272532362491,0.639235025504604,0.0299059110693634,0.791390363825485,0.818463942268863,0.805753854336217,0.0685151207726449,0.653748897369951,NA,2 -0.875396882302278,0.30443329596892,0.593776761321351,0.252662484999746,0.143341688672081,0.202486322494224,0.209879243280739,0.0234908810816705,0.443319902056828,0.865865753497928,0.605711343931034,NA,1 -0.486947118214128,0.524243678199127,0.0881543764844537,0.894031626405194,0.117429533274844,0.153696053894237,0.0199424722231925,0.0166082351934165,0.0904451431706548,0.634674217319116,0.25031628459692,NA,2 -1.52973577099548,0.478108013281599,0.41259130765684,0.922683839220554,0.537282897857949,0.28786034998484,0.518826235085726,0.768439579289407,0.712297242833301,0.471964940661564,0.0293229683302343,NA,1 -0.314146031707177,0.352395879570395,0.994232393568382,0.36731672892347,0.371404710225761,0.226430130191147,0.988522542640567,0.171963014639914,0.590013455832377,0.61276261578314,0.0982697720173746,NA,1 -1.43575656817199,0.343638661550358,0.0898715713992715,0.671082190936431,0.0814627697691321,0.115701446775347,0.404091404750943,0.914198226062581,0.77044213260524,0.724840591661632,0.960911155678332,NA,1 -0.326893619264727,0.748311516828835,0.972395642194897,0.462105612503365,0.849746079416946,0.525140055455267,0.691481846617535,0.64259241358377,0.383345428388566,0.193897266406566,0.471145344432443,NA,2 -5.33184058038735,0.805053810123354,0.694419888081029,0.688123103696853,0.857482564169914,0.412985544884577,0.756523177726194,0.587173974374309,0.698520486243069,0.306543357903138,0.569155315635726,NA,3 --0.080953281849018,0.292309782933444,0.886532769538462,0.590175394201651,0.763742916053161,0.843529825564474,0.444129067705944,0.506667435867712,0.263209822122008,0.25964065361768,0.51411981228739,NA,1 --0.0642885588731738,0.0813713425304741,0.733218841487542,0.645473898854107,0.48057269025594,0.958589600399137,0.70577060454525,0.746377976378426,0.228743602987379,0.102548423456028,0.184051728574559,NA,0.5 --2.52172005222636,0.465416571125388,0.113867231411859,0.423410250106826,0.912441357970238,0.253497969126329,0.957222544355318,0.298387421993539,0.55475458712317,0.320897468365729,0.92878125119023,NA,1 -0.26999574319132,0.0482504856772721,0.870362227084115,0.0656226673163474,0.352718068053946,0.950954237952828,0.0427663850132376,0.726099342806265,0.917875221231952,0.958915420575067,0.913584525929764,NA,0.5 --3.83785409353774,0.645924511365592,0.445315549150109,0.396382685052231,0.90340900560841,0.380887056933716,0.26203626464121,0.901702417759225,0.417292382800952,0.660986073780805,0.41848081164062,NA,2 -1.07793425501999,0.529066368937492,0.877272051293403,0.683232536539435,0.744753323029727,0.3394386311993,0.526980829657987,0.636599146993831,0.445450452622026,0.164547258988023,0.882723631570116,NA,2 -0.62946218194112,0.621317676734179,0.144061228726059,0.754505411023274,0.52779389359057,0.0482079237699509,0.308562529738992,0.983435485279188,0.138843381311744,0.889652899932116,0.731016719480976,NA,2 -0.149334178305908,0.687091745669022,0.0368512177374214,0.239253385690972,0.0125792836770415,0.357229135232046,0.194120459957048,0.189207792049274,0.275957886129618,0.757066001417115,0.0808777082711458,NA,2 -1.08233265888858,0.377081562764943,0.927687839372084,0.388060942757875,0.262930948287249,0.415721071185544,0.595551578560844,0.889767974615097,0.680972298840061,0.78021034412086,0.810993585269898,NA,1 --3.21122703337705,0.507098242873326,0.364786737831309,0.706525600515306,0.0720878560096025,0.39529197756201,0.0853938185609877,0.25313756079413,0.884323446080089,0.164221335435286,0.151563796214759,NA,2 -0.531353908252581,0.237350288545713,0.275147415464744,0.501605095108971,0.0984707665629685,0.615358766168356,0.561869938857853,0.290915884776041,0.80242047063075,0.554143948713318,0.496158995432779,NA,0.5 -0.600730293046287,0.157840367639437,0.632627204060555,0.929030198603868,0.630327652674168,0.730401052860543,0.328352578915656,0.717250395333394,0.666188669623807,0.0548528400249779,0.833407171769068,NA,0.5 -1.01314701175506,0.898484638426453,0.203661296050996,0.846307630417868,0.219352454645559,0.742492161225528,0.517635546857491,0.0817107316106558,0.968180059455335,0.336696855258197,0.662801184924319,NA,3 -0.656553622768969,0.753240135731176,0.343114467570558,0.128794493386522,0.560916883870959,0.668597492389381,0.196455833502114,0.172105104196817,0.634505996946245,0.0560522808227688,0.678417820017785,NA,3 -0.450441176587303,0.0125156014692038,0.882218054495752,0.629003981128335,0.674523104447871,0.600846989545971,0.633095325203612,0.400010914774612,0.962411692133173,0.547679667593911,0.177117998711765,NA,0.5 --3.7895427703596,0.756340250140056,0.397959120338783,0.865862431004643,0.327016743365675,0.974390563322231,0.230026092845947,0.127599307801574,0.0277679341379553,0.932881123386323,0.593616648577154,NA,3 -0.946190693629125,0.235685531515628,0.246450350852683,0.248962875455618,0.362083269981667,0.932028053328395,0.458421003306285,0.347327328752726,0.816151189152151,0.349907604278997,0.481070955749601,NA,0.5 -9.4191799850465,0.849114913959056,0.821041927672923,0.602861605351791,0.340112386737019,0.396447837818414,0.765584346372634,0.106930397683755,0.223846387118101,0.744411860825494,0.93100854405202,NA,3 -1.43916127761723,0.766780707519501,0.445421120384708,0.451271537225693,0.72373710712418,0.184640744002536,0.951642854139209,0.495319120585918,0.00020794733427465,0.887012050719932,0.271078041987494,NA,3 -0.243285053031709,0.172001409810036,0.398085105465725,0.231227887328714,0.164392935112119,0.112842592177913,0.171439406462014,0.502675085794181,0.686919268220663,0.491456889547408,0.454194214195013,NA,0.5 -0.630674862417757,0.932543548289686,0.256847345735878,0.979184590745717,0.631192493718117,0.981027684174478,0.682858986081555,0.169451557798311,0.60733424173668,0.9371962856967,0.350068996660411,NA,3 -0.601104263344641,0.218169360188767,0.321275946451351,0.484401718014851,0.851453062379733,0.309203763958067,0.917364788241684,0.2084669705946,0.245045544346794,0.582212117500603,0.228992780670524,NA,0.5 -0.840336041358887,0.589697151444852,0.685158264124766,0.327503680251539,0.213558244984597,0.763619041070342,0.11949589359574,0.71811868599616,0.408380198059604,0.832140794256702,0.770084762945771,NA,2 --0.154074442892582,0.367701322771609,0.884693963453174,0.0697798361070454,0.434232867090032,0.898273871280253,0.254723551450297,0.111902467207983,0.715196632780135,0.184001997346058,0.933242507977411,NA,1 -0.720696373819241,0.198938927613199,0.606495931511745,0.336577468551695,0.326773484470323,0.370286533143371,0.00082359416410327,0.0124902881216258,0.239021955290809,0.908346865791827,0.0976885615382344,NA,0.5 -2.21060504692689,0.682107795961201,0.287048922851682,0.372942043235525,0.396774863358587,0.819875442888588,0.710184759926051,0.190397032536566,0.480622512986884,0.8761212809477,0.223856529453769,NA,2 --0.00610575291022993,0.067679442698136,0.8684254467953,0.996917813783512,0.089049612171948,0.700067178346217,0.794930153992027,0.930825348943472,0.826459914213046,0.864694126881659,0.0812219593208283,NA,0.5 --0.932580761661321,0.572068417910486,0.0790525975171477,0.1397358069662,0.64970194036141,0.396581323817372,0.13856448000297,0.767174258362502,0.951541883638129,0.0230525995139033,0.823745897971094,NA,2 --1.23014290795405,0.450845072045922,0.671238768612966,0.580475129419938,0.165164324454963,0.766623474657536,0.279710618546233,0.353242868324742,0.870829950086772,0.200475953286514,0.171313758939505,NA,1 --3.16484165617904,0.573105162009597,0.258897884283215,0.213309426326305,0.100215684855357,0.882589143933728,0.585318948607892,0.152414533542469,0.104481044691056,0.666403246112168,0.75058483122848,NA,2 -4.62370372914815,0.73533759615384,0.740031143417582,0.788233703933656,0.918064421508461,0.720053877448663,0.39371900446713,0.0863925719168037,0.545766794355586,0.310717426007614,0.739832061110064,NA,2 -1.06384980222904,0.318392616463825,0.172780332621187,0.879689910449088,0.738148988690227,0.928580733481795,0.954701505368575,0.686478513525799,0.12954113073647,0.81541432836093,0.500270092394203,NA,1 -0.797694547807648,0.528178822714835,0.222754651214927,0.543697185115889,0.210871644085273,0.539274046663195,0.117441890062764,0.171201136894524,0.563767193350941,0.336862071650103,0.51251038024202,NA,2 --1.43612727789891,0.995858771493658,0.930734105641022,0.990805656183511,0.171419916907325,0.054182613035664,0.928726562298834,0.748863881221041,0.940693028038368,0.65918653504923,0.31154797365889,NA,3 --0.662510088784892,0.113952487008646,0.250964630628005,0.410930230980739,0.10554705071263,0.907607411500067,0.386065973900259,0.522172022610903,0.355578461894765,0.730407647090033,0.462258468614891,NA,0.5 -1.38840466048108,0.341306684771553,0.812958737369627,0.613608168903738,0.835363352671266,0.894293455639854,0.379676664015278,0.66161756683141,0.392706391168758,0.0179466558620334,0.213412882061675,NA,1 --0.00107339168469423,0.467106574913487,0.461363392416388,0.494308657478541,0.868279299233109,0.097200368065387,0.598206092603505,0.0933514924254268,0.761134962551296,0.0639440342783928,0.236494644545019,NA,1 --0.781708262436574,0.128517229110003,0.170633271569386,0.299606638494879,0.496000039856881,0.510280701331794,0.300075584789738,0.853872763924301,0.197547832271084,0.155960420612246,0.643106105038896,NA,0.5 -0.341957823006054,0.412145783891901,0.359467171831056,0.964134642854333,0.698599802330136,0.409747923957184,0.337395059876144,0.43656346690841,0.678571406751871,0.0371420315932482,0.82975935889408,NA,1 -0.056551295167318,0.222288713557646,0.49123525666073,0.395386436022818,0.539690778823569,0.574697827454656,0.397546735825017,0.730274567846209,0.652596935629845,0.617730791214854,0.38262702873908,NA,0.5 -0.106784845057808,0.148283806396648,0.0914904365781695,0.869412067811936,0.399036906659603,0.190275267930701,0.531700859311968,0.970284502254799,0.63110744045116,0.105896403081715,0.895602686330676,NA,0.5 --0.416591273016816,0.0369925673585385,0.369142951210961,0.802026466466486,0.83644736954011,0.384842669591308,0.399965344229713,0.497230796376243,0.21596611966379,0.430716055911034,0.814202995272353,NA,0.5 --3.80115746542674,0.896795936627313,0.765040824888274,0.844353931955993,0.740446796640754,0.11038355785422,0.944162652362138,0.749065604060888,0.802961661946028,0.594114093808457,0.982326869154349,NA,3 --0.990579577345106,0.364397091092542,0.75387371936813,0.92530464194715,0.556107219774276,0.472777665359899,0.290147027932107,0.558879483724013,0.734846210572869,0.912317100213841,0.382154474034905,NA,1 --2.24113101562438,0.958884046878666,0.916358025046065,0.622180263744667,0.313185492297634,0.629507611738518,0.766538005322218,0.00391310267150402,0.117573201656342,0.677621009293944,0.86312605952844,NA,3 --1.37915963138867,0.992283673025668,0.1130617109593,0.39246118767187,0.386350285029039,0.618687929585576,0.00955063896253705,0.3101056877058,0.00335579412057996,0.428965806029737,0.334215411450714,NA,3 --0.116138535683225,0.712791633559391,0.826705769635737,0.998709523584694,0.935079188318923,0.999019755981863,0.179841009201482,0.775375233031809,0.565244837198406,0.73825324466452,0.217810853617266,NA,2 -2.8651370208333,0.918227929389104,0.991976723074913,0.556448292685673,0.86033549089916,0.86803180235438,0.593546075979248,0.764600344700739,0.115971907973289,0.874116574181244,0.415115732233971,NA,3 --2.62673850900046,0.958541032392532,0.734085222240537,0.277872398728505,0.903584003215656,0.99169037095271,0.442168229259551,0.931520459009334,0.135586928809062,0.59293475211598,0.717808285029605,NA,3 -0.132124476008359,0.608879194129258,0.450298280920833,0.517520986497402,0.466313583776355,0.215587850660086,0.0298155974596739,0.189444416435435,0.174548792652786,0.0233730303589255,0.0774613090325147,NA,2 -0.293719995172287,0.0288681895472109,0.813382123131305,0.00113681610673666,0.941644994076341,0.476048385258764,0.530871283495799,0.56963656633161,0.737177671631798,0.0965919359587133,0.604818821884692,NA,0.5 --0.103954343912756,0.166011718567461,0.653191721998155,0.923596187727526,0.993117317091674,0.121775610838085,0.0122047434560955,0.826190118445083,0.30410159076564,0.38432927033864,0.433089199708775,NA,0.5 -0.716889016296623,0.424093425739557,0.161204627482221,0.926760187838227,0.688538460293785,0.412991603836417,0.0354802808724344,0.914258436998352,0.532327418681234,0.665652722353116,0.169264691183344,NA,1 --3.89856674282553,0.847578614018857,0.657475670333952,0.999632467282936,0.367575454059988,0.757724079070613,0.600674423156306,0.687886233208701,0.487299331231043,0.809120106277987,0.675410686992109,NA,3 -0.167302299143473,0.0175340240821242,0.519420479191467,0.318057127762586,0.0525879818014801,0.545170752331614,0.493110234150663,0.703221583971754,0.326910747680813,0.955182061996311,0.553046975517645,NA,0.5 -3.41642433081155,0.881483160192147,0.548339178785682,0.000273799290880561,0.615487394854426,0.331853078911081,0.753756532678381,0.619998622685671,0.887993055395782,0.321768817491829,0.723241026746109,NA,3 --0.927780454804531,0.694804448867217,0.636846353765577,0.545869589783251,0.543147100368515,0.651934214634821,0.496286484412849,0.906899358844385,0.585167640121654,0.544923613313586,0.9785473162774,NA,2 --0.119458659437126,0.349967923015356,0.0674725724384189,0.461099995067343,0.208889465779066,0.727099789306521,0.322485944256186,0.0180952423252165,0.586417717626318,0.50533275748603,0.208145300159231,NA,1 --0.388475062884792,0.442952080629766,0.550757031654939,0.0116928820498288,0.919974668882787,0.28176667355001,0.343595980666578,0.00730579695664346,0.0586100376676768,0.816800056258217,0.281768965069205,NA,1 --4.99993600969702,0.63083560904488,0.304470317903906,0.130367321893573,0.00757993618026376,0.672845395281911,0.775554867926985,0.561601209221408,0.809073857031763,0.350174822146073,0.369892556220293,NA,2 --4.0932011291792,0.857258127070963,0.849426775239408,0.048665531212464,0.663260747911409,0.201854383805767,0.374728660099208,0.680570747936144,0.791422239039093,0.155649269232526,0.159891559742391,NA,3 --1.63574784505844,0.475719389040023,0.463194821262732,0.562303752405569,0.269420249154791,0.64509441726841,0.99847622634843,0.944659149041399,0.570544230751693,0.00567547790706158,0.528426631586626,NA,1 --0.11708102669514,0.229862803127617,0.693337155273184,0.0249634576030076,0.472208397230133,0.628053227905184,0.81685835798271,0.652433575829491,0.879498174646869,0.0514549117069691,0.73972195899114,NA,0.5 -0.311440636819556,0.424796005710959,0.588375912979245,0.907698655501008,0.31491112918593,0.173823640448973,0.74140138248913,0.306885188445449,0.665830912999809,0.498456827364862,0.325599604984745,NA,1 --0.523524190351201,0.546965456567705,0.69669742975384,0.907464131480083,0.905781339621171,0.374616772402078,0.97064465098083,0.54574229917489,0.0258567882701755,0.499222912825644,0.301927668973804,NA,2 --0.218557554085927,0.545247800182551,0.532440177630633,0.0326926102861762,0.571362876798958,0.246572606964037,0.62373100570403,0.800766404950991,0.306481578154489,0.010504704201594,0.843400214333087,NA,2 -4.3394388956167,0.816038053482771,0.0158182233572006,0.43390632327646,0.118837420130149,0.119051147252321,0.583459883695468,0.119374568108469,0.339256392093375,0.430539290653542,0.673767470754683,NA,3 --3.26204680706395,0.901603788835928,0.348586000269279,0.311316560022533,0.66575604211539,0.234927158802748,0.4827365425881,0.146938317688182,0.157435120781884,0.520954007748514,0.697913567069918,NA,3 --1.34236542741488,0.620338303502649,0.691506293835118,0.481012200936675,0.556597996968776,0.0661866711452603,0.809830514946952,0.655171810416505,0.967796785058454,0.859031651634723,0.547678065020591,NA,2 -2.62955558280747,0.840353684034199,0.673164369771257,0.499901092844084,0.148192471591756,0.0805800317320973,0.836274195462465,0.529777634423226,0.809363054344431,0.569066174095497,0.855126766953617,NA,3 --0.438281983773096,0.286419975571334,0.0967769827693701,0.10907745687291,0.96515504270792,0.346270593581721,0.213748984504491,0.87659911182709,0.000996117945760489,0.11449680221267,0.900713045150042,NA,1 --0.819668273218654,0.42447182838805,0.305468514794484,0.628598166629672,0.897104371571913,0.292935826117173,0.57036331994459,0.812557882862166,0.743328155716881,0.372404619818553,0.995886287419125,NA,1 --2.15502275916223,0.581904123770073,0.136866772547364,0.686216610949486,0.0187842161394656,0.65494451741688,0.842298676259816,0.63314122450538,0.811359524028376,0.159128638915718,0.789697944186628,NA,2 --4.06652032407565,0.945780334994197,0.103755530901253,0.866593056125566,0.559091728879139,0.662870865082368,0.727699366630986,0.921519288793206,0.788001266308129,0.731000347295776,0.126275058137253,NA,3 --0.928011560521798,0.0429409199859947,0.000712319742888212,0.487844629446045,0.285959116648883,0.604428091552109,0.919549201382324,0.459383252309635,0.727056132396683,0.891209361376241,0.935548935784027,NA,0.5 --0.454744311371444,0.328576610656455,0.304207463981584,0.804804563056678,0.846173255937174,0.218100384576246,0.161341489758343,0.0678012303542346,0.141843453980982,0.453309568809345,0.174222140572965,NA,1 -2.45897650312528,0.572264407062903,0.749721936415881,0.234257691074163,0.789959133137017,0.606415096437559,0.554639648413286,0.0167610717471689,0.64993330463767,0.166575252544135,0.753427488962188,NA,2 -6.3265939587816,0.819421303458512,0.737637117039412,0.894383451668546,0.761343186022714,0.651555245975032,0.303879029816017,0.53969519212842,0.311747952597216,0.319801090052351,0.836284123826772,NA,3 --1.00774699035296,0.27150860009715,0.301227852702141,0.904794576577842,0.526734083425254,0.869649238185957,0.472009405493736,0.629169559106231,0.432649733265862,0.571082683978602,0.996163306059316,NA,1 -0.746872815646722,0.140957412077114,0.958929154090583,0.87989121559076,0.650756301125512,0.234146694419906,0.255590234417468,0.420677730115131,0.597164487466216,0.723511897493154,0.155153512023389,NA,0.5 --0.650446841065496,0.358473981032148,0.862018195679411,0.700934009160846,0.0534398728050292,0.899426093092188,0.338007411221042,0.93593748472631,0.137800792232156,0.827133719110861,0.955730211222544,NA,1 -2.48082038300247,0.961773466318846,0.0908011053688824,0.906130446586758,0.650059416424483,0.286194331245497,0.889659560285509,0.1561406950932,0.112980651902035,0.944480426842347,0.891218760749325,NA,3 --0.0986038818318802,0.0816032441798598,0.00994992698542774,0.0916588774416596,0.106353395618498,0.383079081075266,0.394785706652328,0.774634938454255,0.597865202464163,0.382828871253878,0.0871360464952886,NA,0.5 -1.1520722946698,0.372490294510499,0.642933483468369,0.932021666085348,0.25168374995701,0.656297161011025,0.551260613603517,0.00573913590051234,0.171782038640231,0.776553657837212,0.087719907052815,NA,1 -4.3995392071382,0.889362843940035,0.875090325716883,0.929815792944282,0.647907796083018,0.219110984588042,0.0118211782537401,0.674186609685421,0.940533248940483,0.127728166524321,0.00270650535821915,NA,3 -0.0290089026468435,0.248338227858767,0.142401040298864,0.763918829383329,0.76891547255218,0.181683097267523,0.87568672769703,0.820826095528901,0.332615875406191,0.810349359177053,0.414319772040471,NA,0.5 --0.0716587997708629,0.099602812435478,0.572981350123882,0.823288503335789,0.706261500250548,0.0959598675835878,0.149675080087036,0.507298650220037,0.423980995081365,0.74773308984004,0.686039772350341,NA,0.5 --1.30011461908445,0.42879819823429,0.0999327092431486,0.0399061276111752,0.891762524843216,0.942133626202121,0.152556500863284,0.170978014590219,0.368863882729784,0.0347012423444539,0.505937222158536,NA,1 --0.60014108521951,0.138419456779957,0.436613920144737,0.604420283576474,0.204740515211597,0.765600503887981,0.98930400935933,0.276836613891646,0.0905646500177681,0.362719136057422,0.529882991919294,NA,0.5 --1.21145790381126,0.744622958125547,0.749749571783468,0.360840648878366,0.762236907379702,0.297599644400179,0.55331027880311,0.624136926373467,0.798092990880832,0.851898095104843,0.404137687990442,NA,2 --1.14999787644678,0.708947046194226,0.849501037504524,0.762910400982946,0.842422707471997,0.723715083673596,0.381835736334324,0.65285807219334,0.514471322065219,0.310605535982177,0.957349325995892,NA,2 -0.819526373875359,0.515504683367908,0.370702817803249,0.823700248729438,0.387957931729034,0.642785836942494,0.348471603821963,0.0168538340367377,0.692721860483289,0.232188782887533,0.389768423512578,NA,2 -0.04335007115952,0.00294706947170198,0.814211948076263,0.31387432734482,0.261198920663446,0.432101055746898,0.918104903073981,0.415187336038798,0.976240418851376,0.291730556869879,0.519225212512538,NA,0.5 --0.941306547537896,0.320613505784422,0.445245401933789,0.659601202700287,0.373574797762558,0.309444201411679,0.99856966570951,0.845955910626799,0.385660471161827,0.863749768352136,0.85754614090547,NA,1 --1.18186963828744,0.701132035115734,0.859337297501042,0.348278685705736,0.160000698640943,0.238924831384793,0.72800212725997,0.57119212881662,0.620406114961952,0.0212759706191719,0.0539104880299419,NA,2 -0.357126829159397,0.447363865328953,0.8170770870056,0.286243395647034,0.16769379354082,0.238466019975021,0.662134238751605,0.810090478044003,0.577125230105594,0.075156477978453,0.137602924602106,NA,1 --0.190270625930595,0.114963983185589,0.593716769479215,0.903442053124309,0.815722588449717,0.452215825207531,0.0901011251844466,0.436070464085788,0.167535469401628,0.971105119213462,0.308095294749364,NA,0.5 -2.62078960911659,0.566922665573657,0.568221352994442,0.653094903565943,0.0865914954338223,0.786628634436056,0.917334056925029,0.248513758182526,0.225959765724838,0.70123139442876,0.785259994678199,NA,2 -2.55987015026694,0.833427351433784,0.499585702549666,0.598924180492759,0.138457585126162,0.320701913908124,0.774522821651772,0.431595224887133,0.257062387652695,0.22843798622489,0.532552591059357,NA,3 --6.92858680088879,0.782670683693141,0.0883395997807384,0.773956855759025,0.646301113301888,0.948011690285057,0.624721661442891,0.384072499116883,0.52077893470414,0.161269078962505,0.443590870127082,NA,3 -0.15815266439335,0.246600785059854,0.783747910521924,0.235627511749044,0.865159292006865,0.799394915346056,0.177220857003704,0.427418543258682,0.373554009012878,0.512462994782254,0.892843663692474,NA,0.5 -2.65450595933212,0.957361436681822,0.911815687082708,0.41771270474419,0.507127913879231,0.388280466897413,0.209281904157251,0.512741857906803,0.995567321544513,0.608058721991256,0.980684688081965,NA,3 --0.551655302147764,0.105588861741126,0.765772460727021,0.254358040867373,0.496851657284424,0.270303341327235,0.840589354978874,0.230695051141083,0.313359949970618,0.0843451307155192,0.604350767796859,NA,0.5 --1.8820051243768,0.865534041076899,0.282570220297202,0.518943741451949,0.582809855230153,0.733290414093062,0.464054192649201,0.0420361561700702,0.312909991247579,0.357775181764737,0.217115553095937,NA,3 -1.18644127377823,0.572014078265056,0.47516213497147,0.909601657185704,0.7345625765156,0.159067094791681,0.0758456904441118,0.596963451942429,0.248694570967928,0.609315068461001,0.460826584137976,NA,2 -2.90277751401675,0.714331415947527,0.171784817706794,0.755383554147556,0.103581401053816,0.949289575219154,0.311075175181031,0.892260275082663,0.782703787088394,0.609557135263458,0.0396083393134177,NA,2 -0.452672729165082,0.362838387256488,0.290337451035157,0.464105833088979,0.632470370503142,0.288197767920792,0.256659790175036,0.76495176483877,0.331189053365961,0.102935964474455,0.724172668065876,NA,1 --0.133181544288443,0.697857573861256,0.296673942590132,0.952276430791244,0.360607230570167,0.295894445385784,0.471734570106491,0.0396309678908437,0.965557127026841,0.30653714039363,0.577830789843574,NA,2 -0.520898443292835,0.555650745285675,0.0409085555002093,0.172368012135848,0.690776515984908,0.844665887998417,0.950760766165331,0.402099822415039,0.815108865965158,0.204499331535771,0.857367512304336,NA,2 --4.94863899675115,0.611774453427643,0.924236066872254,0.0708308285102248,0.330309199402109,0.127171179978177,0.491658851737157,0.536354664713144,0.913808302022517,0.952496240613982,0.481220846064389,NA,2 -0.736866988322533,0.438154348172247,0.806642528390512,0.0169571470469236,0.662110002711415,0.0638061012141407,0.217253629118204,0.100625155959278,0.0314080747775733,0.836882830131799,0.524368022568524,NA,1 --2.10491296672012,0.603837736183777,0.607003055978566,0.14466291340068,0.00309824291616678,0.579994248226285,0.499280680436641,0.896025418769568,0.630717721302062,0.990853895898908,0.906316698761657,NA,2 --0.385675311337004,0.0125510159414262,0.516535924980417,0.448626236524433,0.0885730378795415,0.931544049177319,0.0455953879281878,0.932046544039622,0.456221962114796,0.18970851926133,0.15640650712885,NA,0.5 --0.933954992523,0.0433406794909388,0.950225162087008,0.479722076794133,0.840589659288526,0.512415238423273,0.36178726516664,0.267573584103957,0.603488262509927,0.0370614926796407,0.624879651935771,NA,0.5 --0.520842484729798,0.0929510626010597,0.185449638403952,0.0954606782179326,0.354706885525957,0.336916906060651,0.202741828747094,0.128091655438766,0.105544385500252,0.0621682242490351,0.337343172635883,NA,0.5 --1.1902755322199,0.341341719729826,0.27227359986864,0.726273648440838,0.375441956566647,0.0984445584472269,0.24713291740045,0.68887278297916,0.395826757419854,0.263939680531621,0.85243204399012,NA,1 --0.0160453226366918,0.598223075270653,0.515979057410732,0.294184898026288,0.419611923396587,0.451983051374555,0.610102006234229,0.502045850269496,0.0474624985363334,0.312018062686548,0.97183643002063,NA,2 -0.662189390327135,0.269035245059058,0.890601610997692,0.579056739108637,0.532028702553362,0.036886066198349,0.655321523780003,0.785500477533787,0.994296758901328,0.625912741292268,0.597492320695892,NA,1 --1.47635093715271,0.519639003556222,0.0640717146452516,0.487221607938409,0.496073189657182,0.116159395547584,0.006662578554824,0.0628067965153605,0.91214673826471,0.58853417634964,0.529160888632759,NA,2 --0.54737069346927,0.0266456077806652,0.749969109194353,0.375706146936864,0.8472254879307,0.302311453036964,0.892074418487027,0.566268016584218,0.896159957861528,0.375780015485361,0.340594119857997,NA,0.5 -1.45065191789749,0.94707443495281,0.425205362495035,0.337836832040921,0.73194540082477,0.0198679927270859,0.384525762172416,0.486763422377408,0.31629200046882,0.698771668132395,0.0593036743812263,NA,3 --0.36734327686679,0.449535029008985,0.739834612002596,0.998082731151953,0.405127953272313,0.798034045379609,0.383507769554853,0.453684104373679,0.366463227430359,0.449634636053815,0.556320760864764,NA,1 --0.142193432562441,0.0993167369160801,0.780165972188115,0.76557879964821,0.643417255720124,0.0841845641843975,0.952992078615353,0.388045005733147,0.353067258605734,0.494073997018859,0.38412378448993,NA,0.5 --0.0204411745838119,0.182601721491665,0.671439114259556,0.933410462690517,0.063336655497551,0.283508115913719,0.0870876528788358,0.795035666320473,0.261529661715031,0.253830597037449,0.271484842523932,NA,0.5 --0.145046151664827,0.412740072468296,0.922483856091276,0.205560039961711,0.281724784523249,0.308287980500609,0.114223259268329,0.339805831899866,0.59704695106484,0.0895669923629612,0.20351308491081,NA,1 --0.655416133968273,0.132374021690339,0.0385842223186046,0.960168987046927,0.553003982873634,0.0313359131105244,0.998100819531828,0.0783406211994588,0.927472871961072,0.659052587812766,0.785189574118704,NA,0.5 -0.0569453360323003,0.455947425216436,0.609724464127794,0.614067499991506,0.97299245535396,0.99231677222997,0.702998753637075,0.760542156873271,0.224094831617549,0.597768754232675,0.174418788170442,NA,1 --0.407626682681292,0.138098580995575,0.752001489046961,0.289834522875026,0.654601774876937,0.41541023994796,0.00403436576016247,0.770590521162376,0.0943589757662266,0.306590826017782,0.670288912020624,NA,0.5 --1.68139582769653,0.969414204824716,0.606616077246144,0.349470247514546,0.341198356356472,0.052086751209572,0.360128439497203,0.27755997586064,0.483060038648546,0.592371104983613,0.179017000831664,NA,3 --1.41235887891515,0.484311842592433,0.627842574846,0.414641877403483,0.304095499916002,0.200165947666392,0.097875046543777,0.366788988700137,0.650650549912825,0.789857445284724,0.357476759934798,NA,1 -0.452289857944087,0.612423501210287,0.262217645766214,0.175194005016237,0.517215595347807,0.658887927187607,0.398805856704712,0.485592402284965,0.953524999320507,0.391746623674408,0.300889120902866,NA,2 --2.08435011464261,0.759304512990639,0.582466018851846,0.261647397186607,0.554752136347815,0.801065125502646,0.137913859682158,0.877758507616818,0.749789121560752,0.467832755995914,0.0966854214202613,NA,3 --0.250461882730123,0.282862489577383,0.311645941110328,0.60926196561195,0.124141810229048,0.274296503281221,0.622827796963975,0.743674198165536,0.265039020916447,0.945896445307881,0.395986490882933,NA,1 -1.38929004072203,0.581831085495651,0.604764067567885,0.0981524740345776,0.640039406251162,0.59033013205044,0.712040337501094,0.564666945487261,0.492797074839473,0.603931111749262,0.279830668587238,NA,2 -3.37308806241796,0.644188753794879,0.11642846907489,0.388798480387777,0.456940018106252,0.958898683078587,0.534385369159281,0.026345114922151,0.103369149379432,0.836727291345596,0.575546889333054,NA,2 --1.80625513593003,0.665181114571169,0.865635313559324,0.236588786588982,0.625312402611598,0.197498777182773,0.469960128655657,0.853968983516097,0.0933152246288955,0.272778454935178,0.137419609585777,NA,2 --0.598181630839518,0.20983607089147,0.388797290390357,0.747019624104723,0.93496206542477,0.0514743681997061,0.420886440202594,0.680075784912333,0.825326177757233,0.238552219001576,0.866760726552457,NA,0.5 --0.678701228012741,0.367456012638286,0.76701890816912,0.299027033150196,0.638889130437747,0.973676205612719,0.202693411381915,0.047896116040647,0.693167041055858,0.72095926781185,0.85648188367486,NA,1 --0.439448370591987,0.171001069946215,0.7493276654277,0.766403131186962,0.142307656118646,0.68866288033314,0.379330020165071,0.842053881846368,0.453404651256278,0.344623880228028,0.187736973864958,NA,0.5 --0.429402457336857,0.232082100119442,0.912413143087178,0.417194787180051,0.263550848001614,0.0323172744829208,0.70294936420396,0.237633729353547,0.179497604025528,0.155822693370283,0.655707615427673,NA,0.5 -0.260273411620581,0.397116974228993,0.285043719923124,0.626160256098956,0.440539639210328,0.228579026414081,0.950656231259927,0.738680086331442,0.909562820801511,0.836534011876211,0.909986349521205,NA,1 --0.601741747228831,0.00319111370481551,0.801924289902672,0.936224457342178,0.0328773867804557,0.680261990521103,0.817012073937804,0.107413252815604,0.0303064233157784,0.471293790731579,0.67827691603452,NA,0.5 -0.579175363038437,0.382363738259301,0.393253202782944,0.158311059698462,0.802881805924699,0.164180209161714,0.210846305359155,0.0777299583423883,0.271905378438532,0.415683274855837,0.646119900047779,NA,1 -0.233684619832068,0.104158732341602,0.044687672983855,0.755966236116365,0.898829346057028,0.747886597644538,0.785850781947374,0.187898478936404,0.26085401698947,0.919391671195626,0.723531056428328,NA,0.5 --1.18176505505841,0.748055958654732,0.457851321436465,0.75289491051808,0.874799456913024,0.5951328065712,0.474716911558062,0.540802971692756,0.586096824146807,0.983661994105205,0.879120596917346,NA,2 --0.899714329525814,0.598819239530712,0.214800329878926,0.255403825780377,0.657175703207031,0.83159053674899,0.0539117860607803,0.840320016723126,0.30811989097856,0.841044461587444,0.933701136615127,NA,2 -0.646093632296236,0.36379097471945,0.930275573628023,0.413587759248912,0.677580804331228,0.780740413581952,0.0646629584953189,0.859191400231794,0.603091831551865,0.680502528557554,0.776847406988963,NA,1 -1.02643265104361,0.787000266136602,0.130031454842538,0.812010727357119,0.11337968474254,0.226478895405307,0.9479467747733,0.289414885919541,0.826333333505318,0.482536654453725,0.204640087904409,NA,3 --1.60873401630344,0.333655280526727,0.473265617853031,0.913793664192781,0.927346032112837,0.598311467794701,0.802216146606952,0.604303786531091,0.204727914649993,0.256905679125339,0.353850594256073,NA,1 --0.193922357216301,0.309134368319064,0.69424194470048,0.045563388383016,0.0289882610086352,0.702026945538819,0.353906127158552,0.578083222731948,0.998108556028455,0.706539734266698,0.111041986150667,NA,1 -2.25822378003256,0.699067022651434,0.175160866929218,0.675344614312053,0.653209868585691,0.448388746008277,0.639876928878948,0.529742702143267,0.386671084910631,0.871713967528194,0.597260946873575,NA,2 --0.294217127054907,0.0130114420317113,0.745208188192919,0.662316924892366,0.869964884594083,0.346366792684421,0.134350478416309,0.91824146383442,0.628984485287219,0.677548446692526,0.685797044774517,NA,0.5 --5.09182856615671,0.535232451977208,0.519287661882117,0.85244823875837,0.555158670060337,0.251911287195981,0.740879685385153,0.578002710128203,0.67449068557471,0.13997372216545,0.171553277410567,NA,2 --0.957106599062254,0.260365868685767,0.039272784255445,0.205235176021233,0.658755406038836,0.0693191748578101,0.265798757784069,0.553729085018858,0.56509280600585,0.896149475127459,0.0161577914841473,NA,1 --0.74547676965435,0.216713041067123,0.312677611829713,0.271783971460536,0.31258171889931,0.547919424949214,0.143913608510047,0.106598092708737,0.155940489610657,0.147696991218254,0.264999630395323,NA,0.5 --1.7556278223873,0.301559766521677,0.813955902820453,0.378281063167378,0.23688638349995,0.825144364265725,0.807417817413807,0.154609765624627,0.987439922289923,0.27786693489179,0.650420303922147,NA,1 --0.0470943359654031,0.160312748048455,0.732178756734356,0.151784234680235,0.372392004355788,0.374390203272924,0.544937105616555,0.617528083734214,0.677572895074263,0.368166692554951,0.655130254570395,NA,0.5 -0.124658668723908,0.182294715894386,0.995739555219188,0.113104562275112,0.11240058299154,0.208648389670998,0.537716497201473,0.0649266852997243,0.650002721697092,0.0501199315767735,0.357121436158195,NA,0.5 -5.21455443964906,0.779459196841344,0.449038189370185,0.292015180690214,0.143725096713752,0.904136493103579,0.608930496964604,0.354715100722387,0.194870956009254,0.843750018393621,0.0934222042560577,NA,3 -1.75772733370199,0.821147922892123,0.425148322479799,0.171820772113279,0.235684452112764,0.761053857393563,0.315450295107439,0.220257700188085,0.835394476307556,0.984954028623179,0.115882620681077,NA,3 --2.47988607885718,0.629339757142588,0.828968267887831,0.995918438304216,0.607340537942946,0.0347810362000018,0.946215031202883,0.967786622000858,0.126801078440621,0.765530567616224,0.121757410932332,NA,2 --0.894667690433518,0.781423395499587,0.0622167293913662,0.272649483056739,0.198446429567412,0.659495470346883,0.81549444841221,0.649437340209261,0.403396281879395,0.0653008993249387,0.449648660607636,NA,3 --2.59967794033893,0.833517065504566,0.922720778966323,0.283613242208958,0.117495934013277,0.44706165464595,0.493439532117918,0.653966138372198,0.117235817248002,0.946330566890538,0.346774617442861,NA,3 -0.562866817007894,0.066422026604414,0.0124633046798408,0.348506391048431,0.520839661592618,0.444712787633762,0.71667861030437,0.328789311693981,0.026976422406733,0.644805737538263,0.929708110401407,NA,0.5 --0.98353176943091,0.345342023530975,0.675448514288291,0.863113478990272,0.767568795010448,0.745967566501349,0.583446231437847,0.355865387944505,0.183258432662115,0.327033995883539,0.742279504658654,NA,1 -1.39982507996358,0.301927540684119,0.425509258871898,0.963772800983861,0.432708555832505,0.552789221983403,0.862505354220048,0.194295163266361,0.607898182235658,0.666282393038273,0.855093700578436,NA,1 -2.19108366345933,0.784785338677466,0.910598614718765,0.0310040663462132,0.792512105545029,0.913687644060701,0.779540273128077,0.203273280523717,0.580115466378629,0.778640141477808,0.0813250311184675,NA,3 -0.102754799970055,0.0151381115429103,0.303750258870423,0.530344953527674,0.398493434069678,0.551140069263056,0.331004866398871,0.0359833822585642,0.954829844180495,0.751596683170646,0.923089451156557,NA,0.5 --2.71593227987405,0.975154257612303,0.158668037503958,0.417390400543809,0.100585088133812,0.454939389135689,0.426590847782791,0.1122667635791,0.172702635638416,0.360355676151812,0.969007642241195,NA,3 -0.11316914619458,0.144407300045714,0.726451003458351,0.323942760471255,0.0605967913288623,0.890680358512327,0.902225363301113,0.955268873134628,0.0258821186143905,0.516488530673087,0.00521893566474319,NA,0.5 -0.214612463434781,0.382541353581473,0.541727594798431,0.679740986786783,0.000293354038149118,0.129221958573908,0.0420920078177005,0.329942629439756,0.341416005743667,0.02168780984357,0.247941843234003,NA,1 --3.25744451413655,0.878030564635992,0.612484061392024,0.841921371873468,0.389507194282487,0.873570098308846,0.941481810528785,0.270786457462236,0.105030868668109,0.0905804259236902,0.677928207209334,NA,3 --0.612778977564602,0.299138211412355,0.71343940589577,0.362453953828663,0.569098226493225,0.806055163498968,0.984774222364649,0.0434266540687531,0.158101887907833,0.446324746357277,0.10046491608955,NA,1 -0.768254629890036,0.090535826748237,0.598878909833729,0.383121004793793,0.942228753585368,0.728780089877546,0.816637960728258,0.63866300880909,0.776843504980206,0.295215973397717,0.961366922827438,NA,0.5 --0.567603909448942,0.0424070812296122,0.419121650280431,0.861307024024427,0.638702552299947,0.642014807322994,0.769632588373497,0.153189419070259,0.810209463350475,0.170681128744036,0.77459627809003,NA,0.5 --5.9771524195317,0.862377436831594,0.652252439875156,0.904387121321633,0.797982558840886,0.350751545047387,0.78046753979288,0.985694216098636,0.70478395326063,0.913390251807868,0.0379480242263526,NA,3 -3.06750410813849,0.80633880593814,0.679136122344062,0.952588496729732,0.829646679805592,0.881833428051323,0.884054609807208,0.0632740694563836,0.70828592334874,0.175746616674587,0.878892200067639,NA,3 -0.957806121739846,0.112819531699643,0.528766673523933,0.779676516773179,0.985083922278136,0.22144174692221,0.14839206263423,0.560842191567644,0.211447972571477,0.751184774329886,0.262189727742225,NA,0.5 -0.7394269431204,0.0613693608902395,0.0450991443358362,0.469146619085222,0.236004963051528,0.938807192258537,0.464330278337002,0.197913277195767,0.454433776205406,0.08160714013502,0.868174660252407,NA,0.5 -2.89437745658501,0.765211458783597,0.946839487645775,0.702625217847526,0.89387573953718,0.322030391078442,0.635835886467248,0.490745102521032,0.251796741737053,0.57421529921703,0.919654391240329,NA,3 --0.837698154337176,0.911105049075559,0.887382954359055,0.218226345721632,0.503832868300378,0.00183678534813225,0.586421895073727,0.72927557816729,0.53785819420591,0.782691985368729,0.0610043394844979,NA,3 -1.61332292935373,0.767474424093962,0.657503654249012,0.23821485415101,0.3197404849343,0.848195793339983,0.523012418299913,0.220011793542653,0.923964306944981,0.419060569256544,0.20017861854285,NA,3 --0.111281104502172,0.136224425164983,0.408692202530801,0.915172693319619,0.946666706586257,0.544694220647216,0.307216387009248,0.844691653503105,0.402090376475826,0.111828369786963,0.364809120539576,NA,0.5 --0.743268442892708,0.631344011286274,0.416102928575128,0.94413023837842,0.405237703584135,0.0686362243723124,0.856442781630903,0.569145374232903,0.298357689753175,0.543654523789883,0.354759010253474,NA,2 --0.905598862688402,0.429458065889776,0.599177881376818,0.525696570519358,0.556089870398864,0.374567771330476,0.73522638506256,0.165942407678813,0.241547536570579,0.640308366389945,0.910148287890479,NA,1 --0.381371694003123,0.319485255284235,0.775019305292517,0.202426716685295,0.393176882760599,0.383716176263988,0.833298979094252,0.63092312333174,0.456514721736312,0.263742276001722,0.599667183123529,NA,1 --0.423702137252754,0.11936202365905,0.114415199263021,0.187402412761003,0.825105560244992,0.401661385083571,0.598708146018907,0.626203504158184,0.0329728100914508,0.0982055109925568,0.259948245948181,NA,0.5 --0.805397857933071,0.126862512435764,0.561627364018932,0.226535556837916,0.0691305936779827,0.0460475136060268,0.0944148511625826,0.584244670113549,0.75893074576743,0.932498198468238,0.695256362203509,NA,0.5 -0.319155083005407,0.472723427927122,0.787534170318395,0.645320721901953,0.715023176278919,0.209356067935005,0.81911487551406,0.801698182709515,0.70181589666754,0.698942579096183,0.469195405719802,NA,1 --0.0907648903587075,0.463830034947023,0.883693974465132,0.996114502195269,0.933144294423983,0.130785294808447,0.672029396053404,0.39099053083919,0.82090491685085,0.889978097518906,0.481951746391132,NA,1 -0.685509547256731,0.733466965379193,0.0820799909997731,0.0614028854761273,0.193428620696068,0.842947125434875,0.147228190675378,0.63629440870136,0.00185011373832822,0.664431531913579,0.359879235271364,NA,2 --6.88310219990682,0.766328622121364,0.514705869369209,0.0641286331228912,0.71263275318779,0.865839602192864,0.317526940023527,0.472104668151587,0.0943877408280969,0.525016574189067,0.392749238293618,NA,3 --2.30749559493075,0.317944551585242,0.313390593510121,0.785861490992829,0.626033872598782,0.691302561433986,0.909780625021085,0.636675088433549,0.000775783322751522,0.607812081230804,0.350837268168107,NA,1 -0.190249220215772,0.0292525580152869,0.784951515262946,0.726968697272241,0.515227085910738,0.952373256674036,0.985080229351297,0.0235045454464853,0.00416592671535909,0.967333158710971,0.396277377847582,NA,0.5 --1.55456092344735,0.298712126677856,0.329398845555261,0.824424406746402,0.504606149625033,0.462511173449457,0.987116791075096,0.984700593166053,0.738257547840476,0.506412111455575,0.215016127098352,NA,1 --5.71298108096891,0.863856402924284,0.472625948023051,0.0874699417036027,0.909145939163864,0.37907059583813,0.480933768674731,0.478833106346428,0.669854396488518,0.0413195120636374,0.272561500314623,NA,3 -2.72257293704493,0.977674740599468,0.757155697094277,0.207756455754861,0.203775214497,0.78417210560292,0.582602435722947,0.666409130906686,0.686884060502052,0.110600611427799,0.558197030797601,NA,3 -10.2974933407333,0.817447017878294,0.907617374090478,0.650747747858986,0.467643693089485,0.37755541363731,8.57806298881769e-05,0.00122400932013988,0.270136442501098,0.128451892407611,0.193024599459022,NA,3 --0.515443771934163,0.156245535472408,0.951601629611105,0.789141749264672,0.919310998637229,0.852464217459783,0.700894636334851,0.253609915263951,0.141912077786401,0.594247984467074,0.577542662387714,NA,0.5 -3.60755387259999,0.935209455899894,0.168847924564034,0.71698178560473,0.164761062944308,0.820564121939242,0.269563676323742,0.865492029348388,0.0663771124090999,0.757397901965305,0.30812180461362,NA,3 -1.53424165667702,0.495078047504649,0.00743950600735843,0.699956620577723,0.150095566408709,0.820907787652686,0.0250667075160891,0.219115395331755,0.0701660695485771,0.144342829938978,0.661241426831111,NA,1 -2.11945690269233,0.494878040393814,0.335358702111989,0.0859456097241491,0.334762825863436,0.436428788118064,0.871063163038343,0.862540878355503,0.149437770945951,0.172741169808432,0.301684214035049,NA,1 --0.764114777927609,0.472692873096094,0.913892224663869,0.814544472144917,0.345459071453661,0.473493713885546,0.772064296295866,0.964068693341687,0.302707911469042,0.310749824624509,0.326807718491182,NA,1 -1.59341061109396,0.607975764898583,0.0794732617214322,0.974691078998148,0.617390216095373,0.461572230560705,0.722619395470247,0.609765309141949,0.594569276785478,0.192564327269793,0.0523809674195945,NA,2 --1.05008884817102,0.458628129446879,0.276793640572578,0.394021584419534,0.216551521327347,0.656353653175756,0.901549807749689,0.893069363664836,0.394100816454738,0.276124453404918,0.0872335392050445,NA,1 --0.454916415458206,0.100318351062015,0.378878374584019,0.0983213325962424,0.174892703304067,0.711960697313771,0.919599503278732,0.190358696039766,0.269345039268956,0.287819736870006,0.287268604617566,NA,0.5 --0.485510412623777,0.241354882251471,0.949344601249322,0.158448557602242,0.133629100397229,0.866928705945611,0.937315163202584,0.184269865509123,0.327862715348601,0.563999371370301,0.785677004838362,NA,0.5 -0.776775239291567,0.291910991771147,0.0846551819704473,0.77373803849332,0.182577745057642,0.20777322165668,0.364862532587722,0.302542981225997,0.99205335858278,0.101478747557849,0.831732955761254,NA,1 --0.503915118460329,0.0921388799324632,0.69793195463717,0.640171206556261,0.433864553691819,0.868978491751477,0.679959014523774,0.123384471517056,0.638639747863635,0.164615281391889,0.954158689128235,NA,0.5 --2.45868231246521,0.486175883794203,0.204146457370371,0.274052948923782,0.427365210838616,0.0658067418262362,0.238397930283099,0.715437222272158,0.305063903564587,0.979477859567851,0.535615005297586,NA,1 -3.07267136331703,0.672265029046685,0.683067857287824,0.300874865613878,0.990141049260274,0.336872167419642,0.262356182327494,0.166166434762999,0.507457473315299,0.182093014242128,0.988365483703092,NA,2 --1.35725074282868,0.373770251171663,0.989214881556109,0.348896064329892,0.00952328857965767,0.870421917643398,0.0278489568736404,0.971605542581528,0.0709339645691216,0.96520364494063,0.830430411733687,NA,1 --2.09439816462769,0.6810261097271,0.785620857728645,0.798142119776458,0.381587166339159,0.372327584540471,0.625061142724007,0.36749339546077,0.582425137981772,0.641279636882246,0.120793332578614,NA,2 --0.00796859068675282,0.0229878642130643,0.855209650471807,0.982022156473249,0.256182591430843,0.777578515000641,0.825847221305594,0.558763438137248,0.683600728865713,0.737240033224225,0.413099535508081,NA,0.5 --1.4979445469424,0.714179781964049,0.706974609289318,0.173396941041574,0.230190037516877,0.856905392371118,0.897651092847809,0.326673374744132,0.499662009300664,0.31590295699425,0.814334665657952,NA,2 --0.439028039829894,0.69282753020525,0.235939425881952,0.237423748476431,0.719278433360159,0.29460807191208,0.928081472637132,0.536010401323438,0.794380184262991,0.658699807012454,0.520766410976648,NA,2 --0.242951793769155,0.243101859698072,0.307146482169628,0.843962025130168,0.0316330003552139,0.430362788261846,0.351072994992137,0.654738643905148,0.969542238861322,0.60179814370349,0.398643856402487,NA,0.5 --0.630263115026057,0.200559745077044,0.907688920153305,0.334342788904905,0.150655995355919,0.972313160542399,0.542266771197319,0.50963302818127,0.218434902839363,0.748211001046002,0.228817308787256,NA,0.5 -0.696031295956962,0.328547941055149,0.876553914509714,0.164135021390393,0.395798878511414,0.477207755902782,0.283287385711446,0.191637014970183,0.346791804069653,0.0879971194081008,0.774699172470719,NA,1 -1.86218635994138,0.667642474640161,0.175826355349272,0.104500935412943,0.188644989626482,0.935310949804261,0.534588502487168,0.650544292991981,0.280430254060775,0.881018078653142,0.322163599310443,NA,2 -0.0618266822267202,0.966812142170966,0.964485211065039,0.745329377939925,0.468591102631763,0.238328024744987,0.248983833007514,0.304753655567765,0.806436447193846,0.374863224104047,0.268298006616533,NA,3 --0.213469050164475,0.626246140571311,0.371381344739348,0.64032802847214,0.563473800662905,0.100175821688026,0.0235935947857797,0.410097931977361,0.984508452471346,0.908335901331156,0.106187571771443,NA,2 -2.71820413597084,0.342010788153857,0.705079389037564,0.551346761640161,0.598823779728264,0.448870678432286,0.737467909697443,0.322141070617363,0.551015156554058,0.226791883818805,0.913258621236309,NA,1 --1.01015670294422,0.765705335186794,0.826494007604197,0.231260113185272,0.565515470458195,0.248132557841018,0.171093754237518,0.286525617586449,0.441912526264787,0.227831450290978,0.286231678677723,NA,3 --3.14339298170896,0.951644803164527,0.232232666807249,0.908373491605744,0.913206648547202,0.407022410538048,0.953037474770099,0.274196953279898,0.27726714592427,0.842141730245203,0.17599373916164,NA,3 -0.816940374582757,0.00308033730834723,0.223011573543772,0.0552042822819203,0.646219877526164,0.89388394379057,0.35093015129678,0.0684547659475356,0.56139280856587,0.433695368934423,0.780848235357553,NA,0.5 --1.06114260561244,0.488627271028236,0.343027671799064,0.495272737927735,0.920322231249884,0.0372702074237168,0.504528037970886,0.391261521494016,0.130302850389853,0.576689033070579,0.0055089273955673,NA,1 -3.14001425018611,0.936363400658593,0.68013828061521,0.900950629264116,0.699788666563109,0.58299405220896,0.566819867584854,0.539389610989019,0.598359101684764,0.961430250899866,0.322504172567278,NA,3 -3.75658460367109,0.72642320394516,0.340536155505106,0.644040668383241,0.0150993110146374,0.76843476947397,0.496726306155324,0.994598030112684,0.0396893250290304,0.878083458403125,0.48237074050121,NA,2 -1.12695583768729,0.983749621314928,0.788466003490612,0.593351603019983,0.94108819635585,0.306970248464495,0.201373670483008,0.209501806180924,0.305320084095001,0.451809530379251,0.385713089257479,NA,3 -1.52459560077634,0.672651072964072,0.46188158262521,0.209960805252194,0.545736834639683,0.21016333415173,0.909396487753838,0.140858779894188,0.656669340096414,0.829055720707402,0.514228889951482,NA,2 -0.963944520418568,0.0890714421402663,0.796551150269806,0.0773013802245259,0.993034319253638,0.18855186062865,0.839895894518122,0.13716762047261,0.883261514361948,0.305906966794282,0.866397253004834,NA,0.5 -0.129771290079079,0.189816636266187,0.434954048832878,0.614785549463704,0.0503757097758353,0.514615601161495,0.952145036077127,0.637101435568184,0.426087585045025,0.311506387079135,0.234520558966324,NA,0.5 -0.912884979700598,0.683184877503663,0.510933240642771,0.473237108206376,0.597530248807743,0.262252091662958,0.373021523468196,0.595309577649459,0.412889627274126,0.824333656812087,0.681659359484911,NA,2 --2.09362850541451,0.426857191603631,0.896169882267714,0.314928279956803,0.495635512284935,0.482654591090977,0.713752283481881,0.620617669541389,0.778110204497352,0.956990454345942,0.608790414175019,NA,1 --0.0644100579741575,0.721381928073242,0.246724736411124,0.742296724347398,0.949994837632403,0.286190049722791,0.0760854352265596,0.445758877322078,0.949769111117348,0.925280713476241,0.499942892463878,NA,2 -0.10139334086343,0.0579308001324534,0.994543595705181,0.293401257600635,0.00884762965142727,0.889348507625982,0.549911052221432,0.615719662979245,0.280096381437033,0.143207954941317,0.271301218308508,NA,0.5 -2.06653528604937,0.518336966400966,0.506918795639649,0.466876366641372,0.620179544668645,0.486438295803964,0.383961294777691,0.519759878516197,0.0823564794845879,0.572632172610611,0.927550779422745,NA,2 -0.083993183366566,0.792077120160684,0.0227267004083842,0.350600538309664,0.802795631811023,0.271761489799246,0.752974940929562,0.58151887729764,0.448455915320665,0.114585774950683,0.480480066733435,NA,3 -0.871280979378034,0.155268164584413,0.169720521196723,0.209214897826314,0.969444170594215,0.455098179168999,0.162548776483163,0.30731915589422,0.797547695925459,0.937418807530776,0.699685209430754,NA,0.5 --0.993675609227815,0.31696749618277,0.255061091855168,0.329014195129275,0.311154339462519,0.901402737712488,0.881375814089552,0.592734534991905,0.0463438257575035,0.407562130596489,0.94760683667846,NA,1 -3.87620571131626,0.712310006609187,0.150037638610229,0.24496771581471,0.652770188869908,0.754434264730662,0.353124728426337,0.437375633278862,0.317781692370772,0.994260160252452,0.886606745189056,NA,2 --0.184073402751879,0.302567538106814,0.236578963929787,0.529907553223893,0.0782564121764153,0.575638748938218,0.953494219575077,0.74427487119101,0.285246771760285,0.675243055913597,0.700421273009852,NA,1 --0.82056052576034,0.607995933154598,0.269555284408852,0.859022193821147,0.785452287411317,0.556223073042929,0.531818702351302,0.694545754231513,0.475929024163634,0.182002258952707,0.0817325403913856,NA,2 --0.128441943850425,0.689589613582939,0.838425317080691,0.0708509180694818,0.0456330950837582,0.244261574465781,0.352208221564069,0.257190010044724,0.259088592603803,0.793637574184686,0.251339136157185,NA,2 --1.67953017850423,0.513718589907512,0.177739400183782,0.135568092111498,0.754302552901208,0.119685888523236,0.271099909441546,0.881469729356468,0.380955666769296,0.0594409927725792,0.285897349473089,NA,2 -4.01029674994469,0.637662801891565,0.681524514220655,0.0506364298053086,0.195933144073933,0.177333432948217,0.206886797677726,0.685570570174605,0.356021656654775,0.985714256530628,0.845266965450719,NA,2 -0.418247462976808,0.156605343334377,0.341756730806082,0.963353005703539,0.899666843237355,0.4310629342217,0.946215359028429,0.897696364205331,0.776364185847342,0.420296776574105,0.381828723009676,NA,0.5 --1.51268669424462,0.62234056647867,0.577063782373443,0.603003362193704,0.446018574293703,0.340106446295977,0.65706201409921,0.152433987008408,0.260849780868739,0.550666619790718,0.255014094058424,NA,2 -4.89636360756605,0.807543221861124,0.476624209433794,0.202747098170221,0.38368299161084,0.331279368372634,0.282060217810795,0.374299304792657,0.750059184152633,0.606916478602216,0.0227296042721719,NA,3 -2.17393837350942,0.280177400680259,0.634965054458007,0.408001474570483,0.916558355325833,0.714338184334338,0.583447890123352,0.470874862978235,0.0685866239946336,0.717115951701999,0.843359079910442,NA,1 -0.114333681805783,0.438450475456193,0.602708112215623,0.349823107244447,0.414686303352937,0.198583260178566,0.261792612727731,0.163943613413721,0.180993436602876,0.669274596963078,0.0332412673160434,NA,1 --0.0575913747610489,0.180707123363391,0.526421991642565,0.82064001262188,0.521263679722324,0.792229591170326,0.514831050531939,0.655353802023456,0.148887025890872,0.567784914281219,0.786496615968645,NA,0.5 --1.01172391605431,0.463357580360025,0.267548341304064,0.839970801956952,0.421053394908085,0.835206566611305,0.946785928681493,0.722578325076029,0.913413252448663,0.466875985031947,0.519509368343279,NA,1 --1.01622844711912,0.636612149421126,0.569439431419596,0.0444083481561393,0.166187049355358,0.890496745705605,0.972773446468636,0.512027263408527,0.453198195435107,0.185353330336511,0.822659314842895,NA,2 -0.879633812737811,0.937646907521412,0.357851214939728,0.754458939423785,0.757599327713251,0.252516763517633,0.253876760136336,0.213992986362427,0.89202951756306,0.274374803295359,0.646439817501232,NA,3 -3.17475900901958,0.715058491099626,0.757979915942997,0.363833147101104,0.339705190155655,0.342578890034929,0.504253518534824,0.836292343214154,0.440996023127809,0.0837488134857267,0.572989572072402,NA,2 --0.0127358680059262,0.91581250471063,0.954323386540636,0.430059639271349,0.762585739139467,0.368888254277408,0.20797983976081,0.47765283472836,0.167148848297074,0.308806973975152,0.314454952254891,NA,3 --0.8544312017705,0.502859967993572,0.151654090732336,0.396629377966747,0.021326761925593,0.205127923749387,0.539602338802069,0.0801719883456826,0.0248979306779802,0.718125011539087,0.321455252822489,NA,2 --1.82887098666628,0.262003114912659,0.0757940574549139,0.85448711225763,0.564622745616361,0.0720540515612811,0.967587752500549,0.959227720042691,0.638316398952156,0.562585303094238,0.834035288542509,NA,1 -3.2622467572984,0.335096222581342,0.244150118669495,0.183940630406141,0.615293832495809,0.827252987539396,0.267028478439897,0.273504600860178,0.507153154583648,0.388217261061072,0.138010729802772,NA,1 -1.65601016742667,0.838172710733488,0.954755051294342,0.608526331372559,0.836548767983913,0.919543279102072,0.515458332141861,0.598611275665462,0.66656892397441,0.607837118208408,0.493392545264214,NA,3 -0.381933395917132,0.261543265311047,0.95214184653014,0.764482910744846,0.966228008968756,0.493756796000525,0.872538808966056,0.585356997791678,0.242046943400055,0.900154284201562,0.452932744985446,NA,1 -0.061813986449167,0.5902709166985,0.654293139465153,0.18226848798804,0.914126749383286,0.398949036840349,0.773602006956935,0.743641247740015,0.208624595077708,0.4145590539556,0.676244510803372,NA,2 -0.274146489910667,0.560014321003109,0.350247922120616,0.756335157435387,0.416073066880926,0.0586489790584892,0.301071509020403,0.422949285479262,0.210499112727121,0.482112698722631,0.112424957100302,NA,2 --5.75883040392064,0.901148825651035,0.810049108229578,0.878282216377556,0.316235190955922,0.230603245552629,0.250381491146982,0.650639815721661,0.0529408203437924,0.958216429455206,0.889385785441846,NA,3 -0.911749288623184,0.539145407034084,0.596371402265504,0.897788940230384,0.439045100938529,0.16618771571666,0.142886094050482,0.667914129793644,0.535841233795509,0.753273370675743,0.0447190594859421,NA,2 -4.55360390370342,0.945572390919551,0.0924056400544941,0.939797199331224,0.642557488754392,0.0861750931944698,0.632535864366218,0.465771202929318,0.272886190097779,0.772984977345914,0.715489778900519,NA,3 --0.907822599380484,0.423604313516989,0.944986613234505,0.159030285896733,0.922660069772974,0.964939738158137,0.532024955842644,0.634939571144059,0.39645018754527,0.635082328226417,0.593678319593892,NA,1 --0.139344587678803,0.58962686243467,0.590537373209372,0.14912519371137,0.20451657124795,0.0897868368774652,0.738680734764785,0.374435908161104,0.989756791852415,0.584141218336299,0.653865883359686,NA,2 -0.806150843945478,0.678180137882009,0.973715249914676,0.106831170618534,0.300997588550672,0.750364477513358,0.491450151661411,0.0154496582690626,0.259265987202525,0.558927526231855,0.0112981905695051,NA,2 -1.44095642980371,0.917209735373035,0.341985862003639,0.268955789040774,0.109774171141908,0.239792765583843,0.454389738617465,0.522328067803755,0.60719471052289,0.118838940281421,0.388227568939328,NA,3 --3.25161689162916,0.922700179507956,0.0325096684973687,0.352068450069055,0.232990069780499,0.469084792770445,0.981463437899947,0.11463882913813,0.72538296924904,0.246744823874906,0.321137896971777,NA,3 --0.130522430553986,0.401394749758765,0.256456261267886,0.0354134733788669,0.978563639800996,0.239538576453924,0.814784180140123,0.735997994896024,0.749723166925833,0.62831072579138,0.890577062265947,NA,1 -0.669575467655283,0.00989005784504116,0.435264135012403,0.723876874428242,0.849379999563098,0.0575199017766863,0.957934653852135,0.837231518467888,0.46569788781926,0.891659042565152,0.783426602371037,NA,0.5 -2.54014096182668,0.305441607953981,0.0902002388611436,0.118577564833686,0.385022775735706,0.56807669182308,0.334235888672993,0.538188089616597,0.216295424615964,0.158172710798681,0.00234771450050175,NA,1 --0.193944994135837,0.00190780940465629,0.78280071169138,0.764838046627119,0.291868838947266,0.869910318404436,0.341377722565085,0.648056457517669,0.579640600830317,0.215270293876529,0.61375098163262,NA,0.5 -0.0208742133525005,0.863172167912126,0.504580818815157,0.374554778914899,0.836817204486579,0.87124543520622,0.670387550722808,0.737122401129454,0.827918927185237,0.599544936791062,0.942843936849385,NA,3 -0.765792314938424,0.566925200633705,0.849535879679024,0.34304199879989,0.847504819976166,0.892173214582726,0.758812437998131,0.214903841959313,0.45636634901166,0.484204821521416,0.207701456500217,NA,2 -0.352804369431417,0.364061176544055,0.670583426486701,0.117415519431233,0.108332704752684,0.316217254614457,0.767299061873928,0.895374773070216,0.667814850108698,0.922258591279387,0.0134780446533114,NA,1 --1.23698048076279,0.332583913113922,0.612572266953066,0.704983988078311,0.768699376378208,0.385655405465513,0.128826092928648,0.116410831455141,0.520903897006065,0.974680992309004,0.806146663846448,NA,1 --2.56441778784605,0.665565764764324,0.605427384609357,0.17659778567031,0.02090925257653,0.704631736734882,0.73883875226602,0.774222782580182,0.268567186314613,0.696540282806382,0.354841801803559,NA,2 -0.232796925229605,0.569475795142353,0.454856015508994,0.939597767777741,0.744496839120984,0.698145196773112,0.467724752146751,0.0180247377138585,0.297438253415748,0.308250970672816,0.523631568299606,NA,2 --0.916045983780381,0.398941761115566,0.661880399100482,0.385367383016273,0.0125482359435409,0.701195284491405,0.652126201428473,0.324592209421098,0.482771011535078,0.943977946648374,0.711980663239956,NA,1 --3.02511296110412,0.601308670360595,0.492553135380149,0.233871235977858,0.712123312056065,0.867592514958233,0.929739307146519,0.856446724850684,0.754480483010411,0.377171362750232,0.526562972692773,NA,2 -0.08116233752899,0.208068270701915,0.806491583352908,0.138934786897153,0.550345825962722,0.273357491008937,0.738135070772842,0.122066766722128,0.67383626755327,0.776206725975499,0.432750965002924,NA,0.5 --5.27518419045481,0.891456350916997,0.601628924021497,0.281817921902984,0.245754534145817,0.525102650281042,0.350016801385209,0.0317163458094001,0.342192857991904,0.69337944011204,0.535440781386569,NA,3 --0.195996325881689,0.138576569035649,0.187388508580625,0.262100108899176,0.263345927000046,0.493119575316086,0.152671873103827,0.504191888263449,0.913665950298309,0.165937364799902,0.0847142075654119,NA,0.5 -0.315579484938764,0.858387031825259,0.695295455399901,0.467401414411142,0.626965253613889,0.490166572621092,0.617864213185385,0.0769863030873239,0.542200545780361,0.219044300960377,0.0295198371168226,NA,3 -0.0267096393523584,0.266054383246228,0.523457160918042,0.538887435337529,0.597743078833446,0.296432407340035,0.807767712045461,0.84599154163152,0.397196589736268,0.82058397657238,0.093033802928403,NA,1 -1.25261939865097,0.60674773179926,0.706874470924959,0.674645616207272,0.156269168946892,0.858524578856304,0.966854283586144,0.777528064092621,0.547137490706518,0.0508977069985121,0.170303455553949,NA,2 -0.541883133495883,0.299106652615592,0.868335065431893,0.293283038539812,0.337412812747061,0.563003861811012,0.057874902850017,0.206769967218861,0.249845808371902,0.57238865736872,0.431905421428382,NA,1 --0.359359846361593,0.00028673093765974,0.444823984289542,0.267580098472536,0.619911819230765,0.348575697746128,0.445040695136413,0.298632511869073,0.349616114981472,0.0186277241446078,0.727784499060363,NA,0.5 --0.725046148314536,0.738201215397567,0.305303359869868,0.98667918308638,0.875479456968606,0.816655561095104,0.550264997407794,0.78673102799803,0.749039791757241,0.148614945588633,0.50864305277355,NA,2 -2.42662093734972,0.78392618172802,0.779760129982606,0.417675146833062,0.697084854589775,0.642540717031807,0.936560302507132,0.12054528738372,0.307862642919645,0.839434523368254,4.8645306378603e-05,NA,3 --0.841952793676179,0.248953337082639,0.903661191929132,0.214526182506233,0.640308956848457,0.630132407415658,0.557652691844851,0.294213210465387,0.575431503122672,0.117273130686954,0.394496359396726,NA,0.5 --1.03023205439958,0.724063246278092,0.797944512451068,0.531282962532714,0.761441437294707,0.882055644644424,0.695016658399254,0.902350852033123,0.740311119239777,0.755922955228016,0.0983525444753468,NA,2 --0.912771137967775,0.34925606357865,0.676785026676953,0.588860424933955,0.161323110340163,0.957519877469167,0.398865574738011,0.566670776112005,0.161247945856303,0.247955431696028,0.743449069792405,NA,1 -0.517313406285897,0.119058254174888,0.301976219983771,0.2837041572202,0.0562514583580196,0.732049201382324,0.543040792457759,0.178399549797177,0.637692917603999,0.176733382046223,0.913348912494257,NA,0.5 -0.871011941589735,0.308870760258287,0.508238948881626,0.217849156819284,0.755379884736612,0.308171590790153,0.877481517381966,0.958365337457508,0.897053161635995,0.609838765114546,0.357350649312139,NA,1 -0.723881692134752,0.1084195880685,0.285055675543845,0.494687871308997,0.959149545524269,0.867562021827325,0.437210090924054,0.0436040495987982,0.886218967847526,0.134117919020355,0.147130149416625,NA,0.5 --4.1056823910876,0.836923674680293,0.790600413223729,0.218965121312067,0.631649191724136,0.131414104020223,0.986164937028661,0.1794618004933,0.893512472510338,0.397628468694165,0.977409608894959,NA,3 -0.952558604299245,0.281533166533336,0.653331918641925,0.782094214344397,0.486817280761898,0.712452893378213,0.735165594611317,0.22127672447823,0.17275553708896,0.781067404197529,0.301315511809662,NA,1 -0.434834298048182,0.345692228991538,0.579240264836699,0.952218348626047,0.0265811623539776,0.292065045097843,0.140570497373119,0.594512204639614,0.894585548667237,0.163636759854853,0.791753190103918,NA,1 --0.617241139637955,0.24287026701495,0.347456730203703,0.352250093361363,0.567248777253553,0.5325484755449,0.362278310814872,0.104282488347962,0.6345406237524,0.298809389350936,0.061041493434459,NA,0.5 --1.10157854604848,0.589005693560466,0.738731898134574,0.411837223684415,0.0231877483893186,0.510611890815198,0.317921287380159,0.142743433825672,0.362699191318825,0.295025371946394,0.297999273752794,NA,2 -0.0779708729765081,0.61409965949133,0.267073255265132,0.644359648926184,0.435340049676597,0.663986770203337,0.853410524781793,0.439644190017134,0.0916713422629982,0.885482318233699,0.328195882029831,NA,2 --0.790086227163868,0.827004836406559,0.588833221467212,0.510606417665258,0.135837320005521,0.959990479284897,0.622058121720329,0.238031924469396,0.0791387532372028,0.875529285753146,0.385706771630794,NA,3 -0.0121011462499208,0.434019140200689,0.76129729533568,0.565473348135129,0.554319101618603,0.193900709273294,0.253057076362893,0.960415460402146,0.64092730381526,0.927356063388288,0.739957847632468,NA,1 -0.528668478450982,0.00797079806216061,0.620697325794026,0.059939440805465,0.933844800340012,0.282361896475777,0.798007321543992,0.109523738268763,0.205230085877702,0.786281060893089,0.626566218677908,NA,0.5 --1.97146007839586,0.426603920059279,0.232865846948698,0.667433485155925,0.0196325895376503,0.752369344001636,0.0936158304102719,0.0917298034764826,0.746732441009954,0.296932371798903,0.478665569564328,NA,1 -0.917133356519574,0.318299685139209,0.434513642685488,0.549943955615163,0.726315794279799,0.817645830567926,0.811867464100942,0.061435961863026,0.746087650069967,0.447872480610386,0.84255820303224,NA,1 --0.329040167848524,0.118530319537967,0.457290950231254,0.645931467181072,0.893649809062481,0.290868752868846,0.963444818975404,0.287232589442283,0.965041985502467,0.503362905001268,0.700591426342726,NA,0.5 -1.48193500973284,0.625231690006331,0.0553984474390745,0.649881738703698,0.299855779856443,0.496550764422864,0.941366601735353,0.854566443245858,0.833261088933796,0.97451618895866,0.931585984304547,NA,2 --0.173022349447719,0.164724288508296,0.869774119695649,0.643917094916105,0.35930937086232,0.131004156777635,0.581698802066967,0.991578063229099,0.646482613869011,0.973439331632107,0.251832834677771,NA,0.5 --0.798130578450863,0.413501289440319,0.191485705785453,0.66864021285437,0.944299905095249,0.781557427253574,0.426434542285278,0.155345749342814,0.47630034852773,0.368169168476015,0.250597581965849,NA,1 --0.574218404015675,0.242424165830016,0.58770715072751,0.515803768066689,0.381307470612228,0.00103132519870996,0.822859295178205,0.0108886968810111,0.230126534821466,0.557247446617112,0.905343563761562,NA,0.5 -1.20576697217052,0.481467565754429,0.384135053725913,0.250637572957203,0.0487530084792525,0.0894714090973139,0.907941385637969,0.130016501294449,0.108319341670722,0.00571742816828191,0.322211351711303,NA,1 -1.09418799497703,0.428736956324428,0.169671062380075,0.828763456083834,0.351603407878429,0.227558815153316,0.226003634510562,0.346984679810703,0.792112624039873,0.331346849445254,0.432414019247517,NA,1 -1.33295018743162,0.657001668587327,0.637372163590044,0.824361448409036,0.929858485469595,0.478111879900098,0.372787630185485,0.409159137168899,0.687029706314206,0.462824999121949,0.282154636690393,NA,2 -0.767515180489683,0.164000877644867,0.787021941738203,0.480641457485035,0.256527086952701,0.139046657830477,0.260776846436784,0.843071689130738,0.928756222594529,0.0495820019859821,0.623474525287747,NA,0.5 -5.91575131696734,0.968515180051327,0.366838250076398,0.0660362653434277,0.856167251477018,0.179680837085471,0.161809651181102,0.591778725851327,0.48686463618651,0.738822849933058,0.183114260667935,NA,3 --6.07690196330587,0.830175120383501,0.379230955149978,0.115117009729147,0.729175335494801,0.629532365594059,0.782002086285502,0.933919353410602,0.0596355667803437,0.724859771085903,0.652606250252575,NA,3 --0.313191580590458,0.433103942545131,0.486646147910506,0.0679896532092243,0.0214490725193173,0.0353633773047477,0.744490799959749,0.845955922966823,0.958894333336502,0.156010123435408,0.303289767587557,NA,1 --0.540923578326201,0.460490206023678,0.00578354438766837,0.550281241768971,0.788312985096127,0.210271368967369,0.815538283670321,0.28684966918081,0.275849351193756,0.469581529032439,0.647478682687506,NA,1 --0.673999243552475,0.570351638598368,0.868796721566468,0.776297046570107,0.343447036575526,0.135450173402205,0.647280380362645,0.0270581084769219,0.964931801892817,0.511921106604859,0.636638013413176,NA,2 -0.521317196587708,0.0333191037643701,0.283822370693088,0.192718317965046,0.948545307852328,0.824977292912081,0.370365360984579,0.515110278967768,0.214455303736031,0.224180643446743,0.711288265651092,NA,0.5 -0.464127896273477,0.150690599810332,0.122571487911046,0.888559450628236,0.731974716298282,0.0771893607452512,0.541291358415037,0.552816195879132,0.315732726827264,0.151362558128312,0.962228747783229,NA,0.5 -0.0280326256524612,0.218963563675061,0.667505384190008,0.811947909416631,0.687651550397277,0.0374882640317082,0.161606113892049,0.341489754617214,0.114423418417573,0.711193904979154,0.999109187629074,NA,0.5 -1.72545787431865,0.791984313167632,0.713207671185955,0.329182849498466,0.632728930795565,0.319493535207585,0.219288081862032,0.48391812061891,0.699484226293862,0.955756370443851,0.813189222943038,NA,3 -1.15486734783788,0.566656570881605,0.657263582805172,0.788281690096483,0.778249064926058,0.590916589833796,0.385298176435754,0.389463485917076,0.137170025613159,0.138196402462199,0.308245474472642,NA,2 --0.457464364994371,0.203198125585914,0.378648053156212,0.808898270130157,0.906652554869652,0.850757014239207,0.988703542388976,0.178554533980787,0.459032043116167,0.799958050018176,0.297242203727365,NA,0.5 -2.0642220969772,0.574996857671067,0.13869056943804,0.148214791202918,0.0713210175745189,0.281843064585701,0.699210103135556,0.611457530176267,0.412616135785356,0.312300917692482,0.753751856740564,NA,2 -1.79325254054518,0.779666494345292,0.699538435321301,0.961586186429486,0.347297596279532,0.0675467550754547,0.325070535531268,0.176488080061972,0.603379048407078,0.136019740719348,0.473342314828187,NA,3 --0.814168249805148,0.179700289387256,0.290472864639014,0.604973735287786,0.684216808062047,0.566976414760575,0.147398902801797,0.29908396420069,0.0406292479019612,0.475435412954539,0.30359836248681,NA,0.5 -1.05601002941817,0.974190507782623,0.59218838927336,0.234243621351197,0.808535494375974,0.187695320462808,0.772726303897798,0.996423880336806,0.286155195906758,0.0625078319571912,0.898713116999716,NA,3 --0.335678118214373,0.589294733013958,0.64874632912688,0.60927852592431,0.912218964891508,0.116446102270856,0.800978301558644,0.552153271622956,0.498927663313225,0.77924073045142,0.604476742446423,NA,2 --0.297220306313378,0.00104556139558554,0.243074154946953,0.28915421012789,0.584380616666749,0.680923626758158,0.190056581050158,0.489183940691873,0.543326803250238,0.462141246302053,0.0547776529565454,NA,0.5 -3.36381201170981,0.842413400299847,0.836705626221374,0.803220853442326,0.52237598830834,0.0978306015022099,0.44979960215278,0.623058009892702,0.0899901306256652,0.659136149799451,0.511611237656325,NA,3 --1.88385207395372,0.52834680210799,0.734877026872709,0.131855968618765,0.703349165152758,0.0538370700087398,0.573549640597776,0.794007008662447,0.812998212641105,0.404450054978952,0.43243193323724,NA,2 --1.73270176331594,0.779436560813338,0.115024698665366,0.0290894189383835,0.368165988940746,0.811495361849666,0.0700736299622804,0.615472083212808,0.257920050993562,0.869047976098955,0.349914842285216,NA,3 -1.37079153120135,0.570917901350185,0.187404554570094,0.398802037583664,0.846335949376225,0.934954999480397,0.529415836092085,0.910126697039232,0.118137815967202,0.111092046368867,0.248694343725219,NA,2 -0.571363043562507,0.331030112225562,0.450528397923335,0.336339824367315,0.16454178863205,0.430747980251908,0.505673100007698,0.328476364724338,0.21718189259991,0.586265194229782,0.447927470318973,NA,1 --1.79533836540329,0.976926734903827,0.0961627899669111,0.60899573401548,0.599624714581296,0.713024838129058,0.36587701109238,0.517652825685218,0.313225879799575,0.691182613139972,0.946866175159812,NA,3 -2.76102079506608,0.760612237034366,0.802744230488315,0.707125860266387,0.710110249230638,0.0485071069560945,0.337238987907767,0.603377134539187,0.396445233840495,0.336875891312957,0.0135681433603168,NA,3 -1.03587118074943,0.546186395455152,0.0895447288639843,0.102424295851961,0.124316787347198,0.398161188932136,0.0595820427406579,0.0317856874316931,0.711451485287398,0.184543622424826,0.570017953403294,NA,2 -2.64092013401424,0.813234215602279,0.346893396694213,0.999183400999755,0.163131871260703,0.109145125374198,0.742135232314467,0.926391132874414,0.354281723732129,0.518946023192257,0.143736189464107,NA,3 --2.14161378373943,0.623024221044034,0.987261440372095,0.240460366243497,0.0876346873119473,0.137773034861311,0.00308776297606528,0.764373373007402,0.797460315050557,0.443219825392589,0.516431802418083,NA,2 --0.193411209224147,0.0207313429564238,0.243699158774689,0.459949440322816,0.262547085061669,0.695130342617631,0.946504537714645,0.726618548389524,0.864191201748326,0.815897560212761,0.0842976626008749,NA,0.5 -1.04927342445141,0.472177557414398,0.0301995284389704,0.293144194176421,0.790906876791269,0.438358291517943,0.794757833238691,0.85613635671325,0.811791674233973,0.112252983730286,0.169568839715794,NA,1 -0.348988455997724,0.262080669403076,0.927174943732098,0.142983013764024,0.640390169108286,0.366609850898385,0.529223524965346,0.57928417250514,0.206402273848653,0.268847646424547,0.165483654243872,NA,1 -0.696440990671028,0.983129711588845,0.654424325563014,0.83093448006548,0.817050429061055,0.692395427962765,0.485632661730051,0.661732462001964,0.0443739434704185,0.631122295744717,0.821382548194379,NA,3 --1.50946435977094,0.524429301731288,0.394003928638995,0.612721437122673,0.641075712163001,0.8340331655927,0.81034871423617,0.378355852561072,0.118066052906215,0.0152956454548985,0.0607850980013609,NA,2 --3.60722631045085,0.973684054333717,0.862753849243745,0.470457507064566,0.516427976312116,0.571779370773584,0.00344765046611428,0.557507778517902,0.863539802376181,0.618254639208317,0.0438256817869842,NA,3 -0.640760110283842,0.3505244248081,0.388575355056673,0.525333723518997,0.90868836059235,0.656305950833485,0.87575338082388,0.145426131319255,0.142402743920684,0.722210387466475,0.646311188582331,NA,1 -1.0927036128643,0.0793833178468049,0.0104481617454439,0.93192859296687,0.779025187250227,0.358062380226329,0.498329946072772,0.557266015792266,0.794099885271862,0.103802968282253,0.489988661836833,NA,0.5 --0.960320030022167,0.231456351932138,0.585931064095348,0.872712701093405,0.684092679759488,0.747057758970186,0.33389916899614,0.693479184061289,0.234955767868087,0.0913671234156936,0.484689482487738,NA,0.5 -5.55515128084404,0.961973294848576,0.598567689536139,0.797411441337317,0.749971851240844,0.626418237108737,0.475173852639273,0.763921497855335,0.795191584387794,0.464224591152743,0.356469079386443,NA,3 -0.222792689673181,0.216883212793618,0.504161356948316,0.685418124077842,0.138950776075944,0.609323340700939,0.0815524500794709,0.887275154469535,0.673021834110841,0.229804648784921,0.545686026103795,NA,0.5 -0.493718034822244,0.11028397642076,0.503047540783882,0.179511502385139,0.763330884743482,0.77220833580941,0.934788197511807,0.494932206347585,0.731283217668533,0.855659510940313,0.0117787455674261,NA,0.5 -0.461772257413767,0.29381926660426,0.762450747890398,0.234901242889464,0.5249476833269,0.629562719725072,0.200397627661005,0.684597495477647,0.356038145022467,0.802739176666364,0.600334402406588,NA,1 --0.126129491991506,0.463731594849378,0.167885933071375,0.540383089799434,0.282209430588409,0.0468776472844183,0.37753825658001,0.563547776313499,0.0574275727849454,0.122192964889109,0.144121702294797,NA,1 -1.61069779366059,0.353434937307611,0.44708115211688,0.562454687664285,0.681164643960074,0.344466916518286,0.389643874252215,0.47427371609956,0.420238852966577,0.246512484736741,0.0113735438790172,NA,1 -0.455525567857585,0.656678552739322,0.520737201906741,0.440892208134755,0.383611467434093,0.853515934199095,0.312850524671376,0.188586953794584,0.611132161458954,0.416548165027052,0.966824663570151,NA,2 -0.752414003343478,0.368973962496966,0.144383963663131,0.298857402056456,0.603884643642232,0.948024961864576,0.0727454784791917,0.0918680941686034,0.0495259300805628,0.778334744274616,0.237199151189998,NA,1 --1.8071186112866,0.330021998845041,0.0695072312373668,0.405746232718229,0.476144076324999,0.162464414490387,0.143160074716434,0.2627956320066,0.371168155921623,0.219870826462284,0.593052309006453,NA,1 -0.338110168780804,0.0405488722026348,0.971289380686358,0.55798068922013,0.217713665682822,0.626417078776285,0.498581673717126,0.377851612633094,0.218278390355408,0.651240651495755,0.95426289155148,NA,0.5 --3.08507748542801,0.905532934004441,0.329712072620168,0.140193846542388,0.36389043321833,0.374093280639499,0.146446174709126,0.438101051608101,0.795116507913917,0.983287858776748,0.488525949651375,NA,3 --1.72372688590346,0.696682563750073,0.349387434776872,0.505626027705148,0.826131088426337,0.336876499466598,0.579624563921243,0.374704758403823,0.0984764297027141,0.158775766147301,0.86363190994598,NA,2 -0.613683233774074,0.0651458308566362,0.581464517861605,0.907430648803711,0.798097816295922,0.125102420337498,0.432380372891203,0.18545747292228,0.108076399890706,0.328964518383145,0.317265540128574,NA,0.5 --0.133225601541877,0.859262556303293,0.671681129839271,0.68166439072229,0.624605974880978,0.0363332864362746,0.559169779764488,0.0110217127948999,0.197365242987871,0.235231575090438,0.954349690815434,NA,3 -0.174705914799019,0.0484868276398629,0.72244808007963,0.412865192163736,0.177109800977632,0.49629073496908,0.43845453299582,0.883392751449719,0.354612664552405,0.244716308778152,0.773190965875983,NA,0.5 --2.11705848424155,0.688459891825914,0.221930157393217,0.0272207804955542,0.0866714252624661,0.0403791540302336,0.440649472177029,0.0105264887679368,0.100412974366918,0.790013117250055,0.270158890169114,NA,2 -1.1060152777824,0.375678987475112,0.924136486602947,0.474893233040348,0.824483483564109,0.0317611857317388,0.0541152430232614,0.202915545087308,0.701533750165254,0.888495072023943,0.528801799751818,NA,1 --1.40484738126697,0.914472201606259,0.389121263055131,0.115709668258205,0.70249339309521,0.88396455720067,0.0478117982856929,0.939067059895024,0.574597767787054,0.756754127331078,0.176177548244596,NA,3 --1.22766612067787,0.357021145289764,0.348640100797638,0.841263998765498,0.261813498334959,0.820658108917996,0.24100551684387,0.458479619817808,0.121229101205245,0.780505152186379,0.284189038909972,NA,1 --0.137006680090441,0.0419280277565122,0.456591938855127,0.769418158335611,0.3475969475694,0.325503826374188,0.268078948371112,0.801353276707232,0.0109000348020345,0.909261407796293,0.970131604699418,NA,0.5 --0.0472805089628641,0.458404157776386,0.509212339995429,0.472755545051768,0.89819523692131,0.00172450905665755,0.195178902009502,0.491253996500745,0.602235041558743,0.000361819053068757,0.810925672529265,NA,1 --6.69280752542676,0.794113770360127,0.283597531029955,0.415720561984926,0.386682669632137,0.543149482225999,0.528418547241017,0.59079469111748,0.82328638038598,0.7218906681519,0.762335718609393,NA,3 --0.791993317715144,0.94858734565787,0.862443518359214,0.336455043870956,0.178935058414936,0.876441803527996,0.886401900788769,0.0696042117197067,0.201479215407744,0.579562027938664,0.916529270121828,NA,3 --3.5798937723398,0.518584945704788,0.341426242142916,0.980938139371574,0.025910884141922,0.703400827012956,0.611595463240519,0.540800383547321,0.108336265431717,0.802920109592378,0.717154868878424,NA,2 -1.34179692842582,0.256257413187996,0.137025708565488,0.189194887410849,0.356893142918125,0.456141719361767,0.802874229149893,0.286960861412808,0.539246008498594,0.623640240402892,0.202499657869339,NA,1 -0.574692809911355,0.367721185088158,0.189689038787037,0.650288314092904,0.916537459241226,0.189138949615881,0.176210168516263,0.626079858280718,0.105365733848885,0.0375174721702933,0.556417517596856,NA,1 -0.291989692067655,0.374401310924441,0.173651255434379,0.666860758094117,0.863264636369422,0.927415754180402,0.457678058184683,0.499020219780505,0.751850564265624,0.987137568416074,0.329214916797355,NA,1 --1.0948854251991,0.480366722913459,0.382694880710915,0.93061161483638,0.545312937349081,0.46468991599977,0.467403194168583,0.945679706521332,0.918626469327137,0.6478590965271,0.0622258309740573,NA,1 -0.549466325611295,0.832923032343388,0.586885153781623,0.856621402082965,0.0974842766299844,0.193075296701863,0.797113571548834,0.331570270936936,0.681197754340246,0.456899536773562,0.216078756842762,NA,3 --1.18444309638044,0.987224767683074,0.0441201773937792,0.150253226514906,0.924849382368848,0.243760817451403,0.601127431960776,0.941083330428228,0.0376720677595586,0.590444662142545,0.594265560386702,NA,3 --2.6287892127699,0.511187988799065,0.17760677007027,0.0900412548799068,0.707964851288125,0.772246306762099,0.0507865459658206,0.880909067578614,0.565439280355349,0.325563751393929,0.625947890337557,NA,2 diff --git a/include/stochtree/bart.h b/include/stochtree/bart.h new file mode 100644 index 00000000..dc4efb45 --- /dev/null +++ b/include/stochtree/bart.h @@ -0,0 +1,337 @@ +/*! + * Copyright (c) 2026 stochtree authors. All rights reserved. + * Licensed under the MIT License. See LICENSE file in the project root for license information. + */ +#ifndef STOCHTREE_BART_H_ +#define STOCHTREE_BART_H_ + +#include +#include +#include +#include +#include +#include + +namespace StochTree { + +enum class BARTRFXModelSpec { + Custom, + InterceptOnly +}; + +struct BARTData { + // Train set covariates + double* X_train = nullptr; + int n_train = 0; + int p = 0; + + // Test set covariates + double* X_test = nullptr; + int n_test = 0; + + // Train set outcome + double* y_train = nullptr; + + // Basis for leaf regression + double* basis_train = nullptr; + double* basis_test = nullptr; + int basis_dim = 0; + + // Observation weights + double* obs_weights_train = nullptr; + double* obs_weights_test = nullptr; + + // Random effects + int* rfx_group_ids_train = nullptr; + int* rfx_group_ids_test = nullptr; + double* rfx_basis_train = nullptr; + double* rfx_basis_test = nullptr; + int rfx_num_groups = 0; + int rfx_basis_dim = 0; +}; + +struct BARTConfig { + // High level parameters + bool standardize_outcome = true; // whether to standardize the outcome before fitting and unstandardize predictions after + int num_threads = 1; // number of threads to use for sampling + bool verbose = false; // whether to print sampler progress to the console + int cutpoint_grid_size = 100; // number of cutpoints to consider for each covariate when sampling splits + std::vector feature_types; // feature types for each covariate (should be same length as number of covariates in the dataset), where 0 = continuous, 1 = categorical + LinkFunction link_function = LinkFunction::Identity; // link function to use (Identity, Probit, Cloglog) + OutcomeType outcome_type = OutcomeType::Continuous; // type of the outcome variable (Continuous, Binary, Ordinal) + int random_seed = -1; // random seed for reproducibility (if negative, a random seed will be generated) + bool keep_gfr = true; // whether or not to keep GFR samples or simply use them to warm-start an MCMC chain + bool keep_burnin = false; // whether or not to keep "burn-in" MCMC samples (largely a debugging flag) + + // Global error variance parameters + double a_sigma2_global = 0.0; // shape parameter for inverse gamma prior on global error variance + double b_sigma2_global = 0.0; // scale parameter for inverse gamma prior on global error variance + double sigma2_global_init = 1.0; // initial value for global error variance + bool sample_sigma2_global = true; // whether to sample global error variance (if false, it will be fixed at sigma2_global_init) + + // Mean forest parameters + int num_trees_mean = 200; // number of trees in the mean forest + double alpha_mean = 0.95; // alpha parameter for mean forest tree prior + double beta_mean = 2.0; // beta parameter for mean forest tree prior + int min_samples_leaf_mean = 5; // minimum number of samples per leaf for mean forest + int max_depth_mean = -1; // maximum depth for mean forest trees (-1 means no maximum) + bool leaf_constant_mean = true; // whether to use constant leaf model for mean forest + int leaf_dim_mean = 1; // dimension of the leaf for mean forest + bool exponentiated_leaf_mean = false; // whether to exponentiate leaf predictions for mean forest + int num_features_subsample_mean = 0; // number of features to subsample for each mean forest split (0 means no subsampling) + double a_sigma2_mean = 3.0; // shape parameter for inverse gamma prior on mean forest leaf scale + double b_sigma2_mean = -1.0; // scale parameter for inverse gamma prior on mean forest leaf scale (-1 is a sentinel value that triggers a data-informed calibration based on the variance of the outcome and the number of trees) + double sigma2_mean_init = -1.0; // initial value of mean forest leaf scale (-1 is a sentinel value that triggers a data-informed calibration based on the variance of the outcome and the number of trees) + std::vector var_weights_mean; // variable weights for mean forest splits (should be same length as number of covariates in the dataset) + std::vector sigma2_leaf_mean_matrix; // prior covariance matrix Sigma_0 for multivariate leaf regression, stored column-major (size leaf_dim_mean^2); empty = use sigma2_mean_init * I + bool sample_sigma2_leaf_mean = false; // whether to sample mean forest leaf scale (if false, it will be fixed at sigma2_mean_init) + std::vector sweep_update_indices_mean; // indices of trees to update in a given sweep (should be subset of [0, num_trees - 1]) + MeanLeafModelType mean_leaf_model_type = MeanLeafModelType::GaussianConstant; // leaf model type for mean forest + int num_classes_cloglog = 0; // number of classes for cloglog ordinal leaf model (should be set if mean_leaf_model_type = CloglogOrdinal) + double cloglog_leaf_prior_shape = 2.0; // shape parameter for cloglog ordinal leaf model prior + double cloglog_leaf_prior_scale = 2.0; // scale parameter for cloglog ordinal leaf model prior + double cloglog_cutpoint_0 = 0.0; // Fixed value of the first log-scale cutpoint for the cloglog model (defaults to 0 for identifiability) + + // Variance forest parameters + int num_trees_variance = 0; // number of trees in the variance forest + double leaf_prior_calibration_param = 1.5; // calibration parameter for variance forest leaf prior + double shape_variance_forest = -1.0; // shape parameter for variance forest leaf model (calibrated internally based on leaf_prior_calibration_param if set to sentinel value of -1) + double scale_variance_forest = -1.0; // scale parameter for variance forest leaf model (calibrated internally based on leaf_prior_calibration_param if set to sentinel value of -1) + double variance_forest_leaf_init = -1.0; // initial (raw-scale) root value for the variance forest; each leaf starts at log(value)/num_trees_variance. Sentinel <= 0 = calibrate internally + double alpha_variance = 0.5; // alpha parameter for variance forest tree prior + double beta_variance = 2.0; // beta parameter for variance forest tree prior + int min_samples_leaf_variance = 5; // minimum number of samples per leaf for variance forest + int max_depth_variance = -1; // maximum depth for variance forest trees (-1 means no maximum) + bool leaf_constant_variance = true; // whether to use constant leaf model for variance forest + int leaf_dim_variance = 1; // dimension of the leaf for variance forest (should be 1 if leaf_constant_variance=true) + bool exponentiated_leaf_variance = true; // whether to exponentiate leaf predictions for variance forest + int num_features_subsample_variance = 0; // number of features to subsample for each variance forest split (0 means no subsampling) + std::vector var_weights_variance; // variable weights for variance forest splits (should be same length as number of covariates in the dataset) + std::vector sweep_update_indices_variance; // indices of trees to update in a given sweep (should be subset of [0, num_trees - 1]) + + // Random effects parameters + bool has_random_effects = false; // whether or not a model includes a random effects term + BARTRFXModelSpec rfx_model_spec = BARTRFXModelSpec::Custom; // specification for the random effects model; custom relies on a user-provided basis while intercept-only constructs a varying intercept model without needing a user-provided basis + std::vector rfx_working_parameter_mean_prior; // vector of dimension num_basis; empty = use zeros + std::vector rfx_group_parameter_mean_prior; // matrix of dimension num_basis x num_groups, stored column-major; empty = use zeros + std::vector rfx_working_parameter_cov_prior; // matrix of dimension num_basis x num_basis, stored column-major; empty = use identity matrix + std::vector rfx_group_parameter_cov_prior; // matrix of dimension num_basis x num_basis, stored column-major; empty = use identity matrix + double rfx_variance_prior_shape = 1.0; // shape parameter for variance prior in random effects model + double rfx_variance_prior_scale = 1.0; // scale parameter for variance prior in random effects model + + // TODO: Other parameters ... +}; + +struct BARTSamples { + // Posterior samples of training set mean forest predictions (num_samples x n_train, stored column-major) + std::vector mean_forest_predictions_train; + + // Posterior samples of training set variance forest predictions (num_samples x n_train, stored column-major) + std::vector variance_forest_predictions_train; + + // Posterior samples of test set mean forest predictions (num_samples x n_test, stored column-major) + std::vector mean_forest_predictions_test; + + // Posterior samples of test set variance forest predictions (num_samples x n_test, stored column-major) + std::vector variance_forest_predictions_test; + + // Posterior samples of global error variance (num_samples) + std::vector global_error_variance_samples; + + // Posterior samples of leaf scale (num_samples) + std::vector leaf_scale_samples; + + // Pointer to sampled mean forests + std::unique_ptr mean_forests; + + // Pointer to sampled variance forests + std::unique_ptr variance_forests; + + // Posterior samples of cloglog cutpoint parameters (num_samples x num_classes - 1, stored column-major) + std::vector cloglog_cutpoint_samples; + + // Posterior samples of training set RFX predictions (num_samples x n_train, stored column-major) + std::vector rfx_predictions_train; + + // Posterior samples of test set RFX predictions (num_samples x n_test, stored column-major) + std::vector rfx_predictions_test; + + // Pointer to random effects sample container and label mapping + std::unique_ptr rfx_container; + std::unique_ptr rfx_label_mapper; + + // Metadata about the samples (e.g., number of samples, burn-in, etc.) could be added here as needed + int num_samples = 0; + int num_train = 0; + int num_test = 0; + double y_bar = 0.0; + double y_std = 0.0; + + void AppendToJson(nlohmann::json& obj) const { + // Forests, under self-describing named keys, with the num_forests counter + nlohmann::json forests = nlohmann::json::object(); + int num_forests = 0; + if (mean_forests != nullptr) { + forests.emplace("mean_forest", mean_forests->to_json()); + num_forests++; + } + if (variance_forests != nullptr) { + forests.emplace("variance_forest", variance_forests->to_json()); + num_forests++; + } + // NOTE: use operator[] assignment (not emplace) for the top-level keys below. The R CppJson + // envelope pre-creates empty "forests"/"random_effects"/"num_forests"/"num_random_effects" + // keys, and nlohmann::json::emplace is a no-op when the key already exists (like std::map), + // which would silently drop the forests/rfx. Assignment overwrites whether or not the key exists. + obj["forests"] = forests; + obj["num_forests"] = num_forests; + // Parameter traces, under the "parameters" subfolder (presence inferred from non-empty vectors) + nlohmann::json parameters = nlohmann::json::object(); + if (!global_error_variance_samples.empty()) { + parameters.emplace("sigma2_global_samples", global_error_variance_samples); + } + if (!leaf_scale_samples.empty()) { + parameters.emplace("sigma2_leaf_samples", leaf_scale_samples); + } + if (!cloglog_cutpoint_samples.empty()) { + parameters.emplace("cloglog_cutpoint_samples", cloglog_cutpoint_samples); + } + if (!parameters.empty()) { + obj["parameters"] = parameters; + } + // Intrinsic scalars (stored in user-facing scale, matching the existing wire format) + obj["outcome_mean"] = y_bar; + obj["outcome_scale"] = y_std; + obj["num_samples"] = num_samples; + // Random effects + int num_random_effects = 0; + nlohmann::json rfx = nlohmann::json::object(); + if (rfx_container != nullptr && rfx_label_mapper != nullptr) { + rfx.emplace("random_effect_container_0", rfx_container->to_json()); + rfx.emplace("random_effect_label_mapper_0", rfx_label_mapper->to_json()); + rfx.emplace("random_effect_groupids_0", rfx_label_mapper->Keys()); + num_random_effects = 1; + } + obj["random_effects"] = rfx; + obj["num_random_effects"] = num_random_effects; + } + + // Populate this BARTSamples from a parsed JSON object + void FromJson(const nlohmann::json& obj) { + // Unpack forests if present, checking for the expected keys + if (obj.contains("forests")) { + const nlohmann::json& forests = obj.at("forests"); + if (forests.contains("mean_forest")) { + mean_forests = std::make_unique(0, 0, false, false); + mean_forests->from_json(forests.at("mean_forest")); + } + if (forests.contains("variance_forest")) { + variance_forests = std::make_unique(0, 0, false, false); + variance_forests->from_json(forests.at("variance_forest")); + } + } + // Unpack parameters if present, checking for expected keys + if (obj.contains("parameters")) { + const nlohmann::json& parameters = obj.at("parameters"); + if (parameters.contains("sigma2_global_samples")) { + global_error_variance_samples = parameters.at("sigma2_global_samples").get>(); + } + if (parameters.contains("sigma2_leaf_samples")) { + leaf_scale_samples = parameters.at("sigma2_leaf_samples").get>(); + } + if (parameters.contains("cloglog_cutpoint_samples")) { + cloglog_cutpoint_samples = parameters.at("cloglog_cutpoint_samples").get>(); + } + } + // Unpack random effects if present, checking for expected keys + if (obj.contains("num_random_effects") && obj.at("num_random_effects").get() > 0) { + rfx_container = std::make_unique(); + rfx_label_mapper = std::make_unique(); + rfx_container->from_json(obj.at("random_effects").at("random_effect_container_0")); + rfx_label_mapper->from_json(obj.at("random_effects").at("random_effect_label_mapper_0")); + } + // Unpack outcome statistics + if (obj.contains("outcome_mean")) y_bar = obj.at("outcome_mean").get(); + if (obj.contains("outcome_scale")) y_std = obj.at("outcome_scale").get(); + if (obj.contains("num_samples")) num_samples = obj.at("num_samples").get(); + } + + // Append another chain's draws onto this one (multi-chain combine). `this` must already be + // populated (e.g. via FromJson on the first chain); `other` must have matching model structure + // (same forests present, same outcome standardization). Forests are deep-copied sample-by-sample + // and parameter traces concatenated, so draw order is preserved (this's draws, then other's). + void Merge(const BARTSamples& other) { + // Runtime checks for samples objects to be combined + if (y_bar != other.y_bar || y_std != other.y_std) { + Log::Fatal("Cannot merge BARTSamples with different outcome standardization"); + } + if (rfx_container != nullptr && other.rfx_container != nullptr) { + if (rfx_container->NumComponents() != other.rfx_container->NumComponents() || + rfx_container->NumGroups() != other.rfx_container->NumGroups()) { + Log::Fatal("Cannot merge BARTSamples with different random effects structure"); + } + if (rfx_label_mapper->Keys() != other.rfx_label_mapper->Keys()) { + Log::Fatal("Cannot merge BARTSamples with different random effects label mapping"); + } + if (rfx_label_mapper->Map() != other.rfx_label_mapper->Map()) { + Log::Fatal("Cannot merge BARTSamples with different random effects label mapping"); + } + } + // Append forests if they exist in the samples object + AppendForestContainerSamples(mean_forests, other.mean_forests, "mean"); + AppendForestContainerSamples(variance_forests, other.variance_forests, "variance"); + // Append random effects if they exist in the samples object + // Note that the LabelMapper is not sample-specific so we do not need to append to it + AppendRandomEffectsContainerSamples(rfx_container, other.rfx_container); + // Append parameter samples + if ((!global_error_variance_samples.empty() && other.global_error_variance_samples.empty()) || + (global_error_variance_samples.empty() && !other.global_error_variance_samples.empty())) { + Log::Fatal("Cannot merge BARTSamples objects: global error variance samples present in one chain but not the other"); + } + global_error_variance_samples.insert(global_error_variance_samples.end(), + other.global_error_variance_samples.begin(), other.global_error_variance_samples.end()); + if ((!leaf_scale_samples.empty() && other.leaf_scale_samples.empty()) || + (leaf_scale_samples.empty() && !other.leaf_scale_samples.empty())) { + Log::Fatal("Cannot merge BARTSamples objects: leaf scale samples present in one chain but not the other"); + } + leaf_scale_samples.insert(leaf_scale_samples.end(), + other.leaf_scale_samples.begin(), other.leaf_scale_samples.end()); + if ((!cloglog_cutpoint_samples.empty() && other.cloglog_cutpoint_samples.empty()) || + (cloglog_cutpoint_samples.empty() && !other.cloglog_cutpoint_samples.empty())) { + Log::Fatal("Cannot merge BARTSamples objects: cloglog cutpoint samples present in one chain but not the other"); + } + cloglog_cutpoint_samples.insert(cloglog_cutpoint_samples.end(), + other.cloglog_cutpoint_samples.begin(), other.cloglog_cutpoint_samples.end()); + num_samples += other.num_samples; + } + + std::vector OutcomePredictionsTrain() const { + std::vector predictions(num_train * num_samples, 0.0); + for (int i = 0; i < num_train * num_samples; ++i) { + if (!mean_forest_predictions_train.empty()) { + predictions[i] += mean_forest_predictions_train[i]; + } + if (!rfx_predictions_train.empty()) { + predictions[i] += rfx_predictions_train[i]; // Add random effects contribution if present + } + } + return predictions; + } + + std::vector OutcomePredictionsTest() const { + std::vector predictions(num_test * num_samples, 0.0); + for (int i = 0; i < num_test * num_samples; ++i) { + if (!mean_forest_predictions_test.empty()) { + predictions[i] += mean_forest_predictions_test[i]; + } + if (!rfx_predictions_test.empty()) { + predictions[i] += rfx_predictions_test[i]; // Add random effects contribution if present + } + } + return predictions; + } +}; + +} // namespace StochTree + +#endif // STOCHTREE_BART_H_ diff --git a/include/stochtree/bart_sampler.h b/include/stochtree/bart_sampler.h new file mode 100644 index 00000000..fb41742c --- /dev/null +++ b/include/stochtree/bart_sampler.h @@ -0,0 +1,390 @@ +/*! + * Copyright (c) 2026 stochtree authors. All rights reserved. + * Licensed under the MIT License. See LICENSE file in the project root for license information. + */ +#ifndef STOCHTREE_BART_SAMPLER_H_ +#define STOCHTREE_BART_SAMPLER_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace StochTree { + +class BARTSampler { + public: + // If continuation is true, the sampler warm-starts from the last retained sample + // already present in `samples` (rather than initializing forests to root), and the + // existing forest container is preserved so that new samples are appended to it. + BARTSampler(BARTSamples& samples, BARTConfig& config, BARTData& data, bool continuation = false); + + // Main entry point for running the BART GFR sampler + // If num_chains > 0, captures snapshots of the last num_chains GFR states for fork_chains() + void run_gfr(BARTSamples& samples, int num_gfr, bool keep_gfr, int num_chains = 0); + + // Run a single chain of the BART MCMC sampler + void run_mcmc(BARTSamples& samples, int num_burnin, int keep_every, int num_mcmc); + + // Run num_chains independent MCMC chains sequentually based on GFR snapshots captured by run_gfr() or re-initialized from root + void run_mcmc_chains(BARTSamples& samples, int num_chains, int num_burnin, int keep_every, int num_mcmc); + + // Post-process samples by extracting test set predictions and running any necessary transformations + void postprocess_samples(BARTSamples& samples); + + // Serialize the internal RNG state to a string. std::mt19937 round-trips losslessly through + // its stream operators, so this captures the exact position in the random stream -- used to + // persist RNG state across a sample() / continue_sampling() boundary for bit-identical results. + std::string GetRngState() const { + std::ostringstream oss; + oss << rng_; + return oss.str(); + } + + // Restore the internal RNG state from a string produced by GetRngState(). Resumes the random + // stream at exactly the captured position. Must be called after construction (InitializeState + // unconditionally re-seeds rng_) and before any sampling draws. + void SetRngState(const std::string& state) { + std::istringstream iss(state); + iss >> rng_; + } + + private: + /*! Initialize state variables */ + void InitializeState(BARTSamples& samples, bool continuation = false); + bool initialized_ = false; + + /*! Internal function to restore sampler state based on a GFR snapshot */ + void RestoreStateFromGFRSnapshot(BARTSamples& samples, int snapshot_index); + + /*! Internal function to restore sampler state to root / initial values */ + void RestoreStateDefault(); + + /*! Internal sample runner function */ + void RunOneIteration(BARTSamples& samples, bool gfr, bool keep_sample, bool write_snapshot = false); + + /*! Initialization visitor */ + struct MeanForestInitVisitor { + BARTSampler& sampler; + BARTSamples& samples; + void operator()(GaussianConstantLeafModel& model) { + sampler.mean_forest_ = std::make_unique(sampler.config_.num_trees_mean, sampler.config_.leaf_dim_mean, sampler.config_.leaf_constant_mean, sampler.config_.exponentiated_leaf_mean); + samples.mean_forests = std::make_unique(sampler.config_.num_trees_mean, sampler.config_.leaf_dim_mean, sampler.config_.leaf_constant_mean, sampler.config_.exponentiated_leaf_mean); + sampler.mean_forest_tracker_ = std::make_unique(sampler.forest_dataset_->GetCovariates(), sampler.config_.feature_types, sampler.config_.num_trees_mean, sampler.data_.n_train); + sampler.tree_prior_mean_ = std::make_unique(sampler.config_.alpha_mean, sampler.config_.beta_mean, sampler.config_.min_samples_leaf_mean, sampler.config_.max_depth_mean); + sampler.mean_forest_->SetLeafValue(sampler.init_val_mean_ / sampler.config_.num_trees_mean); + UpdateResidualEntireForest(*sampler.mean_forest_tracker_, *sampler.forest_dataset_, *sampler.residual_, sampler.mean_forest_.get(), !sampler.config_.leaf_constant_mean, std::minus()); + sampler.mean_forest_tracker_->UpdatePredictions(sampler.mean_forest_.get(), *sampler.forest_dataset_.get()); + sampler.has_mean_forest_ = true; + } + void operator()(GaussianUnivariateRegressionLeafModel& model) { + sampler.mean_forest_ = std::make_unique(sampler.config_.num_trees_mean, sampler.config_.leaf_dim_mean, sampler.config_.leaf_constant_mean, sampler.config_.exponentiated_leaf_mean); + samples.mean_forests = std::make_unique(sampler.config_.num_trees_mean, sampler.config_.leaf_dim_mean, sampler.config_.leaf_constant_mean, sampler.config_.exponentiated_leaf_mean); + sampler.mean_forest_tracker_ = std::make_unique(sampler.forest_dataset_->GetCovariates(), sampler.config_.feature_types, sampler.config_.num_trees_mean, sampler.data_.n_train); + sampler.tree_prior_mean_ = std::make_unique(sampler.config_.alpha_mean, sampler.config_.beta_mean, sampler.config_.min_samples_leaf_mean, sampler.config_.max_depth_mean); + sampler.mean_forest_->SetLeafValue(sampler.init_val_mean_ / sampler.config_.num_trees_mean); + UpdateResidualEntireForest(*sampler.mean_forest_tracker_, *sampler.forest_dataset_, *sampler.residual_, sampler.mean_forest_.get(), !sampler.config_.leaf_constant_mean, std::minus()); + sampler.mean_forest_tracker_->UpdatePredictions(sampler.mean_forest_.get(), *sampler.forest_dataset_.get()); + sampler.has_mean_forest_ = true; + } + void operator()(GaussianMultivariateRegressionLeafModel& model) { + sampler.mean_forest_ = std::make_unique(sampler.config_.num_trees_mean, sampler.config_.leaf_dim_mean, sampler.config_.leaf_constant_mean, sampler.config_.exponentiated_leaf_mean); + samples.mean_forests = std::make_unique(sampler.config_.num_trees_mean, sampler.config_.leaf_dim_mean, sampler.config_.leaf_constant_mean, sampler.config_.exponentiated_leaf_mean); + sampler.mean_forest_tracker_ = std::make_unique(sampler.forest_dataset_->GetCovariates(), sampler.config_.feature_types, sampler.config_.num_trees_mean, sampler.data_.n_train); + sampler.tree_prior_mean_ = std::make_unique(sampler.config_.alpha_mean, sampler.config_.beta_mean, sampler.config_.min_samples_leaf_mean, sampler.config_.max_depth_mean); + sampler.mean_forest_->SetLeafVector(sampler.init_val_mean_vec_); + UpdateResidualEntireForest(*sampler.mean_forest_tracker_, *sampler.forest_dataset_, *sampler.residual_, sampler.mean_forest_.get(), true, std::minus()); + sampler.mean_forest_tracker_->UpdatePredictions(sampler.mean_forest_.get(), *sampler.forest_dataset_.get()); + sampler.has_mean_forest_ = true; + } + void operator()(CloglogOrdinalLeafModel& model) { + sampler.mean_forest_ = std::make_unique(sampler.config_.num_trees_mean, sampler.config_.leaf_dim_mean, sampler.config_.leaf_constant_mean, sampler.config_.exponentiated_leaf_mean); + samples.mean_forests = std::make_unique(sampler.config_.num_trees_mean, sampler.config_.leaf_dim_mean, sampler.config_.leaf_constant_mean, sampler.config_.exponentiated_leaf_mean); + sampler.mean_forest_tracker_ = std::make_unique(sampler.forest_dataset_->GetCovariates(), sampler.config_.feature_types, sampler.config_.num_trees_mean, sampler.data_.n_train); + sampler.tree_prior_mean_ = std::make_unique(sampler.config_.alpha_mean, sampler.config_.beta_mean, sampler.config_.min_samples_leaf_mean, sampler.config_.max_depth_mean); + sampler.mean_forest_->SetLeafValue(sampler.init_val_mean_ / sampler.config_.num_trees_mean); + UpdateResidualEntireForest(*sampler.mean_forest_tracker_, *sampler.forest_dataset_, *sampler.residual_, sampler.mean_forest_.get(), false, std::minus()); + sampler.mean_forest_tracker_->UpdatePredictions(sampler.mean_forest_.get(), *sampler.forest_dataset_.get()); + sampler.has_mean_forest_ = true; + } + }; + + /*! Initialization visitor */ + struct MeanForestResetVisitor { + BARTSampler& sampler; + BARTSamples& samples; + TreeEnsemble& forest; + void operator()(GaussianConstantLeafModel& model) { + sampler.mean_forest_->ReconstituteFromForest(forest); + sampler.mean_forest_tracker_->ReconstituteFromForest(forest, *sampler.forest_dataset_, *sampler.residual_, true); + sampler.mean_forest_tracker_->UpdatePredictions(sampler.mean_forest_.get(), *sampler.forest_dataset_.get()); + } + void operator()(GaussianUnivariateRegressionLeafModel& model) { + sampler.mean_forest_->ReconstituteFromForest(forest); + sampler.mean_forest_tracker_->ReconstituteFromForest(forest, *sampler.forest_dataset_, *sampler.residual_, true); + sampler.mean_forest_tracker_->UpdatePredictions(sampler.mean_forest_.get(), *sampler.forest_dataset_.get()); + } + void operator()(GaussianMultivariateRegressionLeafModel& model) { + sampler.mean_forest_->ReconstituteFromForest(forest); + sampler.mean_forest_tracker_->ReconstituteFromForest(forest, *sampler.forest_dataset_, *sampler.residual_, true); + sampler.mean_forest_tracker_->UpdatePredictions(sampler.mean_forest_.get(), *sampler.forest_dataset_.get()); + } + void operator()(CloglogOrdinalLeafModel& model) { + sampler.mean_forest_->ReconstituteFromForest(forest); + sampler.mean_forest_tracker_->ReconstituteFromForest(forest, *sampler.forest_dataset_, *sampler.residual_, true); + sampler.mean_forest_tracker_->UpdatePredictions(sampler.mean_forest_.get(), *sampler.forest_dataset_.get()); + } + }; + + /*! Continued-sampling initialization visitor. + * Warm-starts the active mean forest from the last retained sample in the + * existing container (rather than from root), and intentionally does NOT + * re-create samples.mean_forests so that new draws are appended to it. + * Only the basic Gaussian (Identity-link) leaf models are supported for now; + * other leaf models hard-error until their continuation paths are implemented. */ + struct MeanForestContinuationInitVisitor { + BARTSampler& sampler; + BARTSamples& samples; + void WarmStart(bool expand_basis) { + // Step 1: standard root init (mirrors MeanForestInitVisitor) so the forest, + // tracker, and residual reach a consistent state -- but DO NOT re-create + // samples.mean_forests, so continuation appends to the existing container. + sampler.mean_forest_ = std::make_unique(sampler.config_.num_trees_mean, sampler.config_.leaf_dim_mean, sampler.config_.leaf_constant_mean, sampler.config_.exponentiated_leaf_mean); + sampler.mean_forest_tracker_ = std::make_unique(sampler.forest_dataset_->GetCovariates(), sampler.config_.feature_types, sampler.config_.num_trees_mean, sampler.data_.n_train); + sampler.tree_prior_mean_ = std::make_unique(sampler.config_.alpha_mean, sampler.config_.beta_mean, sampler.config_.min_samples_leaf_mean, sampler.config_.max_depth_mean); + sampler.mean_forest_->SetLeafValue(sampler.init_val_mean_ / sampler.config_.num_trees_mean); + UpdateResidualEntireForest(*sampler.mean_forest_tracker_, *sampler.forest_dataset_, *sampler.residual_, sampler.mean_forest_.get(), expand_basis, std::minus()); + sampler.mean_forest_tracker_->UpdatePredictions(sampler.mean_forest_.get(), *sampler.forest_dataset_.get()); + sampler.has_mean_forest_ = true; + // Step 2: reset the now-consistent state to the last retained sample + // (mirrors MeanForestResetVisitor / RestoreStateFromGFRSnapshot). The reset + // requires an already-populated tracker, which step 1 provides. + int last_idx = samples.mean_forests->NumSamples() - 1; + TreeEnsemble& last_forest = *samples.mean_forests->GetEnsemble(last_idx); + sampler.mean_forest_->ReconstituteFromForest(last_forest); + sampler.mean_forest_tracker_->ReconstituteFromForest(last_forest, *sampler.forest_dataset_, *sampler.residual_, true); + sampler.mean_forest_tracker_->UpdatePredictions(sampler.mean_forest_.get(), *sampler.forest_dataset_.get()); + } + void operator()(GaussianConstantLeafModel& model) { + WarmStart(!sampler.config_.leaf_constant_mean); + } + void operator()(GaussianUnivariateRegressionLeafModel& model) { + WarmStart(!sampler.config_.leaf_constant_mean); + } + void operator()(GaussianMultivariateRegressionLeafModel& model) { + Log::Fatal("Continued sampling is not yet supported for multivariate leaf regression models"); + } + void operator()(CloglogOrdinalLeafModel& model) { + Log::Fatal("Continued sampling is not yet supported for cloglog ordinal models"); + } + }; + + /*! GFR iteration visitor */ + struct GFROneIterationVisitor { + BARTSampler& sampler; + BARTSamples& samples; + bool keep_sample; + void operator()(GaussianConstantLeafModel& model) { + GFRSampleOneIter( + *sampler.mean_forest_, *sampler.mean_forest_tracker_, *samples.mean_forests, model, + *sampler.forest_dataset_, *sampler.residual_, *sampler.tree_prior_mean_, sampler.rng_, + sampler.config_.var_weights_mean, sampler.config_.sweep_update_indices_mean, sampler.global_variance_, sampler.config_.feature_types, + sampler.config_.cutpoint_grid_size, /*keep_forest=*/keep_sample, + /*pre_initialized=*/true, /*backfitting=*/true, + /*num_features_subsample=*/sampler.config_.num_features_subsample_mean, sampler.config_.num_threads); + } + void operator()(GaussianUnivariateRegressionLeafModel& model) { + GFRSampleOneIter( + *sampler.mean_forest_, *sampler.mean_forest_tracker_, *samples.mean_forests, model, + *sampler.forest_dataset_, *sampler.residual_, *sampler.tree_prior_mean_, sampler.rng_, + sampler.config_.var_weights_mean, sampler.config_.sweep_update_indices_mean, sampler.global_variance_, sampler.config_.feature_types, + sampler.config_.cutpoint_grid_size, /*keep_forest=*/keep_sample, + /*pre_initialized=*/true, /*backfitting=*/true, + /*num_features_subsample=*/sampler.config_.num_features_subsample_mean, sampler.config_.num_threads); + } + void operator()(GaussianMultivariateRegressionLeafModel& model) { + GFRSampleOneIter( + *sampler.mean_forest_, *sampler.mean_forest_tracker_, *samples.mean_forests, model, + *sampler.forest_dataset_, *sampler.residual_, *sampler.tree_prior_mean_, sampler.rng_, + sampler.config_.var_weights_mean, sampler.config_.sweep_update_indices_mean, sampler.global_variance_, sampler.config_.feature_types, + sampler.config_.cutpoint_grid_size, /*keep_forest=*/keep_sample, + /*pre_initialized=*/true, /*backfitting=*/true, + /*num_features_subsample=*/sampler.config_.num_features_subsample_mean, sampler.config_.num_threads, + sampler.config_.leaf_dim_mean); + } + void operator()(CloglogOrdinalLeafModel& model) { + GFRSampleOneIter( + *sampler.mean_forest_, *sampler.mean_forest_tracker_, *samples.mean_forests, model, + *sampler.forest_dataset_, *sampler.residual_, *sampler.tree_prior_mean_, sampler.rng_, + sampler.config_.var_weights_mean, sampler.config_.sweep_update_indices_mean, sampler.global_variance_, sampler.config_.feature_types, + sampler.config_.cutpoint_grid_size, /*keep_forest=*/keep_sample, + /*pre_initialized=*/true, /*backfitting=*/false, + /*num_features_subsample=*/sampler.config_.num_features_subsample_mean, sampler.config_.num_threads); + } + }; + + /*! MCMC iteration visitor */ + struct MCMCOneIterationVisitor { + BARTSampler& sampler; + BARTSamples& samples; + bool keep_sample; + void operator()(GaussianConstantLeafModel& model) { + MCMCSampleOneIter( + *sampler.mean_forest_, *sampler.mean_forest_tracker_, *samples.mean_forests, model, + *sampler.forest_dataset_, *sampler.residual_, *sampler.tree_prior_mean_, sampler.rng_, + sampler.config_.var_weights_mean, sampler.config_.sweep_update_indices_mean, sampler.global_variance_, /*keep_forest=*/keep_sample, + /*pre_initialized=*/true, /*backfitting=*/true, + /*num_threads=*/sampler.config_.num_threads); + } + void operator()(GaussianUnivariateRegressionLeafModel& model) { + MCMCSampleOneIter( + *sampler.mean_forest_, *sampler.mean_forest_tracker_, *samples.mean_forests, model, + *sampler.forest_dataset_, *sampler.residual_, *sampler.tree_prior_mean_, sampler.rng_, + sampler.config_.var_weights_mean, sampler.config_.sweep_update_indices_mean, sampler.global_variance_, /*keep_forest=*/keep_sample, + /*pre_initialized=*/true, /*backfitting=*/true, + /*num_threads=*/sampler.config_.num_threads); + } + void operator()(GaussianMultivariateRegressionLeafModel& model) { + MCMCSampleOneIter( + *sampler.mean_forest_, *sampler.mean_forest_tracker_, *samples.mean_forests, model, + *sampler.forest_dataset_, *sampler.residual_, *sampler.tree_prior_mean_, sampler.rng_, + sampler.config_.var_weights_mean, sampler.config_.sweep_update_indices_mean, sampler.global_variance_, /*keep_forest=*/keep_sample, + /*pre_initialized=*/true, /*backfitting=*/true, + /*num_threads=*/sampler.config_.num_threads, + sampler.config_.leaf_dim_mean); + } + void operator()(CloglogOrdinalLeafModel& model) { + MCMCSampleOneIter( + *sampler.mean_forest_, *sampler.mean_forest_tracker_, *samples.mean_forests, model, + *sampler.forest_dataset_, *sampler.residual_, *sampler.tree_prior_mean_, sampler.rng_, + sampler.config_.var_weights_mean, sampler.config_.sweep_update_indices_mean, sampler.global_variance_, /*keep_forest=*/keep_sample, + /*pre_initialized=*/true, /*backfitting=*/false, + /*num_threads=*/sampler.config_.num_threads); + } + }; + + /*! Mean forest leaf scale update visitor */ + struct ScaleUpdateVisitor { + BARTSampler& sampler; + double leaf_scale; + void operator()(GaussianConstantLeafModel& model) { + model.SetScale(leaf_scale); + } + void operator()(GaussianUnivariateRegressionLeafModel& model) { + model.SetScale(leaf_scale); + } + void operator()(GaussianMultivariateRegressionLeafModel& model) { + // No-op for multivariate regression leaf model since scale is a vector + } + void operator()(CloglogOrdinalLeafModel& model) { + // No-op for cloglog ordinal leaf model since scale is not a variance parameter + } + }; + + /*! Internal reference to config and data state */ + BARTConfig& config_; + BARTData& data_; + + /*! Leaf model for mean and variance forests */ + std::variant mean_leaf_model_; + LogLinearVarianceLeafModel variance_leaf_model_; + + /*! Mean forest state */ + std::unique_ptr mean_forest_; + std::unique_ptr mean_forest_tracker_; + std::unique_ptr tree_prior_mean_; + bool has_mean_forest_ = false; + double init_val_mean_; + std::vector init_val_mean_vec_; + std::unique_ptr ordinal_sampler_; + + /*! Variance forest state */ + std::unique_ptr variance_forest_; + std::unique_ptr variance_forest_tracker_; + std::unique_ptr tree_prior_variance_; + bool has_variance_forest_ = false; + double init_val_variance_; + + /*! Random effects state */ + std::unique_ptr random_effects_model_; + std::unique_ptr random_effects_tracker_; + std::unique_ptr random_effects_dataset_; + bool has_random_effects_ = false; + + /*! Dataset */ + std::unique_ptr residual_; + std::unique_ptr outcome_raw_; + std::unique_ptr forest_dataset_; + std::unique_ptr forest_dataset_test_; + bool has_test_ = false; + + /*! Random number generator */ + std::mt19937 rng_; + + /*! Model parameters */ + double global_variance_; + double leaf_scale_; + std::vector leaf_scale_multivariate_; + + // Global error scale model + std::unique_ptr var_model_; + bool sample_sigma2_global_ = false; + + // Leaf scale model + std::unique_ptr leaf_scale_model_; + bool sample_sigma2_leaf_ = false; + + /*! Snapshot of sampler state captured at the end of a GFR iteration, used to initialize independent MCMC chains */ + struct GFRSnapshot { + // Forest state + std::unique_ptr mean_forest; // null if no mean forest + std::unique_ptr variance_forest; // null if no variance forest + + // Global parameters + double sigma2; + double leaf_scale; + std::vector leaf_scale_multivariate; + + // Residual (incorporates forest + RFX contributions for a given sampler iteration) + std::vector residual; + + // Heteroskedastic variance model state + std::vector variance_weights; // forest_dataset_ var_weights at snapshot time; only valid when variance_forest != null + + // Cloglog model state + std::vector cloglog_forest_preds; // cached forest predictions; only populated with cloglog link + std::vector cloglog_latent_outcome; // cached latent outcome state; only populated with cloglog link + std::vector cloglog_logscale_cutpoints; // cached logscale cutpoints; only populated with cloglog link and ordinal outcome + + // RFX model state (only populated when has_random_effects_) + Eigen::VectorXd rfx_working_parameter; + Eigen::MatrixXd rfx_group_parameters; + Eigen::MatrixXd rfx_group_parameter_covariance; + Eigen::MatrixXd rfx_working_parameter_covariance; + double rfx_variance_prior_shape; + double rfx_variance_prior_scale; + }; + + /*! GFR snapshots captured during run_gfr() for use by multi-chain sampler */ + std::vector gfr_snapshots_; +}; + +} // namespace StochTree + +#endif // STOCHTREE_BART_SAMPLER_H_ diff --git a/include/stochtree/bcf.h b/include/stochtree/bcf.h new file mode 100644 index 00000000..526eefc7 --- /dev/null +++ b/include/stochtree/bcf.h @@ -0,0 +1,384 @@ +/*! + * Copyright (c) 2026 stochtree authors. All rights reserved. + * Licensed under the MIT License. See LICENSE file in the project root for license information. + */ +#ifndef STOCHTREE_BCF_H_ +#define STOCHTREE_BCF_H_ + +#include +#include +#include +#include +#include +#include + +namespace StochTree { + +enum class BCFRFXModelSpec { + Custom, + InterceptOnly, + InterceptPlusTreatment +}; + +struct BCFData { + // Train set covariates + double* X_train = nullptr; + int n_train = 0; + int p = 0; + + // Test set covariates + double* X_test = nullptr; + int n_test = 0; + + // Treatment + double* treatment_train = nullptr; + double* treatment_test = nullptr; + int treatment_dim = 0; + + // Train set outcome + double* y_train = nullptr; + + // Observation weights + double* obs_weights_train = nullptr; + double* obs_weights_test = nullptr; + + // Random effects + int* rfx_group_ids_train = nullptr; + int* rfx_group_ids_test = nullptr; + double* rfx_basis_train = nullptr; + double* rfx_basis_test = nullptr; + int rfx_num_groups = 0; + int rfx_basis_dim = 0; +}; + +struct BCFConfig { + // High level parameters + bool standardize_outcome = true; // whether to standardize the outcome before fitting and unstandardize predictions after + int num_threads = 1; // number of threads to use for sampling + bool verbose = false; // whether to print sampler progress to the console + int cutpoint_grid_size = 100; // number of cutpoints to consider for each covariate when sampling splits + std::vector feature_types; // feature types for each covariate (should be same length as number of covariates in the dataset), where 0 = continuous, 1 = categorical + LinkFunction link_function = LinkFunction::Identity; // link function to use (Identity, Probit, Cloglog) + OutcomeType outcome_type = OutcomeType::Continuous; // type of the outcome variable (Continuous, Binary, Ordinal) + int random_seed = -1; // random seed for reproducibility (if negative, a random seed will be generated) + bool keep_gfr = true; // whether or not to keep GFR samples or simply use them to warm-start an MCMC chain + bool keep_burnin = false; // whether or not to keep "burn-in" MCMC samples (largely a debugging flag) + bool adaptive_coding = false; // whether or not to use adaptive coding for the BCF model + double b_0_init = 0.0; // initial value for the b_0 parameter in the adaptive coding scheme (only relevant if adaptive_coding=true) + double b_1_init = 1.0; // initial value for the b_1 parameter in the adaptive coding scheme (only relevant if adaptive_coding=true) + + // Global error variance parameters + double a_sigma2_global = 0.0; // shape parameter for inverse gamma prior on global error variance + double b_sigma2_global = 0.0; // scale parameter for inverse gamma prior on global error variance + double sigma2_global_init = 1.0; // initial value for global error variance + bool sample_sigma2_global = true; // whether to sample global error variance (if false, it will be fixed at sigma2_global_init) + + // Prognostic forest parameters + int num_trees_mu = 200; // number of trees in the prognostic forest + double alpha_mu = 0.95; // alpha parameter for prognostic forest tree prior + double beta_mu = 2.0; // beta parameter for prognostic forest tree prior + int min_samples_leaf_mu = 5; // minimum number of samples per leaf for prognostic forest + int max_depth_mu = -1; // maximum depth for prognostic forest trees (-1 means no maximum) + bool leaf_constant_mu = true; // whether to use constant leaf model for prognostic forest + int leaf_dim_mu = 1; // dimension of the leaf for prognostic forest + bool exponentiated_leaf_mu = false; // whether to exponentiate leaf predictions for prognostic forest + int num_features_subsample_mu = 0; // number of features to subsample for each prognostic forest split (0 means no subsampling) + double a_sigma2_mu = 3.0; // shape parameter for inverse gamma prior on prognostic forest leaf scale + double b_sigma2_mu = -1.0; // scale parameter for inverse gamma prior on prognostic forest leaf scale (-1 is a sentinel value that triggers a data-informed calibration based on the variance of the outcome and the number of trees) + double sigma2_mu_init = -1.0; // initial value of prognostic forest leaf scale (-1 is a sentinel value that triggers a data-informed calibration based on the variance of the outcome and the number of trees) + std::vector var_weights_mu; // variable weights for prognostic forest splits (should be same length as number of covariates in the dataset) + bool sample_sigma2_leaf_mu = false; // whether to sample prognostic forest leaf scale (if false, it will be fixed at sigma2_mu_init) + std::vector sweep_update_indices_mu; // indices of trees to update in a given sweep (should be subset of [0, num_trees - 1]) + + // Treatment effect forest parameters + int num_trees_tau = 50; // number of trees in the treatment effect forest + double alpha_tau = 0.95; // alpha parameter for treatment effect forest tree prior + double beta_tau = 2.0; // beta parameter for treatment effect forest tree prior + int min_samples_leaf_tau = 5; // minimum number of samples per leaf for treatment effect forest + int max_depth_tau = -1; // maximum depth for treatment effect forest trees (-1 means no maximum) + bool leaf_constant_tau = false; // whether to use constant leaf model for treatment effect forest (false for univariate/multivariate regression leaf, true for constant leaf) + int leaf_dim_tau = 1; // dimension of the leaf for treatment effect forest + bool exponentiated_leaf_tau = false; // whether to exponentiate leaf predictions for treatment effect forest + int num_features_subsample_tau = 0; // number of features to subsample for each treatment effect forest split (0 means no subsampling) + double a_sigma2_tau = 3.0; // shape parameter for inverse gamma prior on treatment effect forest leaf scale + double b_sigma2_tau = -1.0; // scale parameter for inverse gamma prior on treatment effect forest leaf scale (-1 is a sentinel value that triggers a data-informed calibration based on the variance of the outcome and the number of trees) + double sigma2_tau_init = -1.0; // initial value of treatment effect forest leaf scale (-1 is a sentinel value that triggers a data-informed calibration based on the variance of the outcome and the number of trees) + std::vector var_weights_tau; // variable weights for treatment effect forest splits (should be same length as number of covariates in the dataset) + std::vector sigma2_leaf_tau_matrix; // prior covariance matrix Sigma_0 for multivariate leaf regression, stored column-major (size leaf_dim_tau^2); empty = use sigma2_tau_init * I + bool sample_sigma2_leaf_tau = false; // whether to sample treatment effect forest leaf scale (if false, it will be fixed at sigma2_tau_init) + std::vector sweep_update_indices_tau; // indices of trees to update in a given sweep (should be subset of [0, num_trees - 1]) + MeanLeafModelType tau_leaf_model_type = MeanLeafModelType::GaussianUnivariateRegression; // leaf model type for treatment effect forest + bool sample_tau_0 = true; // whether or not to sample an intercept term on the treatment, additive to the covariate-dependent treatment effect forest + double tau_0_prior_var_scalar = -1.0; // scalar-valued prior variance for treatment intercept (only relevant when sample_tau_0=true; -1 is a sentinel value that triggers a data-informed calibration) + std::vector tau_0_prior_var_multivariate; // vector-valued prior variance for treatment intercept in multivariate treatment case (only relevant when sample_tau_0=true; should be of length treatment_dim; empty = use data-informed calibration) + + // Variance forest parameters + int num_trees_variance = 0; // number of trees in the variance forest + double leaf_prior_calibration_param = 1.5; // calibration parameter for variance forest leaf prior + double shape_variance_forest = -1.0; // shape parameter for variance forest leaf model (calibrated internally based on leaf_prior_calibration_param if set to sentinel value of -1) + double scale_variance_forest = -1.0; // scale parameter for variance forest leaf model (calibrated internally based on leaf_prior_calibration_param if set to sentinel value of -1) + double variance_forest_leaf_init = -1.0; // initial (raw-scale) root value for the variance forest; each leaf starts at log(value)/num_trees_variance. Sentinel <= 0 = calibrate internally + double alpha_variance = 0.5; // alpha parameter for variance forest tree prior + double beta_variance = 2.0; // beta parameter for variance forest tree prior + int min_samples_leaf_variance = 5; // minimum number of samples per leaf for variance forest + int max_depth_variance = -1; // maximum depth for variance forest trees (-1 means no maximum) + bool leaf_constant_variance = true; // whether to use constant leaf model for variance forest + int leaf_dim_variance = 1; // dimension of the leaf for variance forest (should be 1 if leaf_constant_variance=true) + bool exponentiated_leaf_variance = true; // whether to exponentiate leaf predictions for variance forest + int num_features_subsample_variance = 0; // number of features to subsample for each variance forest split (0 means no subsampling) + std::vector var_weights_variance; // variable weights for variance forest splits (should be same length as number of covariates in the dataset) + std::vector sweep_update_indices_variance; // indices of trees to update in a given sweep (should be subset of [0, num_trees - 1]) + + // Random effects parameters + bool has_random_effects = false; // whether or not a model includes a random effects term + BCFRFXModelSpec rfx_model_spec = BCFRFXModelSpec::Custom; // specification for the random effects model; custom relies on a user-provided basis while intercept-only constructs a varying intercept model without needing a user-provided basis + std::vector rfx_working_parameter_mean_prior; // vector of dimension num_basis; empty = use zeros + std::vector rfx_group_parameter_mean_prior; // matrix of dimension num_basis x num_groups, stored column-major; empty = use zeros + std::vector rfx_working_parameter_cov_prior; // matrix of dimension num_basis x num_basis, stored column-major; empty = use identity matrix + std::vector rfx_group_parameter_cov_prior; // matrix of dimension num_basis x num_basis, stored column-major; empty = use identity matrix + double rfx_variance_prior_shape = 1.0; // shape parameter for variance prior in random effects model + double rfx_variance_prior_scale = 1.0; // scale parameter for variance prior in random effects model + + // TODO: Other parameters ... +}; + +struct BCFSamples { + // Posterior samples of training set outcome predictions (n_train x num_samples, stored column-major) + std::vector y_hat_train; + + // Posterior samples of training set prognostic forest predictions (n_train x num_samples, stored column-major) + std::vector mu_forest_predictions_train; + + // Posterior samples of training set treatment effect forest predictions (n_train x num_samples, stored column-major) + std::vector tau_forest_predictions_train; + + // Posterior samples of training set variance forest predictions (n_train x num_samples, stored column-major) + std::vector variance_forest_predictions_train; + + // Posterior samples of test set outcome predictions (n_test x num_samples, stored column-major) + std::vector y_hat_test; + + // Posterior samples of test set prognostic forest predictions (n_test x num_samples, stored column-major) + std::vector mu_forest_predictions_test; + + // Posterior samples of test set treatment effect forest predictions (n_test x num_samples, stored column-major) + std::vector tau_forest_predictions_test; + + // Posterior samples of test set variance forest predictions (n_test x num_samples, stored column-major) + std::vector variance_forest_predictions_test; + + // Posterior samples of global error variance (num_samples) + std::vector global_error_variance_samples; + + // Posterior samples of leaf scale (num_samples) + std::vector leaf_scale_mu_samples; + std::vector leaf_scale_tau_samples; + + // Pointer to sampled prognostic forests + std::unique_ptr mu_forests; + + // Pointer to sampled treatment effect forests + std::unique_ptr tau_forests; + + // Pointer to sampled variance forests + std::unique_ptr variance_forests; + + // Posterior samples of training set RFX predictions (n_train x num_samples, stored column-major) + std::vector rfx_predictions_train; + + // Posterior samples of test set RFX predictions (n_test x num_samples, stored column-major) + std::vector rfx_predictions_test; + + // Treatment intercept samples (treatment_dim x num_samples, stored column-major; only populated when sample_tau_0=true) + std::vector tau_0_samples; + + // Adaptive coding parameter samples + std::vector b0_samples; + std::vector b1_samples; + + // Pointer to random effects sample container and label mapping + std::unique_ptr rfx_container; + std::unique_ptr rfx_label_mapper; + + // Metadata about the samples (e.g., number of samples, burn-in, etc.) could be added here as needed + int num_samples = 0; + int num_train = 0; + int num_test = 0; + int treatment_dim = 0; + double y_bar = 0.0; + double y_std = 0.0; + + // Serialize the samples-owned subtree (forests + parameter traces + intrinsic scalars) into a JSON + // object. BCF mirror of BARTSamples::ToJson -- see that method for the design notes (shared C++ + // source of truth, byte-identical key layout, presence inferred from structure, nlohmann sorts + // keys so insertion order is irrelevant). Forests use the BCF named keys; the extra BCF parameter + // traces (leaf_scale_mu/tau, tau_0, b0/b1) go under the same "parameters" subfolder. + // NOTE: parameter traces are serialized verbatim (identity). The tau_0 user-facing scaling + // (x y_std) and its multivariate (treatment_dim>1) ravel-order are reconciled at the postprocess / + // wiring boundary per the locked scale decision, not here. Random effects are not yet routed + // through this path (guarded to avoid silent drops). + void AppendToJson(nlohmann::json& obj) const { + // Forests, under the BCF self-describing named keys, with the num_forests counter + nlohmann::json forests = nlohmann::json::object(); + int num_forests = 0; + if (mu_forests != nullptr) { + forests.emplace("prognostic_forest", mu_forests->to_json()); + num_forests++; + } + if (tau_forests != nullptr) { + forests.emplace("treatment_forest", tau_forests->to_json()); + num_forests++; + } + if (variance_forests != nullptr) { + forests.emplace("variance_forest", variance_forests->to_json()); + num_forests++; + } + obj["forests"] = forests; + obj["num_forests"] = num_forests; + // Parameter traces, under the "parameters" subfolder (presence inferred from non-empty vectors) + nlohmann::json parameters = nlohmann::json::object(); + if (!global_error_variance_samples.empty()) { + parameters.emplace("sigma2_global_samples", global_error_variance_samples); + } + if (!leaf_scale_mu_samples.empty()) { + parameters.emplace("sigma2_leaf_mu_samples", leaf_scale_mu_samples); + } + if (!leaf_scale_tau_samples.empty()) { + parameters.emplace("sigma2_leaf_tau_samples", leaf_scale_tau_samples); + } + if (!b0_samples.empty()) { + parameters.emplace("b0_samples", b0_samples); + } + if (!b1_samples.empty()) { + parameters.emplace("b1_samples", b1_samples); + } + if (!tau_0_samples.empty()) { + parameters.emplace("tau_0_samples", tau_0_samples); + obj.emplace("tau_0_dim", treatment_dim); + } + if (!parameters.empty()) { + obj["parameters"] = parameters; + } + // Intrinsic scalars (stored in user-facing scale, matching the existing wire format) + obj.emplace("outcome_mean", y_bar); + obj.emplace("outcome_scale", y_std); + obj.emplace("num_samples", num_samples); + obj.emplace("treatment_dim", treatment_dim); + // Random effects + int num_random_effects = 0; + nlohmann::json rfx = nlohmann::json::object(); + if (rfx_container != nullptr && rfx_label_mapper != nullptr) { + rfx.emplace("random_effect_container_0", rfx_container->to_json()); + rfx.emplace("random_effect_label_mapper_0", rfx_label_mapper->to_json()); + rfx.emplace("random_effect_groupids_0", rfx_label_mapper->Keys()); + num_random_effects = 1; + } + obj["random_effects"] = rfx; + obj["num_random_effects"] = num_random_effects; + } + + // Populate this BCFSamples from the samples-owned subtree of a parsed JSON object. Inverse of + // ToJson(); presence inferred from structure rather than envelope booleans. See ToJson() re: scope. + void FromJson(const nlohmann::json& obj) { + // Unpack forests if present, checking for the expected keys + if (obj.contains("forests")) { + const nlohmann::json& forests = obj.at("forests"); + if (forests.contains("prognostic_forest")) { + mu_forests = std::make_unique(0, 0, false, false); + mu_forests->from_json(forests.at("prognostic_forest")); + } + if (forests.contains("treatment_forest")) { + tau_forests = std::make_unique(0, 0, false, false); + tau_forests->from_json(forests.at("treatment_forest")); + } + if (forests.contains("variance_forest")) { + variance_forests = std::make_unique(0, 0, false, false); + variance_forests->from_json(forests.at("variance_forest")); + } + } + // Unpack parameters if present, checking for expected keys + if (obj.contains("parameters")) { + const nlohmann::json& parameters = obj.at("parameters"); + if (parameters.contains("sigma2_global_samples")) { + global_error_variance_samples = parameters.at("sigma2_global_samples").get>(); + } + if (parameters.contains("sigma2_leaf_mu_samples")) { + leaf_scale_mu_samples = parameters.at("sigma2_leaf_mu_samples").get>(); + } + if (parameters.contains("sigma2_leaf_tau_samples")) { + leaf_scale_tau_samples = parameters.at("sigma2_leaf_tau_samples").get>(); + } + if (parameters.contains("b0_samples")) { + b0_samples = parameters.at("b0_samples").get>(); + } + if (parameters.contains("b1_samples")) { + b1_samples = parameters.at("b1_samples").get>(); + } + if (parameters.contains("tau_0_samples")) { + tau_0_samples = parameters.at("tau_0_samples").get>(); + } + } + // Unpack random effects if present, checking for expected keys + if (obj.contains("num_random_effects") && obj.at("num_random_effects").get() > 0) { + rfx_container = std::make_unique(); + rfx_label_mapper = std::make_unique(); + rfx_container->from_json(obj.at("random_effects").at("random_effect_container_0")); + rfx_label_mapper->from_json(obj.at("random_effects").at("random_effect_label_mapper_0")); + } + // Unpack outcome statistics + if (obj.contains("outcome_mean")) y_bar = obj.at("outcome_mean").get(); + if (obj.contains("outcome_scale")) y_std = obj.at("outcome_scale").get(); + if (obj.contains("num_samples")) num_samples = obj.at("num_samples").get(); + if (obj.contains("treatment_dim")) treatment_dim = obj.at("treatment_dim").get(); + } + + // Append another chain's draws onto this one (multi-chain combine). BCF mirror of + // BARTSamples::Merge -- `this` must already be populated, `other` must match model structure + // (same forests present, same standardization, same treatment_dim). Forests are deep-copied + // sample-by-sample and parameter traces concatenated, preserving draw order. + void Merge(const BCFSamples& other) { + // Runtime checks for samples objects to be combined + if (y_bar != other.y_bar || y_std != other.y_std) { + Log::Fatal("Cannot merge BCFSamples with different outcome standardization"); + } + if (rfx_container != nullptr && other.rfx_container != nullptr) { + if (rfx_container->NumComponents() != other.rfx_container->NumComponents() || + rfx_container->NumGroups() != other.rfx_container->NumGroups()) { + Log::Fatal("Cannot merge BARTSamples with different random effects structure"); + } + if (rfx_label_mapper->Keys() != other.rfx_label_mapper->Keys()) { + Log::Fatal("Cannot merge BARTSamples with different random effects label mapping"); + } + if (rfx_label_mapper->Map() != other.rfx_label_mapper->Map()) { + Log::Fatal("Cannot merge BARTSamples with different random effects label mapping"); + } + } + if (treatment_dim != other.treatment_dim) { + Log::Fatal("Cannot merge BCFSamples with different treatment_dim"); + } + // Append forests if they exist in the samples object + AppendForestContainerSamples(mu_forests, other.mu_forests, "prognostic"); + AppendForestContainerSamples(tau_forests, other.tau_forests, "treatment"); + AppendForestContainerSamples(variance_forests, other.variance_forests, "variance"); + // Append random effects if they exist in the samples object + AppendRandomEffectsContainerSamples(rfx_container, other.rfx_container); + // Append parameters samples + auto append = [](std::vector& dst, const std::vector& src, const std::string& name = "") { + if ((!dst.empty() && src.empty()) || (dst.empty() && !src.empty())) { + Log::Fatal("Cannot merge BARTSamples objects: %s samples present in one chain but not the other", name.c_str()); + } + dst.insert(dst.end(), src.begin(), src.end()); + }; + append(global_error_variance_samples, other.global_error_variance_samples, "global error variance"); + append(leaf_scale_mu_samples, other.leaf_scale_mu_samples, "leaf scale mu"); + append(leaf_scale_tau_samples, other.leaf_scale_tau_samples, "leaf scale tau"); + append(tau_0_samples, other.tau_0_samples, "tau_0"); + append(b0_samples, other.b0_samples, "b0"); + append(b1_samples, other.b1_samples, "b1"); + num_samples += other.num_samples; + } +}; + +} // namespace StochTree + +#endif // STOCHTREE_BCF_H_ diff --git a/include/stochtree/bcf_sampler.h b/include/stochtree/bcf_sampler.h new file mode 100644 index 00000000..b108e4fa --- /dev/null +++ b/include/stochtree/bcf_sampler.h @@ -0,0 +1,290 @@ +/*! + * Copyright (c) 2026 stochtree authors. All rights reserved. + * Licensed under the MIT License. See LICENSE file in the project root for license information. + */ +#ifndef STOCHTREE_BCF_SAMPLER_H_ +#define STOCHTREE_BCF_SAMPLER_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace StochTree { + +class BCFSampler { + public: + // If continuation is true, the sampler warm-starts the active mu and tau forests from + // the last retained samples already present in `samples` (rather than initializing them + // to root), and the existing forest containers are preserved so that new samples are + // appended to them. + BCFSampler(BCFSamples& samples, BCFConfig& config, BCFData& data, bool continuation = false); + + // Main entry point for running the BCF GFR sampler + // If num_chains > 0, captures snapshots of the last num_chains GFR states for fork_chains() + void run_gfr(BCFSamples& samples, int num_gfr, bool keep_gfr, int num_chains = 0); + + // Run a single chain of the BCF MCMC sampler + void run_mcmc(BCFSamples& samples, int num_burnin, int keep_every, int num_mcmc); + + // Run num_chains independent MCMC chains sequentually based on GFR snapshots captured by run_gfr() or re-initialized from root + void run_mcmc_chains(BCFSamples& samples, int num_chains, int num_burnin, int keep_every, int num_mcmc); + + // Post-process samples by extracting test set predictions and running any necessary transformations + void postprocess_samples(BCFSamples& samples); + + // Serialize the internal RNG state to a string. std::mt19937 round-trips losslessly through + // its stream operators, so this captures the exact position in the random stream -- used to + // persist RNG state across a sample() / continue_sampling() boundary for bit-identical results. + std::string GetRngState() const { + std::ostringstream oss; + oss << rng_; + return oss.str(); + } + + // Restore the internal RNG state from a string produced by GetRngState(). Resumes the random + // stream at exactly the captured position. Must be called after construction (InitializeState + // unconditionally re-seeds rng_) and before any sampling draws. + void SetRngState(const std::string& state) { + std::istringstream iss(state); + iss >> rng_; + } + + private: + /*! Initialize state variables */ + void InitializeState(BCFSamples& samples, bool continuation = false); + bool initialized_ = false; + + /*! Internal function to restore sampler state based on a GFR snapshot */ + void RestoreStateFromGFRSnapshot(BCFSamples& samples, int snapshot_index); + + /*! Internal function to restore sampler state to root / initial values */ + void RestoreStateDefault(); + + /*! Internal sample runner function */ + void RunOneIteration(BCFSamples& samples, bool gfr, bool keep_sample, bool write_snapshot = false); + + /*! Internal function to sample parametric treatment effect "intercept" term (tau_0 in stochtree nomenclature) */ + void SampleParametricTreatmentEffect(); + + /*! Internal function to sample adaptive coding parameters for binary treatment */ + void SampleAdaptiveCodingParameters(); + + /*! Internal reference to config and data state */ + BCFConfig& config_; + BCFData& data_; + + /*! Leaf model for mean and variance forests */ + GaussianConstantLeafModel mu_leaf_model_; + std::variant tau_leaf_model_; + LogLinearVarianceLeafModel variance_leaf_model_; + + /*! Mean forest state */ + std::unique_ptr mu_forest_; + std::unique_ptr mu_forest_tracker_; + std::unique_ptr tree_prior_mu_; + std::unique_ptr tau_forest_; + std::unique_ptr tau_forest_tracker_; + std::unique_ptr tree_prior_tau_; + double init_val_mu_; + double init_val_tau_; + std::vector init_val_tau_vec_; + + /*! Variance forest state */ + std::unique_ptr variance_forest_; + std::unique_ptr variance_forest_tracker_; + std::unique_ptr tree_prior_variance_; + bool has_variance_forest_ = false; + double init_val_variance_; + + /*! Random effects state */ + std::unique_ptr random_effects_model_; + std::unique_ptr random_effects_tracker_; + std::unique_ptr random_effects_dataset_; + bool has_random_effects_ = false; + + /*! Dataset */ + std::unique_ptr residual_; + std::unique_ptr outcome_raw_; + std::unique_ptr forest_dataset_; + std::unique_ptr forest_dataset_test_; + bool has_test_ = false; + + /*! Random number generator */ + std::mt19937 rng_; + + /*! Model parameters */ + double global_variance_; + double leaf_scale_mu_; + double leaf_scale_tau_; + std::vector leaf_scale_tau_multivariate_; + + /*! Probit terms / helpers */ + std::vector model_preds_; + + /*! Raw tau(x) predictions (sum of leaf values across trees, no z multiplication), maintained each step via leaf-lookup */ + std::vector tau_raw_sum_preds_; + + // Global error scale model + std::unique_ptr var_model_; + bool sample_sigma2_global_ = false; + + // Leaf scale models + std::unique_ptr leaf_scale_model_mu_; + bool sample_sigma2_leaf_mu_ = false; + std::unique_ptr leaf_scale_model_tau_; + bool sample_sigma2_leaf_tau_ = false; + + // Treatment intercept term + double tau_0_scalar_; + std::vector tau_0_vector_; + bool sample_tau_0_ = false; + + // Adaptive coding parameters + double b_0_; + double b_1_; + bool adaptive_coding_ = false; + std::vector tau_basis_vector_train_; + std::vector tau_basis_vector_test_; + + /*! GFR iteration visitor for tau forest */ + struct GFROneIterationVisitorTau { + BCFSampler& sampler; + BCFSamples& samples; + bool keep_sample; + void operator()(GaussianUnivariateRegressionLeafModel& model) { + GFRSampleOneIter( + *sampler.tau_forest_, *sampler.tau_forest_tracker_, *samples.tau_forests, model, + *sampler.forest_dataset_, *sampler.residual_, *sampler.tree_prior_tau_, sampler.rng_, + sampler.config_.var_weights_tau, sampler.config_.sweep_update_indices_tau, sampler.global_variance_, sampler.config_.feature_types, + sampler.config_.cutpoint_grid_size, /*keep_forest=*/keep_sample, + /*pre_initialized=*/true, /*backfitting=*/true, + /*num_features_subsample=*/sampler.config_.num_features_subsample_tau, sampler.config_.num_threads); + } + void operator()(GaussianMultivariateRegressionLeafModel& model) { + GFRSampleOneIter( + *sampler.tau_forest_, *sampler.tau_forest_tracker_, *samples.tau_forests, model, + *sampler.forest_dataset_, *sampler.residual_, *sampler.tree_prior_tau_, sampler.rng_, + sampler.config_.var_weights_tau, sampler.config_.sweep_update_indices_tau, sampler.global_variance_, sampler.config_.feature_types, + sampler.config_.cutpoint_grid_size, /*keep_forest=*/keep_sample, + /*pre_initialized=*/true, /*backfitting=*/true, + /*num_features_subsample=*/sampler.config_.num_features_subsample_tau, sampler.config_.num_threads, + sampler.config_.leaf_dim_tau); + } + }; + + /*! MCMC iteration visitor */ + struct MCMCOneIterationVisitorTau { + BCFSampler& sampler; + BCFSamples& samples; + bool keep_sample; + void operator()(GaussianUnivariateRegressionLeafModel& model) { + MCMCSampleOneIter( + *sampler.tau_forest_, *sampler.tau_forest_tracker_, *samples.tau_forests, model, + *sampler.forest_dataset_, *sampler.residual_, *sampler.tree_prior_tau_, sampler.rng_, + sampler.config_.var_weights_tau, sampler.config_.sweep_update_indices_tau, sampler.global_variance_, /*keep_forest=*/keep_sample, + /*pre_initialized=*/true, /*backfitting=*/true, + /*num_threads=*/sampler.config_.num_threads); + } + void operator()(GaussianMultivariateRegressionLeafModel& model) { + MCMCSampleOneIter( + *sampler.tau_forest_, *sampler.tau_forest_tracker_, *samples.tau_forests, model, + *sampler.forest_dataset_, *sampler.residual_, *sampler.tree_prior_tau_, sampler.rng_, + sampler.config_.var_weights_tau, sampler.config_.sweep_update_indices_tau, sampler.global_variance_, /*keep_forest=*/keep_sample, + /*pre_initialized=*/true, /*backfitting=*/true, + /*num_threads=*/sampler.config_.num_threads, sampler.config_.leaf_dim_tau); + } + }; + + /*! Mu / tau forest leaf scale update visitor */ + struct ScaleUpdateVisitor { + BCFSampler& sampler; + double leaf_scale; + void operator()(GaussianConstantLeafModel& model) { + model.SetScale(leaf_scale); + } + void operator()(GaussianUnivariateRegressionLeafModel& model) { + model.SetScale(leaf_scale); + } + void operator()(GaussianMultivariateRegressionLeafModel& model) { + // No-op for multivariate regression leaf model since scale is a vector + } + void operator()(CloglogOrdinalLeafModel& model) { + // No-op for cloglog ordinal leaf model since scale is not a variance parameter + } + }; + + struct TauForestResetVisitor { + BCFSampler& sampler; + BCFSamples& samples; + TreeEnsemble& forest; + void operator()(GaussianUnivariateRegressionLeafModel& model) { + sampler.tau_forest_->ReconstituteFromForest(forest); + sampler.tau_forest_tracker_->ReconstituteFromForest(forest, *sampler.forest_dataset_, *sampler.residual_, true); + sampler.tau_forest_tracker_->UpdatePredictions(sampler.tau_forest_.get(), *sampler.forest_dataset_.get()); + } + void operator()(GaussianMultivariateRegressionLeafModel& model) { + sampler.tau_forest_->ReconstituteFromForest(forest); + sampler.tau_forest_tracker_->ReconstituteFromForest(forest, *sampler.forest_dataset_, *sampler.residual_, true); + sampler.tau_forest_tracker_->UpdatePredictions(sampler.tau_forest_.get(), *sampler.forest_dataset_.get()); + } + }; + + /*! Snapshot of sampler state captured at the end of a GFR iteration, used to initialize independent MCMC chains */ + struct GFRSnapshot { + // Forest state + std::unique_ptr mu_forest; + std::unique_ptr tau_forest; + std::unique_ptr variance_forest; // null if no variance forest + + // Global parameters + double sigma2; + double leaf_scale_mu; + double leaf_scale_tau; + std::vector leaf_scale_tau_multivariate; + + // Treatment intercept + double tau_0_scalar; + std::vector tau_0_vector; + + // Adaptive coding + double b_0; + double b_1; + + // Residual (incorporates forest + RFX contributions for a given sampler iteration) + std::vector residual; + + // Heteroskedastic variance model state + std::vector variance_weights; // forest_dataset_ var_weights at snapshot time; only valid when variance_forest != null + + // RFX model state (only populated when has_random_effects_) + Eigen::VectorXd rfx_working_parameter; + Eigen::MatrixXd rfx_group_parameters; + Eigen::MatrixXd rfx_group_parameter_covariance; + Eigen::MatrixXd rfx_working_parameter_covariance; + double rfx_variance_prior_shape; + double rfx_variance_prior_scale; + }; + + /*! GFR snapshots captured during run_gfr() for use by multi-chain sampler */ + std::vector gfr_snapshots_; +}; + +} // namespace StochTree + +#endif // STOCHTREE_BCF_SAMPLER_H_ diff --git a/include/stochtree/category_tracker.h b/include/stochtree/category_tracker.h index 1fd86cb0..ea0b3cca 100644 --- a/include/stochtree/category_tracker.h +++ b/include/stochtree/category_tracker.h @@ -1,24 +1,24 @@ /*! * Copyright (c) 2024 stochtree authors. - * + * * General-purpose data structures used for keeping track of categories in a training dataset. - * - * SampleCategoryMapper is a simplified version of SampleNodeMapper, which is not tree-specific - * as it tracks categories loaded into a training dataset, and we do not expect to modify it during + * + * SampleCategoryMapper is a simplified version of SampleNodeMapper, which is not tree-specific + * as it tracks categories loaded into a training dataset, and we do not expect to modify it during * training. - * + * * SampleCategoryMapper is used in two places: * 1. Group random effects: mapping observations to group IDs for the purpose of computing random effects * 2. Heteroskedasticity based on fixed categories (as opposed to partitions as in HBART by Pratola et al 2018) - * - One example of this would be binary treatment causal inference with separate outcome variances + * - One example of this would be binary treatment causal inference with separate outcome variances * for the treated and control groups (as in Krantsevich et al 2023) - * - * CategorySampleTracker is a simplified version of FeatureUnsortedPartition, which as above does + * + * CategorySampleTracker is a simplified version of FeatureUnsortedPartition, which as above does * not vary based on tree / partition and is not expected to change during training. - * - * SampleNodeMapper is inspired by the design of the DataPartition class in LightGBM, + * + * SampleNodeMapper is inspired by the design of the DataPartition class in LightGBM, * released under the MIT license with the following copyright: - * + * * Copyright (c) 2016 Microsoft Corporation. All rights reserved. * Licensed under the MIT License. See LICENSE file in the project root for license information. */ @@ -40,12 +40,20 @@ namespace StochTree { */ class SampleCategoryMapper { public: - SampleCategoryMapper(std::vector& group_indices) { + SampleCategoryMapper(std::vector& group_indices) { num_observations_ = group_indices.size(); observation_indices_ = group_indices; } - - SampleCategoryMapper(SampleCategoryMapper& other){ + + SampleCategoryMapper(int* group_indices, int num_observations) { + num_observations_ = num_observations; + observation_indices_.resize(num_observations_); + for (int i = 0; i < num_observations_; i++) { + observation_indices_[i] = group_indices[i]; + } + } + + SampleCategoryMapper(SampleCategoryMapper& other) { num_observations_ = other.NumObservations(); observation_indices_.resize(num_observations_); for (int i = 0; i < num_observations_; i++) { @@ -62,8 +70,8 @@ class SampleCategoryMapper { CHECK_LT(sample_id, num_observations_); observation_indices_[sample_id] = category_id; } - - inline int NumObservations() {return num_observations_;} + + inline int NumObservations() { return num_observations_; } private: std::vector observation_indices_; @@ -75,21 +83,21 @@ class SampleCategoryMapper { */ class CategorySampleTracker { public: - CategorySampleTracker(const std::vector& group_indices) { + CategorySampleTracker(const std::vector& group_indices) { int n = group_indices.size(); indices_ = std::vector(n); std::iota(indices_.begin(), indices_.end(), 0); - auto comp_op = [&](size_t const &l, size_t const &r) { return std::less{}(group_indices[l], group_indices[r]); }; + auto comp_op = [&](size_t const& l, size_t const& r) { return std::less{}(group_indices[l], group_indices[r]); }; std::stable_sort(indices_.begin(), indices_.end(), comp_op); category_count_ = 0; int observation_count = 0; for (int i = 0; i < n; i++) { bool start_cond = i == 0; - bool end_cond = i == n-1; + bool end_cond = i == n - 1; bool new_group_cond{false}; - if (i > 0) new_group_cond = group_indices[indices_[i]] != group_indices[indices_[i-1]]; + if (i > 0) new_group_cond = group_indices[indices_[i]] != group_indices[indices_[i - 1]]; if (start_cond || new_group_cond) { category_id_map_.insert({group_indices[indices_[i]], category_count_}); unique_category_ids_.push_back(group_indices[indices_[i]]); @@ -103,7 +111,44 @@ class CategorySampleTracker { observation_count = 1; category_count_++; } else if (end_cond) { - category_length_.push_back(observation_count+1); + category_length_.push_back(observation_count + 1); + } else { + observation_count++; + } + // Add the index to the category's node index vector in either case + node_index_vector_[category_count_ - 1].emplace_back(indices_[i]); + } + } + + CategorySampleTracker(int* group_indices, int num_observations) { + int n = num_observations; + indices_ = std::vector(n); + std::iota(indices_.begin(), indices_.end(), 0); + + auto comp_op = [&](size_t const& l, size_t const& r) { return std::less{}(group_indices[l], group_indices[r]); }; + std::stable_sort(indices_.begin(), indices_.end(), comp_op); + + category_count_ = 0; + int observation_count = 0; + for (int i = 0; i < n; i++) { + bool start_cond = i == 0; + bool end_cond = i == n - 1; + bool new_group_cond{false}; + if (i > 0) new_group_cond = group_indices[indices_[i]] != group_indices[indices_[i - 1]]; + if (start_cond || new_group_cond) { + category_id_map_.insert({group_indices[indices_[i]], category_count_}); + unique_category_ids_.push_back(group_indices[indices_[i]]); + node_index_vector_.emplace_back(); + if (i == 0) { + category_begin_.push_back(i); + } else { + category_begin_.push_back(i); + category_length_.push_back(observation_count); + } + observation_count = 1; + category_count_++; + } else if (end_cond) { + category_length_.push_back(observation_count + 1); } else { observation_count++; } @@ -113,16 +158,16 @@ class CategorySampleTracker { } /*! \brief Zero-indexed numeric index that category_id is remapped to internally */ - inline int32_t CategoryNumber(int category_id) { + inline int CategoryNumber(int category_id) { return category_id_map_[category_id]; } /*! \brief First index of data points contained in node_id */ - inline data_size_t CategoryBegin(int category_id) {return category_begin_[category_id_map_[category_id]];} + inline data_size_t CategoryBegin(int category_id) { return category_begin_[category_id_map_[category_id]]; } /*! \brief One past the last index of data points contained in node_id */ inline data_size_t CategoryEnd(int category_id) { - int32_t id = category_id_map_[category_id]; + int id = category_id_map_[category_id]; return category_begin_[id] + category_length_[id]; } @@ -132,37 +177,37 @@ class CategorySampleTracker { } /*! \brief Number of total categories stored */ - inline data_size_t NumCategories() {return category_count_;} + inline data_size_t NumCategories() { return category_count_; } /*! \brief Data indices */ std::vector indices_; /*! \brief Data indices for a given node */ std::vector& NodeIndices(int category_id) { - int32_t id = category_id_map_[category_id]; + int id = category_id_map_[category_id]; return node_index_vector_[id]; } - + /*! \brief Data indices for a given node */ std::vector& NodeIndicesInternalIndex(int internal_category_id) { return node_index_vector_[internal_category_id]; } /*! \brief Returns label index map */ - std::map& GetLabelMap() {return category_id_map_;} + std::map& GetLabelMap() { return category_id_map_; } - std::vector& GetUniqueGroupIds() {return unique_category_ids_;} + std::vector& GetUniqueGroupIds() { return unique_category_ids_; } private: // Vectors tracking indices in each node std::vector category_begin_; std::vector category_length_; - std::map category_id_map_; - std::vector unique_category_ids_; + std::map category_id_map_; + std::vector unique_category_ids_; std::vector> node_index_vector_; - int32_t category_count_; + int category_count_; }; -} // namespace StochTree +} // namespace StochTree -#endif // STOCHTREE_CATEGORY_TRACKER_H_ +#endif // STOCHTREE_CATEGORY_TRACKER_H_ diff --git a/include/stochtree/common.h b/include/stochtree/common.h index cd57eea2..5598926a 100644 --- a/include/stochtree/common.h +++ b/include/stochtree/common.h @@ -53,9 +53,9 @@ namespace StochTree { namespace Common { /*! -* Imbues the stream with the C locale. -*/ -static void C_stringstream(std::stringstream &ss) { + * Imbues the stream with the C locale. + */ +static void C_stringstream(std::stringstream& ss) { ss.imbue(std::locale::classic()); } @@ -190,7 +190,7 @@ inline static std::vector Split(const char* c_str, const char* deli return ret; } -template +template inline static const char* Atoi(const char* p, T* out) { int sign; T value; @@ -214,16 +214,16 @@ inline static const char* Atoi(const char* p, T* out) { return p; } -template +template inline static double Pow(T base, int power) { if (power < 0) { return 1.0 / Pow(base, -power); } else if (power == 0) { return 1; } else if (power % 2 == 0) { - return Pow(base*base, power / 2); + return Pow(base * base, power / 2); } else if (power % 3 == 0) { - return Pow(base*base*base, power / 3); + return Pow(base * base * base, power / 3); } else { return base * Pow(base, power - 1); } @@ -285,18 +285,29 @@ inline static const char* Atof(const char* p, double* out) { } if (expon > 308) expon = 308; // Calculate scaling factor. - while (expon >= 50) { scale *= 1E50; expon -= 50; } - while (expon >= 8) { scale *= 1E8; expon -= 8; } - while (expon > 0) { scale *= 10.0; expon -= 1; } + while (expon >= 50) { + scale *= 1E50; + expon -= 50; + } + while (expon >= 8) { + scale *= 1E8; + expon -= 8; + } + while (expon > 0) { + scale *= 10.0; + expon -= 1; + } } // Return signed and scaled floating point result. *out = sign * (frac ? (value / scale) : (value * scale)); } else { size_t cnt = 0; + // clang-format off while (*(p + cnt) != '\0' && *(p + cnt) != ' ' && *(p + cnt) != '\t' && *(p + cnt) != ',' && *(p + cnt) != '\n' && *(p + cnt) != '\r' && *(p + cnt) != ':') { + // clang-format on ++cnt; } if (cnt > 0) { @@ -331,7 +342,7 @@ inline static const char* AtofPrecise(const char* p, double* out) { // Rare path: Not in RFC 7159 format. Possible "inf", "nan", etc. Fallback to standard library: char* end2; - errno = 0; // This is Required before calling strtod. + errno = 0; // This is Required before calling strtod. *out = std::strtod(p, &end2); // strtod is locale aware. if (end2 == p) { Log::Fatal("no conversion to double for: %s", p); @@ -372,7 +383,7 @@ inline static const char* SkipReturn(const char* p) { return p; } -template +template inline static std::vector ArrayCast(const std::vector& arr) { std::vector ret(arr.size()); for (size_t i = 0; i < arr.size(); ++i) { @@ -381,7 +392,7 @@ inline static std::vector ArrayCast(const std::vector& arr) { return ret; } -template +template struct __StringToTHelper { T operator()(const std::string& str) const { T ret = 0; @@ -390,14 +401,14 @@ struct __StringToTHelper { } }; -template +template struct __StringToTHelper { T operator()(const std::string& str) const { return static_cast(std::stod(str)); } }; -template +template inline static std::vector StringToArray(const std::string& str, char delimiter) { std::vector strs = Split(str.c_str(), delimiter); std::vector ret; @@ -409,7 +420,7 @@ inline static std::vector StringToArray(const std::string& str, char delimite return ret; } -template +template inline static std::vector> StringToArrayofArrays( const std::string& str, char left_bracket, char right_bracket, char delimiter) { std::vector strs = SplitBrackets(str.c_str(), left_bracket, right_bracket); @@ -420,7 +431,7 @@ inline static std::vector> StringToArrayofArrays( return ret; } -template +template inline static std::vector StringToArray(const std::string& str, int n) { if (n == 0) { return std::vector(); @@ -436,16 +447,16 @@ inline static std::vector StringToArray(const std::string& str, int n) { return ret; } -template +template struct __StringToTHelperFast { - const char* operator()(const char*p, T* out) const { + const char* operator()(const char* p, T* out) const { return Atoi(p, out); } }; -template +template struct __StringToTHelperFast { - const char* operator()(const char*p, T* out) const { + const char* operator()(const char* p, T* out) const { double tmp = 0.0f; auto ret = Atof(p, &tmp); *out = static_cast(tmp); @@ -453,7 +464,7 @@ struct __StringToTHelperFast { } }; -template +template inline static std::vector StringToArrayFast(const std::string& str, int n) { if (n == 0) { return std::vector(); @@ -467,7 +478,7 @@ inline static std::vector StringToArrayFast(const std::string& str, int n) { return ret; } -template +template inline static std::string Join(const std::vector& strs, const char* delimiter, const bool force_C_locale = false) { if (strs.empty()) { return std::string(""); @@ -485,7 +496,7 @@ inline static std::string Join(const std::vector& strs, const char* delimiter return str_buf.str(); } -template<> +template <> inline std::string Join(const std::vector& strs, const char* delimiter, const bool force_C_locale) { if (strs.empty()) { return std::string(""); @@ -503,7 +514,7 @@ inline std::string Join(const std::vector& strs, const char* del return str_buf.str(); } -template +template inline static std::string Join(const std::vector& strs, size_t start, size_t end, const char* delimiter, const bool force_C_locale = false) { if (end - start <= 0) { return std::string(""); @@ -539,7 +550,7 @@ inline static int64_t Pow2RoundUp(int64_t x) { * \param p_rec The input/output vector of the values. */ inline static void Softmax(std::vector* p_rec) { - std::vector &rec = *p_rec; + std::vector& rec = *p_rec; double wmax = rec[0]; for (size_t i = 1; i < rec.size(); ++i) { wmax = std::max(rec[i], wmax); @@ -569,16 +580,16 @@ inline static void Softmax(const double* input, double* output, int len) { } } -template +template std::vector ConstPtrInVectorWrapper(const std::vector>& input) { std::vector ret; - for (auto t = input.begin(); t !=input.end(); ++t) { + for (auto t = input.begin(); t != input.end(); ++t) { ret.push_back(t->get()); } return ret; } -template +template inline static void SortForPair(std::vector* keys, std::vector* values, size_t start, bool is_reverse = false) { std::vector> arr; auto& ref_key = *keys; @@ -644,14 +655,14 @@ inline static float AvoidInf(float x) { } } -template inline -static typename std::iterator_traits<_Iter>::value_type* IteratorValType(_Iter) { +template +inline static typename std::iterator_traits<_Iter>::value_type* IteratorValType(_Iter) { return (0); } // Check that all y[] are in interval [ymin, ymax] (end points included); throws error if not template -inline static void CheckElementsIntervalClosed(const T *y, T ymin, T ymax, int ny, const char *callername) { +inline static void CheckElementsIntervalClosed(const T* y, T ymin, T ymax, int ny, const char* callername) { auto fatal_msg = [&y, &ymin, &ymax, &callername](int i) { std::ostringstream os; os << "[%s]: does not tolerate element [#%i = " << y[i] << "] outside [" << ymin << ", " << ymax << "]"; @@ -682,7 +693,7 @@ inline static void CheckElementsIntervalClosed(const T *y, T ymin, T ymax, int n // One-pass scan over array w with nw elements: find min, max and sum of elements; // this is useful for checking weight requirements. template -inline static void ObtainMinMaxSum(const T1 *w, int nw, T1 *mi, T1 *ma, T2 *su) { +inline static void ObtainMinMaxSum(const T1* w, int nw, T1* mi, T1* ma, T2* su) { T1 minw; T1 maxw; T1 sumw; @@ -730,7 +741,7 @@ inline static std::vector EmptyBitset(int n) { return std::vector(size); } -template +template inline static void InsertBitset(std::vector* vec, const T val) { auto& ref_v = *vec; int i1 = val / 32; @@ -741,7 +752,7 @@ inline static void InsertBitset(std::vector* vec, const T val) { ref_v[i1] |= (1 << i2); } -template +template inline static std::vector ConstructBitset(const T* vals, int n) { std::vector ret; for (int i = 0; i < n; ++i) { @@ -755,7 +766,7 @@ inline static std::vector ConstructBitset(const T* vals, int n) { return ret; } -template +template inline static bool FindInBitset(const uint32_t* bits, int n, T pos) { int i1 = pos / 32; if (i1 >= n) { @@ -817,7 +828,7 @@ inline bool CheckAllowedJSON(const std::string& s) { || char_code == 93 // ] || char_code == 123 // { || char_code == 125 // } - ) { + ) { return false; } } @@ -979,19 +990,18 @@ class FunctionTimer { extern Common::Timer global_timer; - /*! -* Provides locale-independent alternatives to Common's methods. -* Essential to make models robust to locale settings. -*/ + * Provides locale-independent alternatives to Common's methods. + * Essential to make models robust to locale settings. + */ namespace CommonC { -template +template inline static std::string Join(const std::vector& strs, const char* delimiter) { return StochTree::Common::Join(strs, delimiter, true); } -template +template inline static std::string Join(const std::vector& strs, size_t start, size_t end, const char* delimiter) { return StochTree::Common::Join(strs, start, end, delimiter, true); } @@ -1000,22 +1010,22 @@ inline static const char* Atof(const char* p, double* out) { return StochTree::Common::Atof(p, out); } -template +template struct __StringToTHelperFast { - const char* operator()(const char*p, T* out) const { + const char* operator()(const char* p, T* out) const { return StochTree::Common::Atoi(p, out); } }; /*! -* \warning Beware that ``Common::Atof`` in ``__StringToTHelperFast``, -* has **less** floating point precision than ``__StringToTHelper``. -* Both versions are kept to maintain bit-for-bit the "legacy" LightGBM behaviour in terms of precision. -* Check ``StringToArrayFast`` and ``StringToArray`` for more details on this. -*/ -template + * \warning Beware that ``Common::Atof`` in ``__StringToTHelperFast``, + * has **less** floating point precision than ``__StringToTHelper``. + * Both versions are kept to maintain bit-for-bit the "legacy" LightGBM behaviour in terms of precision. + * Check ``StringToArrayFast`` and ``StringToArray`` for more details on this. + */ +template struct __StringToTHelperFast { - const char* operator()(const char*p, T* out) const { + const char* operator()(const char* p, T* out) const { double tmp = 0.0f; auto ret = Atof(p, &tmp); *out = static_cast(tmp); @@ -1023,7 +1033,7 @@ struct __StringToTHelperFast { } }; -template +template struct __StringToTHelper { T operator()(const std::string& str) const { T ret = 0; @@ -1033,35 +1043,34 @@ struct __StringToTHelper { }; /*! -* \warning Beware that ``Common::Atof`` in ``__StringToTHelperFast``, -* has **less** floating point precision than ``__StringToTHelper``. -* Both versions are kept to maintain bit-for-bit the "legacy" LightGBM behaviour in terms of precision. -* Check ``StringToArrayFast`` and ``StringToArray`` for more details on this. -* \note It is possible that ``fast_double_parser::parse_number`` is faster than ``Common::Atof``. -*/ -template + * \warning Beware that ``Common::Atof`` in ``__StringToTHelperFast``, + * has **less** floating point precision than ``__StringToTHelper``. + * Both versions are kept to maintain bit-for-bit the "legacy" LightGBM behaviour in terms of precision. + * Check ``StringToArrayFast`` and ``StringToArray`` for more details on this. + * \note It is possible that ``fast_double_parser::parse_number`` is faster than ``Common::Atof``. + */ +template struct __StringToTHelper { T operator()(const std::string& str) const { double tmp; const char* end = Common::AtofPrecise(str.c_str(), &tmp); if (end == str.c_str()) { - Log::Fatal("Failed to parse double: %s", str.c_str()); + Log::Fatal("Failed to parse double: %s", str.c_str()); } return static_cast(tmp); } }; - /*! -* \warning Beware that due to internal use of ``Common::Atof`` in ``__StringToTHelperFast``, -* this method has less precision for floating point numbers than ``StringToArray``, -* which calls ``__StringToTHelper``. -* As such, ``StringToArrayFast`` and ``StringToArray`` are not equivalent! -* Both versions were kept to maintain bit-for-bit the "legacy" LightGBM behaviour in terms of precision. -*/ -template + * \warning Beware that due to internal use of ``Common::Atof`` in ``__StringToTHelperFast``, + * this method has less precision for floating point numbers than ``StringToArray``, + * which calls ``__StringToTHelper``. + * As such, ``StringToArrayFast`` and ``StringToArray`` are not equivalent! + * Both versions were kept to maintain bit-for-bit the "legacy" LightGBM behaviour in terms of precision. + */ +template inline static std::vector StringToArrayFast(const std::string& str, int n) { if (n == 0) { return std::vector(); @@ -1076,11 +1085,11 @@ inline static std::vector StringToArrayFast(const std::string& str, int n) { } /*! -* \warning Do not replace calls to this method by ``StringToArrayFast``. -* This method is more precise for floating point numbers. -* Check ``StringToArrayFast`` for more details. -*/ -template + * \warning Do not replace calls to this method by ``StringToArrayFast``. + * This method is more precise for floating point numbers. + * Check ``StringToArrayFast`` for more details. + */ +template inline static std::vector StringToArray(const std::string& str, int n) { if (n == 0) { return std::vector(); @@ -1097,11 +1106,11 @@ inline static std::vector StringToArray(const std::string& str, int n) { } /*! -* \warning Do not replace calls to this method by ``StringToArrayFast``. -* This method is more precise for floating point numbers. -* Check ``StringToArrayFast`` for more details. -*/ -template + * \warning Do not replace calls to this method by ``StringToArrayFast``. + * This method is more precise for floating point numbers. + * Check ``StringToArrayFast`` for more details. + */ +template inline static std::vector StringToArray(const std::string& str, char delimiter) { std::vector strs = StochTree::Common::Split(str.c_str(), delimiter); std::vector ret; @@ -1114,37 +1123,37 @@ inline static std::vector StringToArray(const std::string& str, char delimite } /*! -* Safely formats a value onto a buffer according to a format string and null-terminates it. -* -* \note It checks that the full value was written or forcefully aborts. -* This safety check serves to prevent incorrect internal API usage. -* Correct usage will never incur in this problem: -* - The received buffer size shall be sufficient at all times for the input format string and value. -*/ + * Safely formats a value onto a buffer according to a format string and null-terminates it. + * + * \note It checks that the full value was written or forcefully aborts. + * This safety check serves to prevent incorrect internal API usage. + * Correct usage will never incur in this problem: + * - The received buffer size shall be sufficient at all times for the input format string and value. + */ template inline static void format_to_buf(char* buffer, const size_t buf_len, const char* format, const T value) { - auto result = fmt::format_to_n(buffer, buf_len, format, value); - if (result.size >= buf_len) { - Log::Fatal("Numerical conversion failed. Buffer is too small."); - } - buffer[result.size] = '\0'; + auto result = fmt::format_to_n(buffer, buf_len, format, value); + if (result.size >= buf_len) { + Log::Fatal("Numerical conversion failed. Buffer is too small."); + } + buffer[result.size] = '\0'; } -template +template struct __TToStringHelper { void operator()(T value, char* buffer, size_t buf_len) const { format_to_buf(buffer, buf_len, "{}", value); } }; -template +template struct __TToStringHelper { void operator()(T value, char* buffer, size_t buf_len) const { format_to_buf(buffer, buf_len, "{:g}", value); } }; -template +template struct __TToStringHelper { void operator()(T value, char* buffer, size_t buf_len) const { format_to_buf(buffer, buf_len, "{:.17g}", value); @@ -1152,14 +1161,14 @@ struct __TToStringHelper { }; /*! -* Converts an array to a string with with values separated by the space character. -* This method replaces Common's ``ArrayToString`` and ``ArrayToStringFast`` functionality -* and is locale-independent. -* -* \note If ``high_precision_output`` is set to true, -* floating point values are output with more digits of precision. -*/ -template + * Converts an array to a string with with values separated by the space character. + * This method replaces Common's ``ArrayToString`` and ``ArrayToStringFast`` functionality + * and is locale-independent. + * + * \note If ``high_precision_output`` is set to true, + * floating point values are output with more digits of precision. + */ +template inline static std::string ArrayToString(const std::vector& arr, size_t n) { if (arr.empty() || n == 0) { return std::string(""); diff --git a/include/stochtree/container.h b/include/stochtree/container.h index 4b75ef2f..7e44d148 100644 --- a/include/stochtree/container.h +++ b/include/stochtree/container.h @@ -1,6 +1,6 @@ /*! * Copyright (c) 2024 stochtree authors. All rights reserved. - * + * * Simple container-like interfaces for samples of common models. */ #ifndef STOCHTREE_CONTAINER_H_ @@ -8,6 +8,7 @@ #include #include +#include #include #include @@ -23,8 +24,8 @@ namespace StochTree { class ForestContainer { public: /*! - * \brief Construct a new ForestContainer object. - * + * \brief Construct a new ForestContainer object. + * * \param num_trees Number of trees in each forest. * \param output_dimension Dimension of the leaf node parameter in each tree of each forest. * \param is_leaf_constant Whether or not the leaves of each tree are treated as "constant." If true, then predicting from an ensemble is simply a matter or determining which leaf node an observation falls into. If false, prediction will multiply a leaf node's parameter(s) for a given observation by a basis vector. @@ -33,7 +34,7 @@ class ForestContainer { ForestContainer(int num_trees, int output_dimension = 1, bool is_leaf_constant = true, bool is_exponentiated = false); /*! * \brief Construct a new ForestContainer object. - * + * * \param num_samples Initial size of a container of forest samples. * \param num_trees Number of trees in each forest. * \param output_dimension Dimension of the leaf node parameter in each tree of each forest. @@ -44,7 +45,7 @@ class ForestContainer { ~ForestContainer() {} /*! * \brief Combine two forests into a single forest by merging their trees - * + * * \param inbound_forest_index Index of the forest that will be appended to * \param outbound_forest_index Index of the forest that will be appended */ @@ -53,7 +54,7 @@ class ForestContainer { } /*! * \brief Add a constant value to every leaf of every tree of a specified forest - * + * * \param forest_index Index of forest whose leaves will be modified * \param constant_value Value to add to every leaf of every tree of the forest at `forest_index` */ @@ -62,7 +63,7 @@ class ForestContainer { } /*! * \brief Multiply every leaf of every tree of a specified forest by a constant value - * + * * \param forest_index Index of forest whose leaves will be modified * \param constant_multiple Value to multiply through by every leaf of every tree of the forest at `forest_index` */ @@ -71,83 +72,83 @@ class ForestContainer { } /*! * \brief Remove a forest from a container of forest samples and delete the corresponding object, freeing its memory. - * + * * \param sample_num Index of forest to be deleted. */ void DeleteSample(int sample_num); /*! * \brief Add a new forest to the container by copying `forest`. - * + * * \param forest Forest to be copied and added to the container of retained forest samples. */ void AddSample(TreeEnsemble& forest); /*! * \brief Initialize a "root" forest of univariate trees as the first element of the container, setting all root node values in every tree to `leaf_value`. - * + * * \param leaf_value Value to assign to the root node of every tree. */ void InitializeRoot(double leaf_value); /*! * \brief Initialize a "root" forest of multivariate trees as the first element of the container, setting all root node values in every tree to `leaf_vector`. - * + * * \param leaf_value Vector of values to assign to the root node of every tree. */ void InitializeRoot(std::vector& leaf_vector); /*! * \brief Pre-allocate space for `num_samples` additional forests in the container. - * + * * \param num_samples Number of (default-constructed) forests to allocated space for in the container. */ void AddSamples(int num_samples); /*! * \brief Copy the forest stored at `previous_sample_id` to the forest stored at `new_sample_id`. - * + * * \param new_sample_id Index of the new forest to be copied from an earlier sample. * \param previous_sample_id Index of the previous forest to copy to `new_sample_id`. */ void CopyFromPreviousSample(int new_sample_id, int previous_sample_id); /*! - * \brief Predict from every forest in the container on every observation in the provided dataset. - * The resulting vector is "column-major", where every forest in a container defines the columns of a - * prediction matrix and every observation in the provided dataset defines the rows. The (`i`,`j`) element - * of this prediction matrix can be read from the `j * num_rows + i` element of the returned `std::vector`, + * \brief Predict from every forest in the container on every observation in the provided dataset. + * The resulting vector is "column-major", where every forest in a container defines the columns of a + * prediction matrix and every observation in the provided dataset defines the rows. The (`i`,`j`) element + * of this prediction matrix can be read from the `j * num_rows + i` element of the returned `std::vector`, * where `num_rows` is equal to the number of observations in `dataset` (i.e. `dataset.NumObservations()`). - * + * * \param dataset Data object containining training data, including covariates, leaf regression bases, and case weights. * \return std::vector Vector of predictions for every forest in the container and every observation in `dataset`. */ std::vector Predict(ForestDataset& dataset); /*! - * \brief Predict from every forest in the container on every observation in the provided dataset. - * The resulting vector stores a possibly three-dimensional array, where the dimensions are arranged as follows - * + * \brief Predict from every forest in the container on every observation in the provided dataset. + * The resulting vector stores a possibly three-dimensional array, where the dimensions are arranged as follows + * * 1. Dimension of the leaf node's raw values (1 for GaussianConstantLeafModel, GaussianUnivariateRegressionLeafModel, and LogLinearVarianceLeafModel, >1 for GaussianMultivariateRegressionLeafModel) * 2. Observations in the provided dataset. * 3. Forest samples in the container. - * + * * If the leaf nodes have univariate values, then the "first dimension" is 1 and the resulting array has the exact same layout as in \ref Predict. - * + * * \param dataset Data object containining training data, including covariates, leaf regression bases, and case weights. * \return std::vector Vector of predictions for every forest in the container and every observation in `dataset`. */ - std::vector PredictRaw(ForestDataset& dataset); + std::vector PredictRaw(ForestDataset& dataset, bool row_major = true); std::vector PredictRaw(ForestDataset& dataset, int forest_num); std::vector PredictRawSingleTree(ForestDataset& dataset, int forest_num, int tree_num); void PredictInPlace(ForestDataset& dataset, std::vector& output); - void PredictRawInPlace(ForestDataset& dataset, std::vector& output); + void PredictRawInPlace(ForestDataset& dataset, std::vector& output, bool row_major = true); void PredictRawInPlace(ForestDataset& dataset, int forest_num, std::vector& output); void PredictRawSingleTreeInPlace(ForestDataset& dataset, int forest_num, int tree_num, std::vector& output); - void PredictLeafIndicesInplace(Eigen::Map>& covariates, - Eigen::Map>& output, + void PredictLeafIndicesInplace(Eigen::Map>& covariates, + Eigen::Map>& output, std::vector& forest_indices, int num_trees, data_size_t n); - inline TreeEnsemble* GetEnsemble(int i) {return forests_[i].get();} - inline int32_t NumSamples() {return num_samples_;} - inline int32_t NumTrees() {return num_trees_;} - inline int32_t NumTrees(int ensemble_num) {return forests_[ensemble_num]->NumTrees();} - inline int32_t NumLeaves(int ensemble_num) {return forests_[ensemble_num]->NumLeaves();} - inline int32_t EnsembleTreeMaxDepth(int ensemble_num, int tree_num) {return forests_[ensemble_num]->TreeMaxDepth(tree_num);} - inline double EnsembleAverageMaxDepth(int ensemble_num) {return forests_[ensemble_num]->AverageMaxDepth();} + inline TreeEnsemble* GetEnsemble(int i) { return forests_[i].get(); } + inline int32_t NumSamples() { return num_samples_; } + inline int32_t NumTrees() { return num_trees_; } + inline int32_t NumTrees(int ensemble_num) { return forests_[ensemble_num]->NumTrees(); } + inline int32_t NumLeaves(int ensemble_num) { return forests_[ensemble_num]->NumLeaves(); } + inline int32_t EnsembleTreeMaxDepth(int ensemble_num, int tree_num) { return forests_[ensemble_num]->TreeMaxDepth(tree_num); } + inline double EnsembleAverageMaxDepth(int ensemble_num) { return forests_[ensemble_num]->AverageMaxDepth(); } inline double AverageMaxDepth() { double numerator = 0.; double denominator = 0.; @@ -159,23 +160,23 @@ class ForestContainer { } return numerator / denominator; } - inline int32_t OutputDimension() {return output_dimension_;} - inline int32_t OutputDimension(int ensemble_num) {return forests_[ensemble_num]->OutputDimension();} - inline bool IsLeafConstant() {return is_leaf_constant_;} - inline bool IsLeafConstant(int ensemble_num) {return forests_[ensemble_num]->IsLeafConstant();} - inline bool IsExponentiated() {return is_exponentiated_;} - inline bool IsExponentiated(int ensemble_num) {return forests_[ensemble_num]->IsExponentiated();} - inline bool AllRoots(int ensemble_num) {return forests_[ensemble_num]->AllRoots();} - inline void SetLeafValue(int ensemble_num, double leaf_value) {forests_[ensemble_num]->SetLeafValue(leaf_value);} - inline void SetLeafVector(int ensemble_num, std::vector& leaf_vector) {forests_[ensemble_num]->SetLeafVector(leaf_vector);} - inline void IncrementSampleCount() {num_samples_++;} + inline int32_t OutputDimension() { return output_dimension_; } + inline int32_t OutputDimension(int ensemble_num) { return forests_[ensemble_num]->OutputDimension(); } + inline bool IsLeafConstant() { return is_leaf_constant_; } + inline bool IsLeafConstant(int ensemble_num) { return forests_[ensemble_num]->IsLeafConstant(); } + inline bool IsExponentiated() { return is_exponentiated_; } + inline bool IsExponentiated(int ensemble_num) { return forests_[ensemble_num]->IsExponentiated(); } + inline bool AllRoots(int ensemble_num) { return forests_[ensemble_num]->AllRoots(); } + inline void SetLeafValue(int ensemble_num, double leaf_value) { forests_[ensemble_num]->SetLeafValue(leaf_value); } + inline void SetLeafVector(int ensemble_num, std::vector& leaf_vector) { forests_[ensemble_num]->SetLeafVector(leaf_vector); } + inline void IncrementSampleCount() { num_samples_++; } void SaveToJsonFile(std::string filename) { nlohmann::json model_json = this->to_json(); std::ofstream output_file(filename); output_file << model_json << std::endl; } - + void LoadFromJsonFile(std::string filename) { std::ifstream f(filename); nlohmann::json file_tree_json = nlohmann::json::parse(f); @@ -219,6 +220,23 @@ class ForestContainer { bool is_leaf_constant_; bool initialized_{false}; }; -} // namespace StochTree -#endif // STOCHTREE_CONTAINER_H_ +/*! + * \brief Append every retained forest sample from `src` onto the end of `dst` (deep copy). + * Used by BARTSamples::Merge / BCFSamples::Merge to combine independently-sampled chains. The forest + * must be present in both samples objects or neither; `name` labels the forest in the mismatch error. + */ +inline void AppendForestContainerSamples(std::unique_ptr& dst, + const std::unique_ptr& src, + const char* name) { + if (src == nullptr && dst == nullptr) return; + if (src == nullptr || dst == nullptr) { + Log::Fatal("Cannot merge samples: %s forest present in one chain but not the other", name); + } + for (int i = 0; i < src->NumSamples(); i++) { + dst->AddSample(*src->GetEnsemble(i)); + } +} +} // namespace StochTree + +#endif // STOCHTREE_CONTAINER_H_ diff --git a/include/stochtree/cutpoint_candidates.h b/include/stochtree/cutpoint_candidates.h index 76f1df4c..a4503c19 100644 --- a/include/stochtree/cutpoint_candidates.h +++ b/include/stochtree/cutpoint_candidates.h @@ -1,39 +1,39 @@ /*! * Copyright (c) 2024 stochtree authors. - * + * * Data structures for enumerating potential cutpoint candidates. - * + * * This is used in the XBART family of algorithms, which samples split rules - * based on the log marginal likelihood of every potential cutpoint. For numeric - * variables with large sample sizes, it is often unnecessary to consider every + * based on the log marginal likelihood of every potential cutpoint. For numeric + * variables with large sample sizes, it is often unnecessary to consider every * unique value, so we allow for an adaptive "grid" of potential cutpoint values. - * - * Algorithms for enumerating cutpoints take Dataset and SortedNodeSampleTracker objects - * as inputs, so that each feature is "pre-sorted" according to its value within a - * given node. The size of the adaptive cutpoint grid is set by the + * + * Algorithms for enumerating cutpoints take Dataset and SortedNodeSampleTracker objects + * as inputs, so that each feature is "pre-sorted" according to its value within a + * given node. The size of the adaptive cutpoint grid is set by the * cutpoint_grid_size configuration parameter. - * + * * Numeric Features * ---------------- - * - * When a node has fewer available observations than cutpoint_grid_size, - * full enumeration of unique available cutpoints is done via the + * + * When a node has fewer available observations than cutpoint_grid_size, + * full enumeration of unique available cutpoints is done via the * `EnumerateNumericCutpointsDeduplication` function - * - * When a node has more available observations than cutpoint_grid_size, - * potential cutpoints are "thinned out" by considering every k-th observation, + * + * When a node has more available observations than cutpoint_grid_size, + * potential cutpoints are "thinned out" by considering every k-th observation, * where k is implied by the number of observations and the target cutpoint_grid_size. - * + * * Ordered Categorical Features * ---------------------------- - * - * In this case, the grid is every unique value of the ordered categorical + * + * In this case, the grid is every unique value of the ordered categorical * feature in ascending order. - * + * * Unordered Categorical Features * ------------------------------ - * - * In this case, the grid is every unique value of the unordered categorical feature, + * + * In this case, the grid is every unique value of the unordered categorical feature, * arranged in an outcome-dependent order, as described in Fisher (1958) */ #ifndef STOCHTREE_CUTPOINT_CANDIDATES_H_ @@ -45,7 +45,7 @@ namespace StochTree { /*! \brief Computing and tracking cutpoints available for a given feature at a given node - * Store cutpoint bins in 0-indexed fashion, so that if a given node has + * Store cutpoint bins in 0-indexed fashion, so that if a given node has */ class FeatureCutpointGrid { public: @@ -66,19 +66,19 @@ class FeatureCutpointGrid { void CalculateStridesUnorderedCategorical(Eigen::MatrixXd& covariates, Eigen::VectorXd& residuals, SortedNodeSampleTracker* feature_node_sort_tracker, int32_t node_id, data_size_t node_begin, data_size_t node_end, int32_t feature_index); /*! \brief Number of potential cutpoints enumerated */ - int32_t NumCutpoints() {return node_stride_begin_.size();} + int32_t NumCutpoints() { return node_stride_begin_.size(); } /*! \brief Beginning index of bin i */ - int32_t BinStartIndex(int i) {return node_stride_begin_.at(i);} + int32_t BinStartIndex(int i) { return node_stride_begin_.at(i); } /*! \brief Size of bin i */ - int32_t BinLength(int i) {return node_stride_length_.at(i);} + int32_t BinLength(int i) { return node_stride_length_.at(i); } /*! \brief Beginning index of bin i */ - int32_t BinEndIndex(int i) {return node_stride_begin_.at(i) + node_stride_length_.at(i);} + int32_t BinEndIndex(int i) { return node_stride_begin_.at(i) + node_stride_length_.at(i); } /*! \brief Value of the upper-bound (cutpoint) implied by bin i */ - double CutpointValue(int i) {return cutpoint_values_.at(i);} + double CutpointValue(int i) { return cutpoint_values_.at(i); } /*! \brief Vector of cutpoint values up to and including bin i * Helper function for converting categorical split "value" (as outlined in Fisher 1958) to a set of categories @@ -135,22 +135,22 @@ class CutpointGridContainer { } /*! \brief Max size of cutpoint grid */ - int32_t CutpointGridSize() {return cutpoint_grid_size_;} + int32_t CutpointGridSize() { return cutpoint_grid_size_; } /*! \brief Number of potential cutpoints enumerated */ - int32_t NumCutpoints(int feature_index) {return feature_cutpoint_grid_[feature_index]->NumCutpoints();} + int32_t NumCutpoints(int feature_index) { return feature_cutpoint_grid_[feature_index]->NumCutpoints(); } /*! \brief Beginning index of bin i */ - int32_t BinStartIndex(int i, int feature_index) {return feature_cutpoint_grid_[feature_index]->BinStartIndex(i);} + int32_t BinStartIndex(int i, int feature_index) { return feature_cutpoint_grid_[feature_index]->BinStartIndex(i); } /*! \brief Size of bin i */ - int32_t BinLength(int i, int feature_index) {return feature_cutpoint_grid_[feature_index]->BinLength(i);} + int32_t BinLength(int i, int feature_index) { return feature_cutpoint_grid_[feature_index]->BinLength(i); } /*! \brief Beginning index of bin i */ - int32_t BinEndIndex(int i, int feature_index) {return feature_cutpoint_grid_[feature_index]->BinEndIndex(i);} + int32_t BinEndIndex(int i, int feature_index) { return feature_cutpoint_grid_[feature_index]->BinEndIndex(i); } /*! \brief Value of the upper-bound (cutpoint) implied by bin i */ - double CutpointValue(int i, int feature_index) {return feature_cutpoint_grid_[feature_index]->CutpointValue(i);} + double CutpointValue(int i, int feature_index) { return feature_cutpoint_grid_[feature_index]->CutpointValue(i); } /*! \brief Vector of cutpoint values up to and including bin i * Helper function for converting categorical split "value" (as outlined in Fisher 1958) to a set of categories @@ -159,7 +159,7 @@ class CutpointGridContainer { return feature_cutpoint_grid_[feature_index]->CutpointVector(i); } - FeatureCutpointGrid* GetFeatureCutpointGrid(int feature_num) {return feature_cutpoint_grid_[feature_num].get(); } + FeatureCutpointGrid* GetFeatureCutpointGrid(int feature_num) { return feature_cutpoint_grid_[feature_num].get(); } private: std::vector> feature_cutpoint_grid_; @@ -184,7 +184,7 @@ class NodeCutpointTracker { void CalculateStridesCategorical(Eigen::MatrixXd& covariates, Eigen::VectorXd& residuals, SortedNodeSampleTracker* feature_node_sort_tracker, data_size_t node_begin, data_size_t node_end, int32_t feature_index); /*! \brief Number of potential cutpoints enumerated */ - int32_t NumCutpoints() {return node_stride_begin_.size();} + int32_t NumCutpoints() { return node_stride_begin_.size(); } /*! \brief Whether a cutpoint grid has been enumerated for a given node */ bool NodeCutpointEvaluated(int32_t node_id) { @@ -192,18 +192,18 @@ class NodeCutpointTracker { } /*! \brief Node id of the node that has been most recently evaluated */ - int32_t CurrentNodeEvaluated() {return current_node_;} + int32_t CurrentNodeEvaluated() { return current_node_; } /*! \brief Vectors of node stride starting points and stride lengths */ std::vector node_stride_begin_; std::vector node_stride_length_; - + private: int32_t cutpoint_grid_size_; std::vector nodes_enumerated_; int32_t current_node_; }; -} // namespace StochTree +} // namespace StochTree -#endif // STOCHTREE_CUTPOINT_CANDIDATES_H_ +#endif // STOCHTREE_CUTPOINT_CANDIDATES_H_ diff --git a/include/stochtree/data.h b/include/stochtree/data.h index 393203b1..10500793 100644 --- a/include/stochtree/data.h +++ b/include/stochtree/data.h @@ -21,10 +21,10 @@ namespace StochTree { */ /*! - * \brief Extract multiple features from the raw data loaded from a file into an `Eigen::MatrixXd`. + * \brief Extract multiple features from the raw data loaded from a file into an `Eigen::MatrixXd`. * Lightly modified from LightGBM's datasetloader interface to support `stochtree`'s use cases. * \internal - * + * * \param text_data Vector of data reads as string from a file. * \param parser Pointer to a parser object (i.e. `CSVParser`). * \param column_indices Integer labels of columns to be extracted from `text_data` into `data`. @@ -32,7 +32,7 @@ namespace StochTree { * \param num_rows Number of observations in the data being loaded. */ static inline void ExtractMultipleFeaturesFromMemory(std::vector* text_data, const Parser* parser, - std::vector& column_indices, Eigen::MatrixXd& data, + std::vector& column_indices, Eigen::MatrixXd& data, data_size_t num_rows) { std::vector> oneline_features; auto& ref_text_data = *text_data; @@ -42,7 +42,7 @@ static inline void ExtractMultipleFeaturesFromMemory(std::vector* t // unpack the vector of textlines read from file into a vector of (int, double) tuples oneline_features.clear(); parser->ParseOneLine(ref_text_data[i].c_str(), &oneline_features); - + // free processed line: ref_text_data[i].clear(); @@ -50,9 +50,8 @@ static inline void ExtractMultipleFeaturesFromMemory(std::vector* t int feature_counter = 0; for (auto& inner_data : oneline_features) { int feature_idx = inner_data.first; - column_matched = (std::find(column_indices.begin(), column_indices.end(), feature_idx) - != column_indices.end()); - if (column_matched){ + column_matched = (std::find(column_indices.begin(), column_indices.end(), feature_idx) != column_indices.end()); + if (column_matched) { data(i, feature_counter) = inner_data.second; feature_counter += 1; } @@ -63,10 +62,10 @@ static inline void ExtractMultipleFeaturesFromMemory(std::vector* t } /*! -* \brief Extract a single feature from the raw data loaded from a file into an `Eigen::VectorXd`. + * \brief Extract a single feature from the raw data loaded from a file into an `Eigen::VectorXd`. * Lightly modified from LightGBM's datasetloader interface to support `stochtree`'s use cases. * \internal - * + * * \param text_data Vector of data reads as string from a file. * \param parser Pointer to a parser object (i.e. `CSVParser`). * \param column_index Integer labels of columns to be extracted from `text_data` into `data`. @@ -74,7 +73,7 @@ static inline void ExtractMultipleFeaturesFromMemory(std::vector* t * \param num_rows Number of observations in the data being loaded. */ static inline void ExtractSingleFeatureFromMemory(std::vector* text_data, const Parser* parser, - int32_t column_index, Eigen::VectorXd& data, data_size_t num_rows) { + int column_index, Eigen::VectorXd& data, data_size_t num_rows) { std::vector> oneline_features; auto& ref_text_data = *text_data; bool column_matched; @@ -82,14 +81,14 @@ static inline void ExtractSingleFeatureFromMemory(std::vector* text // unpack the vector of textlines read from file into a vector of (int, double) tuples oneline_features.clear(); parser->ParseOneLine(ref_text_data[i].c_str(), &oneline_features); - + // free processed line: ref_text_data[i].clear(); // unload the data from oneline_features vector into the dataset variables containers for (auto& inner_data : oneline_features) { int feature_idx = inner_data.first; - if (column_index == feature_idx){ + if (column_index == feature_idx) { data(i) = inner_data.second; } } @@ -106,7 +105,7 @@ static inline std::vector LoadTextDataToMemory(const char* filename return std::move(text_reader.Lines()); } -static inline void FeatureUnpack(std::vector& categorical_variables, const char* var_id) { +static inline void FeatureUnpack(std::vector& categorical_variables, const char* var_id) { std::string var_clean = Common::RemoveQuotationSymbol(Common::Trim(var_id)); int out; bool success = Common::AtoiAndCheck(var_clean.c_str(), &out); @@ -120,7 +119,7 @@ static inline void FeatureUnpack(std::vector& categorical_variables, co static inline std::vector Str2FeatureVec(const char* parameters) { std::vector feature_vec; auto args = Common::Split(parameters, ","); - for (auto arg : args) { + for (const auto& arg : args) { FeatureUnpack(feature_vec, Common::Trim(arg).c_str()); } return feature_vec; @@ -134,7 +133,7 @@ class ColumnMatrix { ColumnMatrix() {} /*! * \brief Construct a new `ColumnMatrix` object from in-memory data buffer. - * + * * \param data_ptr Pointer to first element of a contiguous array of data storing a matrix. * \param num_row Number of rows in the matrix. * \param num_col Number of columns / covariates in the matrix. @@ -143,7 +142,7 @@ class ColumnMatrix { ColumnMatrix(double* data_ptr, data_size_t num_row, int num_col, bool is_row_major); /*! * \brief Construct a new ColumnMatrix object from CSV file - * + * * \param filename Name of the file (including any necessary path prefixes). * \param column_index_string Comma-delimited string listing columns to extract into covariates matrix. * \param header Whether or not the file contains a header of column names / non-data. @@ -153,11 +152,11 @@ class ColumnMatrix { ~ColumnMatrix() {} /*! * \brief Returns the value stored at (`row`, `col`) in the object's internal `Eigen::MatrixXd`. - * + * * \param row Row number to query in the matrix * \param col Column number to query in the matrix */ - double GetElement(data_size_t row_num, int32_t col_num) {return data_(row_num, col_num);} + double GetElement(data_size_t row_num, int col_num) { return data_(row_num, col_num); } /*! * \brief Update an observation in the object's internal `Eigen::MatrixXd` to a new value. * @@ -165,10 +164,10 @@ class ColumnMatrix { * \param col Column number to be overwritten. * \param value New value to write in (`row`, `col`) in the object's internal `Eigen::MatrixXd`. */ - void SetElement(data_size_t row_num, int32_t col_num, double value) {data_(row_num, col_num) = value;} + void SetElement(data_size_t row_num, int col_num, double value) { data_(row_num, col_num) = value; } /*! * \brief Update the data in a `ColumnMatrix` object from an in-memory data buffer. This will erase the existing matrix. - * + * * \param data_ptr Pointer to first element of a contiguous array of data storing a matrix. * \param num_row Number of rows in the matrix. * \param num_col Number of columns / covariates in the matrix. @@ -176,18 +175,19 @@ class ColumnMatrix { */ void LoadData(double* data_ptr, data_size_t num_row, int num_col, bool is_row_major); /*! \brief Number of rows in the object's internal `Eigen::MatrixXd`. */ - inline data_size_t NumRows() {return data_.rows();} + inline data_size_t NumRows() { return data_.rows(); } /*! \brief Number of columns in the object's internal `Eigen::MatrixXd`. */ - inline int NumCols() {return data_.cols();} + inline int NumCols() { return data_.cols(); } /*! \brief Return a reference to the object's internal `Eigen::MatrixXd`, for interfaces that require a raw matrix. */ - inline Eigen::MatrixXd& GetData() {return data_;} + inline Eigen::MatrixXd& GetData() { return data_; } + private: Eigen::MatrixXd data_; }; /*! - * \brief Internal wrapper around `Eigen::VectorXd` interface for univariate floating point data. - * The (frequently updated) full / partial residual used in sampling forests is stored internally + * \brief Internal wrapper around `Eigen::VectorXd` interface for univariate floating point data. + * The (frequently updated) full / partial residual used in sampling forests is stored internally * as a `ColumnVector` by the sampling functions (see \ref sampling_group). */ class ColumnVector { @@ -195,77 +195,78 @@ class ColumnVector { ColumnVector() {} /*! * \brief Construct a new `ColumnVector` object from in-memory data buffer. - * + * * \param data_ptr Pointer to first element of a contiguous array of data storing a vector. * \param num_row Number of rows / elements in the vector. */ ColumnVector(double* data_ptr, data_size_t num_row); /*! * \brief Construct a new ColumnMatrix object from CSV file - * + * * \param filename Name of the file (including any necessary path prefixes). * \param column_index Integer index of the column in `filename` to be unpacked as a vector. * \param header Whether or not the file contains a header of column names / non-data. * \param precise_float_parser Whether floating point numbers in the CSV should be parsed precisely. */ - ColumnVector(std::string filename, int32_t column_index, bool header = true, bool precise_float_parser = false); + ColumnVector(std::string filename, int column_index, bool header = true, bool precise_float_parser = false); ~ColumnVector() {} /*! * \brief Returns the value stored at position `row` in the object's internal `Eigen::VectorXd`. - * + * * \param row Row number to query in the vector */ - double GetElement(data_size_t row) {return data_(row);} + double GetElement(data_size_t row) { return data_(row); } /*! * \brief Returns the value stored at position `row` in the object's internal `Eigen::VectorXd`. - * + * * \param row Row number to query in the vector * \param value New value to write to element `row` of the object's internal `Eigen::VectorXd`. */ - void SetElement(data_size_t row, double value) {data_(row) = value;} + void SetElement(data_size_t row, double value) { data_(row) = value; } /*! * \brief Update the data in a `ColumnVector` object from an in-memory data buffer. This will erase the existing vector. - * + * * \param data_ptr Pointer to first element of a contiguous array of data storing a vector. * \param num_row Number of rows / elements in the vector. */ void LoadData(double* data_ptr, data_size_t num_row); /*! - * \brief Update the data in a `ColumnVector` object from an in-memory data buffer, by adding each value obtained + * \brief Update the data in a `ColumnVector` object from an in-memory data buffer, by adding each value obtained * in `data_ptr` to the existing values in the object's internal `Eigen::VectorXd`. - * + * * \param data_ptr Pointer to first element of a contiguous array of data storing a vector. * \param num_row Number of rows / elements in the vector. */ void AddToData(double* data_ptr, data_size_t num_row); /*! - * \brief Update the data in a `ColumnVector` object from an in-memory data buffer, by subtracting each value obtained + * \brief Update the data in a `ColumnVector` object from an in-memory data buffer, by subtracting each value obtained * in `data_ptr` from the existing values in the object's internal `Eigen::VectorXd`. - * + * * \param data_ptr Pointer to first element of a contiguous array of data storing a vector. * \param num_row Number of rows / elements in the vector. */ void SubtractFromData(double* data_ptr, data_size_t num_row); /*! - * \brief Update the data in a `ColumnVector` object from an in-memory data buffer, by substituting each value obtained + * \brief Update the data in a `ColumnVector` object from an in-memory data buffer, by substituting each value obtained * in `data_ptr` for the existing values in the object's internal `Eigen::VectorXd`. - * + * * \param data_ptr Pointer to first element of a contiguous array of data storing a vector. * \param num_row Number of rows / elements in the vector. */ void OverwriteData(double* data_ptr, data_size_t num_row); /*! \brief Number of rows in the object's internal `Eigen::VectorXd`. */ - inline data_size_t NumRows() {return data_.size();} + inline data_size_t NumRows() { return data_.size(); } /*! \brief Return a reference to the object's internal `Eigen::VectorXd`, for interfaces that require a raw vector. */ - inline Eigen::VectorXd& GetData() {return data_;} + inline Eigen::VectorXd& GetData() { return data_; } + private: Eigen::VectorXd data_; void UpdateData(double* data_ptr, data_size_t num_row, std::function op); }; -/*! - * \brief API for loading and accessing data used to sample tree ensembles - * The covariates / bases / weights used in sampling forests are stored internally +/*! + * \brief API for loading and accessing data used to sample tree ensembles + * The covariates / bases / weights used in sampling forests are stored internally * as a `ForestDataset` by the sampling functions (see \ref sampling_group). */ class ForestDataset { @@ -275,7 +276,7 @@ class ForestDataset { ~ForestDataset() {} /*! * \brief Copy / load covariates from raw memory buffer (often pointer to data in a R matrix or numpy array) - * + * * \param data_ptr Pointer to first element of a contiguous array of data storing a covariate matrix * \param num_row Number of rows in the covariate matrix * \param num_col Number of columns / covariates in the covariate matrix @@ -289,7 +290,7 @@ class ForestDataset { } /*! * \brief Copy / load basis matrix from raw memory buffer (often pointer to data in a R matrix or numpy array) - * + * * \param data_ptr Pointer to first element of a contiguous array of data storing a basis matrix * \param num_row Number of rows in the basis matrix * \param num_col Number of columns in the basis matrix @@ -302,7 +303,7 @@ class ForestDataset { } /*! * \brief Copy / load variance weights from raw memory buffer (often pointer to data in a R vector or numpy array) - * + * * \param data_ptr Pointer to first element of a contiguous array of data storing weights * \param num_row Number of rows in the weight vector */ @@ -312,7 +313,7 @@ class ForestDataset { } /*! * \brief Copy / load covariates from CSV file - * + * * \param filename Name of the file (including any necessary path prefixes) * \param column_index_string Comma-delimited string listing columns to extract into covariates matrix */ @@ -324,7 +325,7 @@ class ForestDataset { } /*! * \brief Copy / load basis matrix from CSV file - * + * * \param filename Name of the file (including any necessary path prefixes) * \param column_index_string Comma-delimited string listing columns to extract into covariates matrix */ @@ -335,67 +336,67 @@ class ForestDataset { } /*! * \brief Copy / load variance / case weights from CSV file - * + * * \param filename Name of the file (including any necessary path prefixes) * \param column_index Integer index of column containing weights */ - void AddVarianceWeightsFromCSV(std::string filename, int32_t column_index, bool header = true, bool precise_float_parser = false) { + void AddVarianceWeightsFromCSV(std::string filename, int column_index, bool header = true, bool precise_float_parser = false) { var_weights_ = ColumnVector(filename, column_index, header, precise_float_parser); has_var_weights_ = true; } /*! \brief Whether or not a `ForestDataset` has (yet) loaded covariate data */ - inline bool HasCovariates() {return has_covariates_;} + inline bool HasCovariates() { return has_covariates_; } /*! \brief Whether or not a `ForestDataset` has (yet) loaded basis data */ - inline bool HasBasis() {return has_basis_;} + inline bool HasBasis() { return has_basis_; } /*! \brief Whether or not a `ForestDataset` has (yet) loaded variance weights */ - inline bool HasVarWeights() {return has_var_weights_;} + inline bool HasVarWeights() { return has_var_weights_; } /*! \brief Number of observations (rows) in the dataset */ - inline data_size_t NumObservations() {return num_observations_;} + inline data_size_t NumObservations() { return num_observations_; } /*! \brief Number of covariate columns in the dataset */ - inline int NumCovariates() {return num_covariates_;} + inline int NumCovariates() { return num_covariates_; } /*! \brief Number of bases in the dataset. This is 0 if the dataset has not been provided a basis matrix. */ - inline int NumBasis() {return num_basis_;} + inline int NumBasis() { return num_basis_; } /*! * \brief Returns a dataset's covariate value stored at (`row`, `col`) - * + * * \param row Row number to query in the covariate matrix * \param col Column number to query in the covariate matrix */ - inline double CovariateValue(data_size_t row, int col) {return covariates_.GetElement(row, col);} + inline double CovariateValue(data_size_t row, int col) { return covariates_.GetElement(row, col); } /*! * \brief Returns a dataset's basis value stored at (`row`, `col`) - * + * * \param row Row number to query in the basis matrix * \param col Column number to query in the basis matrix */ - inline double BasisValue(data_size_t row, int col) {return basis_.GetElement(row, col);} + inline double BasisValue(data_size_t row, int col) { return basis_.GetElement(row, col); } /*! * \brief Returns a dataset's variance weight stored at element `row` - * + * * \param row Index to query in the weight vector */ - inline double VarWeightValue(data_size_t row) {return var_weights_.GetElement(row);} + inline double VarWeightValue(data_size_t row) { return var_weights_.GetElement(row); } /*! * \brief Return a reference to the raw `Eigen::MatrixXd` storing the covariate data - * + * * \return Reference to internal Eigen::MatrixXd */ - inline Eigen::MatrixXd& GetCovariates() {return covariates_.GetData();} + inline Eigen::MatrixXd& GetCovariates() { return covariates_.GetData(); } /*! * \brief Return a reference to the raw `Eigen::MatrixXd` storing the basis data - * + * * \return Reference to internal Eigen::MatrixXd */ - inline Eigen::MatrixXd& GetBasis() {return basis_.GetData();} + inline Eigen::MatrixXd& GetBasis() { return basis_.GetData(); } /*! * \brief Return a reference to the raw `Eigen::VectorXd` storing the variance weights - * + * * \return Reference to internal Eigen::VectorXd */ - inline Eigen::VectorXd& GetVarWeights() {return var_weights_.GetData();} + inline Eigen::VectorXd& GetVarWeights() { return var_weights_.GetData(); } /*! * \brief Update the data in the internal basis matrix to new values stored in a raw double array - * + * * \param data_ptr Pointer to first element of a contiguous array of data storing a basis matrix * \param num_row Number of rows in the basis matrix * \param num_col Number of columns in the basis matrix @@ -408,7 +409,7 @@ class ForestDataset { double temp_value; for (data_size_t i = 0; i < num_row; ++i) { for (int j = 0; j < num_col; ++j) { - if (is_row_major){ + if (is_row_major) { // Numpy 2-d arrays are stored in "row major" order temp_value = static_cast(*(data_ptr + static_cast(num_col) * i + j)); } else { @@ -431,8 +432,10 @@ class ForestDataset { // Copy data from R / Python process memory to Eigen vector double temp_value; for (data_size_t i = 0; i < num_row; ++i) { - if (exponentiate) temp_value = std::exp(static_cast(*(data_ptr + i))); - else temp_value = static_cast(*(data_ptr + i)); + if (exponentiate) + temp_value = std::exp(static_cast(*(data_ptr + i))); + else + temp_value = static_cast(*(data_ptr + i)); var_weights_.SetElement(i, temp_value); } } @@ -466,11 +469,13 @@ class ForestDataset { */ void SetVarWeightValue(data_size_t row_id, double new_value, bool exponentiate = true) { CHECK(has_var_weights_); - if (exponentiate) var_weights_.SetElement(row_id, std::exp(new_value)); - else var_weights_.SetElement(row_id, new_value); + if (exponentiate) + var_weights_.SetElement(row_id, std::exp(new_value)); + else + var_weights_.SetElement(row_id, new_value); } - /*! - * \brief Auxiliary data management methods + /*! + * \brief Auxiliary data management methods * Methods to initialize, get, and set auxiliary data for BART models with more structure than the ``classic`` conjugate-Gaussian leaf BART model */ void AddAuxiliaryDimension(int dim_size) { @@ -506,9 +511,9 @@ class ForestDataset { bool has_basis_{false}; bool has_var_weights_{false}; - /*! - * \brief Vector of vectors to track (potentially jagged) auxiliary data for complex BART models - */ + /*! + * \brief Vector of vectors to track (potentially jagged) auxiliary data for complex BART models + */ std::vector> auxiliary_data_; int num_auxiliary_dims_{0}; bool has_auxiliary_data_{false}; @@ -522,30 +527,30 @@ class RandomEffectsDataset { ~RandomEffectsDataset() {} /*! * \brief Copy / load basis matrix from raw memory buffer (often pointer to data in a R matrix or numpy array) - * + * * \param data_ptr Pointer to first element of a contiguous array of data storing a basis matrix * \param num_row Number of rows in the basis matrix * \param num_col Number of columns in the basis matrix * \param is_row_major Whether or not the data in `data_ptr` are organized in a row-major or column-major fashion */ - void AddBasis(double* data_ptr, data_size_t num_row, int num_col, bool is_row_major) { + void AddBasis(double* data_ptr, data_size_t num_row, int num_col, bool is_row_major) { basis_ = ColumnMatrix(data_ptr, num_row, num_col, is_row_major); num_basis_ = num_col; has_basis_ = true; } /*! * \brief Copy / load variance weights from raw memory buffer (often pointer to data in a R vector or numpy array) - * + * * \param data_ptr Pointer to first element of a contiguous array of data storing weights * \param num_row Number of rows in the weight vector - */ + */ void AddVarianceWeights(double* data_ptr, data_size_t num_row) { var_weights_ = ColumnVector(data_ptr, num_row); has_var_weights_ = true; } /*! * \brief Update the data in the internal basis matrix to new values stored in a raw double array - * + * * \param data_ptr Pointer to first element of a contiguous array of data storing a basis matrix * \param num_row Number of rows in the basis matrix * \param num_col Number of columns in the basis matrix @@ -558,7 +563,7 @@ class RandomEffectsDataset { double temp_value; for (data_size_t i = 0; i < num_row; ++i) { for (int j = 0; j < num_col; ++j) { - if (is_row_major){ + if (is_row_major) { // Numpy 2-d arrays are stored in "row major" order temp_value = static_cast(*(data_ptr + static_cast(num_col) * i + j)); } else { @@ -581,8 +586,10 @@ class RandomEffectsDataset { // Copy data from R / Python process memory to Eigen vector double temp_value; for (data_size_t i = 0; i < num_row; ++i) { - if (exponentiate) temp_value = std::exp(static_cast(*(data_ptr + i))); - else temp_value = static_cast(*(data_ptr + i)); + if (exponentiate) + temp_value = std::exp(static_cast(*(data_ptr + i))); + else + temp_value = static_cast(*(data_ptr + i)); var_weights_.SetElement(i, temp_value); } } @@ -593,7 +600,7 @@ class RandomEffectsDataset { * \param num_row Number of rows in the weight vector * \param exponentiate Whether or not inputs should be exponentiated before being saved to var weight vector */ - void UpdateGroupLabels(std::vector& group_labels, data_size_t num_row) { + void UpdateGroupLabels(std::vector& group_labels, data_size_t num_row) { CHECK(has_group_labels_); CHECK_EQ(this->NumObservations(), num_row) // Copy data from R / Python process memory to internal vector @@ -603,73 +610,87 @@ class RandomEffectsDataset { } /*! * \brief Copy / load group indices for random effects - * - * \param group_labels Vector of integers with as many elements as `num_row` in the basis matrix, + * + * \param group_labels Vector of integers with as many elements as `num_row` in the basis matrix, * where each element corresponds to the group label for a given observation. */ - void AddGroupLabels(std::vector& group_labels) { + void AddGroupLabels(std::vector& group_labels) { group_labels_ = group_labels; has_group_labels_ = true; } + /*! + * \brief Copy / load group indices for random effects + * + * \param group_labels Integer pointer to array with as many elements as `num_row` in the basis matrix, + * where each element corresponds to the group label for a given observation. + */ + void AddGroupLabels(int* group_labels, int num_rows) { + group_labels_.resize(num_rows); + for (int i = 0; i < num_rows; ++i) { + group_labels_[i] = group_labels[i]; + } + has_group_labels_ = true; + } /*! \brief Number of observations (rows) in the dataset */ - inline data_size_t NumObservations() {return basis_.NumRows();} + inline data_size_t NumObservations() { return basis_.NumRows(); } /*! \brief Number of columns of the basis vector in the dataset */ - inline int NumBases() {return basis_.NumCols();} + inline int NumBases() { return basis_.NumCols(); } /*! \brief Whether or not a `RandomEffectsDataset` has (yet) loaded basis data */ - inline bool HasBasis() {return has_basis_;} + inline bool HasBasis() { return has_basis_; } /*! \brief Whether or not a `RandomEffectsDataset` has (yet) loaded variance weights */ - inline bool HasVarWeights() {return has_var_weights_;} + inline bool HasVarWeights() { return has_var_weights_; } /*! \brief Whether or not a `RandomEffectsDataset` has (yet) loaded group labels */ - inline bool HasGroupLabels() {return has_group_labels_;} + inline bool HasGroupLabels() { return has_group_labels_; } /*! * \brief Returns a dataset's basis value stored at (`row`, `col`) - * + * * \param row Row number to query in the basis matrix * \param col Column number to query in the basis matrix */ - inline double BasisValue(data_size_t row, int col) {return basis_.GetElement(row, col);} + inline double BasisValue(data_size_t row, int col) { return basis_.GetElement(row, col); } /*! * \brief Returns a dataset's variance weight stored at element `row` - * + * * \param row Index to query in the weight vector */ - inline double VarWeightValue(data_size_t row) {return var_weights_.GetElement(row);} + inline double VarWeightValue(data_size_t row) { return var_weights_.GetElement(row); } /*! * \brief Returns a dataset's group label stored at element `row` - * + * * \param row Index to query in the group label vector */ - inline int32_t GroupId(data_size_t row) {return group_labels_[row];} + inline int GroupId(data_size_t row) { return group_labels_[row]; } /*! * \brief Return a reference to the raw `Eigen::MatrixXd` storing the basis data - * + * * \return Reference to internal Eigen::MatrixXd */ - inline Eigen::MatrixXd& GetBasis() {return basis_.GetData();} + inline Eigen::MatrixXd& GetBasis() { return basis_.GetData(); } /*! * \brief Return a reference to the raw `Eigen::VectorXd` storing the variance weights - * + * * \return Reference to internal Eigen::VectorXd */ - inline Eigen::VectorXd& GetVarWeights() {return var_weights_.GetData();} + inline Eigen::VectorXd& GetVarWeights() { return var_weights_.GetData(); } /*! * \brief Return a reference to the raw `std::vector` storing the group labels - * + * * \return Reference to internal std::vector */ - inline std::vector& GetGroupLabels() {return group_labels_;} + inline std::vector& GetGroupLabels() { return group_labels_; } + private: ColumnMatrix basis_; ColumnVector var_weights_; - std::vector group_labels_; + std::vector group_labels_; int num_basis_{0}; bool has_basis_{false}; bool has_var_weights_{false}; bool has_group_labels_{false}; }; -/*! \} */ // end of data_group +/*! \} */ // end of data_group -} // namespace StochTree +} // namespace StochTree -#endif // STOCHTREE_DATA_H_ +#endif // STOCHTREE_DATA_H_ diff --git a/include/stochtree/discrete_sampler.h b/include/stochtree/discrete_sampler.h index a513d032..507875d5 100644 --- a/include/stochtree/discrete_sampler.h +++ b/include/stochtree/discrete_sampler.h @@ -12,7 +12,7 @@ namespace StochTree { -/*! \brief Sample without replacement according to a set of probability weights. +/*! \brief Sample without replacement according to a set of probability weights. * This template function is a C++ variant of numpy's implementation: * https://github.com/numpy/numpy/blob/031f44252d613f4524ad181e3eb2ae2791e22187/numpy/random/_generator.pyx#L925 */ @@ -23,19 +23,19 @@ void sample_without_replacement(container_type* output, prob_type* p, container_ std::vector indices(sample_size); std::vector unif_samples(sample_size); std::vector cdf(population_size); - + int fulfilled_sample_count = 0; int remaining_sample_count = sample_size - fulfilled_sample_count; while (fulfilled_sample_count < sample_size) { if (fulfilled_sample_count > 0) { for (int i = 0; i < fulfilled_sample_count; i++) p_copy[indices[i]] = 0.0; } - std::generate(unif_samples.begin(), unif_samples.begin() + remaining_sample_count, [&gen](){ + std::generate(unif_samples.begin(), unif_samples.begin() + remaining_sample_count, [&gen]() { return standard_uniform_draw_53bit(gen); }); std::partial_sum(p_copy.cbegin(), p_copy.cend(), cdf.begin()); for (int i = 0; i < cdf.size(); i++) { - cdf[i] = cdf[i] / cdf[cdf.size()-1]; + cdf[i] = cdf[i] / cdf[cdf.size() - 1]; } std::vector matches(remaining_sample_count); for (int i = 0; i < remaining_sample_count; i++) { @@ -60,6 +60,6 @@ void sample_without_replacement(container_type* output, prob_type* p, container_ } } -} +} // namespace StochTree -#endif // STOCHTREE_DISCRETE_SAMPLER_H_ +#endif // STOCHTREE_DISCRETE_SAMPLER_H_ diff --git a/include/stochtree/distributions.h b/include/stochtree/distributions.h index f946975a..3cf4120a 100644 --- a/include/stochtree/distributions.h +++ b/include/stochtree/distributions.h @@ -2,16 +2,34 @@ #define STOCHTREE_DISTRIBUTIONS_H #include #include +#include /*! * \brief A collection of random number generation utilities. * - * This file is vendored from a broader C++ / R distribution + * This file is vendored from a broader C++ / R distribution * library, where the distributions are subject to rigorous testing. * https://github.com/andrewherren/cpp11_r_rng */ namespace StochTree { +/*! + * Standard normal cumulative distribution function, implemented via the complementary error function. + */ +inline double norm_cdf(double x) { + return 0.5 * boost::math::erfc(-x / std::sqrt(2.0)); +} + +/*! + * Standard normal quantile function (inverse CDF), implemented via the inverse complementary error function. + */ +inline double norm_inv_cdf(double p) { + return -std::sqrt(2.0) * boost::math::erfc_inv(2.0 * p); +} + +/*! Precomputed standard normal CDF at 0 */ +static constexpr double Phi_0 = 0.5; + /*! * Generate a standard uniform random variate to 53 bits of precision via two mersenne twisters, see: * https://github.com/numpy/numpy/blob/0d7986494b39ace565afda3de68be528ddade602/numpy/random/src/mt19937/mt19937.h#L56 @@ -68,9 +86,9 @@ class standard_normal { /*! * Stateless standard normal sampler implementing Marsaglia's polar method. * Without caching, this is half as fast as other methods for repeated normal sampling, - * but this might be acceptable in cases where a relatively small number of + * but this might be acceptable in cases where a relatively small number of * normal draws is desired. - * + * * Reference: https://en.wikipedia.org/wiki/Marsaglia_polar_method */ inline double sample_standard_normal(double mean, double sd, std::mt19937& gen) { @@ -130,10 +148,10 @@ class gamma_sampler { v = v * v * v; double u = standard_uniform_draw_53bit(gen); if (u < 1.0 - 0.0331 * (x * x) * (x * x)) { - return b * v * scale; + return b * v * scale; } if (std::log(u) < 0.5 * x * x + b * (1.0 - v + std::log(v))) { - return b * v * scale; + return b * v * scale; } } } else { @@ -202,23 +220,23 @@ inline double sample_gamma(std::mt19937& gen, double shape, double scale) { while (true) { double x, v; do { - // Marsaglia's polar method for standard normal + // Marsaglia's polar method for standard normal double u1, u2, s; do { u1 = standard_uniform_draw_53bit(gen) * 2.0 - 1.0; u2 = standard_uniform_draw_53bit(gen) * 2.0 - 1.0; s = u1 * u1 + u2 * u2; } while (s >= 1.0 || s == 0.0); - x = u1 * std::sqrt(-2.0 * std::log(s) / s); + x = u1 * std::sqrt(-2.0 * std::log(s) / s); v = 1.0 + c * x; } while (v <= 0.0); v = v * v * v; double u = standard_uniform_draw_53bit(gen); if (u < 1.0 - 0.0331 * (x * x) * (x * x)) { - return b * v * scale; + return b * v * scale; } if (std::log(u) < 0.5 * x * x + b * (1.0 - v + std::log(v))) { - return b * v * scale; + return b * v * scale; } } } else { @@ -228,13 +246,13 @@ inline double sample_gamma(std::mt19937& gen, double shape, double scale) { /*! * Walker-Vose alias method for sampling with replacement from a weighted discrete distribution. - * + * * Simplified from https://github.com/boostorg/random/blob/develop/include/boost/random/discrete_distribution.hpp * Other references: https://en.wikipedia.org/wiki/Alias_method */ class walker_vose { public: - template + template walker_vose(Iterator first, Iterator last) { n_ = std::distance(first, last); probability_.resize(n_); @@ -245,7 +263,7 @@ class walker_vose { for (auto it = first; it != last; ++it) { sum += *it; } - + // Build alias table using Walker's algorithm std::vector p(n_); std::vector below_average, above_average; @@ -258,33 +276,35 @@ class walker_vose { above_average.push_back(i); } } - + while (!below_average.empty() && !above_average.empty()) { - int j = below_average.back(); below_average.pop_back(); - int i = above_average.back(); above_average.pop_back(); - + int j = below_average.back(); + below_average.pop_back(); + int i = above_average.back(); + above_average.pop_back(); + probability_[j] = p[j]; alias_[j] = i; p[i] = (p[i] + p[j]) - 1.0; - + if (p[i] < 1.0) { below_average.push_back(i); } else { above_average.push_back(i); } } - + while (!above_average.empty()) { probability_[above_average.back()] = 1.0; above_average.pop_back(); } - + while (!below_average.empty()) { probability_[below_average.back()] = 1.0; below_average.pop_back(); } } - + int operator()(std::mt19937& gen) { double u = standard_uniform_draw_53bit(gen); int i = static_cast(u * n_); @@ -323,6 +343,22 @@ inline int sample_discrete_stateless(std::mt19937& gen, std::vector& wei return weights.size() - 1; } +/*! + * Generate a single sample from a truncated standard normal distribution, bounded above by 0. + */ +inline double sample_std_truncnorm_upper(std::mt19937& gen) { + double uniform_draw = standard_uniform_draw_53bit(gen); + return norm_inv_cdf(uniform_draw * Phi_0); } -#endif // STOCHTREE_DISTRIBUTIONS_H \ No newline at end of file +/*! + * Generate a single sample from a truncated standard normal distribution, bounded below by 0. + */ +inline double sample_std_truncnorm_lower(std::mt19937& gen) { + double uniform_draw = standard_uniform_draw_53bit(gen); + return norm_inv_cdf(uniform_draw + (1 - uniform_draw) * Phi_0); +} + +} // namespace StochTree + +#endif // STOCHTREE_DISTRIBUTIONS_H \ No newline at end of file diff --git a/include/stochtree/ensemble.h b/include/stochtree/ensemble.h index 4f6ddf42..b9eeb303 100644 --- a/include/stochtree/ensemble.h +++ b/include/stochtree/ensemble.h @@ -1,10 +1,10 @@ /*! * Copyright (c) 2024 stochtree authors. All rights reserved. * Licensed under the MIT License. See LICENSE file in the project root for license information. - * - * Inspired by the design of the Learner, GBTreeModel, and GBTree classes in xgboost, + * + * Inspired by the design of the Learner, GBTreeModel, and GBTree classes in xgboost, * released under the Apache license with the following copyright: - * + * * Copyright 2015-2023 by XGBoost Contributors */ #ifndef STOCHTREE_ENSEMBLE_H_ @@ -32,7 +32,7 @@ class TreeEnsemble { public: /*! * \brief Initialize a new TreeEnsemble - * + * * \param num_trees Number of trees in a forest * \param output_dimension Dimension of the leaf node parameter * \param is_leaf_constant Whether or not the leaves of each tree are treated as "constant." If true, then predicting from an ensemble is simply a matter or determining which leaf node an observation falls into. If false, prediction will multiply a leaf node's parameter(s) for a given observation by a basis vector. @@ -51,10 +51,10 @@ class TreeEnsemble { is_leaf_constant_ = is_leaf_constant; is_exponentiated_ = is_exponentiated; } - + /*! * \brief Initialize an ensemble based on the state of an existing ensemble - * + * * \param ensemble `TreeEnsemble` used to initialize the current ensemble */ TreeEnsemble(TreeEnsemble& ensemble) { @@ -74,12 +74,12 @@ class TreeEnsemble { this->CloneFromExistingTree(j, tree); } } - + ~TreeEnsemble() {} /*! * \brief Combine two forests into a single forest by merging their trees - * + * * \param ensemble Reference to another `TreeEnsemble` that will be merged into the current ensemble */ void MergeForest(TreeEnsemble& ensemble) { @@ -103,7 +103,7 @@ class TreeEnsemble { /*! * \brief Add a constant value to every leaf of every tree in an ensemble. If leaves are multi-dimensional, `constant_value` will be added to every dimension of the leaves. - * + * * \param constant_value Value that will be added to every leaf of every tree */ void AddValueToLeaves(double constant_value) { @@ -115,7 +115,7 @@ class TreeEnsemble { /*! * \brief Multiply every leaf of every tree by a constant value. If leaves are multi-dimensional, `constant_multiple` will be multiplied through every dimension of the leaves. - * + * * \param constant_multiple Value that will be multiplied by every leaf of every tree */ void MultiplyLeavesByValue(double constant_multiple) { @@ -127,9 +127,9 @@ class TreeEnsemble { /*! * \brief Return a pointer to a tree in the forest - * + * * \param i Index (0-based) of a tree to be queried - * \return Tree* + * \return Tree* */ inline Tree* GetTree(int i) { return trees_[i].get(); @@ -147,7 +147,7 @@ class TreeEnsemble { /*! * \brief Reset a single tree in an ensemble * \todo Consider refactoring this and `ResetInitTree` - * + * * \param i Index (0-based) of the tree to be reset */ inline void ResetTree(int i) { @@ -157,7 +157,7 @@ class TreeEnsemble { /*! * \brief Reset a single tree in an ensemble * \todo Consider refactoring this and `ResetTree` - * + * * \param i Index (0-based) of the tree to be reset */ inline void ResetInitTree(int i) { @@ -167,7 +167,7 @@ class TreeEnsemble { /*! * \brief Clone a single tree in an ensemble from an existing tree, overwriting current tree - * + * * \param i Index of the tree to be overwritten * \param tree Pointer to tree used to clone tree `i` */ @@ -177,7 +177,7 @@ class TreeEnsemble { /*! * \brief Reset an ensemble to clone another ensemble - * + * * \param ensemble Reference to an existing `TreeEnsemble` */ inline void ReconstituteFromForest(TreeEnsemble& ensemble) { @@ -207,19 +207,19 @@ class TreeEnsemble { return output; } - std::vector PredictRaw(ForestDataset& dataset) { + std::vector PredictRaw(ForestDataset& dataset, bool row_major = true) { data_size_t n = dataset.NumObservations(); data_size_t total_output_size = n * output_dimension_; std::vector output(total_output_size); - PredictRawInplace(dataset, output, 0); + PredictRawInplace(dataset, output, 0, trees_.size(), 0, row_major); return output; } - - inline void PredictInplace(ForestDataset& dataset, std::vector &output, data_size_t offset = 0) { + + inline void PredictInplace(ForestDataset& dataset, std::vector& output, data_size_t offset = 0) { PredictInplace(dataset, output, 0, trees_.size(), offset); } - inline void PredictInplace(ForestDataset& dataset, std::vector &output, + inline void PredictInplace(ForestDataset& dataset, std::vector& output, int tree_begin, int tree_end, data_size_t offset = 0) { if (is_leaf_constant_) { PredictInplace(dataset.GetCovariates(), output, tree_begin, tree_end, offset); @@ -229,11 +229,11 @@ class TreeEnsemble { } } - inline void PredictInplace(Eigen::MatrixXd& covariates, Eigen::MatrixXd& basis, std::vector &output, data_size_t offset = 0) { + inline void PredictInplace(Eigen::MatrixXd& covariates, Eigen::MatrixXd& basis, std::vector& output, data_size_t offset = 0) { PredictInplace(covariates, basis, output, 0, trees_.size(), offset); } - inline void PredictInplace(Eigen::MatrixXd& covariates, Eigen::MatrixXd& basis, std::vector &output, + inline void PredictInplace(Eigen::MatrixXd& covariates, Eigen::MatrixXd& basis, std::vector& output, int tree_begin, int tree_end, data_size_t offset = 0) { double pred; CHECK_EQ(covariates.rows(), basis.rows()); @@ -247,22 +247,24 @@ class TreeEnsemble { for (data_size_t i = 0; i < n; i++) { pred = 0.0; for (size_t j = tree_begin; j < tree_end; j++) { - auto &tree = *trees_[j]; + auto& tree = *trees_[j]; std::int32_t nidx = EvaluateTree(tree, covariates, i); for (int32_t k = 0; k < output_dimension_; k++) { pred += tree.LeafValue(nidx, k) * basis(i, k); } } - if (is_exponentiated_) output[i + offset] = std::exp(pred); - else output[i + offset] = pred; + if (is_exponentiated_) + output[i + offset] = std::exp(pred); + else + output[i + offset] = pred; } } - inline void PredictInplace(Eigen::MatrixXd& covariates, std::vector &output, data_size_t offset = 0) { + inline void PredictInplace(Eigen::MatrixXd& covariates, std::vector& output, data_size_t offset = 0) { PredictInplace(covariates, output, 0, trees_.size(), offset); } - inline void PredictInplace(Eigen::MatrixXd& covariates, std::vector &output, int tree_begin, int tree_end, data_size_t offset = 0) { + inline void PredictInplace(Eigen::MatrixXd& covariates, std::vector& output, int tree_begin, int tree_end, data_size_t offset = 0) { double pred; data_size_t n = covariates.rows(); data_size_t total_output_size = n; @@ -272,21 +274,23 @@ class TreeEnsemble { for (data_size_t i = 0; i < n; i++) { pred = 0.0; for (size_t j = tree_begin; j < tree_end; j++) { - auto &tree = *trees_[j]; + auto& tree = *trees_[j]; std::int32_t nidx = EvaluateTree(tree, covariates, i); pred += tree.LeafValue(nidx, 0); } - if (is_exponentiated_) output[i + offset] = std::exp(pred); - else output[i + offset] = pred; + if (is_exponentiated_) + output[i + offset] = std::exp(pred); + else + output[i + offset] = pred; } } - inline void PredictRawInplace(ForestDataset& dataset, std::vector &output, data_size_t offset = 0) { - PredictRawInplace(dataset, output, 0, trees_.size(), offset); + inline void PredictRawInplace(ForestDataset& dataset, std::vector& output, data_size_t offset = 0, bool row_major = true) { + PredictRawInplace(dataset, output, 0, trees_.size(), offset, row_major); } - inline void PredictRawInplace(ForestDataset& dataset, std::vector &output, - int tree_begin, int tree_end, data_size_t offset = 0) { + inline void PredictRawInplace(ForestDataset& dataset, std::vector& output, + int tree_begin, int tree_end, data_size_t offset = 0, bool row_major = true) { double pred; Eigen::MatrixXd covariates = dataset.GetCovariates(); CHECK_EQ(output_dimension_, trees_[0]->OutputDimension()); @@ -299,11 +303,15 @@ class TreeEnsemble { for (int32_t k = 0; k < output_dimension_; k++) { pred = 0.0; for (size_t j = tree_begin; j < tree_end; j++) { - auto &tree = *trees_[j]; + auto& tree = *trees_[j]; int32_t nidx = EvaluateTree(tree, covariates, i); pred += tree.LeafValue(nidx, k); } - output[i*output_dimension_ + k + offset] = pred; + if (row_major) { + output[i * output_dimension_ + k + offset] = pred; + } else { + output[k * n + i + offset] = pred; + } } } } @@ -380,30 +388,30 @@ class TreeEnsemble { } /*! - * \brief Obtain a 0-based "maximum" leaf index for an ensemble, which is equivalent to the sum of the + * \brief Obtain a 0-based "maximum" leaf index for an ensemble, which is equivalent to the sum of the * number of leaves in each tree. This is used in conjunction with `PredictLeafIndicesInplace`, * which returns an observation-specific leaf index for every observation-tree pair. */ int GetMaxLeafIndex() { int max_leaf = 0; for (int j = 0; j < num_trees_; j++) { - auto &tree = *trees_[j]; + auto& tree = *trees_[j]; max_leaf += tree.NumLeaves(); } return max_leaf; } /*! - * \brief Obtain a 0-based leaf index for every tree in an ensemble and for each - * observation in a ForestDataset. Internally, trees are stored as essentially - * vectors of node information, and the leaves_ vector gives us node IDs for every - * leaf in the tree. Here, we would like to know, for every observation in a dataset, - * which leaf number it is mapped to. Since the leaf numbers themselves - * do not carry any information, we renumber them from 0 to `leaves_.size()-1`. - * We compute this at the tree-level and coordinate this computation at the + * \brief Obtain a 0-based leaf index for every tree in an ensemble and for each + * observation in a ForestDataset. Internally, trees are stored as essentially + * vectors of node information, and the leaves_ vector gives us node IDs for every + * leaf in the tree. Here, we would like to know, for every observation in a dataset, + * which leaf number it is mapped to. Since the leaf numbers themselves + * do not carry any information, we renumber them from 0 to `leaves_.size()-1`. + * We compute this at the tree-level and coordinate this computation at the * ensemble level. * - * Note: this assumes the creation of a vector of column indices of size + * Note: this assumes the creation of a vector of column indices of size * `dataset.NumObservations()` x `ensemble.NumTrees()` * \param ForestDataset Dataset with which to predict leaf indices from the tree * \param output Vector of length num_trees*n which stores the leaf node prediction @@ -415,16 +423,16 @@ class TreeEnsemble { } /*! - * \brief Obtain a 0-based leaf index for every tree in an ensemble and for each - * observation in a ForestDataset. Internally, trees are stored as essentially - * vectors of node information, and the leaves_ vector gives us node IDs for every - * leaf in the tree. Here, we would like to know, for every observation in a dataset, - * which leaf number it is mapped to. Since the leaf numbers themselves - * do not carry any information, we renumber them from 0 to `leaves_.size()-1`. - * We compute this at the tree-level and coordinate this computation at the + * \brief Obtain a 0-based leaf index for every tree in an ensemble and for each + * observation in a ForestDataset. Internally, trees are stored as essentially + * vectors of node information, and the leaves_ vector gives us node IDs for every + * leaf in the tree. Here, we would like to know, for every observation in a dataset, + * which leaf number it is mapped to. Since the leaf numbers themselves + * do not carry any information, we renumber them from 0 to `leaves_.size()-1`. + * We compute this at the tree-level and coordinate this computation at the * ensemble level. * - * Note: this assumes the creation of a vector of column indices of size + * Note: this assumes the creation of a vector of column indices of size * `dataset.NumObservations()` x `ensemble.NumTrees()` * \param covariates Matrix of covariates * \param output Vector of length num_trees*n which stores the leaf node prediction @@ -432,11 +440,11 @@ class TreeEnsemble { * \param n Size of dataset */ void PredictLeafIndicesInplace(Eigen::Map>& covariates, std::vector& output, int num_trees, data_size_t n) { - CHECK_GE(output.size(), num_trees*n); + CHECK_GE(output.size(), num_trees * n); int offset = 0; int max_leaf = 0; for (int j = 0; j < num_trees; j++) { - auto &tree = *trees_[j]; + auto& tree = *trees_[j]; int num_leaves = tree.NumLeaves(); tree.PredictLeafIndexInplace(covariates, output, offset, max_leaf); offset += n; @@ -445,13 +453,13 @@ class TreeEnsemble { } /*! - * \brief Obtain a 0-based leaf index for every tree in an ensemble and for each - * observation in a ForestDataset. Internally, trees are stored as essentially - * vectors of node information, and the leaves_ vector gives us node IDs for every - * leaf in the tree. Here, we would like to know, for every observation in a dataset, - * which leaf number it is mapped to. Since the leaf numbers themselves - * do not carry any information, we renumber them from 0 to `leaves_.size()-1`. - * We compute this at the tree-level and coordinate this computation at the + * \brief Obtain a 0-based leaf index for every tree in an ensemble and for each + * observation in a ForestDataset. Internally, trees are stored as essentially + * vectors of node information, and the leaves_ vector gives us node IDs for every + * leaf in the tree. Here, we would like to know, for every observation in a dataset, + * which leaf number it is mapped to. Since the leaf numbers themselves + * do not carry any information, we renumber them from 0 to `leaves_.size()-1`. + * We compute this at the tree-level and coordinate this computation at the * ensemble level. * * Note: this assumes the creation of a matrix of column indices with `num_trees*n` rows @@ -462,14 +470,14 @@ class TreeEnsemble { * \param num_trees Number of trees in an ensemble * \param n Size of dataset */ - void PredictLeafIndicesInplace(Eigen::Map>& covariates, - Eigen::Map>& output, + void PredictLeafIndicesInplace(Eigen::Map>& covariates, + Eigen::Map>& output, int column_ind, int num_trees, data_size_t n) { - CHECK_GE(output.size(), num_trees*n); + CHECK_GE(output.size(), num_trees * n); int offset = 0; int max_leaf = 0; for (int j = 0; j < num_trees; j++) { - auto &tree = *trees_[j]; + auto& tree = *trees_[j]; int num_leaves = tree.NumLeaves(); tree.PredictLeafIndexInplace(covariates, output, column_ind, offset, max_leaf); offset += n; @@ -478,16 +486,16 @@ class TreeEnsemble { } /*! - * \brief Obtain a 0-based leaf index for every tree in an ensemble and for each - * observation in a ForestDataset. Internally, trees are stored as essentially - * vectors of node information, and the leaves_ vector gives us node IDs for every - * leaf in the tree. Here, we would like to know, for every observation in a dataset, - * which leaf number it is mapped to. Since the leaf numbers themselves - * do not carry any information, we renumber them from 0 to `leaves_.size()-1`. - * We compute this at the tree-level and coordinate this computation at the + * \brief Obtain a 0-based leaf index for every tree in an ensemble and for each + * observation in a ForestDataset. Internally, trees are stored as essentially + * vectors of node information, and the leaves_ vector gives us node IDs for every + * leaf in the tree. Here, we would like to know, for every observation in a dataset, + * which leaf number it is mapped to. Since the leaf numbers themselves + * do not carry any information, we renumber them from 0 to `leaves_.size()-1`. + * We compute this at the tree-level and coordinate this computation at the * ensemble level. * - * Note: this assumes the creation of a vector of column indices of size + * Note: this assumes the creation of a vector of column indices of size * `dataset.NumObservations()` x `ensemble.NumTrees()` * \param ForestDataset Dataset with which to predict leaf indices from the tree * \param output Vector of length num_trees*n which stores the leaf node prediction @@ -495,11 +503,11 @@ class TreeEnsemble { * \param n Size of dataset */ void PredictLeafIndicesInplace(Eigen::MatrixXd& covariates, std::vector& output, int num_trees, data_size_t n) { - CHECK_GE(output.size(), num_trees*n); + CHECK_GE(output.size(), num_trees * n); int offset = 0; int max_leaf = 0; for (int j = 0; j < num_trees; j++) { - auto &tree = *trees_[j]; + auto& tree = *trees_[j]; int num_leaves = tree.NumLeaves(); tree.PredictLeafIndexInplace(covariates, output, offset, max_leaf); offset += n; @@ -514,7 +522,7 @@ class TreeEnsemble { std::vector PredictLeafIndices(ForestDataset* dataset) { int num_trees = num_trees_; data_size_t n = dataset->NumObservations(); - std::vector output(n*num_trees); + std::vector output(n * num_trees); PredictLeafIndicesInplace(dataset, output, num_trees, n); return output; } @@ -532,10 +540,10 @@ class TreeEnsemble { tree_label = "tree_" + std::to_string(i); result_obj.emplace(tree_label, trees_[i]->to_json()); } - + return result_obj; } - + /*! \brief Load from JSON */ void from_json(const json& ensemble_json) { this->num_trees_ = ensemble_json.at("num_trees"); @@ -561,8 +569,8 @@ class TreeEnsemble { bool is_exponentiated_; }; -/*! \} */ // end of forest_group +/*! \} */ // end of forest_group -} // namespace StochTree +} // namespace StochTree -#endif // STOCHTREE_ENSEMBLE_H_ +#endif // STOCHTREE_ENSEMBLE_H_ diff --git a/include/stochtree/export.h b/include/stochtree/export.h index 4b651749..fb5b339b 100644 --- a/include/stochtree/export.h +++ b/include/stochtree/export.h @@ -1,11 +1,11 @@ /*! - * Export macros ensure that the C++ code can be used as a library cross-platform - * (declspec needed to load names from a DLL on windows) and can be wrapped in a + * Export macros ensure that the C++ code can be used as a library cross-platform + * (declspec needed to load names from a DLL on windows) and can be wrapped in a * C program. - * - * This code modifies (changing names of) the export macros in LightGBM, which carries + * + * This code modifies (changing names of) the export macros in LightGBM, which carries * the following copyright information: - * + * * Copyright (c) 2017 Microsoft Corporation. All rights reserved. * Licensed under the MIT License. See LICENSE file in the project root for license information. */ @@ -24,8 +24,8 @@ #define STOCHTREE_EXPORT __declspec(dllexport) #define STOCHTREE_C_EXPORT STOCHTREE_EXTERN_C __declspec(dllexport) #else -#define STOCHTREE_EXPORT __attribute__ ((visibility ("default"))) -#define STOCHTREE_C_EXPORT STOCHTREE_EXTERN_C __attribute__ ((visibility ("default"))) +#define STOCHTREE_EXPORT __attribute__((visibility("default"))) +#define STOCHTREE_C_EXPORT STOCHTREE_EXTERN_C __attribute__((visibility("default"))) #endif #endif /** STOCHTREE_EXPORT_H_ **/ diff --git a/include/stochtree/gamma_sampler.h b/include/stochtree/gamma_sampler.h index 1c524f4f..53d2f332 100644 --- a/include/stochtree/gamma_sampler.h +++ b/include/stochtree/gamma_sampler.h @@ -12,11 +12,11 @@ class GammaSampler { GammaSampler() {} ~GammaSampler() {} double Sample(double a, double b, std::mt19937& gen, bool rate_param = true) { - double scale = rate_param ? 1./b : b; + double scale = rate_param ? 1. / b : b; return sample_gamma(gen, a, scale); } }; -} // namespace StochTree +} // namespace StochTree -#endif // STOCHTREE_GAMMA_SAMPLER_H_ \ No newline at end of file +#endif // STOCHTREE_GAMMA_SAMPLER_H_ \ No newline at end of file diff --git a/include/stochtree/ig_sampler.h b/include/stochtree/ig_sampler.h index 669b9e56..6da28fad 100644 --- a/include/stochtree/ig_sampler.h +++ b/include/stochtree/ig_sampler.h @@ -13,15 +13,15 @@ class InverseGammaSampler { ~InverseGammaSampler() {} double Sample(double a, double b, std::mt19937& gen, bool scale_param = true) { // C++ standard library provides a gamma distribution with scale - // parameter, but the correspondence between gamma and IG is that + // parameter, but the correspondence between gamma and IG is that // 1 / gamma(a,b) ~ IG(a,b) when b is a __rate__ parameter. - // Before sampling, we convert ig_scale to a gamma scale parameter by + // Before sampling, we convert ig_scale to a gamma scale parameter by // taking its multiplicative inverse. - double gamma_scale = scale_param ? 1./b : b; - return (1/sample_gamma(gen, a, gamma_scale)); + double gamma_scale = scale_param ? 1. / b : b; + return (1 / sample_gamma(gen, a, gamma_scale)); } }; -} // namespace StochTree +} // namespace StochTree -#endif // STOCHTREE_IG_SAMPLER_H_ \ No newline at end of file +#endif // STOCHTREE_IG_SAMPLER_H_ \ No newline at end of file diff --git a/include/stochtree/io.h b/include/stochtree/io.h index 55963946..eeb259e7 100644 --- a/include/stochtree/io.h +++ b/include/stochtree/io.h @@ -7,8 +7,8 @@ * parser.h * pipeline_reader.h * text_reader.h - * - * LightGBM is MIT licensed and released with the following copyright header + * + * LightGBM is MIT licensed and released with the following copyright header * (with different copyright years in different files): * * Copyright (c) 2016 Microsoft Corporation. All rights reserved. @@ -42,9 +42,9 @@ namespace StochTree { const size_t kGbs = size_t(1024) * 1024 * 1024; /*! -* \brief Contains some operation for an array, e.g. ArgMax, TopK. -*/ -template + * \brief Contains some operation for an array, e.g. ArgMax, TopK. + */ +template class ArrayArgs { public: inline static size_t ArgMax(const std::vector& array) { @@ -116,18 +116,35 @@ class ArrayArgs { std::vector& ref = *arr; VAL_T v = ref[end - 1]; for (;;) { - while (ref[++i] > v) {} - while (v > ref[--j]) { if (j == start) { break; } } - if (i >= j) { break; } + while (ref[++i] > v) { + } + while (v > ref[--j]) { + if (j == start) { + break; + } + } + if (i >= j) { + break; + } std::swap(ref[i], ref[j]); - if (ref[i] == v) { p++; std::swap(ref[p], ref[i]); } - if (v == ref[j]) { q--; std::swap(ref[j], ref[q]); } + if (ref[i] == v) { + p++; + std::swap(ref[p], ref[i]); + } + if (v == ref[j]) { + q--; + std::swap(ref[j], ref[q]); + } } std::swap(ref[i], ref[end - 1]); j = i - 1; i = i + 1; - for (int k = start; k <= p; k++, j--) { std::swap(ref[k], ref[j]); } - for (int k = end - 2; k >= q; k--, i++) { std::swap(ref[i], ref[k]); } + for (int k = start; k <= p; k++, j--) { + std::swap(ref[k], ref[j]); + } + for (int k = end - 2; k >= q; k--, i++) { + std::swap(ref[i], ref[k]); + } *l = j; *r = i; } @@ -193,24 +210,24 @@ class ArrayArgs { }; /*! - * \brief An interface for serializing binary data to a buffer - */ + * \brief An interface for serializing binary data to a buffer + */ struct BinaryWriter { /*! - * \brief Append data to this binary target - * \param data Buffer to write from - * \param bytes Number of bytes to write from buffer - * \return Number of bytes written - */ + * \brief Append data to this binary target + * \param data Buffer to write from + * \param bytes Number of bytes to write from buffer + * \return Number of bytes written + */ virtual size_t Write(const void* data, size_t bytes) = 0; /*! - * \brief Append data to this binary target aligned on a given byte size boundary - * \param data Buffer to write from - * \param bytes Number of bytes to write from buffer - * \param alignment The size of bytes to align to in whole increments - * \return Number of bytes written - */ + * \brief Append data to this binary target aligned on a given byte size boundary + * \param data Buffer to write from + * \param bytes Number of bytes to write from buffer + * \param alignment The size of bytes to align to in whole increments + * \return Number of bytes written + */ size_t AlignedWrite(const void* data, size_t bytes, size_t alignment = 8) { auto ret = Write(data, bytes); if (bytes % alignment != 0) { @@ -222,11 +239,11 @@ struct BinaryWriter { } /*! - * \brief The aligned size of a buffer length. - * \param bytes The number of bytes in a buffer - * \param alignment The size of bytes to align to in whole increments - * \return Number of aligned bytes - */ + * \brief The aligned size of a buffer length. + * \param bytes The number of bytes in a buffer + * \param alignment The size of bytes to align to in whole increments + * \return Number of aligned bytes + */ static size_t AlignedSize(size_t bytes, size_t alignment = 8) { if (bytes % alignment == 0) { return bytes; @@ -301,42 +318,42 @@ class Parser { Parser() {} /*! - * \brief Constructor for customized parser. The constructor accepts content not path because need to save/load the config along with model string - */ + * \brief Constructor for customized parser. The constructor accepts content not path because need to save/load the config along with model string + */ explicit Parser(std::string) {} /*! \brief virtual destructor */ virtual ~Parser() {} /*! - * \brief Parse one line with label - * \param str One line record, string format, should end with '\0' - * \param out_features Output columns, store in (column_idx, values) - */ + * \brief Parse one line with label + * \param str One line record, string format, should end with '\0' + * \param out_features Output columns, store in (column_idx, values) + */ virtual void ParseOneLine(const char* str, std::vector>* out_features) const = 0; virtual int NumFeatures() const = 0; /*! - * \brief Create an object of parser, will auto choose the format depend on file - * \param filename One Filename of data - * \param header whether input file contains header - * \param num_features Pass num_features of this data file if you know, <=0 means don't know - * \param precise_float_parser using precise floating point number parsing if true - * \return Object of parser - */ + * \brief Create an object of parser, will auto choose the format depend on file + * \param filename One Filename of data + * \param header whether input file contains header + * \param num_features Pass num_features of this data file if you know, <=0 means don't know + * \param precise_float_parser using precise floating point number parsing if true + * \return Object of parser + */ static Parser* CreateParser(const char* filename, bool header, int num_features, bool precise_float_parser); }; -class CSVParser: public Parser { +class CSVParser : public Parser { public: explicit CSVParser(int total_columns, AtofFunc atof) - :total_columns_(total_columns), atof_(atof) { + : total_columns_(total_columns), atof_(atof) { } inline void ParseOneLine(const char* str, - std::vector>* out_features) const override { + std::vector>* out_features) const override { int idx = 0; double val = 0.0f; int offset = 0; @@ -365,22 +382,22 @@ class CSVParser: public Parser { }; /*! -* \brief A pipeline file reader, use 2 threads, one read block from file, the other process the block -*/ + * \brief A pipeline file reader, use 2 threads, one read block from file, the other process the block + */ class PipelineReader { public: /*! - * \brief Read data from a file, use pipeline methods - * \param filename Filename of data - * \process_fun Process function - */ + * \brief Read data from a file, use pipeline methods + * \param filename Filename of data + * \process_fun Process function + */ static size_t Read(const char* filename, int skip_bytes, const std::function& process_fun) { auto reader = VirtualFileReader::Make(filename); if (!reader->Init()) { return 0; } size_t cnt = 0; - const size_t buffer_size = 16 * 1024 * 1024; + const size_t buffer_size = 16 * 1024 * 1024; // buffer used for the process_fun auto buffer_process = std::vector(buffer_size); // buffer used for the file reading @@ -397,9 +414,9 @@ class PipelineReader { while (read_cnt > 0) { // start read thread std::thread read_worker = std::thread( - [=, &last_read_cnt, &reader, &buffer_read] { - last_read_cnt = reader->Read(buffer_read.data(), buffer_size); - }); + [=, &last_read_cnt, &reader, &buffer_read] { + last_read_cnt = reader->Read(buffer_read.data(), buffer_size); + }); // start process cnt += process_fun(buffer_process.data(), read_cnt); // wait for read thread @@ -413,18 +430,17 @@ class PipelineReader { }; /*! -* \brief Read text data from file -*/ -template + * \brief Read text data from file + */ +template class TextReader { public: /*! - * \brief Constructor - * \param filename Filename of data - * \param is_skip_first_line True if need to skip header - */ - TextReader(const char* filename, bool is_skip_first_line, size_t progress_interval_bytes = SIZE_MAX): - filename_(filename), is_skip_first_line_(is_skip_first_line), read_progress_interval_bytes_(progress_interval_bytes) { + * \brief Constructor + * \param filename Filename of data + * \param is_skip_first_line True if need to skip header + */ + TextReader(const char* filename, bool is_skip_first_line, size_t progress_interval_bytes = SIZE_MAX) : filename_(filename), is_skip_first_line_(is_skip_first_line), read_progress_interval_bytes_(progress_interval_bytes) { if (is_skip_first_line_) { auto reader = VirtualFileReader::Make(filename); if (!reader->Init()) { @@ -454,33 +470,33 @@ class TextReader { } } /*! - * \brief Destructor - */ + * \brief Destructor + */ ~TextReader() { Clear(); } /*! - * \brief Clear cached data - */ + * \brief Clear cached data + */ inline void Clear() { lines_.clear(); lines_.shrink_to_fit(); } /*! - * \brief return first line of data - */ + * \brief return first line of data + */ inline std::string first_line() { return first_line_; } /*! - * \brief Get text data that read from file - * \return Text data, store in std::vector by line - */ + * \brief Get text data that read from file + * \return Text data, store in std::vector by line + */ inline std::vector& Lines() { return lines_; } /*! - * \brief Get joined text data that read from file - * \return Text data, store in std::string, joined all lines by delimiter - */ + * \brief Get joined text data that read from file + * \return Text data, store in std::string, joined all lines by delimiter + */ inline std::string JoinedLines(std::string delimiter = "\n") { std::stringstream ss; for (auto line : lines_) { @@ -494,47 +510,48 @@ class TextReader { INDEX_T total_cnt = 0; size_t bytes_read = 0; PipelineReader::Read(filename_, skip_bytes_, - [&process_fun, &bytes_read, &total_cnt, this] - (const char* buffer_process, size_t read_cnt) { - size_t cnt = 0; - size_t i = 0; - size_t last_i = 0; - // skip the break between \r and \n - if (last_line_.size() == 0 && buffer_process[0] == '\n') { - i = 1; - last_i = i; - } - while (i < read_cnt) { - if (buffer_process[i] == '\n' || buffer_process[i] == '\r') { - if (last_line_.size() > 0) { - last_line_.append(buffer_process + last_i, i - last_i); - process_fun(total_cnt, last_line_.c_str(), last_line_.size()); - last_line_ = ""; - } else { - process_fun(total_cnt, buffer_process + last_i, i - last_i); - } - ++cnt; - ++i; - ++total_cnt; - // skip end of line - while ((buffer_process[i] == '\n' || buffer_process[i] == '\r') && i < read_cnt) { ++i; } - last_i = i; - } else { - ++i; - } - } - if (last_i != read_cnt) { - last_line_.append(buffer_process + last_i, read_cnt - last_i); - } - - size_t prev_bytes_read = bytes_read; - bytes_read += read_cnt; - if (prev_bytes_read / read_progress_interval_bytes_ < bytes_read / read_progress_interval_bytes_) { - Log::Debug("Read %.1f GBs from %s.", 1.0 * bytes_read / kGbs, filename_); - } - - return cnt; - }); + [&process_fun, &bytes_read, &total_cnt, this](const char* buffer_process, size_t read_cnt) { + size_t cnt = 0; + size_t i = 0; + size_t last_i = 0; + // skip the break between \r and \n + if (last_line_.size() == 0 && buffer_process[0] == '\n') { + i = 1; + last_i = i; + } + while (i < read_cnt) { + if (buffer_process[i] == '\n' || buffer_process[i] == '\r') { + if (last_line_.size() > 0) { + last_line_.append(buffer_process + last_i, i - last_i); + process_fun(total_cnt, last_line_.c_str(), last_line_.size()); + last_line_ = ""; + } else { + process_fun(total_cnt, buffer_process + last_i, i - last_i); + } + ++cnt; + ++i; + ++total_cnt; + // skip end of line + while ((buffer_process[i] == '\n' || buffer_process[i] == '\r') && i < read_cnt) { + ++i; + } + last_i = i; + } else { + ++i; + } + } + if (last_i != read_cnt) { + last_line_.append(buffer_process + last_i, read_cnt - last_i); + } + + size_t prev_bytes_read = bytes_read; + bytes_read += read_cnt; + if (prev_bytes_read / read_progress_interval_bytes_ < bytes_read / read_progress_interval_bytes_) { + Log::Debug("Read %.1f GBs from %s.", 1.0 * bytes_read / kGbs, filename_); + } + + return cnt; + }); // if last line of file doesn't contain end of line if (last_line_.size() > 0) { Log::Info("Warning: last line of %s has no end of line, still using this line", filename_); @@ -546,14 +563,14 @@ class TextReader { } /*! - * \brief Read all text data from file in memory - * \return number of lines of text data - */ + * \brief Read all text data from file in memory + * \return number of lines of text data + */ INDEX_T ReadAllLines() { return ReadAllAndProcess( - [=](INDEX_T, const char* buffer, size_t size) { - lines_.emplace_back(buffer, size); - }); + [=](INDEX_T, const char* buffer, size_t size) { + lines_.emplace_back(buffer, size); + }); } std::vector ReadContent(size_t* out_len) { @@ -577,8 +594,7 @@ class TextReader { INDEX_T SampleFromFile(Random* random, INDEX_T sample_cnt, std::vector* out_sampled_data) { INDEX_T cur_sample_cnt = 0; return ReadAllAndProcess([=, &random, &cur_sample_cnt, - &out_sampled_data] - (INDEX_T line_idx, const char* buffer, size_t size) { + &out_sampled_data](INDEX_T line_idx, const char* buffer, size_t size) { if (cur_sample_cnt < sample_cnt) { out_sampled_data->emplace_back(buffer, size); ++cur_sample_cnt; @@ -591,54 +607,52 @@ class TextReader { }); } /*! - * \brief Read part of text data from file in memory, use filter_fun to filter data - * \param filter_fun Function that perform data filter - * \param out_used_data_indices Store line indices that read text data - * \return The number of total data - */ + * \brief Read part of text data from file in memory, use filter_fun to filter data + * \param filter_fun Function that perform data filter + * \param out_used_data_indices Store line indices that read text data + * \return The number of total data + */ INDEX_T ReadAndFilterLines(const std::function& filter_fun, std::vector* out_used_data_indices) { out_used_data_indices->clear(); INDEX_T total_cnt = ReadAllAndProcess( - [&filter_fun, &out_used_data_indices, this] - (INDEX_T line_idx , const char* buffer, size_t size) { - bool is_used = filter_fun(line_idx); - if (is_used) { - out_used_data_indices->push_back(line_idx); - lines_.emplace_back(buffer, size); - } - }); + [&filter_fun, &out_used_data_indices, this](INDEX_T line_idx, const char* buffer, size_t size) { + bool is_used = filter_fun(line_idx); + if (is_used) { + out_used_data_indices->push_back(line_idx); + lines_.emplace_back(buffer, size); + } + }); return total_cnt; } INDEX_T SampleAndFilterFromFile(const std::function& filter_fun, std::vector* out_used_data_indices, - Random* random, INDEX_T sample_cnt, std::vector* out_sampled_data) { + Random* random, INDEX_T sample_cnt, std::vector* out_sampled_data) { INDEX_T cur_sample_cnt = 0; out_used_data_indices->clear(); INDEX_T total_cnt = ReadAllAndProcess( [=, &filter_fun, &out_used_data_indices, &random, &cur_sample_cnt, - &out_sampled_data] - (INDEX_T line_idx, const char* buffer, size_t size) { - bool is_used = filter_fun(line_idx); - if (is_used) { - out_used_data_indices->push_back(line_idx); - if (cur_sample_cnt < sample_cnt) { - out_sampled_data->emplace_back(buffer, size); - ++cur_sample_cnt; - } else { - const size_t idx = static_cast(random->NextInt(0, static_cast(out_used_data_indices->size()))); - if (idx < static_cast(sample_cnt)) { - out_sampled_data->operator[](idx) = std::string(buffer, size); + &out_sampled_data](INDEX_T line_idx, const char* buffer, size_t size) { + bool is_used = filter_fun(line_idx); + if (is_used) { + out_used_data_indices->push_back(line_idx); + if (cur_sample_cnt < sample_cnt) { + out_sampled_data->emplace_back(buffer, size); + ++cur_sample_cnt; + } else { + const size_t idx = static_cast(random->NextInt(0, static_cast(out_used_data_indices->size()))); + if (idx < static_cast(sample_cnt)) { + out_sampled_data->operator[](idx) = std::string(buffer, size); + } + } } - } - } - }); + }); return total_cnt; } INDEX_T CountLine() { return ReadAllAndProcess( - [=](INDEX_T, const char*, size_t) { - }); + [=](INDEX_T, const char*, size_t) { + }); } INDEX_T ReadAllAndProcessParallelWithFilter(const std::function&)>& process_fun, const std::function& filter_fun) { @@ -647,56 +661,57 @@ class TextReader { size_t bytes_read = 0; INDEX_T used_cnt = 0; PipelineReader::Read(filename_, skip_bytes_, - [&process_fun, &filter_fun, &total_cnt, &bytes_read, &used_cnt, this] - (const char* buffer_process, size_t read_cnt) { - size_t cnt = 0; - size_t i = 0; - size_t last_i = 0; - INDEX_T start_idx = used_cnt; - // skip the break between \r and \n - if (last_line_.size() == 0 && buffer_process[0] == '\n') { - i = 1; - last_i = i; - } - while (i < read_cnt) { - if (buffer_process[i] == '\n' || buffer_process[i] == '\r') { - if (last_line_.size() > 0) { - last_line_.append(buffer_process + last_i, i - last_i); - if (filter_fun(used_cnt, total_cnt)) { - lines_.push_back(last_line_); - ++used_cnt; - } - last_line_ = ""; - } else { - if (filter_fun(used_cnt, total_cnt)) { - lines_.emplace_back(buffer_process + last_i, i - last_i); - ++used_cnt; - } - } - ++cnt; - ++i; - ++total_cnt; - // skip end of line - while ((buffer_process[i] == '\n' || buffer_process[i] == '\r') && i < read_cnt) { ++i; } - last_i = i; - } else { - ++i; - } - } - process_fun(start_idx, lines_); - lines_.clear(); - if (last_i != read_cnt) { - last_line_.append(buffer_process + last_i, read_cnt - last_i); - } - - size_t prev_bytes_read = bytes_read; - bytes_read += read_cnt; - if (prev_bytes_read / read_progress_interval_bytes_ < bytes_read / read_progress_interval_bytes_) { - Log::Debug("Read %.1f GBs from %s.", 1.0 * bytes_read / kGbs, filename_); - } - - return cnt; - }); + [&process_fun, &filter_fun, &total_cnt, &bytes_read, &used_cnt, this](const char* buffer_process, size_t read_cnt) { + size_t cnt = 0; + size_t i = 0; + size_t last_i = 0; + INDEX_T start_idx = used_cnt; + // skip the break between \r and \n + if (last_line_.size() == 0 && buffer_process[0] == '\n') { + i = 1; + last_i = i; + } + while (i < read_cnt) { + if (buffer_process[i] == '\n' || buffer_process[i] == '\r') { + if (last_line_.size() > 0) { + last_line_.append(buffer_process + last_i, i - last_i); + if (filter_fun(used_cnt, total_cnt)) { + lines_.push_back(last_line_); + ++used_cnt; + } + last_line_ = ""; + } else { + if (filter_fun(used_cnt, total_cnt)) { + lines_.emplace_back(buffer_process + last_i, i - last_i); + ++used_cnt; + } + } + ++cnt; + ++i; + ++total_cnt; + // skip end of line + while ((buffer_process[i] == '\n' || buffer_process[i] == '\r') && i < read_cnt) { + ++i; + } + last_i = i; + } else { + ++i; + } + } + process_fun(start_idx, lines_); + lines_.clear(); + if (last_i != read_cnt) { + last_line_.append(buffer_process + last_i, read_cnt - last_i); + } + + size_t prev_bytes_read = bytes_read; + bytes_read += read_cnt; + if (prev_bytes_read / read_progress_interval_bytes_ < bytes_read / read_progress_interval_bytes_) { + Log::Debug("Read %.1f GBs from %s.", 1.0 * bytes_read / kGbs, filename_); + } + + return cnt; + }); // if last line of file doesn't contain end of line if (last_line_.size() > 0) { Log::Info("Warning: last line of %s has no end of line, still using this line", filename_); @@ -718,13 +733,13 @@ class TextReader { INDEX_T ReadPartAndProcessParallel(const std::vector& used_data_indices, const std::function&)>& process_fun) { return ReadAllAndProcessParallelWithFilter(process_fun, - [&used_data_indices](INDEX_T used_cnt, INDEX_T total_cnt) { - if (static_cast(used_cnt) < used_data_indices.size() && total_cnt == used_data_indices[used_cnt]) { - return true; - } else { - return false; - } - }); + [&used_data_indices](INDEX_T used_cnt, INDEX_T total_cnt) { + if (static_cast(used_cnt) < used_data_indices.size() && total_cnt == used_data_indices[used_cnt]) { + return true; + } else { + return false; + } + }); } private: @@ -745,4 +760,4 @@ class TextReader { } // namespace StochTree -#endif // STOCHTREE_IO_H_ +#endif // STOCHTREE_IO_H_ diff --git a/include/stochtree/leaf_model.h b/include/stochtree/leaf_model.h index 6f42c110..b40156e6 100644 --- a/include/stochtree/leaf_model.h +++ b/include/stochtree/leaf_model.h @@ -383,8 +383,8 @@ class GaussianConstantSuffStat { void IncrementSuffStat(ForestDataset& dataset, Eigen::VectorXd& outcome, ForestTracker& tracker, data_size_t row_idx, int tree_idx) { n += 1; if (dataset.HasVarWeights()) { - sum_w += 1/dataset.VarWeightValue(row_idx); - sum_yw += outcome(row_idx, 0)/dataset.VarWeightValue(row_idx); + sum_w += 1 / dataset.VarWeightValue(row_idx); + sum_yw += outcome(row_idx, 0) / dataset.VarWeightValue(row_idx); } else { sum_w += 1.0; sum_yw += outcome(row_idx, 0); @@ -462,7 +462,10 @@ class GaussianConstantLeafModel { * * \param tau Leaf node prior scale parameter */ - GaussianConstantLeafModel(double tau) {tau_ = tau; normal_sampler_ = UnivariateNormalSampler();} + GaussianConstantLeafModel(double tau) { + tau_ = tau; + normal_sampler_ = UnivariateNormalSampler(); + } ~GaussianConstantLeafModel() {} /*! * \brief Log marginal likelihood for a proposed split, evaluated only for observations that fall into the node being split. @@ -511,11 +514,12 @@ class GaussianConstantLeafModel { * * \param tau Leaf node prior scale parameter */ - void SetScale(double tau) {tau_ = tau;} + void SetScale(double tau) { tau_ = tau; } /*! * \brief Whether this model requires a basis vector for posterior inference and prediction */ - inline bool RequiresBasis() {return false;} + inline bool RequiresBasis() { return false; } + private: double tau_; UnivariateNormalSampler normal_sampler_; @@ -547,11 +551,11 @@ class GaussianUnivariateRegressionSuffStat { void IncrementSuffStat(ForestDataset& dataset, Eigen::VectorXd& outcome, ForestTracker& tracker, data_size_t row_idx, int tree_idx) { n += 1; if (dataset.HasVarWeights()) { - sum_xxw += dataset.BasisValue(row_idx, 0)*dataset.BasisValue(row_idx, 0)/dataset.VarWeightValue(row_idx); - sum_yxw += outcome(row_idx, 0)*dataset.BasisValue(row_idx, 0)/dataset.VarWeightValue(row_idx); + sum_xxw += dataset.BasisValue(row_idx, 0) * dataset.BasisValue(row_idx, 0) / dataset.VarWeightValue(row_idx); + sum_yxw += outcome(row_idx, 0) * dataset.BasisValue(row_idx, 0) / dataset.VarWeightValue(row_idx); } else { - sum_xxw += dataset.BasisValue(row_idx, 0)*dataset.BasisValue(row_idx, 0); - sum_yxw += outcome(row_idx, 0)*dataset.BasisValue(row_idx, 0); + sum_xxw += dataset.BasisValue(row_idx, 0) * dataset.BasisValue(row_idx, 0); + sum_yxw += outcome(row_idx, 0) * dataset.BasisValue(row_idx, 0); } } /*! @@ -562,7 +566,7 @@ class GaussianUnivariateRegressionSuffStat { sum_xxw = 0.0; sum_yxw = 0.0; } - /*! + /*! * \brief Increment the value of each sufficient statistic by the values provided by `suff_stat` * * \param suff_stat Sufficient statistic to be added to the current sufficient statistics @@ -621,7 +625,10 @@ class GaussianUnivariateRegressionSuffStat { /*! \brief Marginal likelihood and posterior computation for gaussian homoskedastic constant leaf outcome model */ class GaussianUnivariateRegressionLeafModel { public: - GaussianUnivariateRegressionLeafModel(double tau) {tau_ = tau; normal_sampler_ = UnivariateNormalSampler();} + GaussianUnivariateRegressionLeafModel(double tau) { + tau_ = tau; + normal_sampler_ = UnivariateNormalSampler(); + } ~GaussianUnivariateRegressionLeafModel() {} /*! * \brief Log marginal likelihood for a proposed split, evaluated only for observations that fall into the node being split. @@ -665,8 +672,9 @@ class GaussianUnivariateRegressionLeafModel { */ void SampleLeafParameters(ForestDataset& dataset, ForestTracker& tracker, ColumnVector& residual, Tree* tree, int tree_num, double global_variance, std::mt19937& gen); void SetEnsembleRootPredictedValue(ForestDataset& dataset, TreeEnsemble* ensemble, double root_pred_value); - void SetScale(double tau) {tau_ = tau;} - inline bool RequiresBasis() {return true;} + void SetScale(double tau) { tau_ = tau; } + inline bool RequiresBasis() { return true; } + private: double tau_; UnivariateNormalSampler normal_sampler_; @@ -703,16 +711,16 @@ class GaussianMultivariateRegressionSuffStat { n += 1; if (dataset.HasVarWeights()) { for (int i = 0; i < p; i++) { - ytWX(0,i) += outcome(row_idx, 0) * dataset.BasisValue(row_idx, i) / dataset.VarWeightValue(row_idx); + ytWX(0, i) += outcome(row_idx, 0) * dataset.BasisValue(row_idx, i) / dataset.VarWeightValue(row_idx); for (int j = 0; j < p; j++) { - XtWX(i,j) += dataset.BasisValue(row_idx, i) * dataset.BasisValue(row_idx, j) / dataset.VarWeightValue(row_idx); + XtWX(i, j) += dataset.BasisValue(row_idx, i) * dataset.BasisValue(row_idx, j) / dataset.VarWeightValue(row_idx); } } } else { for (int i = 0; i < p; i++) { - ytWX(0,i) += outcome(row_idx, 0) * dataset.BasisValue(row_idx, i); + ytWX(0, i) += outcome(row_idx, 0) * dataset.BasisValue(row_idx, i); for (int j = 0; j < p; j++) { - XtWX(i,j) += dataset.BasisValue(row_idx, i) * dataset.BasisValue(row_idx, j); + XtWX(i, j) += dataset.BasisValue(row_idx, i) * dataset.BasisValue(row_idx, j); } } } @@ -729,7 +737,7 @@ class GaussianMultivariateRegressionSuffStat { } } } - /*! + /*! * \brief Increment the value of each sufficient statistic by the values provided by `suff_stat` * * \param suff_stat Sufficient statistic to be added to the current sufficient statistics @@ -793,7 +801,10 @@ class GaussianMultivariateRegressionLeafModel { * * \param Sigma_0 Prior covariance, must have the same number of rows and columns as dimensions of the basis vector for the multivariate regression problem */ - GaussianMultivariateRegressionLeafModel(Eigen::MatrixXd& Sigma_0) {Sigma_0_ = Sigma_0; multivariate_normal_sampler_ = MultivariateNormalSampler();} + GaussianMultivariateRegressionLeafModel(Eigen::MatrixXd& Sigma_0) { + Sigma_0_ = Sigma_0; + multivariate_normal_sampler_ = MultivariateNormalSampler(); + } ~GaussianMultivariateRegressionLeafModel() {} /*! * \brief Log marginal likelihood for a proposed split, evaluated only for observations that fall into the node being split. @@ -837,8 +848,9 @@ class GaussianMultivariateRegressionLeafModel { */ void SampleLeafParameters(ForestDataset& dataset, ForestTracker& tracker, ColumnVector& residual, Tree* tree, int tree_num, double global_variance, std::mt19937& gen); void SetEnsembleRootPredictedValue(ForestDataset& dataset, TreeEnsemble* ensemble, double root_pred_value); - void SetScale(Eigen::MatrixXd& Sigma_0) {Sigma_0_ = Sigma_0;} - inline bool RequiresBasis() {return true;} + void SetScale(Eigen::MatrixXd& Sigma_0) { Sigma_0_ = Sigma_0; } + inline bool RequiresBasis() { return true; } + private: Eigen::MatrixXd Sigma_0_; MultivariateNormalSampler multivariate_normal_sampler_; @@ -864,7 +876,7 @@ class LogLinearVarianceSuffStat { */ void IncrementSuffStat(ForestDataset& dataset, Eigen::VectorXd& outcome, ForestTracker& tracker, data_size_t row_idx, int tree_idx) { n += 1; - weighted_sum_ei += std::exp(std::log(outcome(row_idx)*outcome(row_idx)) - tracker.GetSamplePrediction(row_idx) + tracker.GetTreeSamplePrediction(row_idx, tree_idx)); + weighted_sum_ei += std::exp(std::log(outcome(row_idx) * outcome(row_idx)) - tracker.GetSamplePrediction(row_idx) + tracker.GetTreeSamplePrediction(row_idx, tree_idx)); } /*! * \brief Reset all of the sufficient statistics to zero @@ -929,7 +941,11 @@ class LogLinearVarianceSuffStat { /*! \brief Marginal likelihood and posterior computation for heteroskedastic log-linear variance model */ class LogLinearVarianceLeafModel { public: - LogLinearVarianceLeafModel(double a, double b) {a_ = a; b_ = b; gamma_sampler_ = GammaSampler();} + LogLinearVarianceLeafModel(double a, double b) { + a_ = a; + b_ = b; + gamma_sampler_ = GammaSampler(); + } ~LogLinearVarianceLeafModel() {} /*! * \brief Log marginal likelihood for a proposed split, evaluated only for observations that fall into the node being split. @@ -974,16 +990,16 @@ class LogLinearVarianceLeafModel { */ void SampleLeafParameters(ForestDataset& dataset, ForestTracker& tracker, ColumnVector& residual, Tree* tree, int tree_num, double global_variance, std::mt19937& gen); void SetEnsembleRootPredictedValue(ForestDataset& dataset, TreeEnsemble* ensemble, double root_pred_value); - void SetPriorShape(double a) {a_ = a;} - void SetPriorRate(double b) {b_ = b;} - inline bool RequiresBasis() {return false;} + void SetPriorShape(double a) { a_ = a; } + void SetPriorRate(double b) { b_ = b; } + inline bool RequiresBasis() { return false; } + private: double a_; double b_; GammaSampler gamma_sampler_; }; - /*! \brief Sufficient statistic and associated operations for complementary log-log ordinal BART model */ class CloglogOrdinalSuffStat { public: @@ -1016,20 +1032,20 @@ class CloglogOrdinalSuffStat { unsigned int y = static_cast(outcome(row_idx)); // Get auxiliary data from tracker (assuming types: 0=latents Z, 1=forest predictions, 2=cutpoints gamma, 3=cumsum exp of gamma) - double Z = dataset.GetAuxiliaryDataValue(0, row_idx); // latent variables Z + double Z = dataset.GetAuxiliaryDataValue(0, row_idx); // latent variables Z double lambda_minus = dataset.GetAuxiliaryDataValue(1, row_idx); // forest predictions excluding current tree // Get cutpoints gamma and cumulative sum of exp(gamma) const std::vector& gamma = dataset.GetAuxiliaryDataVectorConst(2); // cutpoints gamma const std::vector& seg = dataset.GetAuxiliaryDataVectorConst(3); // cumsum exp of gamma - + int K = gamma.size() + 1; // Number of ordinal categories if (y == K - 1) { other_sum += std::exp(lambda_minus) * seg[y]; // checked and it's correct } else { sum_Y_less_K += 1.0; - other_sum += std::exp(lambda_minus) * (Z * std::exp(gamma[y]) + seg[y]); // checked and it's correct + other_sum += std::exp(lambda_minus) * (Z * std::exp(gamma[y]) + seg[y]); // checked and it's correct } } @@ -1150,7 +1166,7 @@ class CloglogOrdinalLeafModel { * Samples from log-gamma: sample from gamma, then take log. */ void SampleLeafParameters(ForestDataset& dataset, ForestTracker& tracker, ColumnVector& residual, Tree* tree, int tree_num, double global_variance, std::mt19937& gen); - inline bool RequiresBasis() {return false;} + inline bool RequiresBasis() { return false; } private: double a_; @@ -1185,12 +1201,12 @@ using LeafModelVariant = std::variant; -template +template static inline SuffStatVariant createSuffStat(SuffStatConstructorArgs... leaf_suff_stat_args) { return SuffStatType(leaf_suff_stat_args...); } -template +template static inline LeafModelVariant createLeafModel(LeafModelConstructorArgs... leaf_model_args) { return LeafModelType(leaf_model_args...); } @@ -1238,12 +1254,11 @@ static inline LeafModelVariant leafModelFactory(ModelType model_type, double tau } } -template +template static inline void AccumulateSuffStatProposed( - SuffStatType& node_suff_stat, SuffStatType& left_suff_stat, SuffStatType& right_suff_stat, ForestDataset& dataset, ForestTracker& tracker, - ColumnVector& residual, double global_variance, TreeSplit& split, int tree_num, int leaf_num, int split_feature, int num_threads, - SuffStatConstructorArgs&... suff_stat_args -) { + SuffStatType& node_suff_stat, SuffStatType& left_suff_stat, SuffStatType& right_suff_stat, ForestDataset& dataset, ForestTracker& tracker, + ColumnVector& residual, double global_variance, TreeSplit& split, int tree_num, int leaf_num, int split_feature, int num_threads, + SuffStatConstructorArgs&... suff_stat_args) { // Determine the position of the node's indices in the forest tracking data structure int node_begin_index = tracker.UnsortedNodeBegin(tree_num, leaf_num); int node_end_index = tracker.UnsortedNodeEnd(tree_num, leaf_num); @@ -1309,9 +1324,9 @@ static inline void AccumulateSuffStatProposed( } } -template +template static inline void AccumulateSuffStatExisting(SuffStatType& node_suff_stat, SuffStatType& left_suff_stat, SuffStatType& right_suff_stat, ForestDataset& dataset, ForestTracker& tracker, - ColumnVector& residual, double global_variance, int tree_num, int split_node_id, int left_node_id, int right_node_id) { + ColumnVector& residual, double global_variance, int tree_num, int split_node_id, int left_node_id, int right_node_id) { // Acquire iterators auto left_node_begin_iter = tracker.UnsortedNodeBeginIterator(tree_num, left_node_id); auto left_node_end_iter = tracker.UnsortedNodeEndIterator(tree_num, left_node_id); @@ -1333,7 +1348,7 @@ static inline void AccumulateSuffStatExisting(SuffStatType& node_suff_stat, Suff } } -template +template static inline void AccumulateSingleNodeSuffStat(SuffStatType& node_suff_stat, ForestDataset& dataset, ForestTracker& tracker, ColumnVector& residual, int tree_num, int node_id) { // Acquire iterators std::vector::iterator node_begin_iter; @@ -1354,10 +1369,10 @@ static inline void AccumulateSingleNodeSuffStat(SuffStatType& node_suff_stat, Fo } } -template +template static inline void AccumulateCutpointBinSuffStat(SuffStatType& left_suff_stat, ForestTracker& tracker, CutpointGridContainer& cutpoint_grid_container, - ForestDataset& dataset, ColumnVector& residual, double global_variance, int tree_num, int node_id, - int feature_num, int cutpoint_num) { + ForestDataset& dataset, ColumnVector& residual, double global_variance, int tree_num, int node_id, + int feature_num, int cutpoint_num) { // Acquire iterators auto node_begin_iter = tracker.SortedNodeBeginIterator(node_id, feature_num); auto node_end_iter = tracker.SortedNodeEndIterator(node_id, feature_num); @@ -1382,8 +1397,8 @@ static inline void AccumulateCutpointBinSuffStat(SuffStatType& left_suff_stat, F } } -/*! \} */ // end of leaf_model_group +/*! \} */ // end of leaf_model_group -} // namespace StochTree +} // namespace StochTree -#endif // STOCHTREE_LEAF_MODEL_H_ +#endif // STOCHTREE_LEAF_MODEL_H_ diff --git a/include/stochtree/linear_regression.h b/include/stochtree/linear_regression.h new file mode 100644 index 00000000..4e30dfd4 --- /dev/null +++ b/include/stochtree/linear_regression.h @@ -0,0 +1,163 @@ +/*! + * Copyright (c) 2026 stochtree authors. All rights reserved. + * Licensed under the MIT License. See LICENSE file in the project root for license information. + */ +#ifndef STOCHTREE_REGRESSION_H_ +#define STOCHTREE_REGRESSION_H_ + +#include +#include +#include + +#include + +namespace StochTree { + +/*! + * \brief Sample a regression coefficient from the posterior distribution of a univariate Gaussian regression model with known error variance and known prior variance. + * + * \param y Pointer to outcome array of length n + * \param x Pointer to covariate array of length n + * \param error_variance Known error variance (sigma^2) + * \param prior_variance Known prior variance (tau^2) + * \param n Number of observations + * \param gen Random number generator + * \return double + */ +static double sample_univariate_gaussian_regression_coefficient(double* y, double* x, double error_variance, double prior_variance, int n, std::mt19937& gen) { + double sum_xx = 0.0; + double sum_yx = 0.0; + for (int i = 0; i < n; i++) { + sum_xx += x[i] * x[i]; + sum_yx += y[i] * x[i]; + } + double post_var = (prior_variance * error_variance) / (sum_xx * prior_variance + error_variance); + double post_mean = post_var * (sum_yx / error_variance); + return sample_standard_normal(post_mean, std::sqrt(post_var), gen); +} + +/*! + * \brief Sample regression coefficients from the posterior distribution of a bivariate Gaussian regression model with known error variance and known prior variance. + * + * \param output Pointer to sampled coefficient array of length 2, where the sampled coefficients will be stored + * \param y Pointer to outcome array of length n + * \param x1 Pointer to first covariate array of length n + * \param x2 Pointer to second covariate array of length n + * \param error_variance Known error variance (sigma^2) + * \param prior_variance_11 First diagonal element of the known prior covariance matrix (tau^2 for the first coefficient) + * \param prior_variance_12 Off-diagonal element of the known prior covariance matrix (covariance between the two coefficients) -- note that this is the same as prior_variance_21 so only one off-diagonal element is needed + * \param prior_variance_22 Second diagonal element of the known prior covariance matrix (tau^2 for the second coefficient) + * \param n Number of observations + * \param gen Random number generator + */ +static void sample_general_bivariate_gaussian_regression_coefficients(double* output, double* y, double* x1, double* x2, double error_variance, double prior_variance_11, double prior_variance_12, double prior_variance_22, int n, std::mt19937& gen) { + double det_prior_var = prior_variance_11 * prior_variance_22 - prior_variance_12 * prior_variance_12; + double inv_prior_var_11 = prior_variance_22 / det_prior_var; + double inv_prior_var_12 = -prior_variance_12 / det_prior_var; + double inv_prior_var_22 = prior_variance_11 / det_prior_var; + double sum_x1x1 = 0.0; + double sum_x1x2 = 0.0; + double sum_x2x2 = 0.0; + double sum_yx1 = 0.0; + double sum_yx2 = 0.0; + for (int i = 0; i < n; i++) { + sum_x1x1 += x1[i] * x1[i]; + sum_x1x2 += x1[i] * x2[i]; + sum_x2x2 += x2[i] * x2[i]; + sum_yx1 += y[i] * x1[i]; + sum_yx2 += y[i] * x2[i]; + } + double post_var_pre_inv_11 = inv_prior_var_11 + sum_x1x1 / error_variance; + double post_var_pre_inv_12 = inv_prior_var_12 + sum_x1x2 / error_variance; + double post_var_pre_inv_22 = inv_prior_var_22 + sum_x2x2 / error_variance; + double det_post_var_pre_inv = post_var_pre_inv_11 * post_var_pre_inv_22 - post_var_pre_inv_12 * post_var_pre_inv_12; + double post_var_11 = post_var_pre_inv_22 / det_post_var_pre_inv; + double post_var_12 = -post_var_pre_inv_12 / det_post_var_pre_inv; + double post_var_22 = post_var_pre_inv_11 / det_post_var_pre_inv; + double post_mean_1 = post_var_11 * (sum_yx1 / error_variance) + post_var_12 * (sum_yx2 / error_variance); + double post_mean_2 = post_var_12 * (sum_yx1 / error_variance) + post_var_22 * (sum_yx2 / error_variance); + double chol_var_11 = std::sqrt(post_var_11); + double chol_var_12 = post_var_12 / chol_var_11; + double chol_var_22 = std::sqrt(post_var_22 - chol_var_12 * chol_var_12); + double z1 = sample_standard_normal(0.0, 1.0, gen); + double z2 = sample_standard_normal(0.0, 1.0, gen); + output[0] = post_mean_1 + chol_var_11 * z1; + output[1] = post_mean_2 + chol_var_12 * z1 + chol_var_22 * z2; +} + +/*! + * \brief Sample regression coefficients from the posterior distribution of a bivariate Gaussian regression model with known error variance and known diagonal prior variance. + * + * \param output Pointer to sampled coefficient array of length 2, where the sampled coefficients will be stored + * \param y Pointer to outcome array of length n + * \param x1 Pointer to first covariate array of length n + * \param x2 Pointer to second covariate array of length n + * \param error_variance Known error variance (sigma^2) + * \param prior_variance_11 First diagonal element of the known prior covariance matrix (tau^2 for the first coefficient) + * \param prior_variance_22 Second diagonal element of the known prior covariance matrix (tau^2 for the second coefficient) + * \param n Number of observations + * \param gen Random number generator + */ +static void sample_diagonal_bivariate_gaussian_regression_coefficients(double* output, double* y, double* x1, double* x2, double error_variance, double prior_variance_11, double prior_variance_22, int n, std::mt19937& gen) { + double inv_prior_var_11 = 1.0 / prior_variance_11; + double inv_prior_var_22 = 1.0 / prior_variance_22; + double sum_x1x1 = 0.0; + double sum_x1x2 = 0.0; + double sum_x2x2 = 0.0; + double sum_yx1 = 0.0; + double sum_yx2 = 0.0; + for (int i = 0; i < n; i++) { + sum_x1x1 += x1[i] * x1[i]; + sum_x1x2 += x1[i] * x2[i]; + sum_x2x2 += x2[i] * x2[i]; + sum_yx1 += y[i] * x1[i]; + sum_yx2 += y[i] * x2[i]; + } + double post_var_pre_inv_11 = inv_prior_var_11 + sum_x1x1 / error_variance; + double post_var_pre_inv_12 = sum_x1x2 / error_variance; + double post_var_pre_inv_22 = inv_prior_var_22 + sum_x2x2 / error_variance; + double det_post_var_pre_inv = post_var_pre_inv_11 * post_var_pre_inv_22 - post_var_pre_inv_12 * post_var_pre_inv_12; + double post_var_11 = post_var_pre_inv_22 / det_post_var_pre_inv; + double post_var_12 = -post_var_pre_inv_12 / det_post_var_pre_inv; + double post_var_22 = post_var_pre_inv_11 / det_post_var_pre_inv; + double post_mean_1 = post_var_11 * (sum_yx1 / error_variance) + post_var_12 * (sum_yx2 / error_variance); + double post_mean_2 = post_var_12 * (sum_yx1 / error_variance) + post_var_22 * (sum_yx2 / error_variance); + double chol_var_11 = std::sqrt(post_var_11); + double chol_var_12 = post_var_12 / chol_var_11; + double chol_var_22 = std::sqrt(post_var_22 - chol_var_12 * chol_var_12); + double z1 = sample_standard_normal(0.0, 1.0, gen); + double z2 = sample_standard_normal(0.0, 1.0, gen); + output[0] = post_mean_1 + chol_var_11 * z1; + output[1] = post_mean_2 + chol_var_12 * z1 + chol_var_22 * z2; +} + +/*! + * \brief Sample regression coefficients from the posterior distribution of a bivariate Gaussian regression model with known error variance and known diagonal prior variance. + * + * \param y Eigen::VectorXd of outcomes of length n + * \param X Eigen::MatrixXd of covariates with n rows and p columns + * \param error_variance Known error variance (sigma^2) + * \param prior_variance Eigen::MatrixXd of known prior covariance matrix (tau^2 for the coefficients) of dimension p x p + * \param n Number of observations + * \param gen Random number generator + */ +static Eigen::VectorXd sample_general_gaussian_regression_coefficients(const Eigen::Ref& y, const Eigen::Ref& X, double error_variance, const Eigen::Ref& prior_variance, int n, std::mt19937& gen) { + int p = X.cols(); + Eigen::MatrixXd inv_prior_var = prior_variance.inverse(); + Eigen::MatrixXd XtX = X.transpose() * X; + Eigen::VectorXd Xty = X.transpose() * y; + Eigen::MatrixXd post_var_pre_inv = inv_prior_var + XtX / error_variance; + Eigen::MatrixXd post_var = post_var_pre_inv.inverse(); + Eigen::VectorXd post_mean = post_var * (Xty / error_variance); + Eigen::LLT chol(post_var); + Eigen::MatrixXd L = chol.matrixL(); + Eigen::VectorXd z(p); + for (int i = 0; i < p; i++) { + z(i) = sample_standard_normal(0.0, 1.0, gen); + } + return post_mean + L * z; +} + +} // namespace StochTree + +#endif // STOCHTREE_REGRESSION_H_ diff --git a/include/stochtree/log.h b/include/stochtree/log.h index 9f64c31b..8ce87f79 100644 --- a/include/stochtree/log.h +++ b/include/stochtree/log.h @@ -1,9 +1,9 @@ /*! * Logging and runtime value checking utilities. - * - * This code is largely included as-is from LightGBM, which carries + * + * This code is largely included as-is from LightGBM, which carries * the following copyright information: - * + * * Copyright (c) 2016 Microsoft Corporation. All rights reserved. * Licensed under the MIT License. See LICENSE file in the project root for * license information. @@ -65,10 +65,10 @@ namespace StochTree { #endif #ifndef CHECK_NOTNULL -#define CHECK_NOTNULL(pointer) \ - if ((pointer) == nullptr) \ +#define CHECK_NOTNULL(pointer) \ + if ((pointer) == nullptr) \ StochTree::Log::Fatal(#pointer " Can't be NULL at %s, line %d .\n", \ - __FILE__, __LINE__); + __FILE__, __LINE__); #endif enum class LogLevel : int { @@ -83,7 +83,7 @@ enum class LogLevel : int { */ class Log { public: - using Callback = void (*)(const char *); + using Callback = void (*)(const char*); /*! * \brief Resets the minimal log level. It is INFO by default. * \param level The new minimal log level. @@ -92,25 +92,25 @@ class Log { static void ResetCallBack(Callback callback) { GetLogCallBack() = callback; } - static void Debug(const char *format, ...) { + static void Debug(const char* format, ...) { va_list val; va_start(val, format); Write(LogLevel::Debug, "Debug", format, val); va_end(val); } - static void Info(const char *format, ...) { + static void Info(const char* format, ...) { va_list val; va_start(val, format); Write(LogLevel::Info, "Info", format, val); va_end(val); } - static void Warning(const char *format, ...) { + static void Warning(const char* format, ...) { va_list val; va_start(val, format); Write(LogLevel::Warning, "Warning", format, val); va_end(val); } - static void Fatal(const char *format, ...) { + static void Fatal(const char* format, ...) { va_list val; const size_t kBufSize = 1024; char str_buf[kBufSize]; @@ -135,7 +135,7 @@ class Log { } private: - static void Write(LogLevel level, const char *level_str, const char *format, + static void Write(LogLevel level, const char* level_str, const char* format, va_list val) { if (level <= GetLevel()) { // omit the message with low level // R code should write back to R's output stream, @@ -166,12 +166,12 @@ class Log { // a trick to use static variable in header file. // May be not good, but avoid to use an additional cpp file - static LogLevel &GetLevel() { + static LogLevel& GetLevel() { static THREAD_LOCAL LogLevel level = LogLevel::Info; return level; } - static Callback &GetLogCallBack() { + static Callback& GetLogCallBack() { static THREAD_LOCAL Callback callback = nullptr; return callback; } diff --git a/include/stochtree/mainpage.h b/include/stochtree/mainpage.h index dc39f162..09cfdc28 100644 --- a/include/stochtree/mainpage.h +++ b/include/stochtree/mainpage.h @@ -3,80 +3,80 @@ /*! * \mainpage stochtree C++ Documentation - * + * * \section getting-started Getting Started - * + * * `stochtree` can be built and run as a standalone C++ program directly from source using `cmake`: - * + * * \subsection cloning-repo Cloning the Repository - * - * To clone the repository, you must have git installed, which you can do following these instructions. - * - * Once git is available at the command line, navigate to the folder that will store this project (in bash / zsh, this is done by running `cd` followed by the path to the directory). + * + * To clone the repository, you must have git installed, which you can do following these instructions. + * + * Once git is available at the command line, navigate to the folder that will store this project (in bash / zsh, this is done by running `cd` followed by the path to the directory). * Then, clone the `stochtree` repo as a subfolder by running * \code{.sh} * git clone --recursive https://github.com/StochasticTree/stochtree.git * \endcode - * - * NOTE: this project incorporates several dependencies as git submodules, - * which is why the `--recursive` flag is necessary (some systems may perform a recursive clone without this flag, but - * `--recursive` ensures this behavior on all platforms). If you have already cloned the repo without the `--recursive` flag, + * + * NOTE: this project incorporates several dependencies as git submodules, + * which is why the `--recursive` flag is necessary (some systems may perform a recursive clone without this flag, but + * `--recursive` ensures this behavior on all platforms). If you have already cloned the repo without the `--recursive` flag, * you can retrieve the submodules recursively by running `git submodule update --init --recursive` in the main repo directory. - * + * * \section key-components Key Components - * + * * The stochtree C++ core consists of thousands of lines of C++ code, but it can organized and understood through several components (see [topics](topics.html) for more detail): - * + * * - Trees: the most important "primitive" of decision tree algorithms is the \ref tree_group "decision tree itself", which in stochtree is defined by a \ref StochTree::Tree "Tree" class as well as a series of static helper functions for prediction. * - Forest: individual trees are combined into a \ref forest_group "forest", or ensemble, which in stochtree is defined by the \ref StochTree::TreeEnsemble "TreeEnsemble" class and a container of forests is defined by the \ref StochTree::ForestContainer "ForestContainer" class. * - Dataset: data can be loaded from a variety of sources into a `stochtree` \ref data_group "data layer". * - Leaf Model: `stochtree`'s data structures are generalized to support a wide range of models, which are defined via specialized classes in the \ref leaf_model_group "leaf model layer". * - Sampler: helper functions that sample forests from training data comprise the \ref sampling_group "sampling layer" of `stochtree`. - * + * * \section extending-stochtree Extending stochtree - * + * * \subsection custom-leaf-models Custom Leaf Models - * - * The \ref leaf_model_group "leaf model documentation" details the key components of new decision tree models: - * custom `LeafModel` and `SuffStat` classes that implement a model's log marginal likelihood and posterior computations. - * - * Adding a new leaf model will consist largely of implementing new versions of each of these classes which track the - * API of the existing classes. Once these classes exist, they need to be reflected in several places. - * + * + * The \ref leaf_model_group "leaf model documentation" details the key components of new decision tree models: + * custom `LeafModel` and `SuffStat` classes that implement a model's log marginal likelihood and posterior computations. + * + * Adding a new leaf model will consist largely of implementing new versions of each of these classes which track the + * API of the existing classes. Once these classes exist, they need to be reflected in several places. + * * Suppose, for the sake of illustration, that the newest custom leaf model is a multinomial logit model. - * + * * First, add an entry to the \ref StochTree::ModelType "ModelType" enumeration for this new model type - * + * * \code{.cpp} * enum ModelType { - * kConstantLeafGaussian, - * kUnivariateRegressionLeafGaussian, - * kMultivariateRegressionLeafGaussian, - * kLogLinearVariance, - * kMultinomialLogit, + * kConstantLeafGaussian, + * kUnivariateRegressionLeafGaussian, + * kMultivariateRegressionLeafGaussian, + * kLogLinearVariance, + * kMultinomialLogit, * }; - * \endcode - * + * \endcode + * * Next, add entries to the `std::variants` that bundle related `SuffStat` and `LeafModel` classes - * + * * \code{.cpp} - * using SuffStatVariant = std::variant; - * \endcode - * + * \endcode + * * \code{.cpp} - * using LeafModelVariant = std::variant; - * \endcode - * + * \endcode + * * Finally, update the \ref StochTree::suffStatFactory "suffStatFactory" and \ref StochTree::leafModelFactory "leafModelFactory" functions to add a logic branch registering these new objects - * + * * \code{.cpp} * static inline SuffStatVariant suffStatFactory(ModelType model_type, int basis_dim = 0) { * if (model_type == kConstantLeafGaussian) { @@ -93,8 +93,8 @@ * Log::Fatal("Incompatible model type provided to suff stat factory"); * } * } - * \endcode - * + * \endcode + * * \code{.cpp} * static inline LeafModelVariant leafModelFactory(ModelType model_type, double tau, Eigen::MatrixXd& Sigma0, double a, double b) { * if (model_type == kConstantLeafGaussian) { @@ -111,8 +111,8 @@ * Log::Fatal("Incompatible model type provided to leaf model factory"); * } * } - * \endcode - * + * \endcode + * */ #endif // STOCHTREE_MAINPAGE_H_ diff --git a/include/stochtree/meta.h b/include/stochtree/meta.h index d0aa4049..84647938 100644 --- a/include/stochtree/meta.h +++ b/include/stochtree/meta.h @@ -1,9 +1,9 @@ /*! * Macros, constants, and type definitions used elsewhere in the codebase - * - * This code is largely included as-is from LightGBM, which carries + * + * This code is largely included as-is from LightGBM, which carries * the following copyright information: - * + * * Copyright (c) 2016 Microsoft Corporation. All rights reserved. * Licensed under the MIT License. See LICENSE file in the project root for license information. */ @@ -20,20 +20,42 @@ #include #if (defined(_MSC_VER) && (defined(_M_IX86) || defined(_M_AMD64))) || defined(__INTEL_COMPILER) || MM_PREFETCH - #include - #define PREFETCH_T0(addr) _mm_prefetch(reinterpret_cast(addr), _MM_HINT_T0) +#include +#define PREFETCH_T0(addr) _mm_prefetch(reinterpret_cast(addr), _MM_HINT_T0) #elif defined(__GNUC__) - #define PREFETCH_T0(addr) __builtin_prefetch(reinterpret_cast(addr), 0, 3) +#define PREFETCH_T0(addr) __builtin_prefetch(reinterpret_cast(addr), 0, 3) #else - #define PREFETCH_T0(addr) do {} while (0) +#define PREFETCH_T0(addr) \ + do { \ + } while (0) #endif namespace StochTree { +enum class LinkFunction { + Identity, + Probit, + Cloglog +}; + +enum class OutcomeType { + Continuous, + Binary, + Ordinal +}; + +enum class MeanLeafModelType { + GaussianConstant, + GaussianUnivariateRegression, + GaussianMultivariateRegression, + LogLinearVariance, + CloglogOrdinal +}; + /*! \brief Integer encoding of feature types */ enum FeatureType { - kNumeric, /*!< Numeric feature */ - kOrderedCategorical, /*!< Ordered categorical feature */ + kNumeric, /*!< Numeric feature */ + kOrderedCategorical, /*!< Ordered categorical feature */ kUnorderedCategorical /*!< Unordered categorical feature */ }; @@ -123,31 +145,29 @@ typedef int32_t node_t; typedef double split_cond_t; using PredictFunction = -std::function>&, double* output)>; + std::function>&, double* output)>; using PredictSparseFunction = -std::function>&, std::vector>* output)>; - -typedef void(*ReduceFunction)(const char* input, char* output, int type_size, comm_size_t array_size); - + std::function>&, std::vector>* output)>; -typedef void(*ReduceScatterFunction)(char* input, comm_size_t input_size, int type_size, - const comm_size_t* block_start, const comm_size_t* block_len, int num_block, char* output, comm_size_t output_size, - const ReduceFunction& reducer); +typedef void (*ReduceFunction)(const char* input, char* output, int type_size, comm_size_t array_size); -typedef void(*AllgatherFunction)(char* input, comm_size_t input_size, const comm_size_t* block_start, - const comm_size_t* block_len, int num_block, char* output, comm_size_t output_size); +typedef void (*ReduceScatterFunction)(char* input, comm_size_t input_size, int type_size, + const comm_size_t* block_start, const comm_size_t* block_len, int num_block, char* output, comm_size_t output_size, + const ReduceFunction& reducer); +typedef void (*AllgatherFunction)(char* input, comm_size_t input_size, const comm_size_t* block_start, + const comm_size_t* block_len, int num_block, char* output, comm_size_t output_size); #define NO_SPECIFIC (-1) const int kAlignedSize = 32; -#define SIZE_ALIGNED(t) ((t) + kAlignedSize - 1) / kAlignedSize * kAlignedSize +#define SIZE_ALIGNED(t) ((t) + kAlignedSize - 1) / kAlignedSize* kAlignedSize // Refer to https://docs.microsoft.com/en-us/cpp/error-messages/compiler-warnings/compiler-warning-level-4-c4127?view=vs-2019 #ifdef _MSC_VER - #pragma warning(disable : 4127) +#pragma warning(disable : 4127) #endif } // namespace StochTree diff --git a/include/stochtree/normal_sampler.h b/include/stochtree/normal_sampler.h index 60b8a550..bca02079 100644 --- a/include/stochtree/normal_sampler.h +++ b/include/stochtree/normal_sampler.h @@ -12,11 +12,12 @@ namespace StochTree { class UnivariateNormalSampler { public: - UnivariateNormalSampler() {std_normal_dist_ = standard_normal();} + UnivariateNormalSampler() { std_normal_dist_ = standard_normal(); } ~UnivariateNormalSampler() {} double Sample(double mean, double variance, std::mt19937& gen) { return mean + std::sqrt(variance) * std_normal_dist_(gen); } + private: /*! \brief Standard normal distribution */ standard_normal std_normal_dist_; @@ -24,7 +25,7 @@ class UnivariateNormalSampler { class MultivariateNormalSampler { public: - MultivariateNormalSampler() {std_normal_dist_ = standard_normal();} + MultivariateNormalSampler() { std_normal_dist_ = standard_normal(); } ~MultivariateNormalSampler() {} std::vector Sample(Eigen::VectorXd& mean, Eigen::MatrixXd& covariance, std::mt19937& gen) { // Dimension extraction and checks @@ -32,7 +33,7 @@ class MultivariateNormalSampler { int cov_rows = covariance.rows(); int cov_cols = covariance.cols(); CHECK_EQ(mean_cols, cov_cols); - + // Variance cholesky decomposition Eigen::LLT decomposition(covariance); Eigen::MatrixXd covariance_chol = decomposition.matrixL(); @@ -57,7 +58,7 @@ class MultivariateNormalSampler { int cov_rows = covariance.rows(); int cov_cols = covariance.cols(); CHECK_EQ(mean_cols, cov_cols); - + // Variance cholesky decomposition Eigen::LLT decomposition(covariance); Eigen::MatrixXd covariance_chol = decomposition.matrixL(); @@ -71,11 +72,12 @@ class MultivariateNormalSampler { // Compute and return the sampled value return mean + covariance_chol * std_norm_vec; } + private: /*! \brief Standard normal distribution */ standard_normal std_normal_dist_; }; -} // namespace StochTree +} // namespace StochTree -#endif // STOCHTREE_NORMAL_SAMPLER_H_ \ No newline at end of file +#endif // STOCHTREE_NORMAL_SAMPLER_H_ \ No newline at end of file diff --git a/include/stochtree/openmp_utils.h b/include/stochtree/openmp_utils.h index 78c17234..e3add212 100644 --- a/include/stochtree/openmp_utils.h +++ b/include/stochtree/openmp_utils.h @@ -13,43 +13,43 @@ namespace StochTree { // OpenMP thread management inline int get_max_threads() { - return omp_get_max_threads(); + return omp_get_max_threads(); } inline int get_thread_num() { - return omp_get_thread_num(); + return omp_get_thread_num(); } inline int get_num_threads() { - return omp_get_num_threads(); + return omp_get_num_threads(); } inline void set_num_threads(int num_threads) { - omp_set_num_threads(num_threads); + omp_set_num_threads(num_threads); } - + #define STOCHTREE_PARALLEL_FOR(num_threads) \ - _Pragma("omp parallel for num_threads(num_threads)") + _Pragma("omp parallel for num_threads(num_threads)") #define STOCHTREE_REDUCTION_ADD(var) \ - _Pragma("omp reduction(+:var)") + _Pragma("omp reduction(+:var)") #define STOCHTREE_CRITICAL \ - _Pragma("omp critical") + _Pragma("omp critical") #else #define STOCHTREE_HAS_OPENMP 0 -inline int get_max_threads() {return 1;} +inline int get_max_threads() { return 1; } -inline int get_thread_num() {return 0;} +inline int get_thread_num() { return 0; } -inline int get_num_threads() {return 1;} +inline int get_num_threads() { return 1; } inline void set_num_threads(int num_threads) {} - + #define STOCHTREE_PARALLEL_FOR(num_threads) - + #define STOCHTREE_REDUCTION_ADD(var) #define STOCHTREE_CRITICAL @@ -57,57 +57,57 @@ inline void set_num_threads(int num_threads) {} #endif static int GetMaxThreads() { - return get_max_threads(); + return get_max_threads(); } static int GetCurrentThreadNum() { - return get_thread_num(); + return get_thread_num(); } - + static int GetNumThreads() { - return get_num_threads(); + return get_num_threads(); } - + static void SetNumThreads(int num_threads) { - set_num_threads(num_threads); + set_num_threads(num_threads); } - + static bool IsOpenMPAvailable() { - return STOCHTREE_HAS_OPENMP; + return STOCHTREE_HAS_OPENMP; } - + static int GetOptimalThreadCount(int workload_size, int min_work_per_thread = 1000) { - if (!IsOpenMPAvailable()) { - return 1; - } - - int max_threads = GetMaxThreads(); - int optimal_threads = workload_size / min_work_per_thread; - - return std::min(optimal_threads, max_threads); + if (!IsOpenMPAvailable()) { + return 1; + } + + int max_threads = GetMaxThreads(); + int optimal_threads = workload_size / min_work_per_thread; + + return std::min(optimal_threads, max_threads); } // Parallel execution utilities -template +template void ParallelFor(int start, int end, int num_threads, Func func) { - if (num_threads <= 0) { - num_threads = GetOptimalThreadCount(end - start); + if (num_threads <= 0) { + num_threads = GetOptimalThreadCount(end - start); + } + + if (num_threads == 1 || !STOCHTREE_HAS_OPENMP) { + // Sequential execution + for (int i = start; i < end; ++i) { + func(i); } - - if (num_threads == 1 || !STOCHTREE_HAS_OPENMP) { - // Sequential execution - for (int i = start; i < end; ++i) { - func(i); - } - } else { - // Parallel execution - STOCHTREE_PARALLEL_FOR(num_threads) - for (int i = start; i < end; ++i) { - func(i); - } + } else { + // Parallel execution + STOCHTREE_PARALLEL_FOR(num_threads) + for (int i = start; i < end; ++i) { + func(i); } + } } -} // namespace StochTree +} // namespace StochTree -#endif // STOCHTREE_OPENMP_UTILS_H \ No newline at end of file +#endif // STOCHTREE_OPENMP_UTILS_H \ No newline at end of file diff --git a/include/stochtree/ordinal_sampler.h b/include/stochtree/ordinal_sampler.h index d67563e2..4d5777b8 100644 --- a/include/stochtree/ordinal_sampler.h +++ b/include/stochtree/ordinal_sampler.h @@ -13,29 +13,28 @@ #include #include -#include namespace StochTree { static double sample_truncated_exponential_low_high(double u, double rate, double low, double high) { - return -std::log((1-u)*std::exp(-rate*low) + u*std::exp(-rate*high))/rate; + return -std::log((1 - u) * std::exp(-rate * low) + u * std::exp(-rate * high)) / rate; } static double sample_truncated_exponential_low(double u, double rate, double low) { - return -std::log((1-u)*std::exp(-rate*low))/rate; + return -std::log((1 - u) * std::exp(-rate * low)) / rate; } static double sample_truncated_exponential_high(double u, double rate, double high) { - return -std::log1p(u*std::expm1(-high*rate))/rate; + return -std::log1p(u * std::expm1(-high * rate)) / rate; } static double sample_exponential(double u, double rate) { - return -std::log1p(-u)/rate; + return -std::log1p(-u) / rate; } /*! * \brief Sampler for ordinal model hyperparameters - * + * * This class handles MCMC sampling for ordinal-specific parameters: * - Truncated exponential latent variables (Z) * - Cutpoint parameters (gamma) @@ -50,9 +49,9 @@ class OrdinalSampler { /*! * \brief Sample from truncated exponential distribution - * + * * Samples from exponential distribution truncated to [low,high] - * + * * \param gen Random number generator * \param rate Rate parameter for exponential distribution * \param low Lower truncation bound @@ -63,7 +62,7 @@ class OrdinalSampler { /*! * \brief Update truncated exponential latent variables (Z) - * + * * \param dataset Forest dataset containing training data (covariates) and auxiliary data needed for sampling * \param outcome Vector of outcome values * \param gen Random number generator @@ -72,7 +71,7 @@ class OrdinalSampler { /*! * \brief Update gamma cutpoint parameters - * + * * \param dataset Forest dataset containing training data (covariates) and auxiliary data needed for sampling * \param outcome Vector of outcome values * \param alpha_gamma Shape parameter for log-gamma prior on cutpoints gamma @@ -80,13 +79,13 @@ class OrdinalSampler { * \param gamma_0 Fixed value for first cutpoint parameter (for identifiability) * \param gen Random number generator */ - void UpdateGammaParams(ForestDataset& dataset, Eigen::VectorXd& outcome, - double alpha_gamma, double beta_gamma, + void UpdateGammaParams(ForestDataset& dataset, Eigen::VectorXd& outcome, + double alpha_gamma, double beta_gamma, double gamma_0, std::mt19937& gen); /*! * \brief Update cumulative exponential sums (seg) - * + * * \param dataset Forest dataset containing training data (covariates) and auxiliary data needed for sampling */ void UpdateCumulativeExpSums(ForestDataset& dataset); @@ -95,6 +94,6 @@ class OrdinalSampler { GammaSampler gamma_sampler_; }; -} // namespace StochTree +} // namespace StochTree -#endif // STOCHTREE_ORDINAL_SAMPLER_H_ +#endif // STOCHTREE_ORDINAL_SAMPLER_H_ diff --git a/include/stochtree/partition_tracker.h b/include/stochtree/partition_tracker.h index f25c875c..164ffae0 100644 --- a/include/stochtree/partition_tracker.h +++ b/include/stochtree/partition_tracker.h @@ -1,25 +1,25 @@ /*! * Copyright (c) 2024 stochtree authors. - * + * * Data structures used for tracking dataset through the tree building process. - * + * * The first category of data structure tracks observations available in nodes of a tree. - * a. UnsortedNodeSampleTracker tracks the observations available in every leaf of every tree in an ensemble, + * a. UnsortedNodeSampleTracker tracks the observations available in every leaf of every tree in an ensemble, * in no feature-specific sort order. It is primarily designed for use in BART-based algorithms. - * b. SortedNodeSampleTracker tracks the observations available in a every leaf of a tree, pre-sorted + * b. SortedNodeSampleTracker tracks the observations available in a every leaf of a tree, pre-sorted * separately for each feature. It is primarily designed for use in XBART-based algorithms. - * + * * The second category, SampleNodeMapper, maps observations from a dataset to leaf nodes. - * - * SampleNodeMapper is inspired by the design of the DataPartition class in LightGBM, + * + * SampleNodeMapper is inspired by the design of the DataPartition class in LightGBM, * released under the MIT license with the following copyright: - * + * * Copyright (c) 2016 Microsoft Corporation. All rights reserved. * Licensed under the MIT License. See LICENSE file in the project root for license information. - * - * SortedNodeSampleTracker is inspired by the "approximate" split finding method in xgboost, released + * + * SortedNodeSampleTracker is inspired by the "approximate" split finding method in xgboost, released * under the Apache license with the following copyright: - * + * * Copyright 2015~2023 by XGBoost Contributors */ #ifndef STOCHTREE_PARTITION_TRACKER_H_ @@ -47,7 +47,7 @@ class ForestTracker { public: /*! * \brief Construct a new `ForestTracker` object - * + * * \param covariates Matrix of covariate data * \param feature_types Type of each feature (column) in `covariates`. This is represented by the enum `StochTree::FeatureType` * \param num_trees Number of trees in an ensemble to be sampled @@ -83,14 +83,16 @@ class ForestTracker { std::vector::iterator UnsortedNodeEndIterator(int tree_id, int node_id); std::vector::iterator SortedNodeBeginIterator(int node_id, int feature_id); std::vector::iterator SortedNodeEndIterator(int node_id, int feature_id); - SamplePredMapper* GetSamplePredMapper() {return sample_pred_mapper_.get();} - SampleNodeMapper* GetSampleNodeMapper() {return sample_node_mapper_.get();} - UnsortedNodeSampleTracker* GetUnsortedNodeSampleTracker() {return unsorted_node_sample_tracker_.get();} - SortedNodeSampleTracker* GetSortedNodeSampleTracker() {return sorted_node_sample_tracker_.get();} - int GetNumObservations() {return num_observations_;} - int GetNumTrees() {return num_trees_;} - int GetNumFeatures() {return num_features_;} - bool Initialized() {return initialized_;} + SamplePredMapper* GetSamplePredMapper() { return sample_pred_mapper_.get(); } + SampleNodeMapper* GetSampleNodeMapper() { return sample_node_mapper_.get(); } + UnsortedNodeSampleTracker* GetUnsortedNodeSampleTracker() { return unsorted_node_sample_tracker_.get(); } + SortedNodeSampleTracker* GetSortedNodeSampleTracker() { return sorted_node_sample_tracker_.get(); } + double* GetSumPredictions() { return sum_predictions_.data(); } + int GetNumObservations() { return num_observations_; } + int GetNumTrees() { return num_trees_; } + int GetNumFeatures() { return num_features_; } + bool Initialized() { return initialized_; } + private: /*! \brief Mapper from observations to predicted values summed over every tree in a forest */ std::vector sum_predictions_; @@ -102,7 +104,7 @@ class ForestTracker { * Primarily used in MCMC algorithms */ std::unique_ptr unsorted_node_sample_tracker_; - /*! \brief Data structure tracking / updating observations available in each node for each feature (pre-sorted) for a given tree in a forest + /*! \brief Data structure tracking / updating observations available in each node for each feature (pre-sorted) for a given tree in a forest * Primarily used in GFR algorithms */ std::unique_ptr presort_container_; @@ -145,10 +147,10 @@ class SamplePredMapper { CHECK_LT(tree_id, num_trees_); tree_preds_[tree_id][sample_id] = value; } - - inline int NumTrees() {return num_trees_;} - - inline int NumObservations() {return num_observations_;} + + inline int NumTrees() { return num_trees_; } + + inline int NumObservations() { return num_observations_; } inline void AssignAllSamplesToConstantPrediction(int tree_id, double value) { for (data_size_t i = 0; i < num_observations_; i++) { @@ -174,8 +176,8 @@ class SampleNodeMapper { tree_observation_indices_[j].resize(num_observations_); } } - - SampleNodeMapper(SampleNodeMapper& other){ + + SampleNodeMapper(SampleNodeMapper& other) { num_trees_ = other.NumTrees(); num_observations_ = other.NumObservations(); // Initialize the vector of vectors of leaf indices for each tree @@ -214,10 +216,10 @@ class SampleNodeMapper { CHECK_LT(tree_id, num_trees_); tree_observation_indices_[tree_id][sample_id] = node_id; } - - inline int NumTrees() {return num_trees_;} - - inline int NumObservations() {return num_observations_;} + + inline int NumTrees() { return num_trees_; } + + inline int NumObservations() { return num_observations_; } inline void AssignAllSamplesToRoot(int tree_id) { for (data_size_t i = 0; i < num_observations_; i++) { @@ -333,10 +335,11 @@ class UnsortedNodeSampleTracker { void PartitionTreeNode(Eigen::MatrixXd& covariates, int tree_id, int node_id, int left_node_id, int right_node_id, int feature_split, std::vector const& category_list) { return feature_partitions_[tree_id]->PartitionNode(covariates, node_id, left_node_id, right_node_id, feature_split, category_list); } - + /*! \brief Convert a tree to root */ void ResetTreeToRoot(int tree_id, data_size_t n) { - feature_partitions_[tree_id].reset(new FeatureUnsortedPartition(n));; + feature_partitions_[tree_id].reset(new FeatureUnsortedPartition(n)); + ; } /*! \brief Convert a (currently split) node to a leaf */ @@ -418,7 +421,7 @@ class UnsortedNodeSampleTracker { /*! \brief Update SampleNodeMapper for all the observations in tree */ void UpdateObservationMapping(Tree* tree, int tree_id, SampleNodeMapper* sample_node_mapper) { - std::vector leaves = tree->GetLeaves(); + std::vector const& leaves = tree->GetLeaves(); int leaf; for (int i = 0; i < leaves.size(); i++) { leaf = leaves[i]; @@ -447,15 +450,15 @@ class NodeOffsetSize { ~NodeOffsetSize() {} - void SetSorted() {presorted_ = true;} + void SetSorted() { presorted_ = true; } - bool IsSorted() {return presorted_;} + bool IsSorted() { return presorted_; } - data_size_t Begin() {return node_begin_;} + data_size_t Begin() { return node_begin_; } - data_size_t End() {return node_end_;} + data_size_t End() { return node_end_; } - data_size_t Size() {return node_size_;} + data_size_t Size() { return node_size_; } private: data_size_t node_begin_; @@ -468,16 +471,17 @@ class NodeOffsetSize { class FeaturePresortPartition; /*! \brief Data structure for presorting a feature by its values - * - * This class is intended to be run *once* on a dataset as it + * + * This class is intended to be run *once* on a dataset as it * pre-sorts each feature across the entire dataset. - * + * * FeaturePresortPartition is intended for use in recursive construction - * of new trees, and each new tree's FeaturePresortPartition is initialized + * of new trees, and each new tree's FeaturePresortPartition is initialized * from a FeaturePresortRoot class so that features are only arg-sorted one time. */ class FeaturePresortRoot { - friend FeaturePresortPartition; + friend FeaturePresortPartition; + public: FeaturePresortRoot(Eigen::MatrixXd& covariates, int32_t feature_index, FeatureType feature_type) { feature_index_ = feature_index; @@ -488,17 +492,17 @@ class FeaturePresortRoot { void ArgsortRoot(Eigen::MatrixXd& covariates) { data_size_t num_obs = covariates.rows(); - + // Make a vector of indices from 0 to num_obs - 1 - if (feature_sort_indices_.size() != num_obs){ + if (feature_sort_indices_.size() != num_obs) { feature_sort_indices_.resize(num_obs, 0); } std::iota(feature_sort_indices_.begin(), feature_sort_indices_.end(), 0); // Define a custom comparator to be used with stable_sort: - // For every two indices l and r store as elements of `data_sort_indices_`, + // For every two indices l and r store as elements of `data_sort_indices_`, // compare them for sorting purposes by indexing the covariate's raw data with both l and r - auto comp_op = [&](size_t const &l, size_t const &r) { return std::less{}(covariates(l, feature_index_), covariates(r, feature_index_)); }; + auto comp_op = [&](size_t const& l, size_t const& r) { return std::less{}(covariates(l, feature_index_), covariates(r, feature_index_)); }; std::stable_sort(feature_sort_indices_.begin(), feature_sort_indices_.end(), comp_op); } @@ -520,21 +524,21 @@ class FeaturePresortRootContainer { ~FeaturePresortRootContainer() {} - FeaturePresortRoot* GetFeaturePresort(int feature_num) {return feature_presort_[feature_num].get(); } + FeaturePresortRoot* GetFeaturePresort(int feature_num) { return feature_presort_[feature_num].get(); } private: std::vector> feature_presort_; int num_features_; }; -/*! \brief Data structure that tracks pre-sorted feature values +/*! \brief Data structure that tracks pre-sorted feature values * through a tree's split lifecycle - * - * This class is initialized from a FeaturePresortRoot which has computed the - * sort indices for a given feature over the entire dataset, so that sorting + * + * This class is initialized from a FeaturePresortRoot which has computed the + * sort indices for a given feature over the entire dataset, so that sorting * is not necessary for each new tree. - * - * When a split is made, this class handles sifting for each feature, so that + * + * When a split is made, this class handles sifting for each feature, so that * the presorted feature values available at each node are easily queried. */ class FeaturePresortPartition { @@ -563,28 +567,29 @@ class FeaturePresortPartition { void SplitFeatureCategorical(Eigen::MatrixXd& covariates, int32_t node_id, int32_t feature_index, std::vector const& category_list); /*! \brief Start position of node indexed by node_id */ - data_size_t NodeBegin(int32_t node_id) {return node_offset_sizes_[node_id].Begin();} + data_size_t NodeBegin(int32_t node_id) { return node_offset_sizes_[node_id].Begin(); } /*! \brief End position of node indexed by node_id */ - data_size_t NodeEnd(int32_t node_id) {return node_offset_sizes_[node_id].End();} + data_size_t NodeEnd(int32_t node_id) { return node_offset_sizes_[node_id].End(); } /*! \brief Size (in observations) of node indexed by node_id */ - data_size_t NodeSize(int32_t node_id) {return node_offset_sizes_[node_id].Size();} + data_size_t NodeSize(int32_t node_id) { return node_offset_sizes_[node_id].Size(); } /*! \brief Data indices for a given node */ std::vector NodeIndices(int node_id); /*! \brief Feature sort index j */ - data_size_t SortIndex(data_size_t j) {return feature_sort_indices_[j];} + data_size_t SortIndex(data_size_t j) { return feature_sort_indices_[j]; } /*! \brief Feature type */ - FeatureType GetFeatureType() {return feature_type_;} + FeatureType GetFeatureType() { return feature_type_; } /*! \brief Update SampleNodeMapper for all the observations in node_id */ void UpdateObservationMapping(int node_id, int tree_id, SampleNodeMapper* sample_node_mapper); /*! \brief Feature sort indices */ std::vector feature_sort_indices_; + private: /*! \brief Add left and right nodes */ void AddLeftRightNodes(data_size_t left_node_begin, data_size_t left_node_size, data_size_t right_node_begin, data_size_t right_node_size); @@ -663,7 +668,7 @@ class SortedNodeSampleTracker { } /*! \brief Feature sort index j for feature_index */ - data_size_t SortIndex(data_size_t j, int feature_index) {return feature_partitions_[feature_index]->SortIndex(j); } + data_size_t SortIndex(data_size_t j, int feature_index) { return feature_partitions_[feature_index]->SortIndex(j); } /*! \brief Update SampleNodeMapper for all the observations in node_id */ void UpdateObservationMapping(int node_id, int tree_id, SampleNodeMapper* sample_node_mapper, int feature_index = 0) { @@ -675,6 +680,6 @@ class SortedNodeSampleTracker { int num_features_; }; -} // namespace StochTree +} // namespace StochTree -#endif // STOCHTREE_PARTITION_TRACKER_H_ +#endif // STOCHTREE_PARTITION_TRACKER_H_ diff --git a/include/stochtree/prediction.h b/include/stochtree/prediction.h new file mode 100644 index 00000000..760bb55a --- /dev/null +++ b/include/stochtree/prediction.h @@ -0,0 +1,184 @@ +/*! + * Copyright (c) 2026 stochtree authors. All rights reserved. + * Licensed under the MIT License. See LICENSE file in the project root for license information. + */ +#ifndef STOCHTREE_PREDICTION_H_ +#define STOCHTREE_PREDICTION_H_ + +#include +#include +#include +#include +#include +#include + +namespace StochTree { + +/*! \brief Determines whether posterior predictions are returned as-is or pre-aggregated. */ +enum class PredType { + kPosterior, + kMean +}; + +/*! \brief Determines the scale of predictions (i.e. whether a probability / class transformation is applied) + * + * \details Options are + * + * - linear (i.e. raw) scale, + * - probability scale, + * - class predictions. + * + * Only valid for binary / ordinal outcome models. + */ +enum class PredScale { + kLinear, + kProbability, + kClass +}; + +/*! \brief Selector for model terms that should be predicted. */ +struct BARTPredTerms { + bool y_hat = true; + bool mean_forest = false; + bool variance_forest = false; + bool random_effects = false; +}; + +/*! \brief Struct returning BART model predictions + * + * \details The BART prediction routine can return predictions of one or more model terms, + * and this struct serves as a "container" for model predictions. All observation-specific + * prediction terms can be: + * 1. Pre-aggregated (type = mean) or contain the entire posterior (type = posterior) + * 2. Linear scale, probability scale or class predictions (for binary / ordinal outcomes) + */ +struct BARTPredictionResult { + // Outcome conditional mean + std::vector y_hat; + + // Covariate-dependent prognostic term (mu(x)) + std::vector mean_forest_predictions; + + // Conditional variance term + std::vector variance_forest_predictions; + + // Random effects predictions + std::vector rfx_predictions; +}; + +/*! \brief Metadata for the BART prediction routine + * + * \details Stores details about the underlying model as well as prediction specifications needed for the prediction routine. + */ +struct BARTPredictionMetadata { + // Metadata about the samples / model (e.g., number of samples, burn-in, etc.) + int num_samples = 0; + int num_obs = 0; + int num_basis = 0; + double y_bar = 0.0; + double y_std = 0.0; + bool has_variance_forest = false; + bool has_rfx = false; + BARTRFXModelSpec rfx_model_spec; + PredType pred_type = PredType::kPosterior; + BARTPredTerms pred_terms; + PredScale pred_scale = PredScale::kLinear; + LinkFunction link_function = LinkFunction::Identity; + OutcomeType outcome_type = OutcomeType::Continuous; + int cloglog_num_classes = 0; +}; + +/*! \brief BART prediction function + * + * \details Accepts BARTData, BARTSamples, and a struct of metadata, which dictates the model terms for which + * predictions are computed / returned and any transformations done before returning (i.e. pre-aggregation, + * probit function transformation). + * + * \param data Struct wrapping pointers to prediction data from R / Python + * \param samples Object storing BART samples + * \param metadata Struct containing prediction metadata + * \return BARTPredictionResult struct containing prediction vectors + */ +BARTPredictionResult predict_bart_model(BARTData& data, BARTSamples& samples, BARTPredictionMetadata& metadata); + +/*! \brief Selector for model terms that should be predicted. */ +struct BCFPredTerms { + bool y_hat = true; + bool mu_x = false; + bool tau_x = false; + bool prognostic_function = false; + bool cate = false; + bool conditional_variance = false; + bool random_effects = false; +}; + +/*! \brief Struct returning BCF model predictions + * + * \details The BCF prediction routine can return predictions of one or more model terms, + * and this struct serves as a "container" for model predictions. All observation-specific + * prediction terms can be: + * 1. Pre-aggregated (type = mean) or contain the entire posterior (type = posterior) + * 2. Linear scale, probability scale or class predictions (for binary / ordinal outcomes) + */ +struct BCFPredictionResult { + // Outcome conditional mean + std::vector y_hat; + + // Covariate-dependent prognostic term (mu(x)) + std::vector mu_x; + + // Covariate-dependent treatment effect term (tau(x)) + std::vector tau_x; + + // Prognostic function (includes mu(x) and any random intercepts, provided random effects + // were estimated with `intercept_only` or `intercept_plus_treatment` specification) + std::vector prognostic_function; + + // CATE function (includes tau(x) and any random slopes on treatment, provided + // random effects were estimated with `intercept_plus_treatment` specification) + std::vector cate; + + // Conditional variance term + std::vector conditional_variance; + + // Random effects predictions + std::vector random_effects; +}; + +/*! \brief Metadata for the BCF prediction routine + * + * \details Stores details about the underlying model as well as prediction specifications needed for the prediction routine. + */ +struct BCFPredictionMetadata { + // Metadata about the samples / model (e.g., number of samples, burn-in, etc.) could be added here as needed + int num_samples = 0; + int num_obs = 0; + int treatment_dim = 0; + double y_bar = 0.0; + double y_std = 0.0; + bool has_variance_forest = false; + bool has_rfx = false; + BCFRFXModelSpec rfx_model_spec; + bool adaptive_coding = false; + bool sample_tau_0 = false; + PredType pred_type = PredType::kPosterior; + BCFPredTerms pred_terms; + PredScale pred_scale = PredScale::kLinear; +}; + +/*! \brief BCF prediction function + * + * \details Accepts BCFData, BCFSamples, and a struct of metadata, which dictates the model terms for which + * predictions are computed / returned and any transformations done before returning (i.e. pre-aggregation, + * probit function transformation). + * + * \param data Struct wrapping pointers to prediction data from R / Python + * \param samples Object storing BCF samples + * \param metadata Struct containing prediction metadata + * \return BCFPredictionResult struct containing prediction vectors + */ +BCFPredictionResult predict_bcf_model(BCFData& data, BCFSamples& samples, BCFPredictionMetadata& metadata); + +} // namespace StochTree + +#endif // STOCHTREE_PREDICTION_H_ diff --git a/include/stochtree/prior.h b/include/stochtree/prior.h index 5d8686f7..3e1117e7 100644 --- a/include/stochtree/prior.h +++ b/include/stochtree/prior.h @@ -25,18 +25,19 @@ class RandomEffectsRegressionGaussianPrior : public RandomEffectsGaussianPrior { num_groups_ = num_groups; } ~RandomEffectsRegressionGaussianPrior() {} - double GetPriorVarianceShape() {return a_;} - double GetPriorVarianceScale() {return b_;} - int32_t GetNumComponents() {return num_components_;} - int32_t GetNumGroups() {return num_groups_;} - void SetPriorVarianceShape(double a) {a_ = a;} - void SetPriorVarianceScale(double b) {b_ = b;} - void SetNumComponents(int32_t num_components) {num_components_ = num_components;} - void SetNumGroups(int32_t num_groups) {num_groups_ = num_groups;} + double GetPriorVarianceShape() { return a_; } + double GetPriorVarianceScale() { return b_; } + int32_t GetNumComponents() { return num_components_; } + int32_t GetNumGroups() { return num_groups_; } + void SetPriorVarianceShape(double a) { a_ = a; } + void SetPriorVarianceScale(double b) { b_ = b; } + void SetNumComponents(int32_t num_components) { num_components_ = num_components; } + void SetNumGroups(int32_t num_groups) { num_groups_ = num_groups; } + private: double a_; double b_; - int32_t num_components_; + int32_t num_components_; int32_t num_groups_; }; @@ -49,14 +50,15 @@ class TreePrior { max_depth_ = max_depth; } ~TreePrior() {} - double GetAlpha() {return alpha_;} - double GetBeta() {return beta_;} - int32_t GetMinSamplesLeaf() {return min_samples_in_leaf_;} - int32_t GetMaxDepth() {return max_depth_;} - void SetAlpha(double alpha) {alpha_ = alpha;} - void SetBeta(double beta) {beta_ = beta;} - void SetMinSamplesLeaf(int32_t min_samples_in_leaf) {min_samples_in_leaf_ = min_samples_in_leaf;} - void SetMaxDepth(int32_t max_depth) {max_depth_ = max_depth;} + double GetAlpha() { return alpha_; } + double GetBeta() { return beta_; } + int32_t GetMinSamplesLeaf() { return min_samples_in_leaf_; } + int32_t GetMaxDepth() { return max_depth_; } + void SetAlpha(double alpha) { alpha_ = alpha; } + void SetBeta(double beta) { beta_ = beta; } + void SetMinSamplesLeaf(int32_t min_samples_in_leaf) { min_samples_in_leaf_ = min_samples_in_leaf; } + void SetMaxDepth(int32_t max_depth) { max_depth_ = max_depth; } + private: double alpha_; double beta_; @@ -71,15 +73,16 @@ class IGVariancePrior { scale_ = scale; } ~IGVariancePrior() {} - double GetShape() {return shape_;} - double GetScale() {return scale_;} - void SetShape(double shape) {shape_ = shape;} - void SetScale(double scale) {scale_ = scale;} + double GetShape() { return shape_; } + double GetScale() { return scale_; } + void SetShape(double shape) { shape_ = shape; } + void SetScale(double scale) { scale_ = scale; } + private: double shape_; double scale_; }; -} // namespace StochTree +} // namespace StochTree -#endif // STOCHTREE_PRIOR_H_ \ No newline at end of file +#endif // STOCHTREE_PRIOR_H_ \ No newline at end of file diff --git a/include/stochtree/probit.h b/include/stochtree/probit.h new file mode 100644 index 00000000..24a49252 --- /dev/null +++ b/include/stochtree/probit.h @@ -0,0 +1,35 @@ +/*! + * Copyright (c) 2024 stochtree authors. All rights reserved. + * Licensed under the MIT License. See LICENSE file in the project root for license information. + */ +#ifndef STOCHTREE_PROBIT_H_ +#define STOCHTREE_PROBIT_H_ + +#include + +namespace StochTree { + +static void sample_probit_latent_outcome(std::mt19937& gen, double* outcome, double* conditional_mean, double* partial_residual, double y_bar, int n) { + double uniform_draw_std; + double uniform_draw_trunc; + double quantile; + double cond_mean; + double latent_outcome; + for (int i = 0; i < n; i++) { + cond_mean = conditional_mean[i] + y_bar; + uniform_draw_std = standard_uniform_draw_53bit(gen); + quantile = norm_cdf(0 - cond_mean); + if (outcome[i] == 1.0) { + uniform_draw_trunc = quantile + uniform_draw_std * (1.0 - quantile); + latent_outcome = norm_inv_cdf(uniform_draw_trunc) + cond_mean; + } else { + uniform_draw_trunc = uniform_draw_std * quantile; + latent_outcome = norm_inv_cdf(uniform_draw_trunc) + cond_mean; + } + partial_residual[i] = latent_outcome - cond_mean; + } +} + +} // namespace StochTree + +#endif // STOCHTREE_PROBIT_H_ \ No newline at end of file diff --git a/include/stochtree/random.h b/include/stochtree/random.h index 3d39b647..6f5931e1 100644 --- a/include/stochtree/random.h +++ b/include/stochtree/random.h @@ -12,13 +12,13 @@ namespace StochTree { /*! -* \brief A wrapper for random generator -*/ + * \brief A wrapper for random generator + */ class Random { public: /*! - * \brief Constructor, with random seed - */ + * \brief Constructor, with random seed + */ Random() { std::random_device rd; auto genrator = std::mt19937(rd()); @@ -26,45 +26,45 @@ class Random { x = distribution(genrator); } /*! - * \brief Constructor, with specific seed - */ + * \brief Constructor, with specific seed + */ explicit Random(int seed) { x = seed; } /*! - * \brief Generate random integer, int16 range. `[0, 65536]` - * \param lower_bound lower bound - * \param upper_bound upper bound - * \return The random integer between [lower_bound, upper_bound) - */ + * \brief Generate random integer, int16 range. `[0, 65536]` + * \param lower_bound lower bound + * \param upper_bound upper bound + * \return The random integer between [lower_bound, upper_bound) + */ inline int NextShort(int lower_bound, int upper_bound) { return (RandInt16()) % (upper_bound - lower_bound) + lower_bound; } /*! - * \brief Generate random integer, int32 range - * \param lower_bound lower bound - * \param upper_bound upper bound - * \return The random integer between [lower_bound, upper_bound) - */ + * \brief Generate random integer, int32 range + * \param lower_bound lower bound + * \param upper_bound upper bound + * \return The random integer between [lower_bound, upper_bound) + */ inline int NextInt(int lower_bound, int upper_bound) { return (RandInt32()) % (upper_bound - lower_bound) + lower_bound; } /*! - * \brief Generate random float data - * \return The random float between `[0.0, 1.0)` - */ + * \brief Generate random float data + * \return The random float between `[0.0, 1.0)` + */ inline float NextFloat() { // get random float in `[0,1)` return static_cast(RandInt16()) / (32768.0f); } /*! - * \brief Sample K data from `{0,1,...,N-1}` - * \param N - * \param K - * \return K Ordered sampled data from `{0,1,...,N-1}` - */ + * \brief Sample K data from `{0,1,...,N-1}` + * \param N + * \param K + * \return K Ordered sampled data from `{0,1,...,N-1}` + */ inline std::vector Sample(int N, int K) { std::vector ret; ret.reserve(K); @@ -110,7 +110,6 @@ class Random { unsigned int x = 123456789; }; - } // namespace StochTree -#endif // STOCHTREE_RANDOM_H_ +#endif // STOCHTREE_RANDOM_H_ diff --git a/include/stochtree/random_effects.h b/include/stochtree/random_effects.h index b322a560..a1b02e7c 100644 --- a/include/stochtree/random_effects.h +++ b/include/stochtree/random_effects.h @@ -34,31 +34,33 @@ class RandomEffectsContainer; /*! \brief Wrapper around data structures for random effects sampling algorithms */ class RandomEffectsTracker { public: - RandomEffectsTracker(std::vector& group_indices); + RandomEffectsTracker(std::vector& group_indices); + RandomEffectsTracker(int* group_indices, int num_observations); ~RandomEffectsTracker() {} - inline data_size_t GetCategoryId(int observation_num) {return sample_category_mapper_->GetCategoryId(observation_num);} - inline data_size_t CategoryBegin(int category_id) {return category_sample_tracker_->CategoryBegin(category_id);} - inline data_size_t CategoryEnd(int category_id) {return category_sample_tracker_->CategoryEnd(category_id);} - inline data_size_t CategorySize(int category_id) {return category_sample_tracker_->CategorySize(category_id);} - inline int32_t NumCategories() {return num_categories_;} - inline int32_t CategoryNumber(int32_t category_id) {return category_sample_tracker_->CategoryNumber(category_id);} - SampleCategoryMapper* GetSampleCategoryMapper() {return sample_category_mapper_.get();} - CategorySampleTracker* GetCategorySampleTracker() {return category_sample_tracker_.get();} + inline data_size_t GetCategoryId(int observation_num) { return sample_category_mapper_->GetCategoryId(observation_num); } + inline data_size_t CategoryBegin(int category_id) { return category_sample_tracker_->CategoryBegin(category_id); } + inline data_size_t CategoryEnd(int category_id) { return category_sample_tracker_->CategoryEnd(category_id); } + inline data_size_t CategorySize(int category_id) { return category_sample_tracker_->CategorySize(category_id); } + inline int NumCategories() { return num_categories_; } + inline int CategoryNumber(int category_id) { return category_sample_tracker_->CategoryNumber(category_id); } + SampleCategoryMapper* GetSampleCategoryMapper() { return sample_category_mapper_.get(); } + CategorySampleTracker* GetCategorySampleTracker() { return category_sample_tracker_.get(); } std::vector::iterator UnsortedNodeBeginIterator(int category_id); std::vector::iterator UnsortedNodeEndIterator(int category_id); - std::map& GetLabelMap() {return category_sample_tracker_->GetLabelMap();} - std::vector& GetUniqueGroupIds() {return category_sample_tracker_->GetUniqueGroupIds();} - std::vector& NodeIndices(int category_id) {return category_sample_tracker_->NodeIndices(category_id);} - std::vector& NodeIndicesInternalIndex(int internal_category_id) {return category_sample_tracker_->NodeIndicesInternalIndex(internal_category_id);} - double GetPrediction(data_size_t observation_num) {return rfx_predictions_.at(observation_num);} - void SetPrediction(data_size_t observation_num, double pred) {rfx_predictions_.at(observation_num) = pred;} + std::map& GetLabelMap() { return category_sample_tracker_->GetLabelMap(); } + std::vector& GetUniqueGroupIds() { return category_sample_tracker_->GetUniqueGroupIds(); } + std::vector& NodeIndices(int category_id) { return category_sample_tracker_->NodeIndices(category_id); } + std::vector& NodeIndicesInternalIndex(int internal_category_id) { return category_sample_tracker_->NodeIndicesInternalIndex(internal_category_id); } + double* GetPredictions() { return rfx_predictions_.data(); } + double GetPrediction(data_size_t observation_num) { return rfx_predictions_.at(observation_num); } + void SetPrediction(data_size_t observation_num, double pred) { rfx_predictions_.at(observation_num) = pred; } /*! \brief Resets RFX tracker based on a specific sample. Assumes tracker already exists in main memory. */ - void ResetFromSample(MultivariateRegressionRandomEffectsModel& rfx_model, + void ResetFromSample(MultivariateRegressionRandomEffectsModel& rfx_model, RandomEffectsDataset& rfx_dataset, ColumnVector& residual); - /*! \brief Resets RFX tracker to initial default. Assumes tracker already exists in main memory. + /*! \brief Resets RFX tracker to initial default. Assumes tracker already exists in main memory. * Assumes that the initial "clean slate" prediction of a random effects model is 0. */ - void RootReset(MultivariateRegressionRandomEffectsModel& rfx_model, + void RootReset(MultivariateRegressionRandomEffectsModel& rfx_model, RandomEffectsDataset& rfx_dataset, ColumnVector& residual); private: @@ -77,22 +79,26 @@ class RandomEffectsTracker { class LabelMapper { public: LabelMapper() {} - LabelMapper(std::map label_map) { + LabelMapper(std::map label_map) { label_map_ = label_map; for (const auto& [key, value] : label_map) keys_.push_back(key); } ~LabelMapper() {} - void LoadFromLabelMap(std::map label_map) { + void LoadFromLabelMap(std::map label_map) { label_map_ = label_map; for (const auto& [key, value] : label_map) keys_.push_back(key); } - bool ContainsLabel(int32_t category_id) { + bool ContainsLabel(int category_id) { auto pos = label_map_.find(category_id); return pos != label_map_.end(); } - int32_t CategoryNumber(int32_t category_id) { + int CategoryNumber(int category_id) { return label_map_[category_id]; } + void CopyFromOther(LabelMapper& other) { + keys_ = other.Keys(); + label_map_ = other.Map(); + } void SaveToJsonFile(std::string filename) { nlohmann::json model_json = this->to_json(); std::ofstream output_file(filename); @@ -113,14 +119,18 @@ class LabelMapper { this->Reset(); this->from_json(rfx_label_mapper_json); } - std::vector& Keys() {return keys_;} - std::map& Map() {return label_map_;} - void Reset() {label_map_.clear(); keys_.clear();} + std::vector& Keys() { return keys_; } + std::map& Map() { return label_map_; } + void Reset() { + label_map_.clear(); + keys_.clear(); + } nlohmann::json to_json(); void from_json(const nlohmann::json& rfx_label_mapper_json); + private: - std::map label_map_; - std::vector keys_; + std::map label_map_; + std::vector keys_; }; /*! \brief Posterior computation and sampling and state storage for random effects model with a group-level multivariate basis regression */ @@ -140,7 +150,7 @@ class MultivariateRegressionRandomEffectsModel { /*! \brief Reconstruction from serialized model parameter samples */ void ResetFromSample(RandomEffectsContainer& rfx_container, int sample_num); - + /*! \brief Samplers */ void SampleRandomEffects(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& tracker, double global_variance, std::mt19937& gen); void SampleWorkingParameter(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& tracker, double global_variance, std::mt19937& gen); @@ -154,7 +164,7 @@ class MultivariateRegressionRandomEffectsModel { void SetGroupParameters(Eigen::MatrixXd& group_parameters) { group_parameters_ = group_parameters; } - void SetGroupParameter(Eigen::VectorXd& group_parameter, int32_t group_id) { + void SetGroupParameter(Eigen::VectorXd& group_parameter, int group_id) { group_parameters_(Eigen::all, group_id) = group_parameter; } void SetWorkingParameterCovariance(Eigen::MatrixXd& working_parameter_covariance) { @@ -163,7 +173,7 @@ class MultivariateRegressionRandomEffectsModel { void SetGroupParameterCovariance(Eigen::MatrixXd& group_parameter_covariance) { group_parameter_covariance_ = group_parameter_covariance; } - void SetGroupParameterVarianceComponent(double value, int32_t component_id) { + void SetGroupParameterVarianceComponent(double value, int component_id) { group_parameter_covariance_(component_id, component_id) = value; } void SetVariancePriorShape(double value) { @@ -192,9 +202,9 @@ class MultivariateRegressionRandomEffectsModel { double GetVariancePriorScale() { return variance_prior_scale_; } - inline int32_t NumComponents() {return num_components_;} - inline int32_t NumGroups() {return num_groups_;} - + inline int NumComponents() { return num_components_; } + inline int NumGroups() { return num_groups_; } + std::vector Predict(RandomEffectsDataset& dataset, RandomEffectsTracker& tracker) { std::vector output(dataset.NumObservations()); PredictInplace(dataset, tracker, output); @@ -203,12 +213,12 @@ class MultivariateRegressionRandomEffectsModel { void PredictInplace(RandomEffectsDataset& dataset, RandomEffectsTracker& tracker, std::vector& output) { Eigen::MatrixXd X = dataset.GetBasis(); - std::vector group_labels = dataset.GetGroupLabels(); + std::vector group_labels = dataset.GetGroupLabels(); CHECK_EQ(X.rows(), group_labels.size()); int n = X.rows(); CHECK_EQ(n, output.size()); Eigen::MatrixXd alpha_diag = working_parameter_.asDiagonal().toDenseMatrix(); - std::int32_t group_ind; + int group_ind; for (int i = 0; i < n; i++) { group_ind = tracker.CategoryNumber(group_labels[i]); output[i] = X(i, Eigen::all) * alpha_diag * group_parameters_(Eigen::all, group_ind); @@ -229,13 +239,13 @@ class MultivariateRegressionRandomEffectsModel { void SubtractNewPredictionFromResidual(RandomEffectsDataset& dataset, RandomEffectsTracker& tracker, ColumnVector& residual) { Eigen::MatrixXd X = dataset.GetBasis(); - std::vector group_labels = dataset.GetGroupLabels(); + std::vector group_labels = dataset.GetGroupLabels(); CHECK_EQ(X.rows(), group_labels.size()); int n = X.rows(); double new_pred; double new_resid; Eigen::MatrixXd alpha_diag = working_parameter_.asDiagonal().toDenseMatrix(); - std::int32_t group_ind; + int group_ind; for (int i = 0; i < n; i++) { group_ind = tracker.CategoryNumber(group_labels[i]); new_pred = X(i, Eigen::all) * alpha_diag * group_parameters_(Eigen::all, group_ind); @@ -250,13 +260,13 @@ class MultivariateRegressionRandomEffectsModel { /*! \brief Compute the posterior covariance of the working parameter, conditional on the group parameters and the variance components */ Eigen::MatrixXd WorkingParameterVariance(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance); /*! \brief Compute the posterior mean of a group parameter, conditional on the working parameter and the variance components */ - Eigen::VectorXd GroupParameterMean(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, int32_t group_id); + Eigen::VectorXd GroupParameterMean(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, int group_id); /*! \brief Compute the posterior covariance of a group parameter, conditional on the working parameter and the variance components */ - Eigen::MatrixXd GroupParameterVariance(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, int32_t group_id); + Eigen::MatrixXd GroupParameterVariance(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, int group_id); /*! \brief Compute the posterior shape of the group variance component, conditional on the working and group parameters */ - double VarianceComponentShape(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, int32_t component_id); + double VarianceComponentShape(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, int component_id); /*! \brief Compute the posterior scale of the group variance component, conditional on the working and group parameters */ - double VarianceComponentScale(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, int32_t component_id); + double VarianceComponentScale(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, int component_id); private: /*! \brief Samplers */ @@ -266,7 +276,7 @@ class MultivariateRegressionRandomEffectsModel { /*! \brief Random effects structure details */ int num_components_; int num_groups_; - + /*! \brief Group mean parameters, decomposed into "working parameter" and individual parameters * under the "redundant" parameterization of Gelman et al (2008) */ @@ -275,7 +285,7 @@ class MultivariateRegressionRandomEffectsModel { /*! \brief Variance components for the group parameters */ Eigen::MatrixXd group_parameter_covariance_; - + /*! \brief Variance components for the working parameter */ Eigen::MatrixXd working_parameter_covariance_; @@ -317,15 +327,25 @@ class RandomEffectsContainer { this->Reset(); this->from_json(rfx_container_json); } + void CopyFromOther(RandomEffectsContainer& other) { + this->Reset(); + num_samples_ = other.NumSamples(); + num_components_ = other.NumComponents(); + num_groups_ = other.NumGroups(); + beta_ = other.GetBeta(); + alpha_ = other.GetAlpha(); + xi_ = other.GetXi(); + sigma_xi_ = other.GetSigma(); + } void AddSample(MultivariateRegressionRandomEffectsModel& model); void DeleteSample(int sample_num); void Predict(RandomEffectsDataset& dataset, LabelMapper& label_mapper, std::vector& output); - inline int NumSamples() {return num_samples_;} - inline int NumComponents() {return num_components_;} - inline int NumGroups() {return num_groups_;} - inline void SetNumSamples(int num_samples) {num_samples_ = num_samples;} - inline void SetNumComponents(int num_components) {num_components_ = num_components;} - inline void SetNumGroups(int num_groups) {num_groups_ = num_groups;} + inline int NumSamples() { return num_samples_; } + inline int NumComponents() { return num_components_; } + inline int NumGroups() { return num_groups_; } + inline void SetNumSamples(int num_samples) { num_samples_ = num_samples; } + inline void SetNumComponents(int num_components) { num_components_ = num_components; } + inline void SetNumGroups(int num_groups) { num_groups_ = num_groups; } void Reset() { num_samples_ = 0; num_components_ = 0; @@ -335,13 +355,14 @@ class RandomEffectsContainer { xi_.clear(); sigma_xi_.clear(); } - std::vector& GetBeta() {return beta_;} - std::vector& GetAlpha() {return alpha_;} - std::vector& GetXi() {return xi_;} - std::vector& GetSigma() {return sigma_xi_;} + std::vector& GetBeta() { return beta_; } + std::vector& GetAlpha() { return alpha_; } + std::vector& GetXi() { return xi_; } + std::vector& GetSigma() { return sigma_xi_; } nlohmann::json to_json(); void from_json(const nlohmann::json& rfx_container_json); void append_from_json(const nlohmann::json& rfx_container_json); + private: int num_samples_; int num_components_; @@ -355,6 +376,39 @@ class RandomEffectsContainer { void AddSigma(MultivariateRegressionRandomEffectsModel& model); }; -} // namespace StochTree +/*! + * \brief Append every retained random effects sample from `src` onto the end of `dst` (deep copy). + * Used by BARTSamples::Merge / BCFSamples::Merge to combine independently-sampled chains. The random effects + * must be present in both samples objects or neither. + */ +inline void AppendRandomEffectsContainerSamples(std::unique_ptr& dst, + const std::unique_ptr& src) { + if (src == nullptr && dst == nullptr) return; + if (src == nullptr || dst == nullptr) { + Log::Fatal("Cannot merge samples: random effects container present in one chain but not the other"); + } + // Check that group and basis dimensions match between the two containers before appending + if (src->NumComponents() != dst->NumComponents()) { + Log::Fatal("Cannot merge samples: random effects container has %d components in one chain but %d in the other", src->NumComponents(), dst->NumComponents()); + } + if (src->NumGroups() != dst->NumGroups()) { + Log::Fatal("Cannot merge samples: random effects container has %d groups in one chain but %d in the other", src->NumGroups(), dst->NumGroups()); + } + dst->SetNumSamples(dst->NumSamples() + src->NumSamples()); + std::vector& dst_beta = dst->GetBeta(); + std::vector& dst_alpha = dst->GetAlpha(); + std::vector& dst_xi = dst->GetXi(); + std::vector& dst_sigma_xi = dst->GetSigma(); + std::vector& src_beta = src->GetBeta(); + std::vector& src_alpha = src->GetAlpha(); + std::vector& src_xi = src->GetXi(); + std::vector& src_sigma_xi = src->GetSigma(); + dst_beta.insert(dst_beta.end(), src_beta.begin(), src_beta.end()); + dst_alpha.insert(dst_alpha.end(), src_alpha.begin(), src_alpha.end()); + dst_xi.insert(dst_xi.end(), src_xi.begin(), src_xi.end()); + dst_sigma_xi.insert(dst_sigma_xi.end(), src_sigma_xi.begin(), src_sigma_xi.end()); +} + +} // namespace StochTree -#endif // STOCHTREE_RANDOM_EFFECTS_H_ +#endif // STOCHTREE_RANDOM_EFFECTS_H_ diff --git a/include/stochtree/tree.h b/include/stochtree/tree.h index 3810e3cb..61bf6005 100644 --- a/include/stochtree/tree.h +++ b/include/stochtree/tree.h @@ -52,13 +52,13 @@ enum FeatureSplitType { /*! \brief Forward declaration of TreeSplit class */ class TreeSplit; -/*! +/*! * \defgroup tree_group Tree API - * + * * \brief Classes / functions for creating and modifying decision trees. - * + * * \section tree_design Design - * + * * \{ */ @@ -68,7 +68,7 @@ class Tree { static constexpr std::int32_t kInvalidNodeId{-1}; static constexpr std::int32_t kDeletedNodeMarker = std::numeric_limits::max(); static constexpr std::int32_t kRoot{0}; - + Tree() = default; // ~Tree() = default; Tree(Tree const&) = delete; @@ -76,9 +76,9 @@ class Tree { Tree(Tree&&) noexcept = default; Tree& operator=(Tree&&) noexcept = default; /*! - * \brief Copy the structure and parameters of another tree. If the `Tree` object calling this method already + * \brief Copy the structure and parameters of another tree. If the `Tree` object calling this method already * has a non-root tree structure / parameters, this will be erased and replaced with a copy of `tree`. - * + * * \param tree Tree to be cloned */ void CloneFromTree(Tree* tree); @@ -102,19 +102,19 @@ class Tree { void ExpandNode(std::int32_t nid, int split_index, double split_value, std::vector left_value_vector, std::vector right_value_vector); /*! \brief Expand a node based on a categorical split rule */ void ExpandNode(std::int32_t nid, int split_index, std::vector const& categorical_indices, std::vector left_value_vector, std::vector right_value_vector); - /*! \brief Expand a node based on a generic split rule */ + /*! \brief Expand a node based on a generic split rule */ void ExpandNode(std::int32_t nid, int split_index, TreeSplit& split, double left_value, double right_value); /*! \brief Expand a node based on a generic split rule */ void ExpandNode(std::int32_t nid, int split_index, TreeSplit& split, std::vector left_value_vector, std::vector right_value_vector); /*! \brief Whether or not a tree is a "stump" consisting of a single root node */ - inline bool IsRoot() {return leaves_.size() == 1;} - + inline bool IsRoot() { return leaves_.size() == 1; } + /*! \brief Convert tree to JSON and return JSON in-memory */ json to_json(); - /*! - * \brief Load from JSON - * + /*! + * \brief Load from JSON + * * \param tree_json In-memory json object (of type `nlohmann::json`) */ void from_json(const json& tree_json); @@ -135,7 +135,7 @@ class Tree { // TODO refactor and add this to the multivariate case as well if (!IsRoot(nid)) { int parent_id = Parent(nid); - if ((IsLeaf(LeftChild(parent_id))) && (IsLeaf(RightChild(parent_id)))){ + if ((IsLeaf(LeftChild(parent_id))) && (IsLeaf(RightChild(parent_id)))) { leaf_parents_.push_back(parent_id); } } @@ -174,12 +174,12 @@ class Tree { // TODO refactor and add this to the multivariate case as well if (!IsRoot(nid)) { int parent_id = Parent(nid); - if ((IsLeaf(LeftChild(parent_id))) && (IsLeaf(RightChild(parent_id)))){ + if ((IsLeaf(LeftChild(parent_id))) && (IsLeaf(RightChild(parent_id)))) { leaf_parents_.push_back(parent_id); } } } - + /*! * \brief Collapse an internal node to a leaf node, deleting its children from the tree * \param nid Node id of the new leaf node @@ -200,7 +200,7 @@ class Tree { /*! * \brief Add a constant value to every leaf of a tree. If leaves are multi-dimensional, `constant_value` will be added to every dimension of the leaves. - * + * * \param constant_value Value that will be added to every leaf of a tree */ void AddValueToLeaves(double constant_value) { @@ -217,7 +217,7 @@ class Tree { /*! * \brief Multiply every leaf of a tree by a constant value. If leaves are multi-dimensional, `constant_value` will be multiplied through every dimension of the leaves. - * + * * \param constant_multiple Value that will be multiplied by every leaf of a tree */ void MultiplyLeavesByValue(double constant_multiple) { @@ -234,14 +234,15 @@ class Tree { /*! * \brief Iterate through all nodes in this tree. - * + * * \tparam Func Function object type, must map `std::int32_t` to `bool`. * \param func Function that accepts a node index and returns `False` when iteration through a given branch of the tree should stop and `True` otherwise. */ - template void WalkTree(Func func) const { + template + void WalkTree(Func func) const { std::stack nodes; nodes.push(kRoot); - auto &self = *this; + auto& self = *this; while (!nodes.empty()) { auto nidx = nodes.top(); nodes.pop(); @@ -285,7 +286,7 @@ class Tree { bool IsLogScale() const { return is_log_scale_; } - + /*! * \brief Index of the node's parent * \param nid ID of node being queried @@ -293,7 +294,7 @@ class Tree { std::int32_t Parent(std::int32_t nid) const { return parent_[nid]; } - + /*! * \brief Index of the node's left child * \param nid ID of node being queried @@ -301,7 +302,7 @@ class Tree { std::int32_t LeftChild(std::int32_t nid) const { return cleft_[nid]; } - + /*! * \brief Index of the node's right child * \param nid ID of node being queried @@ -309,7 +310,7 @@ class Tree { std::int32_t RightChild(std::int32_t nid) const { return cright_[nid]; } - + /*! * \brief Index of the node's "default" child (potentially used in the case of a missing feature at prediction time) * \param nid ID of node being queried @@ -317,7 +318,7 @@ class Tree { std::int32_t DefaultChild(std::int32_t nid) const { return cleft_[nid]; } - + /*! * \brief Feature index defining the node's split rule * \param nid ID of node being queried @@ -325,7 +326,7 @@ class Tree { std::int32_t SplitIndex(std::int32_t nid) const { return split_index_[nid]; } - + /*! * \brief Whether the node is a leaf node * \param nid ID of node being queried @@ -333,7 +334,7 @@ class Tree { bool IsLeaf(std::int32_t nid) const { return cleft_[nid] == kInvalidNodeId; } - + /*! * \brief Whether the node is root * \param nid ID of node being queried @@ -357,7 +358,7 @@ class Tree { double LeafValue(std::int32_t nid) const { return leaf_value_[nid]; } - + /*! * \brief Get parameter value of a node (typically though not necessarily a leaf node) at a given output dimension * \param nid ID of node being queried @@ -386,7 +387,7 @@ class Tree { std::stack node_depths; nodes.push(kRoot); node_depths.push(0); - auto &self = *this; + auto& self = *this; while (!nodes.empty()) { auto nidx = nodes.top(); nodes.pop(); @@ -399,11 +400,11 @@ class Tree { auto right = self.RightChild(nidx); if (left != Tree::kInvalidNodeId) { nodes.push(left); - node_depths.push(node_depth+1); + node_depths.push(node_depth + 1); } if (right != Tree::kInvalidNodeId) { nodes.push(right); - node_depths.push(node_depth+1); + node_depths.push(node_depth + 1); } } } @@ -457,7 +458,7 @@ class Tree { } return result; } - + /*! * \brief Tests whether the leaf node has a non-empty leaf vector * \param nid ID of node being queried @@ -476,7 +477,7 @@ class Tree { /*! * \brief Get list of all categories belonging to the left child node. - * Categories are integers ranging from 0 to (n-1), where n is the number of categories in that particular feature. + * Categories are integers ranging from 0 to (n-1), where n is the number of categories in that particular feature. * This list is assumed to be in ascending order. * * \param nid ID of node being queried @@ -538,7 +539,7 @@ class Tree { bool is_right_leaf = false; // Check if node nidx is a leaf, if so, return false bool is_leaf = this->IsLeaf(nid); - if (is_leaf){ + if (is_leaf) { return false; } else { // If nidx is not a leaf, it must have left and right nodes @@ -571,7 +572,7 @@ class Tree { [[nodiscard]] std::vector const& GetLeafParents() const { return leaf_parents_; } - + /*! * \brief Get indices of all valid (non-deleted) nodes. */ @@ -579,11 +580,11 @@ class Tree { std::vector output; auto const& self = *this; this->WalkTree([&output, &self](std::int32_t nidx) { - if (!self.IsDeleted(nidx)) { - output.push_back(nidx); - } - return true; - }); + if (!self.IsDeleted(nidx)) { + output.push_back(nidx); + } + return true; + }); return output; } @@ -604,12 +605,12 @@ class Tree { * \brief Get the total number of nodes including deleted ones in this tree. */ [[nodiscard]] std::int32_t NumNodes() const noexcept { return num_nodes; } - + /** * \brief Get the total number of deleted nodes in this tree. */ [[nodiscard]] std::int32_t NumDeletedNodes() const noexcept { return num_deleted_nodes; } - + /** * \brief Get the total number of valid nodes in this tree. */ @@ -675,7 +676,7 @@ class Tree { */ void SetNumericSplit( std::int32_t nid, std::int32_t split_index, double threshold); - + /*! * \brief Create a categorical split * \param nid ID of node being updated @@ -685,8 +686,8 @@ class Tree { * which node the category list should represent. */ void SetCategoricalSplit(std::int32_t nid, std::int32_t split_index, - std::vector const& category_list); - + std::vector const& category_list); + /*! * \brief Set the leaf value of the node * \param nid ID of node being updated @@ -703,18 +704,18 @@ class Tree { /*! * \brief Obtain a 0-based leaf index for each observation in a ForestDataset. - * Internally, trees are stored as vectors of node information, + * Internally, trees are stored as vectors of node information, * and the `leaves_` vector gives us node IDs for every leaf in the tree. - * Here, we would like to know, for every observation in a dataset, - * which leaf number it is mapped to. Since the leaf numbers themselves - * do not carry any information, we renumber them from `0` to `leaves_.size()-1`. + * Here, we would like to know, for every observation in a dataset, + * which leaf number it is mapped to. Since the leaf numbers themselves + * do not carry any information, we renumber them from `0` to `leaves_.size()-1`. * - * Note: this is a tree-level helper function for an ensemble-level function. - * It assumes the creation of: + * Note: this is a tree-level helper function for an ensemble-level function. + * It assumes the creation of: * -# a vector of column indices of size `dataset.NumObservations()` x `ensemble.NumTrees()`, stored in "tree-major" order - * -# a running counter of the number of tree-observations already indexed in the ensemble + * -# a running counter of the number of tree-observations already indexed in the ensemble * (used as offsets for the leaf number computed and returned here) - * Users running this function for a single tree may simply pre-allocate an output vector as + * Users running this function for a single tree may simply pre-allocate an output vector as * `std::vector output(dataset->NumObservations())` and set the offset to 0. * \param dataset Dataset with which to predict leaf indices from the tree * \param output Pre-allocated output vector storing a matrix of column indices, with "rows" corresponding to observations in `dataset` and "columns" corresponding to trees in an ensemble @@ -725,18 +726,18 @@ class Tree { /*! * \brief Obtain a 0-based leaf index for each observation in a ForestDataset. - * Internally, trees are stored as vectors of node information, + * Internally, trees are stored as vectors of node information, * and the `leaves_` vector gives us node IDs for every leaf in the tree. - * Here, we would like to know, for every observation in a dataset, - * which leaf number it is mapped to. Since the leaf numbers themselves - * do not carry any information, we renumber them from `0` to `leaves_.size()-1`. + * Here, we would like to know, for every observation in a dataset, + * which leaf number it is mapped to. Since the leaf numbers themselves + * do not carry any information, we renumber them from `0` to `leaves_.size()-1`. * - * Note: this is a tree-level helper function for an ensemble-level function. - * It assumes the creation of: + * Note: this is a tree-level helper function for an ensemble-level function. + * It assumes the creation of: * -# a vector of column indices of size `dataset.NumObservations()` x `ensemble.NumTrees()`, stored in "tree-major" order - * -# a running counter of the number of tree-observations already indexed in the ensemble + * -# a running counter of the number of tree-observations already indexed in the ensemble * (used as offsets for the leaf number computed and returned here) - * Users running this function for a single tree may simply pre-allocate an output vector as + * Users running this function for a single tree may simply pre-allocate an output vector as * `std::vector output(dataset->NumObservations())` and set the offset to 0. * \param covariates Eigen matrix with which to predict leaf indices * \param output Pre-allocated output vector storing a matrix of column indices, with "rows" corresponding to observations in `covariates` and "columns" corresponding to trees in an ensemble @@ -747,18 +748,18 @@ class Tree { /*! * \brief Obtain a 0-based leaf index for each observation in a ForestDataset. - * Internally, trees are stored as vectors of node information, + * Internally, trees are stored as vectors of node information, * and the `leaves_` vector gives us node IDs for every leaf in the tree. - * Here, we would like to know, for every observation in a dataset, - * which leaf number it is mapped to. Since the leaf numbers themselves - * do not carry any information, we renumber them from `0` to `leaves_.size()-1`. + * Here, we would like to know, for every observation in a dataset, + * which leaf number it is mapped to. Since the leaf numbers themselves + * do not carry any information, we renumber them from `0` to `leaves_.size()-1`. * - * Note: this is a tree-level helper function for an ensemble-level function. - * It assumes the creation of: + * Note: this is a tree-level helper function for an ensemble-level function. + * It assumes the creation of: * -# a vector of column indices of size `dataset.NumObservations()` x `ensemble.NumTrees()`, stored in "tree-major" order - * -# a running counter of the number of tree-observations already indexed in the ensemble + * -# a running counter of the number of tree-observations already indexed in the ensemble * (used as offsets for the leaf number computed and returned here) - * Users running this function for a single tree may simply pre-allocate an output vector as + * Users running this function for a single tree may simply pre-allocate an output vector as * `std::vector output(dataset->NumObservations())` and set the offset to 0. * \param covariates Eigen matrix with which to predict leaf indices * \param output Pre-allocated output vector storing a matrix of column indices, with "rows" corresponding to observations in `covariates` and "columns" corresponding to trees in an ensemble @@ -767,8 +768,8 @@ class Tree { */ void PredictLeafIndexInplace(Eigen::Map>& covariates, std::vector& output, int32_t offset, int32_t max_leaf); - void PredictLeafIndexInplace(Eigen::Map>& covariates, - Eigen::Map>& output, + void PredictLeafIndexInplace(Eigen::Map>& covariates, + Eigen::Map>& output, int column_ind, int32_t offset, int32_t max_leaf); // Node info @@ -784,7 +785,7 @@ class Tree { std::vector leaves_; std::vector leaf_parents_; std::vector deleted_nodes_; - + // Leaf vector std::vector leaf_vector_; std::vector leaf_vector_begin_; @@ -803,27 +804,26 @@ class Tree { /*! \brief Comparison operator for trees */ inline bool operator==(const Tree& lhs, const Tree& rhs) { return ( - (lhs.has_categorical_split_ == rhs.has_categorical_split_) && - (lhs.output_dimension_ == rhs.output_dimension_) && - (lhs.is_log_scale_ == rhs.is_log_scale_) && - (lhs.node_type_ == rhs.node_type_) && - (lhs.parent_ == rhs.parent_) && - (lhs.cleft_ == rhs.cleft_) && - (lhs.cright_ == rhs.cright_) && - (lhs.split_index_ == rhs.split_index_) && - (lhs.leaf_value_ == rhs.leaf_value_) && - (lhs.threshold_ == rhs.threshold_) && - (lhs.internal_nodes_ == rhs.internal_nodes_) && - (lhs.leaves_ == rhs.leaves_) && - (lhs.leaf_parents_ == rhs.leaf_parents_) && - (lhs.deleted_nodes_ == rhs.deleted_nodes_) && - (lhs.leaf_vector_ == rhs.leaf_vector_) && - (lhs.leaf_vector_begin_ == rhs.leaf_vector_begin_) && - (lhs.leaf_vector_end_ == rhs.leaf_vector_end_) && - (lhs.category_list_ == rhs.category_list_) && - (lhs.category_list_begin_ == rhs.category_list_begin_) && - (lhs.category_list_end_ == rhs.category_list_end_) - ); + (lhs.has_categorical_split_ == rhs.has_categorical_split_) && + (lhs.output_dimension_ == rhs.output_dimension_) && + (lhs.is_log_scale_ == rhs.is_log_scale_) && + (lhs.node_type_ == rhs.node_type_) && + (lhs.parent_ == rhs.parent_) && + (lhs.cleft_ == rhs.cleft_) && + (lhs.cright_ == rhs.cright_) && + (lhs.split_index_ == rhs.split_index_) && + (lhs.leaf_value_ == rhs.leaf_value_) && + (lhs.threshold_ == rhs.threshold_) && + (lhs.internal_nodes_ == rhs.internal_nodes_) && + (lhs.leaves_ == rhs.leaves_) && + (lhs.leaf_parents_ == rhs.leaf_parents_) && + (lhs.deleted_nodes_ == rhs.deleted_nodes_) && + (lhs.leaf_vector_ == rhs.leaf_vector_) && + (lhs.leaf_vector_begin_ == rhs.leaf_vector_begin_) && + (lhs.leaf_vector_end_ == rhs.leaf_vector_end_) && + (lhs.category_list_ == rhs.category_list_) && + (lhs.category_list_begin_ == rhs.category_list_begin_) && + (lhs.category_list_end_ == rhs.category_list_end_)); } /*! \brief Determine whether an observation produces a "true" value in a numeric split node @@ -847,15 +847,13 @@ inline bool SplitTrueCategorical(double fvalue, std::vector const // A valid (integer) category must satisfy two criteria: // 1) it must be exactly representable as double // 2) it must fit into uint32_t - auto max_representable_int - = std::min(static_cast(std::numeric_limits::max()), - static_cast(std::uint64_t(1) << std::numeric_limits::digits)); + auto max_representable_int = std::min(static_cast(std::numeric_limits::max()), + static_cast(std::uint64_t(1) << std::numeric_limits::digits)); if (fvalue < 0 || std::fabs(fvalue) > max_representable_int) { category_matched = false; } else { auto const category_value = static_cast(fvalue); - category_matched = (std::find(category_list.begin(), category_list.end(), category_value) - != category_list.end()); + category_matched = (std::find(category_list.begin(), category_list.end(), category_value) != category_list.end()); } return category_matched; } @@ -880,9 +878,9 @@ inline int NextNodeCategorical(double fvalue, std::vector const& return SplitTrueCategorical(fvalue, category_list) ? left_child : right_child; } -/*! +/*! * Determine the node at which a tree places a given observation - * + * * \param tree Tree object used for prediction * \param data Dataset used for prediction * \param row Row indexing the prediction observation @@ -897,7 +895,7 @@ inline int EvaluateTree(Tree const& tree, Eigen::MatrixXd& data, int row) { } else { if (tree.NodeType(node_id) == StochTree::TreeNodeType::kCategoricalSplitNode) { node_id = NextNodeCategorical(fvalue, tree.CategoryList(node_id), - tree.LeftChild(node_id), tree.RightChild(node_id)); + tree.LeftChild(node_id), tree.RightChild(node_id)); } else { node_id = NextNodeNumeric(fvalue, tree.Threshold(node_id), tree.LeftChild(node_id), tree.RightChild(node_id)); } @@ -906,9 +904,9 @@ inline int EvaluateTree(Tree const& tree, Eigen::MatrixXd& data, int row) { return node_id; } -/*! +/*! * Determine the node at which a tree places a given observation - * + * * \param tree Tree object used for prediction * \param data Dataset used for prediction * \param row Row indexing the prediction observation @@ -923,7 +921,7 @@ inline int EvaluateTree(Tree const& tree, Eigen::Map& split_categories) { @@ -979,22 +977,25 @@ class TreeSplit { split_set_ = true; } ~TreeSplit() {} - bool SplitSet() {return split_set_;} + bool SplitSet() { return split_set_; } /*! \brief Whether or not a `TreeSplit` rule is numeric */ - bool NumericSplit() {return numeric_;} + bool NumericSplit() { return numeric_; } /*! * \brief Whether a given covariate value is `True` or `False` on the rule defined by a `TreeSplit` object - * + * * \param fvalue Value of the covariate */ bool SplitTrue(double fvalue) { - if (numeric_) return SplitTrueNumeric(fvalue, split_value_); - else return SplitTrueCategorical(fvalue, split_categories_); + if (numeric_) + return SplitTrueNumeric(fvalue, split_value_); + else + return SplitTrueCategorical(fvalue, split_categories_); } /*! \brief Numeric cutoff value defining a `TreeSplit` object */ - double SplitValue() {return split_value_;} + double SplitValue() { return split_value_; } /*! \brief Categories defining a `TreeSplit` object */ - std::vector SplitCategories() {return split_categories_;} + std::vector SplitCategories() { return split_categories_; } + private: bool split_set_{false}; bool numeric_; @@ -1002,8 +1003,8 @@ class TreeSplit { std::vector split_categories_; }; -/*! \} */ // end of tree_group +/*! \} */ // end of tree_group -} // namespace StochTree +} // namespace StochTree -#endif // STOCHTREE_TREE_H_ +#endif // STOCHTREE_TREE_H_ diff --git a/include/stochtree/tree_sampler.h b/include/stochtree/tree_sampler.h index 8a12f81f..b7ab5381 100644 --- a/include/stochtree/tree_sampler.h +++ b/include/stochtree/tree_sampler.h @@ -23,12 +23,12 @@ namespace StochTree { /*! * \defgroup sampling_group Forest Sampler API * - * \brief Functions for sampling from a forest. The core interface of these functions, - * as used by the R, Python, and standalone C++ program, is defined by - * \ref MCMCSampleOneIter, which runs one iteration of the MCMC sampler for a - * given forest, and \ref GFRSampleOneIter, which runs one iteration of the - * grow-from-root (GFR) algorithm for a given forest. All other functions are - * essentially helpers used in a sampling function, which are documented here + * \brief Functions for sampling from a forest. The core interface of these functions, + * as used by the R, Python, and standalone C++ program, is defined by + * \ref MCMCSampleOneIter, which runs one iteration of the MCMC sampler for a + * given forest, and \ref GFRSampleOneIter, which runs one iteration of the + * grow-from-root (GFR) algorithm for a given forest. All other functions are + * essentially helpers used in a sampling function, which are documented here * to make extending the C++ codebase more straightforward. * * \{ @@ -36,7 +36,7 @@ namespace StochTree { /*! * \brief Computer the range of available split values for a continuous variable, given the current structure of a tree. - * + * * \param tracker Tracking data structures that speed up sampler operations. * \param dataset Data object containining training data, including covariates, leaf regression bases, and case weights. * \param tree_num Index of the tree for which a split is proposed. @@ -49,16 +49,18 @@ static inline void VarSplitRange(ForestTracker& tracker, ForestDataset& dataset, var_min = std::numeric_limits::max(); var_max = std::numeric_limits::min(); double feature_value; - + std::vector::iterator node_begin_iter = tracker.UnsortedNodeBeginIterator(tree_num, leaf_split); std::vector::iterator node_end_iter = tracker.UnsortedNodeEndIterator(tree_num, leaf_split); - + for (auto i = node_begin_iter; i != node_end_iter; i++) { auto idx = *i; feature_value = dataset.CovariateValue(idx, feature_split); + // Compute var_min and var_max for the feature being split on, which are used to determine the range of valid cutpoints for this split if (feature_value < var_min) { var_min = feature_value; - } else if (feature_value > var_max) { + } + if (feature_value > var_max) { var_max = feature_value; } } @@ -66,14 +68,14 @@ static inline void VarSplitRange(ForestTracker& tracker, ForestDataset& dataset, /*! * \brief Determines whether a proposed split creates two leaf nodes with constant values for every feature (thus ensuring that the tree cannot split further). - * + * * \param dataset Data object containining training data, including covariates, leaf regression bases, and case weights. * \param tracker Tracking data structures that speed up sampler operations. * \param split Proposed split of tree `tree_num` at node `leaf_split`. * \param tree_num Index of the tree for which a split is proposed. * \param leaf_split Index of the leaf in `tree_num` for which a split is proposed. * \param feature_split Index of the feature to which `split` will be applied - * \return `true` if `split` creates two nodes with constant values for every feature in `dataset`, `false` otherwise. + * \return `true` if `split` creates two nodes with constant values for every feature in `dataset`, `false` otherwise. */ static inline bool NodesNonConstantAfterSplit(ForestDataset& dataset, ForestTracker& tracker, TreeSplit& split, int tree_num, int leaf_split, int feature_split) { int p = dataset.GetCovariates().cols(); @@ -84,7 +86,7 @@ static inline bool NodesNonConstantAfterSplit(ForestDataset& dataset, ForestTrac double var_min_left; double var_max_right; double var_min_right; - + for (int j = 0; j < p; j++) { auto node_begin_iter = tracker.UnsortedNodeBeginIterator(tree_num, leaf_split); auto node_end_iter = tracker.UnsortedNodeEndIterator(tree_num, leaf_split); @@ -124,7 +126,7 @@ static inline bool NodeNonConstant(ForestDataset& dataset, ForestTracker& tracke double feature_value; double var_max; double var_min; - + for (int j = 0; j < p; j++) { auto node_begin_iter = tracker.UnsortedNodeBeginIterator(tree_num, node_id); auto node_end_iter = tracker.UnsortedNodeEndIterator(tree_num, node_id); @@ -147,7 +149,7 @@ static inline bool NodeNonConstant(ForestDataset& dataset, ForestTracker& tracke return false; } -static inline void AddSplitToModel(ForestTracker& tracker, ForestDataset& dataset, TreePrior& tree_prior, TreeSplit& split, std::mt19937& gen, Tree* tree, +static inline void AddSplitToModel(ForestTracker& tracker, ForestDataset& dataset, TreePrior& tree_prior, TreeSplit& split, std::mt19937& gen, Tree* tree, int tree_num, int leaf_node, int feature_split, bool keep_sorted = false, int num_threads = -1) { // Use zeros as a "temporary" leaf values since we draw leaf parameters after tree sampling is complete if (tree->OutputDimension() > 1) { @@ -164,7 +166,7 @@ static inline void AddSplitToModel(ForestTracker& tracker, ForestDataset& datase tracker.AddSplit(dataset.GetCovariates(), split, feature_split, tree_num, leaf_node, left_node, right_node, keep_sorted, num_threads); } -static inline void RemoveSplitFromModel(ForestTracker& tracker, ForestDataset& dataset, TreePrior& tree_prior, std::mt19937& gen, Tree* tree, +static inline void RemoveSplitFromModel(ForestTracker& tracker, ForestDataset& dataset, TreePrior& tree_prior, std::mt19937& gen, Tree* tree, int tree_num, int leaf_node, int left_node, int right_node, bool keep_sorted = false) { // Use zeros as a "temporary" leaf values since we draw leaf parameters after tree sampling is complete if (tree->OutputDimension() > 1) { @@ -203,7 +205,7 @@ static inline double ComputeVarianceOutcome(ColumnVector& residual) { return sum_y_sq / static_cast(n) - (sum_y * sum_y) / (static_cast(n) * static_cast(n)); } -static inline void UpdateModelVarianceForest(ForestTracker& tracker, ForestDataset& dataset, ColumnVector& residual, +static inline void UpdateModelVarianceForest(ForestTracker& tracker, ForestDataset& dataset, ColumnVector& residual, TreeEnsemble* forest, bool requires_basis, std::function op) { data_size_t n = dataset.GetCovariates().rows(); double tree_pred = 0.; @@ -222,7 +224,7 @@ static inline void UpdateModelVarianceForest(ForestTracker& tracker, ForestDatas tracker.SetTreeSamplePrediction(i, j, tree_pred); pred_value += tree_pred; } - + // Run op (either plus or minus) on the residual and the new prediction new_resid = op(residual.GetElement(i), pred_value); residual.SetElement(i, new_resid); @@ -230,7 +232,7 @@ static inline void UpdateModelVarianceForest(ForestTracker& tracker, ForestDatas tracker.SyncPredictions(); } -static inline void UpdateResidualNoTrackerUpdate(ForestTracker& tracker, ForestDataset& dataset, ColumnVector& residual, TreeEnsemble* forest, +static inline void UpdateResidualNoTrackerUpdate(ForestTracker& tracker, ForestDataset& dataset, ColumnVector& residual, TreeEnsemble* forest, bool requires_basis, std::function op) { data_size_t n = dataset.GetCovariates().rows(); double tree_pred = 0.; @@ -248,14 +250,14 @@ static inline void UpdateResidualNoTrackerUpdate(ForestTracker& tracker, ForestD } pred_value += tree_pred; } - + // Run op (either plus or minus) on the residual and the new prediction new_resid = op(residual.GetElement(i), pred_value); residual.SetElement(i, new_resid); } } -static inline void UpdateResidualEntireForest(ForestTracker& tracker, ForestDataset& dataset, ColumnVector& residual, TreeEnsemble* forest, +static inline void UpdateResidualEntireForest(ForestTracker& tracker, ForestDataset& dataset, ColumnVector& residual, TreeEnsemble* forest, bool requires_basis, std::function op) { data_size_t n = dataset.GetCovariates().rows(); double tree_pred = 0.; @@ -274,7 +276,7 @@ static inline void UpdateResidualEntireForest(ForestTracker& tracker, ForestData tracker.SetTreeSamplePrediction(i, j, tree_pred); pred_value += tree_pred; } - + // Run op (either plus or minus) on the residual and the new prediction new_resid = op(residual.GetElement(i), pred_value); residual.SetElement(i, new_resid); @@ -296,8 +298,8 @@ static inline void UpdateResidualNewOutcome(ForestTracker& tracker, ColumnVector } } -static inline void UpdateMeanModelTree(ForestTracker& tracker, ForestDataset& dataset, ColumnVector& residual, Tree* tree, int tree_num, - bool requires_basis, std::function op, bool tree_new) { +static inline void UpdateMeanModelTree(ForestTracker& tracker, ForestDataset& dataset, ColumnVector& residual, Tree* tree, int tree_num, + bool requires_basis, std::function op, bool tree_new) { data_size_t n = dataset.GetCovariates().rows(); double pred_value; int32_t leaf_pred; @@ -305,7 +307,7 @@ static inline void UpdateMeanModelTree(ForestTracker& tracker, ForestDataset& da double pred_delta; for (data_size_t i = 0; i < n; i++) { if (tree_new) { - // If the tree has been newly sampled or adjusted, we must rerun the prediction + // If the tree has been newly sampled or adjusted, we must rerun the prediction // method and update the SamplePredMapper stored in tracker leaf_pred = tracker.GetNodeId(i, tree_num); if (requires_basis) { @@ -317,7 +319,7 @@ static inline void UpdateMeanModelTree(ForestTracker& tracker, ForestDataset& da tracker.SetTreeSamplePrediction(i, tree_num, pred_value); tracker.SetSamplePrediction(i, tracker.GetSamplePrediction(i) + pred_delta); } else { - // If the tree has not yet been modified via a sampling step, + // If the tree has not yet been modified via a sampling step, // we can query its prediction directly from the SamplePredMapper stored in tracker pred_value = tracker.GetTreeSamplePrediction(i, tree_num); } @@ -345,13 +347,13 @@ static inline void UpdateResidualNewBasis(ForestTracker& tracker, ForestDataset& // Compute new prediction based on updated basis leaf_pred = tracker.GetNodeId(i, tree_num); new_tree_pred = tree->PredictFromNode(leaf_pred, dataset.GetBasis(), i); - + // Cache the new prediction in the tracker tracker.SetTreeSamplePrediction(i, tree_num, new_tree_pred); // Subtract out the updated tree prediction new_resid -= new_tree_pred; - + // Propagate the change back to the residual residual.SetElement(i, new_resid); } @@ -359,7 +361,7 @@ static inline void UpdateResidualNewBasis(ForestTracker& tracker, ForestDataset& tracker.SyncPredictions(); } -static inline void UpdateVarModelTree(ForestTracker& tracker, ForestDataset& dataset, ColumnVector& residual, Tree* tree, +static inline void UpdateVarModelTree(ForestTracker& tracker, ForestDataset& dataset, ColumnVector& residual, Tree* tree, int tree_num, bool requires_basis, std::function op, bool tree_new) { data_size_t n = dataset.GetCovariates().rows(); double pred_value; @@ -370,7 +372,7 @@ static inline void UpdateVarModelTree(ForestTracker& tracker, ForestDataset& dat double prev_pred; for (data_size_t i = 0; i < n; i++) { if (tree_new) { - // If the tree has been newly sampled or adjusted, we must rerun the prediction + // If the tree has been newly sampled or adjusted, we must rerun the prediction // method and update the SamplePredMapper stored in tracker leaf_pred = tracker.GetNodeId(i, tree_num); if (requires_basis) { @@ -386,7 +388,7 @@ static inline void UpdateVarModelTree(ForestTracker& tracker, ForestDataset& dat new_weight = std::log(dataset.VarWeightValue(i)) + pred_value; dataset.SetVarWeightValue(i, new_weight, true); } else { - // If the tree has not yet been modified via a sampling step, + // If the tree has not yet been modified via a sampling step, // we can query its prediction directly from the SamplePredMapper stored in tracker pred_value = tracker.GetTreeSamplePrediction(i, tree_num); new_weight = std::log(dataset.VarWeightValue(i)) - pred_value; @@ -395,8 +397,8 @@ static inline void UpdateVarModelTree(ForestTracker& tracker, ForestDataset& dat } } -static inline void UpdateCLogLogModelTree(ForestTracker& tracker, ForestDataset& dataset, ColumnVector& residual, Tree* tree, int tree_num, - bool requires_basis, bool tree_new) { +static inline void UpdateCLogLogModelTree(ForestTracker& tracker, ForestDataset& dataset, ColumnVector& residual, Tree* tree, int tree_num, + bool requires_basis, bool tree_new) { data_size_t n = dataset.GetCovariates().rows(); double pred_value; @@ -404,7 +406,7 @@ static inline void UpdateCLogLogModelTree(ForestTracker& tracker, ForestDataset& double pred_delta; for (data_size_t i = 0; i < n; i++) { if (tree_new) { - // If the tree has been newly sampled or adjusted, we must rerun the prediction + // If the tree has been newly sampled or adjusted, we must rerun the prediction // method and update the SamplePredMapper stored in tracker leaf_pred = tracker.GetNodeId(i, tree_num); if (requires_basis) { @@ -418,7 +420,7 @@ static inline void UpdateCLogLogModelTree(ForestTracker& tracker, ForestDataset& // Set auxiliary data slot 1 to forest predictions excluding the current tree (tree_num) dataset.SetAuxiliaryDataValue(1, i, tracker.GetSamplePrediction(i) - pred_value); } else { - // If the tree has not yet been modified via a sampling step, + // If the tree has not yet been modified via a sampling step, // we can query its prediction directly from the SamplePredMapper stored in tracker pred_value = tracker.GetTreeSamplePrediction(i, tree_num); // Set auxiliary data slot 1 to forest predictions excluding the current tree (tree_num): needed? since tree not changed? @@ -431,10 +433,9 @@ static inline void UpdateCLogLogModelTree(ForestTracker& tracker, ForestDataset& template static inline std::tuple EvaluateProposedSplit( - ForestDataset& dataset, ForestTracker& tracker, ColumnVector& residual, LeafModel& leaf_model, - TreeSplit& split, int tree_num, int leaf_num, int split_feature, double global_variance, - int num_threads, LeafSuffStatConstructorArgs&... leaf_suff_stat_args -) { + ForestDataset& dataset, ForestTracker& tracker, ColumnVector& residual, LeafModel& leaf_model, + TreeSplit& split, int tree_num, int leaf_num, int split_feature, double global_variance, + int num_threads, LeafSuffStatConstructorArgs&... leaf_suff_stat_args) { // Initialize sufficient statistics LeafSuffStat node_suff_stat = LeafSuffStat(leaf_suff_stat_args...); LeafSuffStat left_suff_stat = LeafSuffStat(leaf_suff_stat_args...); @@ -442,10 +443,9 @@ static inline std::tuple EvaluatePropo // Accumulate sufficient statistics AccumulateSuffStatProposed( - node_suff_stat, left_suff_stat, right_suff_stat, dataset, tracker, - residual, global_variance, split, tree_num, leaf_num, split_feature, num_threads, - leaf_suff_stat_args... - ); + node_suff_stat, left_suff_stat, right_suff_stat, dataset, tracker, + residual, global_variance, split, tree_num, leaf_num, split_feature, num_threads, + leaf_suff_stat_args...); data_size_t left_n = left_suff_stat.n; data_size_t right_n = right_suff_stat.n; @@ -458,17 +458,16 @@ static inline std::tuple EvaluatePropo template static inline std::tuple EvaluateExistingSplit( - ForestDataset& dataset, ForestTracker& tracker, ColumnVector& residual, LeafModel& leaf_model, - double global_variance, int tree_num, int split_node_id, int left_node_id, int right_node_id, - LeafSuffStatConstructorArgs&... leaf_suff_stat_args -) { + ForestDataset& dataset, ForestTracker& tracker, ColumnVector& residual, LeafModel& leaf_model, + double global_variance, int tree_num, int split_node_id, int left_node_id, int right_node_id, + LeafSuffStatConstructorArgs&... leaf_suff_stat_args) { // Initialize sufficient statistics LeafSuffStat node_suff_stat = LeafSuffStat(leaf_suff_stat_args...); LeafSuffStat left_suff_stat = LeafSuffStat(leaf_suff_stat_args...); LeafSuffStat right_suff_stat = LeafSuffStat(leaf_suff_stat_args...); // Accumulate sufficient statistics - AccumulateSuffStatExisting(node_suff_stat, left_suff_stat, right_suff_stat, dataset, tracker, + AccumulateSuffStatExisting(node_suff_stat, left_suff_stat, right_suff_stat, dataset, tracker, residual, global_variance, tree_num, split_node_id, left_node_id, right_node_id); data_size_t left_n = left_suff_stat.n; data_size_t right_n = right_suff_stat.n; @@ -481,8 +480,8 @@ static inline std::tuple EvaluateExist } template -static inline void AdjustStateBeforeTreeSampling(ForestTracker& tracker, LeafModel& leaf_model, ForestDataset& dataset, - ColumnVector& residual, TreePrior& tree_prior, bool backfitting, Tree* tree, int tree_num) { +static inline void AdjustStateBeforeTreeSampling(ForestTracker& tracker, LeafModel& leaf_model, ForestDataset& dataset, + ColumnVector& residual, TreePrior& tree_prior, bool backfitting, Tree* tree, int tree_num) { if constexpr (std::is_same_v) { UpdateCLogLogModelTree(tracker, dataset, residual, tree, tree_num, leaf_model.RequiresBasis(), false); } else if (backfitting) { @@ -494,8 +493,8 @@ static inline void AdjustStateBeforeTreeSampling(ForestTracker& tracker, LeafMod } template -static inline void AdjustStateAfterTreeSampling(ForestTracker& tracker, LeafModel& leaf_model, ForestDataset& dataset, - ColumnVector& residual, TreePrior& tree_prior, bool backfitting, Tree* tree, int tree_num) { +static inline void AdjustStateAfterTreeSampling(ForestTracker& tracker, LeafModel& leaf_model, ForestDataset& dataset, + ColumnVector& residual, TreePrior& tree_prior, bool backfitting, Tree* tree, int tree_num) { if constexpr (std::is_same_v) { UpdateCLogLogModelTree(tracker, dataset, residual, tree, tree_num, leaf_model.RequiresBasis(), true); } else if (backfitting) { @@ -507,10 +506,10 @@ static inline void AdjustStateAfterTreeSampling(ForestTracker& tracker, LeafMode } template -static inline void SampleSplitRule(Tree* tree, ForestTracker& tracker, LeafModel& leaf_model, ForestDataset& dataset, ColumnVector& residual, - TreePrior& tree_prior, std::mt19937& gen, int tree_num, double global_variance, int cutpoint_grid_size, - std::unordered_map>& node_index_map, std::deque& split_queue, - int node_id, data_size_t node_begin, data_size_t node_end, std::vector& variable_weights, +static inline void SampleSplitRule(Tree* tree, ForestTracker& tracker, LeafModel& leaf_model, ForestDataset& dataset, ColumnVector& residual, + TreePrior& tree_prior, std::mt19937& gen, int tree_num, double global_variance, int cutpoint_grid_size, + std::unordered_map>& node_index_map, std::deque& split_queue, + int node_id, data_size_t node_begin, data_size_t node_end, std::vector& variable_weights, std::vector& feature_types, std::vector feature_subset, int num_threads, LeafSuffStatConstructorArgs&... leaf_suff_stat_args) { // Leaf depth int leaf_depth = tree->GetDepth(node_id); @@ -519,15 +518,14 @@ static inline void SampleSplitRule(Tree* tree, ForestTracker& tracker, LeafModel int32_t max_depth = tree_prior.GetMaxDepth(); if ((max_depth == -1) || (leaf_depth < max_depth)) { - // Vector of vectors to store results for each feature int p = dataset.NumCovariates(); - std::vector> feature_log_cutpoint_evaluations(p+1); - std::vector> feature_cutpoint_values(p+1); - std::vector feature_cutpoint_counts(p+1, 0); + std::vector> feature_log_cutpoint_evaluations(p + 1); + std::vector> feature_cutpoint_values(p + 1); + std::vector feature_cutpoint_counts(p + 1, 0); StochTree::data_size_t valid_cutpoint_count; - - // Evaluate all possible cutpoints according to the leaf node model, + + // Evaluate all possible cutpoints according to the leaf node model, // recording their log-likelihood and other split information in a series of vectors. // Initialize node sufficient statistics @@ -545,7 +543,7 @@ static inline void SampleSplitRule(Tree* tree, ForestTracker& tracker, LeafModel Eigen::VectorXd var_weights; bool has_weights = dataset.HasVarWeights(); if (has_weights) var_weights = dataset.GetVarWeights(); - + // Minimum size of newly created leaf nodes (used to rule out invalid splits) int32_t min_samples_in_leaf = tree_prior.GetMinSamplesLeaf(); @@ -557,13 +555,13 @@ static inline void SampleSplitRule(Tree* tree, ForestTracker& tracker, LeafModel // Initialize cutpoint grid container CutpointGridContainer cutpoint_grid_container(covariates, outcome, cutpoint_grid_size); - + // Evaluate all possible splits for each feature in parallel StochTree::ParallelFor(0, covariates.cols(), num_threads, [&](int j) { if ((std::abs(variable_weights.at(j)) > kEpsilon) && (feature_subset[j])) { // Enumerate cutpoint strides cutpoint_grid_container.CalculateStrides(covariates, outcome, tracker.GetSortedNodeSampleTracker(), node_id, node_begin, node_end, j, feature_types); - + // Left and right node sufficient statistics LeafSuffStat left_suff_stat = LeafSuffStat(leaf_suff_stat_args...); LeafSuffStat right_suff_stat = LeafSuffStat(leaf_suff_stat_args...); @@ -584,18 +582,18 @@ static inline void SampleSplitRule(Tree* tree, ForestTracker& tracker, LeafModel // Compute the corresponding right node sufficient statistics right_suff_stat.SubtractSuffStat(node_suff_stat, left_suff_stat); - // Store the bin index as the "cutpoint value" - we can use this to query the actual split + // Store the bin index as the "cutpoint value" - we can use this to query the actual split // value or the set of split categories later on once a split is chose double cutoff_value = cutpoint_idx; // Only include cutpoint for consideration if it defines a valid split in the training data - bool valid_split = (left_suff_stat.SampleGreaterThanEqual(min_samples_in_leaf) && + bool valid_split = (left_suff_stat.SampleGreaterThanEqual(min_samples_in_leaf) && right_suff_stat.SampleGreaterThanEqual(min_samples_in_leaf)); if (valid_split) { feature_cutpoint_counts[j]++; // Add to split rule vector feature_cutpoint_values[j].push_back(cutoff_value); - // Add the log marginal likelihood of the split to the split eval vector + // Add the log marginal likelihood of the split to the split eval vector double split_log_ml = leaf_model.SplitLogMarginalLikelihood(left_suff_stat, right_suff_stat, global_variance); feature_log_cutpoint_evaluations[j].push_back(split_log_ml); } @@ -608,29 +606,29 @@ static inline void SampleSplitRule(Tree* tree, ForestTracker& tracker, LeafModel // Add the log marginal likelihood of the "no-split" option (adjusted for tree prior and cutpoint size per the XBART paper) feature_log_cutpoint_evaluations[covariates.cols()].push_back(no_split_log_ml); - + // Compute an adjustment to reflect the no split prior probability and the number of cutpoints double bart_prior_no_split_adj; double alpha = tree_prior.GetAlpha(); double beta = tree_prior.GetBeta(); int node_depth = tree->GetDepth(node_id); if (valid_cutpoint_count == 0) { - bart_prior_no_split_adj = std::log(((std::pow(1+node_depth, beta))/alpha) - 1.0); + bart_prior_no_split_adj = std::log(((std::pow(1 + node_depth, beta)) / alpha) - 1.0); } else { - bart_prior_no_split_adj = std::log(((std::pow(1+node_depth, beta))/alpha) - 1.0) + std::log(valid_cutpoint_count); + bart_prior_no_split_adj = std::log(((std::pow(1 + node_depth, beta)) / alpha) - 1.0) + std::log(valid_cutpoint_count); } feature_log_cutpoint_evaluations[covariates.cols()][0] += bart_prior_no_split_adj; - // Convert log marginal likelihood to marginal likelihood, normalizing by the maximum log-likelihood double largest_ml = -std::numeric_limits::infinity(); for (int j = 0; j < p + 1; j++) { if (feature_log_cutpoint_evaluations[j].size() > 0) { - double feature_max_ml = *std::max_element(feature_log_cutpoint_evaluations[j].begin(), feature_log_cutpoint_evaluations[j].end());; + double feature_max_ml = *std::max_element(feature_log_cutpoint_evaluations[j].begin(), feature_log_cutpoint_evaluations[j].end()); + ; largest_ml = std::max(largest_ml, feature_max_ml); } } - std::vector> feature_cutpoint_evaluations(p+1); + std::vector> feature_cutpoint_evaluations(p + 1); for (int j = 0; j < p + 1; j++) { if (feature_log_cutpoint_evaluations[j].size() > 0) { feature_cutpoint_evaluations[j].resize(feature_log_cutpoint_evaluations[j].size()); @@ -641,7 +639,7 @@ static inline void SampleSplitRule(Tree* tree, ForestTracker& tracker, LeafModel } // Compute sum of marginal likelihoods for each feature - std::vector feature_total_cutpoint_evaluations(p+1, 0.0); + std::vector feature_total_cutpoint_evaluations(p + 1, 0.0); for (int j = 0; j < p + 1; j++) { if (feature_log_cutpoint_evaluations[j].size() > 0) { feature_total_cutpoint_evaluations[j] = std::accumulate(feature_cutpoint_evaluations[j].begin(), feature_cutpoint_evaluations[j].end(), 0.0); @@ -655,8 +653,8 @@ static inline void SampleSplitRule(Tree* tree, ForestTracker& tracker, LeafModel // Then, sample a cutpoint according to feature_cutpoint_evaluations[feature_chosen] int cutpoint_chosen = sample_discrete_stateless(gen, feature_cutpoint_evaluations[feature_chosen], feature_total_cutpoint_evaluations[feature_chosen]); - - if (feature_chosen == p){ + + if (feature_chosen == p) { // "No split" sampled, don't split or add any nodes to split queue return; } else { @@ -665,14 +663,14 @@ static inline void SampleSplitRule(Tree* tree, ForestTracker& tracker, LeafModel FeatureType feature_type = feature_types[feature_split]; double split_value = feature_cutpoint_values[feature_split][cutpoint_chosen]; // Perform all of the relevant "split" operations in the model, tree and training dataset - + // Compute node sample size data_size_t node_n = node_end - node_begin; - + // Actual numeric cutpoint used for ordered categorical and numeric features double split_value_numeric; TreeSplit tree_split; - + // We will use these later in the model expansion data_size_t left_n = 0; data_size_t right_n = 0; @@ -696,7 +694,7 @@ static inline void SampleSplitRule(Tree* tree, ForestTracker& tracker, LeafModel } else { Log::Fatal("Invalid split type"); } - + // Add split to tree and trackers AddSplitToModel(tracker, dataset, tree_prior, tree_split, gen, tree, tree_num, node_id, feature_split, true, num_threads); @@ -715,15 +713,15 @@ static inline void SampleSplitRule(Tree* tree, ForestTracker& tracker, LeafModel // Add the left and right nodes to the split tracker split_queue.push_front(right_node); - split_queue.push_front(left_node); + split_queue.push_front(left_node); } } } template static inline void GFRSampleTreeOneIter(Tree* tree, ForestTracker& tracker, ForestContainer& forests, LeafModel& leaf_model, ForestDataset& dataset, - ColumnVector& residual, TreePrior& tree_prior, std::mt19937& gen, std::vector& variable_weights, - int tree_num, double global_variance, std::vector& feature_types, int cutpoint_grid_size, + ColumnVector& residual, TreePrior& tree_prior, std::mt19937& gen, std::vector& variable_weights, + int tree_num, double global_variance, std::vector& feature_types, int cutpoint_grid_size, int num_features_subsample, int num_threads, LeafSuffStatConstructorArgs&... leaf_suff_stat_args) { int root_id = Tree::kRoot; int curr_node_id; @@ -748,9 +746,8 @@ static inline void GFRSampleTreeOneIter(Tree* tree, ForestTracker& tracker, Fore std::iota(feature_indices.begin(), feature_indices.end(), 0); std::vector features_selected(num_features_subsample); sample_without_replacement( - features_selected.data(), variable_weights.data(), feature_indices.data(), - p, num_features_subsample, gen - ); + features_selected.data(), variable_weights.data(), feature_indices.data(), + p, num_features_subsample, gen); for (int i = 0; i < p; i++) { feature_subset.at(i) = false; } @@ -778,21 +775,20 @@ static inline void GFRSampleTreeOneIter(Tree* tree, ForestTracker& tracker, Fore curr_node_end = begin_end.second; // Draw a split rule at random SampleSplitRule( - tree, tracker, leaf_model, dataset, residual, tree_prior, gen, tree_num, global_variance, cutpoint_grid_size, - node_index_map, split_queue, curr_node_id, curr_node_begin, curr_node_end, variable_weights, feature_types, - feature_subset, num_threads, leaf_suff_stat_args... - ); + tree, tracker, leaf_model, dataset, residual, tree_prior, gen, tree_num, global_variance, cutpoint_grid_size, + node_index_map, split_queue, curr_node_id, curr_node_begin, curr_node_end, variable_weights, feature_types, + feature_subset, num_threads, leaf_suff_stat_args...); } } -/*! +/*! * Runs one iteration of the "grow-from-root" (GFR) sampler for a tree ensemble model, which consists of two steps for every tree in a forest: * 1. Grow a tree by recursively sampling cutpoint via the GFR algorithm * 2. Sampling leaf node parameters, conditional on an updated tree, via a Gibbs sampler - * + * * \tparam LeafModel Leaf model type (i.e. `GaussianConstantLeafModel`, `GaussianUnivariateRegressionLeafModel`, etc...) * \tparam LeafSuffStat Leaf sufficient statistic type (i.e. `GaussianConstantSuffStat`, `GaussianUnivariateRegressionSuffStat`, etc...) - * \tparam LeafSuffStatConstructorArgs Type of constructor arguments used to initialize `LeafSuffStat` class. For `GaussianMultivariateRegressionSuffStat`, + * \tparam LeafSuffStatConstructorArgs Type of constructor arguments used to initialize `LeafSuffStat` class. For `GaussianMultivariateRegressionSuffStat`, * this is `int`, while each of the other three sufficient statistic classes do not take a constructor argument. * \param active_forest Current state of an ensemble from the sampler's perspective. This is managed through an "active forest" class, as distinct from a "forest container" class * of stored ensemble samples because we often wish to update model state without saving the result (e.g. during burn-in or thinning of an MCMC sampler). @@ -817,38 +813,37 @@ static inline void GFRSampleTreeOneIter(Tree* tree, ForestTracker& tracker, Fore * \param leaf_suff_stat_args Any arguments which must be supplied to initialize a `LeafSuffStat` object. */ template -static inline void GFRSampleOneIter(TreeEnsemble& active_forest, ForestTracker& tracker, ForestContainer& forests, LeafModel& leaf_model, ForestDataset& dataset, - ColumnVector& residual, TreePrior& tree_prior, std::mt19937& gen, std::vector& variable_weights, - std::vector& sweep_update_indices, double global_variance, std::vector& feature_types, int cutpoint_grid_size, +static inline void GFRSampleOneIter(TreeEnsemble& active_forest, ForestTracker& tracker, ForestContainer& forests, LeafModel& leaf_model, ForestDataset& dataset, + ColumnVector& residual, TreePrior& tree_prior, std::mt19937& gen, std::vector& variable_weights, + std::vector& sweep_update_indices, double global_variance, std::vector& feature_types, int cutpoint_grid_size, bool keep_forest, bool pre_initialized, bool backfitting, int num_features_subsample, int num_threads, LeafSuffStatConstructorArgs&... leaf_suff_stat_args) { // Run the GFR algorithm for each tree int num_trees = forests.NumTrees(); for (const int& i : sweep_update_indices) { // Adjust any model state needed to run a tree sampler - // For models that involve Bayesian backfitting, this amounts to adding tree i's + // For models that involve Bayesian backfitting, this amounts to adding tree i's // predictions back to the residual (thus, training a model on the "partial residual") // For more general "blocked MCMC" models, this might require changes to a ForestTracker or Dataset object Tree* tree = active_forest.GetTree(i); AdjustStateBeforeTreeSampling(tracker, leaf_model, dataset, residual, tree_prior, backfitting, tree, i); - + // Reset the tree and sample trackers active_forest.ResetInitTree(i); tracker.ResetRoot(dataset.GetCovariates(), feature_types, i); tree = active_forest.GetTree(i); - + // Sample tree i GFRSampleTreeOneIter( - tree, tracker, forests, leaf_model, dataset, residual, tree_prior, gen, - variable_weights, i, global_variance, feature_types, cutpoint_grid_size, - num_features_subsample, num_threads, leaf_suff_stat_args... - ); - + tree, tracker, forests, leaf_model, dataset, residual, tree_prior, gen, + variable_weights, i, global_variance, feature_types, cutpoint_grid_size, + num_features_subsample, num_threads, leaf_suff_stat_args...); + // Sample leaf parameters for tree i tree = active_forest.GetTree(i); leaf_model.SampleLeafParameters(dataset, tracker, residual, tree, i, global_variance, gen); - + // Adjust any model state needed to run a tree sampler - // For models that involve Bayesian backfitting, this amounts to subtracting tree i's + // For models that involve Bayesian backfitting, this amounts to subtracting tree i's // predictions back out of the residual (thus, using an updated "partial residual" in the following interation). // For more general "blocked MCMC" models, this might require changes to a ForestTracker or Dataset object AdjustStateAfterTreeSampling(tracker, leaf_model, dataset, residual, tree_prior, backfitting, tree, i); @@ -860,8 +855,8 @@ static inline void GFRSampleOneIter(TreeEnsemble& active_forest, ForestTracker& } template -static inline void MCMCGrowTreeOneIter(Tree* tree, ForestTracker& tracker, LeafModel& leaf_model, ForestDataset& dataset, ColumnVector& residual, - TreePrior& tree_prior, std::mt19937& gen, int tree_num, std::vector& variable_weights, +static inline void MCMCGrowTreeOneIter(Tree* tree, ForestTracker& tracker, LeafModel& leaf_model, ForestDataset& dataset, ColumnVector& residual, + TreePrior& tree_prior, std::mt19937& gen, int tree_num, std::vector& variable_weights, double global_variance, double prob_grow_old, int num_threads, LeafSuffStatConstructorArgs&... leaf_suff_stat_args) { // Extract dataset information data_size_t n = dataset.GetCovariates().rows(); @@ -870,7 +865,7 @@ static inline void MCMCGrowTreeOneIter(Tree* tree, ForestTracker& tracker, LeafM int num_leaves = tree->NumLeaves(); std::vector leaves = tree->GetLeaves(); std::vector leaf_weights(num_leaves); - std::fill(leaf_weights.begin(), leaf_weights.end(), 1.0/num_leaves); + std::fill(leaf_weights.begin(), leaf_weights.end(), 1.0 / num_leaves); walker_vose leaf_dist(leaf_weights.begin(), leaf_weights.end()); int leaf_chosen = leaves[leaf_dist(gen)]; int leaf_depth = tree->GetDepth(leaf_chosen); @@ -883,7 +878,6 @@ static inline void MCMCGrowTreeOneIter(Tree* tree, ForestTracker& tracker, LeafM if ((leaf_depth >= max_depth) && (max_depth != -1)) { accept = false; } else { - // Select a split variable at random int p = dataset.GetCovariates().cols(); CHECK_EQ(variable_weights.size(), p); @@ -897,7 +891,7 @@ static inline void MCMCGrowTreeOneIter(Tree* tree, ForestTracker& tracker, LeafM if (var_max <= var_min) { return; } - + // Split based on var_min to var_max in a given node double split_point_chosen = standard_uniform_draw_53bit(gen) * (var_max - var_min) + var_min; @@ -906,8 +900,7 @@ static inline void MCMCGrowTreeOneIter(Tree* tree, ForestTracker& tracker, LeafM // Compute the marginal likelihood of split and no split, given the leaf prior std::tuple split_eval = EvaluateProposedSplit( - dataset, tracker, residual, leaf_model, split, tree_num, leaf_chosen, var_chosen, global_variance, num_threads, leaf_suff_stat_args... - ); + dataset, tracker, residual, leaf_model, split, tree_num, leaf_chosen, var_chosen, global_variance, num_threads, leaf_suff_stat_args...); double split_log_marginal_likelihood = std::get<0>(split_eval); double no_split_log_marginal_likelihood = std::get<1>(split_eval); int32_t left_n = std::get<2>(split_eval); @@ -917,18 +910,17 @@ static inline void MCMCGrowTreeOneIter(Tree* tree, ForestTracker& tracker, LeafM bool left_node_sample_cutoff = left_n >= tree_prior.GetMinSamplesLeaf(); bool right_node_sample_cutoff = right_n >= tree_prior.GetMinSamplesLeaf(); if ((left_node_sample_cutoff) && (right_node_sample_cutoff)) { - // Determine probability of growing the split node and its two new left and right nodes - double pg = tree_prior.GetAlpha() * std::pow(1+leaf_depth, -tree_prior.GetBeta()); - double pgl = tree_prior.GetAlpha() * std::pow(1+leaf_depth+1, -tree_prior.GetBeta()); - double pgr = tree_prior.GetAlpha() * std::pow(1+leaf_depth+1, -tree_prior.GetBeta()); + double pg = tree_prior.GetAlpha() * std::pow(1 + leaf_depth, -tree_prior.GetBeta()); + double pgl = tree_prior.GetAlpha() * std::pow(1 + leaf_depth + 1, -tree_prior.GetBeta()); + double pgr = tree_prior.GetAlpha() * std::pow(1 + leaf_depth + 1, -tree_prior.GetBeta()); // Determine whether a "grow" move is possible from the newly formed tree // in order to compute the probability of choosing "prune" from the new tree // (which is always possible by construction) bool non_constant = NodesNonConstantAfterSplit(dataset, tracker, split, tree_num, leaf_chosen, var_chosen); - bool min_samples_left_check = left_n >= 2*tree_prior.GetMinSamplesLeaf(); - bool min_samples_right_check = right_n >= 2*tree_prior.GetMinSamplesLeaf(); + bool min_samples_left_check = left_n >= 2 * tree_prior.GetMinSamplesLeaf(); + bool min_samples_right_check = right_n >= 2 * tree_prior.GetMinSamplesLeaf(); double prob_prune_new; if (non_constant && (min_samples_left_check || min_samples_right_check)) { prob_prune_new = 0.5; @@ -938,14 +930,12 @@ static inline void MCMCGrowTreeOneIter(Tree* tree, ForestTracker& tracker, LeafM // Determine the number of leaves in the current tree and leaf parents in the proposed tree int num_leaf_parents = tree->NumLeafParents(); - double p_leaf = 1/static_cast(num_leaves); - double p_leaf_parent = 1/static_cast(num_leaf_parents+1); + double p_leaf = 1 / static_cast(num_leaves); + double p_leaf_parent = 1 / static_cast(num_leaf_parents + 1); // Compute the final MH ratio - double log_mh_ratio = ( - std::log(pg) + std::log(1-pgl) + std::log(1-pgr) - std::log(1-pg) + std::log(prob_prune_new) + - std::log(p_leaf_parent) - std::log(prob_grow_old) - std::log(p_leaf) - no_split_log_marginal_likelihood + split_log_marginal_likelihood - ); + double log_mh_ratio = (std::log(pg) + std::log(1 - pgl) + std::log(1 - pgr) - std::log(1 - pg) + std::log(prob_prune_new) + + std::log(p_leaf_parent) - std::log(prob_grow_old) - std::log(p_leaf) - no_split_log_marginal_likelihood + split_log_marginal_likelihood); // Threshold at 0 if (log_mh_ratio > 0) { log_mh_ratio = 0; @@ -967,35 +957,34 @@ static inline void MCMCGrowTreeOneIter(Tree* tree, ForestTracker& tracker, LeafM } template -static inline void MCMCPruneTreeOneIter(Tree* tree, ForestTracker& tracker, LeafModel& leaf_model, ForestDataset& dataset, ColumnVector& residual, - TreePrior& tree_prior, std::mt19937& gen, int tree_num, double global_variance, int num_threads, +static inline void MCMCPruneTreeOneIter(Tree* tree, ForestTracker& tracker, LeafModel& leaf_model, ForestDataset& dataset, ColumnVector& residual, + TreePrior& tree_prior, std::mt19937& gen, int tree_num, double global_variance, int num_threads, LeafSuffStatConstructorArgs&... leaf_suff_stat_args) { // Choose a "leaf parent" node at random int num_leaves = tree->NumLeaves(); int num_leaf_parents = tree->NumLeafParents(); std::vector leaf_parents = tree->GetLeafParents(); std::vector leaf_parent_weights(num_leaf_parents); - std::fill(leaf_parent_weights.begin(), leaf_parent_weights.end(), 1.0/num_leaf_parents); + std::fill(leaf_parent_weights.begin(), leaf_parent_weights.end(), 1.0 / num_leaf_parents); walker_vose leaf_parent_dist(leaf_parent_weights.begin(), leaf_parent_weights.end()); int leaf_parent_chosen = leaf_parents[leaf_parent_dist(gen)]; int leaf_parent_depth = tree->GetDepth(leaf_parent_chosen); int left_node = tree->LeftChild(leaf_parent_chosen); int right_node = tree->RightChild(leaf_parent_chosen); int feature_split = tree->SplitIndex(leaf_parent_chosen); - + // Compute the marginal likelihood for the leaf parent and its left and right nodes std::tuple split_eval = EvaluateExistingSplit( - dataset, tracker, residual, leaf_model, global_variance, tree_num, leaf_parent_chosen, left_node, right_node, leaf_suff_stat_args... - ); + dataset, tracker, residual, leaf_model, global_variance, tree_num, leaf_parent_chosen, left_node, right_node, leaf_suff_stat_args...); double split_log_marginal_likelihood = std::get<0>(split_eval); double no_split_log_marginal_likelihood = std::get<1>(split_eval); int32_t left_n = std::get<2>(split_eval); int32_t right_n = std::get<3>(split_eval); - + // Determine probability of growing the split node and its two new left and right nodes - double pg = tree_prior.GetAlpha() * std::pow(1+leaf_parent_depth, -tree_prior.GetBeta()); - double pgl = tree_prior.GetAlpha() * std::pow(1+leaf_parent_depth+1, -tree_prior.GetBeta()); - double pgr = tree_prior.GetAlpha() * std::pow(1+leaf_parent_depth+1, -tree_prior.GetBeta()); + double pg = tree_prior.GetAlpha() * std::pow(1 + leaf_parent_depth, -tree_prior.GetBeta()); + double pgl = tree_prior.GetAlpha() * std::pow(1 + leaf_parent_depth + 1, -tree_prior.GetBeta()); + double pgr = tree_prior.GetAlpha() * std::pow(1 + leaf_parent_depth + 1, -tree_prior.GetBeta()); // Determine whether a "prune" move is possible from the new tree, // in order to compute the probability of choosing "grow" from the new tree @@ -1020,14 +1009,12 @@ static inline void MCMCPruneTreeOneIter(Tree* tree, ForestTracker& tracker, Leaf } // Determine the number of leaves in the current tree and leaf parents in the proposed tree - double p_leaf = 1/static_cast(num_leaves-1); - double p_leaf_parent = 1/static_cast(num_leaf_parents); + double p_leaf = 1 / static_cast(num_leaves - 1); + double p_leaf_parent = 1 / static_cast(num_leaf_parents); // Compute the final MH ratio - double log_mh_ratio = ( - std::log(1-pg) - std::log(pg) - std::log(1-pgl) - std::log(1-pgr) + std::log(prob_prune_old) + - std::log(p_leaf) - std::log(prob_grow_new) - std::log(p_leaf_parent) + no_split_log_marginal_likelihood - split_log_marginal_likelihood - ); + double log_mh_ratio = (std::log(1 - pg) - std::log(pg) - std::log(1 - pgl) - std::log(1 - pgr) + std::log(prob_prune_old) + + std::log(p_leaf) - std::log(prob_grow_new) - std::log(p_leaf_parent) + no_split_log_marginal_likelihood - split_log_marginal_likelihood); // Threshold at 0 if (log_mh_ratio > 0) { log_mh_ratio = 0; @@ -1046,12 +1033,12 @@ static inline void MCMCPruneTreeOneIter(Tree* tree, ForestTracker& tracker, Leaf template static inline void MCMCSampleTreeOneIter(Tree* tree, ForestTracker& tracker, ForestContainer& forests, LeafModel& leaf_model, ForestDataset& dataset, - ColumnVector& residual, TreePrior& tree_prior, std::mt19937& gen, std::vector& variable_weights, + ColumnVector& residual, TreePrior& tree_prior, std::mt19937& gen, std::vector& variable_weights, int tree_num, double global_variance, int num_threads, LeafSuffStatConstructorArgs&... leaf_suff_stat_args) { // Determine whether it is possible to grow any of the leaves bool grow_possible = false; std::vector leaves = tree->GetLeaves(); - for (auto& leaf: leaves) { + for (auto& leaf : leaves) { if (tracker.UnsortedNodeSize(tree_num, leaf) > 2 * tree_prior.GetMinSamplesLeaf()) { grow_possible = true; break; @@ -1084,15 +1071,13 @@ static inline void MCMCSampleTreeOneIter(Tree* tree, ForestTracker& tracker, For // Draw a split rule at random data_size_t step_chosen = step_dist(gen); bool accept; - + if (step_chosen == 0) { MCMCGrowTreeOneIter( - tree, tracker, leaf_model, dataset, residual, tree_prior, gen, tree_num, variable_weights, global_variance, prob_grow, num_threads, leaf_suff_stat_args... - ); + tree, tracker, leaf_model, dataset, residual, tree_prior, gen, tree_num, variable_weights, global_variance, prob_grow, num_threads, leaf_suff_stat_args...); } else { MCMCPruneTreeOneIter( - tree, tracker, leaf_model, dataset, residual, tree_prior, gen, tree_num, global_variance, num_threads, leaf_suff_stat_args... - ); + tree, tracker, leaf_model, dataset, residual, tree_prior, gen, tree_num, global_variance, num_threads, leaf_suff_stat_args...); } } @@ -1100,10 +1085,10 @@ static inline void MCMCSampleTreeOneIter(Tree* tree, ForestTracker& tracker, For * \brief Runs one iteration of the MCMC sampler for a tree ensemble model, which consists of two steps for every tree in a forest: * 1. Sampling "birth-death" tree modifications via the Metropolis-Hastings algorithm * 2. Sampling leaf node parameters, conditional on a (possibly-updated) tree, via a Gibbs sampler - * + * * \tparam LeafModel Leaf model type (i.e. `GaussianConstantLeafModel`, `GaussianUnivariateRegressionLeafModel`, etc...) * \tparam LeafSuffStat Leaf sufficient statistic type (i.e. `GaussianConstantSuffStat`, `GaussianUnivariateRegressionSuffStat`, etc...) - * \tparam LeafSuffStatConstructorArgs Type of constructor arguments used to initialize `LeafSuffStat` class. For `GaussianMultivariateRegressionSuffStat`, + * \tparam LeafSuffStatConstructorArgs Type of constructor arguments used to initialize `LeafSuffStat` class. For `GaussianMultivariateRegressionSuffStat`, * this is `int`, while each of the other three sufficient statistic classes do not take a constructor argument. * \param active_forest Current state of an ensemble from the sampler's perspective. This is managed through an "active forest" class, as distinct from a "forest container" class * of stored ensemble samples because we often wish to update model state without saving the result (e.g. during burn-in or thinning of an MCMC sampler). @@ -1125,33 +1110,32 @@ static inline void MCMCSampleTreeOneIter(Tree* tree, ForestTracker& tracker, For * \param leaf_suff_stat_args Any arguments which must be supplied to initialize a `LeafSuffStat` object. */ template -static inline void MCMCSampleOneIter(TreeEnsemble& active_forest, ForestTracker& tracker, ForestContainer& forests, LeafModel& leaf_model, ForestDataset& dataset, - ColumnVector& residual, TreePrior& tree_prior, std::mt19937& gen, std::vector& variable_weights, - std::vector& sweep_update_indices, double global_variance, bool keep_forest, bool pre_initialized, bool backfitting, int num_threads, +static inline void MCMCSampleOneIter(TreeEnsemble& active_forest, ForestTracker& tracker, ForestContainer& forests, LeafModel& leaf_model, ForestDataset& dataset, + ColumnVector& residual, TreePrior& tree_prior, std::mt19937& gen, std::vector& variable_weights, + std::vector& sweep_update_indices, double global_variance, bool keep_forest, bool pre_initialized, bool backfitting, int num_threads, LeafSuffStatConstructorArgs&... leaf_suff_stat_args) { // Run the MCMC algorithm for each tree int num_trees = forests.NumTrees(); for (const int& i : sweep_update_indices) { // Adjust any model state needed to run a tree sampler - // For models that involve Bayesian backfitting, this amounts to adding tree i's + // For models that involve Bayesian backfitting, this amounts to adding tree i's // predictions back to the residual (thus, training a model on the "partial residual") // For more general "blocked MCMC" models, this might require changes to a ForestTracker or Dataset object Tree* tree = active_forest.GetTree(i); AdjustStateBeforeTreeSampling(tracker, leaf_model, dataset, residual, tree_prior, backfitting, tree, i); - + // Sample tree i tree = active_forest.GetTree(i); MCMCSampleTreeOneIter( - tree, tracker, forests, leaf_model, dataset, residual, tree_prior, gen, variable_weights, i, - global_variance, num_threads, leaf_suff_stat_args... - ); - + tree, tracker, forests, leaf_model, dataset, residual, tree_prior, gen, variable_weights, i, + global_variance, num_threads, leaf_suff_stat_args...); + // Sample leaf parameters for tree i tree = active_forest.GetTree(i); leaf_model.SampleLeafParameters(dataset, tracker, residual, tree, i, global_variance, gen); - + // Adjust any model state needed to run a tree sampler - // For models that involve Bayesian backfitting, this amounts to subtracting tree i's + // For models that involve Bayesian backfitting, this amounts to subtracting tree i's // predictions back out of the residual (thus, using an updated "partial residual" in the following interation). // For more general "blocked MCMC" models, this might require changes to a ForestTracker or Dataset object AdjustStateAfterTreeSampling(tracker, leaf_model, dataset, residual, tree_prior, backfitting, tree, i); @@ -1162,8 +1146,8 @@ static inline void MCMCSampleOneIter(TreeEnsemble& active_forest, ForestTracker& } } -/*! \} */ // end of sampling_group +/*! \} */ // end of sampling_group -} // namespace StochTree +} // namespace StochTree -#endif // STOCHTREE_TREE_SAMPLER_H_ +#endif // STOCHTREE_TREE_SAMPLER_H_ diff --git a/include/stochtree/variance_model.h b/include/stochtree/variance_model.h index af6bbd0d..c738391d 100644 --- a/include/stochtree/variance_model.h +++ b/include/stochtree/variance_model.h @@ -19,7 +19,7 @@ namespace StochTree { /*! \brief Marginal likelihood and posterior computation for gaussian homoskedastic constant leaf outcome model */ class GlobalHomoskedasticVarianceModel { public: - GlobalHomoskedasticVarianceModel() {ig_sampler_ = InverseGammaSampler();} + GlobalHomoskedasticVarianceModel() { ig_sampler_ = InverseGammaSampler(); } ~GlobalHomoskedasticVarianceModel() {} double PosteriorShape(Eigen::VectorXd& residuals, double a, double b) { data_size_t n = residuals.rows(); @@ -55,6 +55,7 @@ class GlobalHomoskedasticVarianceModel { double ig_scale = PosteriorScale(residuals, weights, a, b); return ig_sampler_.Sample(ig_shape, ig_scale, gen); } + private: InverseGammaSampler ig_sampler_; }; @@ -62,25 +63,26 @@ class GlobalHomoskedasticVarianceModel { /*! \brief Marginal likelihood and posterior computation for gaussian homoskedastic constant leaf outcome model */ class LeafNodeHomoskedasticVarianceModel { public: - LeafNodeHomoskedasticVarianceModel() {ig_sampler_ = InverseGammaSampler();} + LeafNodeHomoskedasticVarianceModel() { ig_sampler_ = InverseGammaSampler(); } ~LeafNodeHomoskedasticVarianceModel() {} double PosteriorShape(TreeEnsemble* ensemble, double a, double b) { data_size_t num_leaves = ensemble->NumLeaves(); - return (a/2.0) + (num_leaves/2.0); + return (a / 2.0) + (num_leaves / 2.0); } double PosteriorScale(TreeEnsemble* ensemble, double a, double b) { double mu_sq = ensemble->SumLeafSquared(); - return (b/2.0) + (mu_sq/2.0); + return (b / 2.0) + (mu_sq / 2.0); } double SampleVarianceParameter(TreeEnsemble* ensemble, double a, double b, std::mt19937& gen) { double ig_shape = PosteriorShape(ensemble, a, b); double ig_scale = PosteriorScale(ensemble, a, b); return ig_sampler_.Sample(ig_shape, ig_scale, gen); } + private: InverseGammaSampler ig_sampler_; }; -} // namespace StochTree +} // namespace StochTree -#endif // STOCHTREE_VARIANCE_MODEL_H_ \ No newline at end of file +#endif // STOCHTREE_VARIANCE_MODEL_H_ \ No newline at end of file diff --git a/man/CppJson.Rd b/man/CppJson.Rd index 7819c293..46e02b3e 100644 --- a/man/CppJson.Rd +++ b/man/CppJson.Rd @@ -48,6 +48,13 @@ For tutorials on the "proper" usage of the stochtree's advanced workflow, we pro \item \href{#method-CppJson-get_integer}{\code{CppJson$get_integer()}} \item \href{#method-CppJson-get_boolean}{\code{CppJson$get_boolean()}} \item \href{#method-CppJson-get_string}{\code{CppJson$get_string()}} +\item \href{#method-CppJson-contains}{\code{CppJson$contains()}} +\item \href{#method-CppJson-get_scalar_or_default}{\code{CppJson$get_scalar_or_default()}} +\item \href{#method-CppJson-get_integer_or_default}{\code{CppJson$get_integer_or_default()}} +\item \href{#method-CppJson-get_boolean_or_default}{\code{CppJson$get_boolean_or_default()}} +\item \href{#method-CppJson-get_string_or_default}{\code{CppJson$get_string_or_default()}} +\item \href{#method-CppJson-rename_field}{\code{CppJson$rename_field()}} +\item \href{#method-CppJson-erase_field}{\code{CppJson$erase_field()}} \item \href{#method-CppJson-get_vector}{\code{CppJson$get_vector()}} \item \href{#method-CppJson-get_integer_vector}{\code{CppJson$get_integer_vector()}} \item \href{#method-CppJson-get_string_vector}{\code{CppJson$get_string_vector()}} @@ -78,13 +85,16 @@ A new \code{CppJson} object. \subsection{Method \code{add_forest()}}{ Convert a forest container to json and add to the current \code{CppJson} object \subsection{Usage}{ -\if{html}{\out{
}}\preformatted{CppJson$add_forest(forest_samples)}\if{html}{\out{
}} +\if{html}{\out{
}}\preformatted{CppJson$add_forest(forest_samples, forest_label = "")}\if{html}{\out{
}} } \subsection{Arguments}{ \if{html}{\out{
}} \describe{ \item{\code{forest_samples}}{\code{ForestSamples} R class} + +\item{\code{forest_label}}{Key under the \code{forests} subfolder to store this +container. Defaults to an auto-numbered \verb{forest_} label when empty.} } \if{html}{\out{
}} } @@ -413,6 +423,172 @@ None } } \if{html}{\out{
}} +\if{html}{\out{}} +\if{latex}{\out{\hypertarget{method-CppJson-contains}{}}} +\subsection{Method \code{contains()}}{ +Whether the json object contains a field "field_name" (with optional subfolder "subfolder_name") +\subsection{Usage}{ +\if{html}{\out{
}}\preformatted{CppJson$contains(field_name, subfolder_name = NULL)}\if{html}{\out{
}} +} + +\subsection{Arguments}{ +\if{html}{\out{
}} +\describe{ +\item{\code{field_name}}{The name of the field to check for} + +\item{\code{subfolder_name}}{(Optional) Name of the subfolder / hierarchy} +} +\if{html}{\out{
}} +} +\subsection{Returns}{ +Logical, \code{TRUE} if the field is present +} +} +\if{html}{\out{
}} +\if{html}{\out{}} +\if{latex}{\out{\hypertarget{method-CppJson-get_scalar_or_default}{}}} +\subsection{Method \code{get_scalar_or_default()}}{ +Retrieve a scalar value, returning \code{default} if the field is absent. Used to +"augment" older JSON within a schema version (see RFC 0005). +\subsection{Usage}{ +\if{html}{\out{
}}\preformatted{CppJson$get_scalar_or_default(field_name, default, subfolder_name = NULL)}\if{html}{\out{
}} +} + +\subsection{Arguments}{ +\if{html}{\out{
}} +\describe{ +\item{\code{field_name}}{The name of the field} + +\item{\code{default}}{Value returned if the field is absent} + +\item{\code{subfolder_name}}{(Optional) Name of the subfolder / hierarchy} +} +\if{html}{\out{
}} +} +\subsection{Returns}{ +The stored value, or \code{default} if absent +} +} +\if{html}{\out{
}} +\if{html}{\out{}} +\if{latex}{\out{\hypertarget{method-CppJson-get_integer_or_default}{}}} +\subsection{Method \code{get_integer_or_default()}}{ +Retrieve an integer value, returning \code{default} if the field is absent. +\subsection{Usage}{ +\if{html}{\out{
}}\preformatted{CppJson$get_integer_or_default(field_name, default, subfolder_name = NULL)}\if{html}{\out{
}} +} + +\subsection{Arguments}{ +\if{html}{\out{
}} +\describe{ +\item{\code{field_name}}{The name of the field} + +\item{\code{default}}{Value returned if the field is absent} + +\item{\code{subfolder_name}}{(Optional) Name of the subfolder / hierarchy} +} +\if{html}{\out{
}} +} +\subsection{Returns}{ +The stored value, or \code{default} if absent +} +} +\if{html}{\out{
}} +\if{html}{\out{}} +\if{latex}{\out{\hypertarget{method-CppJson-get_boolean_or_default}{}}} +\subsection{Method \code{get_boolean_or_default()}}{ +Retrieve a boolean value, returning \code{default} if the field is absent. +\subsection{Usage}{ +\if{html}{\out{
}}\preformatted{CppJson$get_boolean_or_default(field_name, default, subfolder_name = NULL)}\if{html}{\out{
}} +} + +\subsection{Arguments}{ +\if{html}{\out{
}} +\describe{ +\item{\code{field_name}}{The name of the field} + +\item{\code{default}}{Value returned if the field is absent} + +\item{\code{subfolder_name}}{(Optional) Name of the subfolder / hierarchy} +} +\if{html}{\out{
}} +} +\subsection{Returns}{ +The stored value, or \code{default} if absent +} +} +\if{html}{\out{
}} +\if{html}{\out{}} +\if{latex}{\out{\hypertarget{method-CppJson-get_string_or_default}{}}} +\subsection{Method \code{get_string_or_default()}}{ +Retrieve a string value, returning \code{default} if the field is absent. +\subsection{Usage}{ +\if{html}{\out{
}}\preformatted{CppJson$get_string_or_default(field_name, default, subfolder_name = NULL)}\if{html}{\out{
}} +} + +\subsection{Arguments}{ +\if{html}{\out{
}} +\describe{ +\item{\code{field_name}}{The name of the field} + +\item{\code{default}}{Value returned if the field is absent} + +\item{\code{subfolder_name}}{(Optional) Name of the subfolder / hierarchy} +} +\if{html}{\out{
}} +} +\subsection{Returns}{ +The stored value, or \code{default} if absent +} +} +\if{html}{\out{
}} +\if{html}{\out{}} +\if{latex}{\out{\hypertarget{method-CppJson-rename_field}{}}} +\subsection{Method \code{rename_field()}}{ +Rename a field "old_name" to "new_name" (with optional subfolder). No-op if +"old_name" is absent. Used by JSON schema migrations (RFC 0005). +\subsection{Usage}{ +\if{html}{\out{
}}\preformatted{CppJson$rename_field(old_name, new_name, subfolder_name = NULL)}\if{html}{\out{
}} +} + +\subsection{Arguments}{ +\if{html}{\out{
}} +\describe{ +\item{\code{old_name}}{Current field name} + +\item{\code{new_name}}{New field name} + +\item{\code{subfolder_name}}{(Optional) Name of the subfolder / hierarchy} +} +\if{html}{\out{
}} +} +\subsection{Returns}{ +None +} +} +\if{html}{\out{
}} +\if{html}{\out{}} +\if{latex}{\out{\hypertarget{method-CppJson-erase_field}{}}} +\subsection{Method \code{erase_field()}}{ +Erase a field "field_name" (with optional subfolder). No-op if absent. +\subsection{Usage}{ +\if{html}{\out{
}}\preformatted{CppJson$erase_field(field_name, subfolder_name = NULL)}\if{html}{\out{
}} +} + +\subsection{Arguments}{ +\if{html}{\out{
}} +\describe{ +\item{\code{field_name}}{The name of the field to erase} + +\item{\code{subfolder_name}}{(Optional) Name of the subfolder / hierarchy} +} +\if{html}{\out{
}} +} +\subsection{Returns}{ +None +} +} +\if{html}{\out{
}} \if{html}{\out{}} \if{latex}{\out{\hypertarget{method-CppJson-get_vector}{}}} \subsection{Method \code{get_vector()}}{ diff --git a/man/ForestModel.Rd b/man/ForestModel.Rd index 5ff3d208..0b4b2f14 100644 --- a/man/ForestModel.Rd +++ b/man/ForestModel.Rd @@ -96,7 +96,7 @@ Run a single iteration of the forest sampling algorithm (MCMC or GFR) rng, forest_model_config, global_model_config, - num_threads = -1, + num_threads = 1, keep_forest = TRUE, gfr = TRUE )}\if{html}{\out{}} @@ -119,7 +119,7 @@ Run a single iteration of the forest sampling algorithm (MCMC or GFR) \item{\code{global_model_config}}{GlobalModelConfig object containing global model parameters and settings} -\item{\code{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 \code{1}, otherwise to the maximum number of available threads.} +\item{\code{num_threads}}{Number of threads to use in the GFR and MCMC algorithms, as well as prediction. Defaults to \code{1} (single-threaded). Set to \code{-1} to use the maximum number of available threads, or a positive integer for a specific count. OpenMP must be available for values other than \code{1}.} \item{\code{keep_forest}}{(Optional) Whether the updated forest sample should be saved to \code{forest_samples}. Default: \code{TRUE}.} diff --git a/man/ForestSamples.Rd b/man/ForestSamples.Rd index e2fe6486..ff8bd998 100644 --- a/man/ForestSamples.Rd +++ b/man/ForestSamples.Rd @@ -199,7 +199,7 @@ Create a new \code{ForestContainer} object from a json object \describe{ \item{\code{json_object}}{Object of class \code{CppJson}} -\item{\code{json_forest_label}}{Label referring to a particular forest (i.e. "forest_0") in the overall json hierarchy} +\item{\code{json_forest_label}}{Label referring to a particular forest (i.e. "mean_forest") in the overall json hierarchy} } \if{html}{\out{}} } @@ -221,7 +221,7 @@ Append to a \code{ForestContainer} object from a json object \describe{ \item{\code{json_object}}{Object of class \code{CppJson}} -\item{\code{json_forest_label}}{Label referring to a particular forest (i.e. "forest_0") in the overall json hierarchy} +\item{\code{json_forest_label}}{Label referring to a particular forest (i.e. "mean_forest") in the overall json hierarchy} } \if{html}{\out{}} } @@ -243,7 +243,7 @@ Create a new \code{ForestContainer} object from a json object \describe{ \item{\code{json_string}}{JSON string which parses into object of class \code{CppJson}} -\item{\code{json_forest_label}}{Label referring to a particular forest (i.e. "forest_0") in the overall json hierarchy} +\item{\code{json_forest_label}}{Label referring to a particular forest (i.e. "mean_forest") in the overall json hierarchy} } \if{html}{\out{}} } @@ -265,7 +265,7 @@ Append to a \code{ForestContainer} object from a json object \describe{ \item{\code{json_string}}{JSON string which parses into object of class \code{CppJson}} -\item{\code{json_forest_label}}{Label referring to a particular forest (i.e. "forest_0") in the overall json hierarchy} +\item{\code{json_forest_label}}{Label referring to a particular forest (i.e. "mean_forest") in the overall json hierarchy} } \if{html}{\out{}} } diff --git a/man/ForestSamplesSerialization.Rd b/man/ForestSamplesSerialization.Rd index af460f19..a29e6c37 100644 --- a/man/ForestSamplesSerialization.Rd +++ b/man/ForestSamplesSerialization.Rd @@ -16,7 +16,7 @@ loadForestContainerCombinedJsonString(json_string_list, json_forest_label) \arguments{ \item{json_object}{Object of class \code{CppJson}} -\item{json_forest_label}{Label referring to a particular forest (i.e. "forest_0") in the overall json hierarchy (must exist in every json object in a list if a list is provided)} +\item{json_forest_label}{Label referring to a particular forest (i.e. "mean_forest") in the overall json hierarchy (must exist in every json object in a list if a list is provided)} \item{json_object_list}{List of objects of class \code{CppJson}} @@ -31,7 +31,7 @@ entire \code{bartmodel} and \code{bcfmodel} objects, this function group provide loading a single \link{ForestSamples} container from a broader BART / BCF model (which may include multiple forests and other parametric terms). \code{loadForestContainerJson} converts a \link{CppJson} object representing a BART or BCF model into a \link{ForestSamples} container -by extracting the JSON indexed by a forest label (i.e. \code{"forest_0"}) and deserializing it into a \link{ForestSamples} object. +by extracting the JSON indexed by a forest label (i.e. \code{"mean_forest"}) and deserializing it into a \link{ForestSamples} object. Both \code{loadForestContainerJson} and \code{loadForestContainerCombinedJson} operate similarly, but on a list of \link{CppJson} or JSON string representations of BART / BCF models with the same structure. @@ -48,8 +48,8 @@ bart_json <- saveBARTModelToJson(bart_model) bart_json_string <- saveBARTModelToJsonString(bart_model) bart_json_list <- list(bart_json) bart_json_string_list <- list(bart_json_string) -mean_forest <- loadForestContainerJson(bart_json, "forest_0") -mean_forest <- loadForestContainerCombinedJson(bart_json_list, "forest_0") -mean_forest <- loadForestContainerCombinedJsonString(bart_json_string_list, "forest_0") +mean_forest <- loadForestContainerJson(bart_json, "mean_forest") +mean_forest <- loadForestContainerCombinedJson(bart_json_list, "mean_forest") +mean_forest <- loadForestContainerCombinedJsonString(bart_json_string_list, "mean_forest") } diff --git a/man/bart.Rd b/man/bart.Rd index b0d6caab..e5da70e9 100644 --- a/man/bart.Rd +++ b/man/bart.Rd @@ -14,6 +14,7 @@ bart( leaf_basis_test = NULL, rfx_group_ids_test = NULL, rfx_basis_test = NULL, + observation_weights_train = NULL, observation_weights = NULL, num_gfr = 5, num_burnin = 0, @@ -60,12 +61,14 @@ that were not in the training set.} \item{rfx_basis_test}{(Optional) Test set basis for "random-slope" regression in additive random effects model.} -\item{observation_weights}{(Optional) Numeric vector of observation weights of length \code{nrow(X_train)}. Weights are +\item{observation_weights_train}{(Optional) Numeric vector of observation weights of length \code{nrow(X_train)}. Weights are applied as \code{y_i | - ~ N(mu(X_i), sigma^2 / w_i)}, so larger weights increase an observation's influence on the fit. All weights must be non-negative. Default: \code{NULL} (all observations equally weighted). Compatible with Gaussian (continuous/identity) and probit outcome models; not compatible with cloglog link functions. Note: these are referred to internally in the C++ layer as "variance weights" (\code{var_weights}), since they scale the residual variance.} +\item{observation_weights}{Deprecated alias for \code{observation_weights_train}; will be removed in a future release.} + \item{num_gfr}{Number of "warm-start" iterations run using the grow-from-root algorithm (He and Hahn, 2021). Default: 5.} \item{num_burnin}{Number of "burn-in" iterations of the MCMC sampler. Default: 0.} @@ -93,7 +96,7 @@ referred to internally in the C++ layer as "variance weights" (\code{var_weights \item \code{verbose} Whether or not to print progress during the sampling loops. Default: \code{FALSE}. \item \code{outcome_model} A structured \code{OutcomeModel} object that specifies the outcome type and desired link function. This argument pre-empts the legacy (deprecated) \code{probit_outcome_model} option. Default: \code{OutcomeModel(outcome='continuous', link='identity')}. \item \code{probit_outcome_model} Deprecated in favor of \code{outcome_model}. Whether or not the outcome should be modeled as explicitly binary via a probit link. If \code{TRUE}, \code{y} must only contain the values \code{0} and \code{1}. Default: \code{FALSE}. -\item \code{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 \code{1}, otherwise to the maximum number of available threads. +\item \code{num_threads} Number of threads to use in the GFR and MCMC algorithms, as well as prediction. Defaults to \code{1} (single-threaded). Set to \code{-1} to use the maximum number of available threads, or a positive integer for a specific count. OpenMP must be available for values other than \code{1}. }} \item{mean_forest_params}{(Optional) A list of mean forest model parameters, each of which has a default value processed internally, so this argument list is optional. diff --git a/man/bcf.Rd b/man/bcf.Rd index 7e0fb442..6c831664 100644 --- a/man/bcf.Rd +++ b/man/bcf.Rd @@ -16,6 +16,7 @@ bcf( propensity_test = NULL, rfx_group_ids_test = NULL, rfx_basis_test = NULL, + observation_weights_train = NULL, observation_weights = NULL, num_gfr = 5, num_burnin = 0, @@ -64,12 +65,14 @@ that were not in the training set.} \item{rfx_basis_test}{(Optional) Test set basis for "random-slope" regression in additive random effects model.} -\item{observation_weights}{(Optional) Numeric vector of observation weights of length \code{nrow(X_train)}. Weights are +\item{observation_weights_train}{(Optional) Numeric vector of observation weights of length \code{nrow(X_train)}. Weights are applied as \code{y_i | - ~ N(mu(X_i), sigma^2 / w_i)}, so larger weights increase an observation's influence on the fit. All weights must be non-negative. Default: \code{NULL} (all observations equally weighted). Applied to both the prognostic and treatment effect forests. Compatible with Gaussian (continuous/identity) and probit outcome models; not compatible with cloglog link functions.} +\item{observation_weights}{Deprecated alias for \code{observation_weights_train}; will be removed in a future release.} + \item{num_gfr}{Number of "warm-start" iterations run using the grow-from-root algorithm (He and Hahn, 2021). Default: 5.} \item{num_burnin}{Number of "burn-in" iterations of the MCMC sampler. Default: 0.} @@ -94,7 +97,7 @@ Default: \code{NULL}.} \item \code{sigma2_global_shape} Shape parameter in the \code{IG(sigma2_global_shape, sigma2_global_scale)} global error variance model. Default: \code{0}. \item \code{sigma2_global_scale} Scale parameter in the \code{IG(sigma2_global_shape, sigma2_global_scale)} global error variance model. Default: \code{0}. \item \code{variable_weights} Numeric weights reflecting the relative probability of splitting on each variable. Does not need to sum to 1 but cannot be negative. Defaults to \code{rep(1/ncol(X_train), ncol(X_train))} if not set here. Note that if the propensity score is included as a covariate in either forest, its weight will default to \code{1/ncol(X_train)}. A workaround if you wish to provide a custom weight for the propensity score is to include it as a column in \code{X_train} and then set \code{propensity_covariate} to \code{'none'} adjust \code{keep_vars} accordingly for the \code{prognostic} or \code{treatment_effect} forests. -\item \code{propensity_covariate} Whether to include the propensity score as a covariate in either or both of the forests. Enter \code{"none"} for neither, \code{"prognostic"} for the prognostic forest, \code{"treatment_effect"} for the treatment forest, and \code{"both"} for both forests. If this is not \code{"none"} and a propensity score is not provided, it will be estimated from (\code{X_train}, \code{Z_train}) using \code{stochtree::bart()}. Default: \code{"mu"}. +\item \code{propensity_covariate} Whether to include the propensity score as a covariate in either or both of the forests. Enter \code{"none"} for neither, \code{"prognostic"} for the prognostic forest, \code{"treatment_effect"} for the treatment forest, and \code{"both"} for both forests. If this is not \code{"none"} and a propensity score is not provided, it will be estimated from (\code{X_train}, \code{Z_train}) using \code{stochtree::bart()}. Default: \code{"prognostic"}. \item \code{adaptive_coding} Whether or not to use an "adaptive coding" scheme in which a binary treatment variable is not coded manually as (0,1) or (-1,1) but learned via parameters \code{b_0} and \code{b_1} that attach to the outcome model \verb{[b_0 (1-Z) + b_1 Z] tau(X)}. This is ignored when Z is not binary. Default: \code{FALSE}. \item \code{control_coding_init} Initial value of the "control" group coding parameter. This is ignored when Z is not binary. Default: \code{-0.5}. \item \code{treated_coding_init} Initial value of the "treatment" group coding parameter. This is ignored when Z is not binary. Default: \code{0.5}. @@ -107,7 +110,7 @@ Default: \code{NULL}.} \item \code{verbose} Whether or not to print progress during the sampling loops. Default: \code{FALSE}. \item \code{outcome_model} A structured \code{OutcomeModel} object that specifies the outcome type and desired link function. This argument pre-empts the legacy (deprecated) \code{probit_outcome_model} option. Default: \code{OutcomeModel(outcome='continuous', link='identity')}. \item \code{probit_outcome_model} Deprecated in favor of \code{outcome_model}. Whether or not the outcome should be modeled as explicitly binary via a probit link. If \code{TRUE}, \code{y} must only contain the values \code{0} and \code{1}. Default: \code{FALSE}. -\item \code{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 \code{1}, otherwise to the maximum number of available threads. +\item \code{num_threads} Number of threads to use in the GFR and MCMC algorithms, as well as prediction. Defaults to \code{1} (single-threaded). Set to \code{-1} to use the maximum number of available threads, or a positive integer for a specific count. OpenMP must be available for values other than \code{1}. }} \item{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. diff --git a/src/Makevars.in b/src/Makevars.in index de1457c9..ff6facc8 100644 --- a/src/Makevars.in +++ b/src/Makevars.in @@ -23,6 +23,11 @@ PKG_LIBS = \ OBJECTS = \ forest.o \ kernel.o \ + bart_sampler.o \ + bcf_sampler.o \ + R_bart.o \ + R_bcf.o \ + R_samples.o \ R_data.o \ R_random_effects.o \ R_utils.o \ @@ -35,6 +40,7 @@ OBJECTS = \ io.o \ leaf_model.o \ ordinal_sampler.o \ + prediction.o \ partition_tracker.o \ random_effects.o \ tree.o diff --git a/src/Makevars.win.in b/src/Makevars.win.in index af8d6f83..a7b045f1 100644 --- a/src/Makevars.win.in +++ b/src/Makevars.win.in @@ -24,6 +24,11 @@ PKG_LIBS = \ OBJECTS = \ forest.o \ kernel.o \ + bart_sampler.o \ + bcf_sampler.o \ + R_bart.o \ + R_bcf.o \ + R_samples.o \ R_data.o \ R_random_effects.o \ R_utils.o \ @@ -37,5 +42,6 @@ OBJECTS = \ leaf_model.o \ ordinal_sampler.o \ partition_tracker.o \ + prediction.o \ random_effects.o \ tree.o diff --git a/src/R_bart.cpp b/src/R_bart.cpp new file mode 100644 index 00000000..e387faaa --- /dev/null +++ b/src/R_bart.cpp @@ -0,0 +1,359 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include "stochtree_types.h" + +StochTree::BARTConfig convert_list_to_bart_config(cpp11::list config) { + StochTree::BARTConfig output; + + // Global model parameters + output.standardize_outcome = get_config_scalar_default(config, "standardize_outcome", true); + output.num_threads = get_config_scalar_default(config, "num_threads", 1); + output.verbose = get_config_scalar_default(config, "verbose", false); + output.cutpoint_grid_size = get_config_scalar_default(config, "cutpoint_grid_size", 100); + output.link_function = static_cast(get_config_scalar_default(config, "link_function", 0)); + output.outcome_type = static_cast(get_config_scalar_default(config, "outcome_type", 0)); + output.random_seed = get_config_scalar_default(config, "random_seed", 1); + output.keep_gfr = get_config_scalar_default(config, "keep_gfr", true); + output.keep_burnin = get_config_scalar_default(config, "keep_burnin", false); + + // Global error variance parameters + output.a_sigma2_global = get_config_scalar_default(config, "a_sigma2_global", 0.0); + output.b_sigma2_global = get_config_scalar_default(config, "b_sigma2_global", 0.0); + output.sigma2_global_init = get_config_scalar_default(config, "sigma2_global_init", 1.0); + output.sample_sigma2_global = get_config_scalar_default(config, "sample_sigma2_global", true); + + // Mean forest parameters + output.num_trees_mean = get_config_scalar_default(config, "num_trees_mean", 200); + output.alpha_mean = get_config_scalar_default(config, "alpha_mean", 0.95); + output.beta_mean = get_config_scalar_default(config, "beta_mean", 2.0); + output.min_samples_leaf_mean = get_config_scalar_default(config, "min_samples_leaf_mean", 5); + output.max_depth_mean = get_config_scalar_default(config, "max_depth_mean", -1); + output.leaf_constant_mean = get_config_scalar_default(config, "leaf_constant_mean", true); + output.leaf_dim_mean = get_config_scalar_default(config, "leaf_dim_mean", 1); + output.exponentiated_leaf_mean = get_config_scalar_default(config, "exponentiated_leaf_mean", false); + output.num_features_subsample_mean = get_config_scalar_default(config, "num_features_subsample_mean", 0); + output.a_sigma2_mean = get_config_scalar_default(config, "a_sigma2_mean", 3.0); + output.b_sigma2_mean = get_config_scalar_default(config, "b_sigma2_mean", -1.0); + output.sigma2_mean_init = get_config_scalar_default(config, "sigma2_mean_init", -1.0); + output.sample_sigma2_leaf_mean = get_config_scalar_default(config, "sample_sigma2_leaf_mean", false); + output.mean_leaf_model_type = static_cast(get_config_scalar_default(config, "mean_leaf_model_type", 0)); + output.num_classes_cloglog = get_config_scalar_default(config, "num_classes_cloglog", 0); + output.cloglog_leaf_prior_shape = get_config_scalar_default(config, "cloglog_leaf_prior_shape", 2.0); + output.cloglog_leaf_prior_scale = get_config_scalar_default(config, "cloglog_leaf_prior_scale", 2.0); + output.cloglog_cutpoint_0 = get_config_scalar_default(config, "cloglog_cutpoint_0", 0.0); + + // Variance forest parameters + output.num_trees_variance = get_config_scalar_default(config, "num_trees_variance", 0); + output.leaf_prior_calibration_param = get_config_scalar_default(config, "leaf_prior_calibration_param", 1.5); + output.shape_variance_forest = get_config_scalar_default(config, "shape_variance_forest", -1.0); + output.scale_variance_forest = get_config_scalar_default(config, "scale_variance_forest", -1.0); + output.variance_forest_leaf_init = get_config_scalar_default(config, "variance_forest_leaf_init", -1.0); + output.alpha_variance = get_config_scalar_default(config, "alpha_variance", 0.5); + output.beta_variance = get_config_scalar_default(config, "beta_variance", 2.0); + output.min_samples_leaf_variance = get_config_scalar_default(config, "min_samples_leaf_variance", 5); + output.max_depth_variance = get_config_scalar_default(config, "max_depth_variance", -1); + output.leaf_constant_variance = get_config_scalar_default(config, "leaf_constant_variance", true); + output.leaf_dim_variance = get_config_scalar_default(config, "leaf_dim_variance", 1); + output.exponentiated_leaf_variance = get_config_scalar_default(config, "exponentiated_leaf_variance", true); + output.num_features_subsample_variance = get_config_scalar_default(config, "num_features_subsample_variance", 0); + + // Random effect parameters + output.has_random_effects = get_config_scalar_default(config, "has_random_effects", false); + output.rfx_model_spec = static_cast(get_config_scalar_default(config, "rfx_model_spec", 0)); + output.rfx_variance_prior_shape = get_config_scalar_default(config, "rfx_variance_prior_shape", 1.0); + output.rfx_variance_prior_scale = get_config_scalar_default(config, "rfx_variance_prior_scale", 1.0); + + // Handle vector conversions separately + SEXP feature_type_raw = static_cast(config["feature_types"]); + if (!Rf_isNull(feature_type_raw)) { + cpp11::integers feature_types_r_vec(feature_type_raw); + for (auto i : feature_types_r_vec) { + output.feature_types.push_back(static_cast(i)); + } + } + SEXP sweep_update_indices_mean_raw = static_cast(config["sweep_update_indices_mean"]); + if (!Rf_isNull(sweep_update_indices_mean_raw)) { + cpp11::integers sweep_update_indices_mean_r_vec(sweep_update_indices_mean_raw); + output.sweep_update_indices_mean.assign(sweep_update_indices_mean_r_vec.begin(), sweep_update_indices_mean_r_vec.end()); + } + SEXP sweep_update_indices_variance_raw = static_cast(config["sweep_update_indices_variance"]); + if (!Rf_isNull(sweep_update_indices_variance_raw)) { + cpp11::integers sweep_update_indices_variance_r_vec(sweep_update_indices_variance_raw); + output.sweep_update_indices_variance.assign(sweep_update_indices_variance_r_vec.begin(), sweep_update_indices_variance_r_vec.end()); + } + SEXP var_weights_mean_raw = static_cast(config["var_weights_mean"]); + if (!Rf_isNull(var_weights_mean_raw)) { + cpp11::doubles var_weights_mean_r_vec(var_weights_mean_raw); + output.var_weights_mean.assign(var_weights_mean_r_vec.begin(), var_weights_mean_r_vec.end()); + } + SEXP sigma2_leaf_mean_matrix_raw = static_cast(config["sigma2_leaf_mean_matrix"]); + if (!Rf_isNull(sigma2_leaf_mean_matrix_raw)) { + cpp11::doubles sigma2_leaf_mean_matrix_r_vec(sigma2_leaf_mean_matrix_raw); + output.sigma2_leaf_mean_matrix.assign(sigma2_leaf_mean_matrix_r_vec.begin(), sigma2_leaf_mean_matrix_r_vec.end()); + } + SEXP var_weights_variance_raw = static_cast(config["var_weights_variance"]); + if (!Rf_isNull(var_weights_variance_raw)) { + cpp11::doubles var_weights_variance_r_vec(var_weights_variance_raw); + output.var_weights_variance.assign(var_weights_variance_r_vec.begin(), var_weights_variance_r_vec.end()); + } + SEXP rfx_working_parameter_mean_prior_raw = static_cast(config["rfx_working_parameter_mean_prior"]); + if (!Rf_isNull(rfx_working_parameter_mean_prior_raw)) { + cpp11::doubles rfx_working_parameter_mean_prior_r_vec(rfx_working_parameter_mean_prior_raw); + output.rfx_working_parameter_mean_prior.assign(rfx_working_parameter_mean_prior_r_vec.begin(), rfx_working_parameter_mean_prior_r_vec.end()); + } + SEXP rfx_group_parameter_mean_prior_raw = static_cast(config["rfx_group_parameter_mean_prior"]); + if (!Rf_isNull(rfx_group_parameter_mean_prior_raw)) { + cpp11::doubles rfx_group_parameter_mean_prior_r_vec(rfx_group_parameter_mean_prior_raw); + output.rfx_group_parameter_mean_prior.assign(rfx_group_parameter_mean_prior_r_vec.begin(), rfx_group_parameter_mean_prior_r_vec.end()); + } + SEXP rfx_working_parameter_cov_prior_raw = static_cast(config["rfx_working_parameter_cov_prior"]); + if (!Rf_isNull(rfx_working_parameter_cov_prior_raw)) { + cpp11::doubles rfx_working_parameter_cov_prior_r_vec(rfx_working_parameter_cov_prior_raw); + output.rfx_working_parameter_cov_prior.assign(rfx_working_parameter_cov_prior_r_vec.begin(), rfx_working_parameter_cov_prior_r_vec.end()); + } + SEXP rfx_group_parameter_cov_prior_raw = static_cast(config["rfx_group_parameter_cov_prior"]); + if (!Rf_isNull(rfx_group_parameter_cov_prior_raw)) { + cpp11::doubles rfx_group_parameter_cov_prior_r_vec(rfx_group_parameter_cov_prior_raw); + output.rfx_group_parameter_cov_prior.assign(rfx_group_parameter_cov_prior_r_vec.begin(), rfx_group_parameter_cov_prior_r_vec.end()); + } + return output; +} + +cpp11::writable::list create_bart_metadata(StochTree::BARTConfig& config) { + // Unpack metadata about the model that was sampled + cpp11::writable::list model_metadata; + model_metadata.push_back(cpp11::named_arg("sigma2_global_init") = config.sigma2_global_init); + model_metadata.push_back(cpp11::named_arg("sigma2_mean_init") = config.sigma2_mean_init); + model_metadata.push_back(cpp11::named_arg("b_sigma2_mean") = config.b_sigma2_mean); + model_metadata.push_back(cpp11::named_arg("shape_variance_forest") = config.shape_variance_forest); + model_metadata.push_back(cpp11::named_arg("scale_variance_forest") = config.scale_variance_forest); + return model_metadata; +} + +[[cpp11::register]] +cpp11::writable::list bart_sample_cpp( + cpp11::external_pointer samples, + cpp11::sexp X_train, + cpp11::sexp y_train, + cpp11::sexp X_test, + int n_train, + int n_test, + int p, + cpp11::sexp basis_train, + cpp11::sexp basis_test, + int basis_dim, + cpp11::sexp obs_weights_train, + cpp11::sexp obs_weights_test, + cpp11::sexp rfx_group_ids_train, + cpp11::sexp rfx_group_ids_test, + cpp11::sexp rfx_basis_train, + cpp11::sexp rfx_basis_test, + int rfx_num_groups, + int rfx_basis_dim, + int num_gfr, + int num_burnin, + int keep_every, + int num_mcmc, + int num_chains, + cpp11::list config_input) { + // Extract pointers to raw data + int protect_count = 0; + double* X_train_ptr = extract_numeric_pointer(X_train, "X_train", protect_count); + double* y_train_ptr = extract_numeric_pointer(y_train, "y_train", protect_count); + double* X_test_ptr = extract_numeric_pointer(X_test, "X_test", protect_count); + double* basis_train_ptr = extract_numeric_pointer(basis_train, "basis_train", protect_count); + double* basis_test_ptr = extract_numeric_pointer(basis_test, "basis_test", protect_count); + double* obs_weights_train_ptr = extract_numeric_pointer(obs_weights_train, "obs_weights_train", protect_count); + double* obs_weights_test_ptr = extract_numeric_pointer(obs_weights_test, "obs_weights_test", protect_count); + int* rfx_group_ids_train_ptr = extract_integer_pointer(rfx_group_ids_train, "rfx_group_ids_train", protect_count); + int* rfx_group_ids_test_ptr = extract_integer_pointer(rfx_group_ids_test, "rfx_group_ids_test", protect_count); + double* rfx_basis_train_ptr = extract_numeric_pointer(rfx_basis_train, "rfx_basis_train", protect_count); + double* rfx_basis_test_ptr = extract_numeric_pointer(rfx_basis_test, "rfx_basis_test", protect_count); + + // Load the BARTData struct + // Consider reading directly from the R objects or at least checking for matches with the R object dimensions) + StochTree::BARTData data; + data.X_train = X_train_ptr; + data.y_train = y_train_ptr; + data.X_test = X_test_ptr; + data.n_train = n_train; + data.p = p; + data.n_test = n_test; + data.basis_train = basis_train_ptr; + data.basis_test = basis_test_ptr; + data.basis_dim = basis_dim; + data.obs_weights_train = obs_weights_train_ptr; + data.obs_weights_test = obs_weights_test_ptr; + data.rfx_group_ids_train = rfx_group_ids_train_ptr; + data.rfx_group_ids_test = rfx_group_ids_test_ptr; + data.rfx_basis_train = rfx_basis_train_ptr; + data.rfx_basis_test = rfx_basis_test_ptr; + data.rfx_num_groups = rfx_num_groups; + data.rfx_basis_dim = rfx_basis_dim; + + // Create the BARTConfig object + StochTree::BARTConfig config = convert_list_to_bart_config(config_input); + + // Initialize a BART sampler + StochTree::BARTSampler bart_sampler(*samples, config, data); + + // Run the sampler + bart_sampler.run_gfr(*samples, num_gfr, config.keep_gfr, num_chains); + if (num_chains > 1) { + bart_sampler.run_mcmc_chains(*samples, num_chains, num_burnin, keep_every, num_mcmc); + } else { + bart_sampler.run_mcmc(*samples, num_burnin, keep_every, num_mcmc); + } + bart_sampler.postprocess_samples(*samples); + + // Unprotect protected R objects + UNPROTECT(protect_count); + + // Unpack metadata + cpp11::writable::list metadata_list = create_bart_metadata(config); + return metadata_list; +} + +cpp11::writable::list convert_bart_preds_to_list(StochTree::BARTPredictionResult& bart_preds) { + cpp11::writable::list output; + + // Predictions + SEXP y_hat_sexp = !bart_preds.y_hat.empty() + ? static_cast(cpp11::writable::doubles(bart_preds.y_hat.begin(), bart_preds.y_hat.end())) + : R_NilValue; + output.push_back(cpp11::named_arg("y_hat") = y_hat_sexp); + + SEXP mean_forest_pred_sexp = !bart_preds.mean_forest_predictions.empty() + ? static_cast(cpp11::writable::doubles(bart_preds.mean_forest_predictions.begin(), bart_preds.mean_forest_predictions.end())) + : R_NilValue; + output.push_back(cpp11::named_arg("mean_forest_predictions") = mean_forest_pred_sexp); + + SEXP variance_forest_pred_sexp = !bart_preds.variance_forest_predictions.empty() + ? static_cast(cpp11::writable::doubles(bart_preds.variance_forest_predictions.begin(), bart_preds.variance_forest_predictions.end())) + : R_NilValue; + output.push_back(cpp11::named_arg("variance_forest_predictions") = variance_forest_pred_sexp); + + SEXP rfx_predictions_sexp = !bart_preds.rfx_predictions.empty() + ? static_cast(cpp11::writable::doubles(bart_preds.rfx_predictions.begin(), bart_preds.rfx_predictions.end())) + : R_NilValue; + output.push_back(cpp11::named_arg("rfx_predictions") = rfx_predictions_sexp); + + return output; +} + +[[cpp11::register]] +cpp11::writable::list bart_predict_cpp( + cpp11::external_pointer bart_samples_ptr, + cpp11::list bart_model_metadata, + cpp11::sexp X, + cpp11::sexp leaf_basis, + int n, + int p, + int num_basis, + cpp11::sexp obs_weights, + cpp11::sexp rfx_group_ids, + cpp11::sexp rfx_basis, + int rfx_num_groups, + int rfx_basis_dim, + bool posterior, + int scale, + bool predict_y_hat, + bool predict_mean_forest, + bool predict_variance_forest, + bool predict_random_effects) { + // Extract pointers to raw data + int protect_count = 0; + double* X_ptr = extract_numeric_pointer(X, "X", protect_count); + double* leaf_basis_ptr = extract_numeric_pointer(leaf_basis, "leaf_basis", protect_count); + double* obs_weights_ptr = extract_numeric_pointer(obs_weights, "obs_weights", protect_count); + int* rfx_group_ids_ptr = extract_integer_pointer(rfx_group_ids, "rfx_group_ids", protect_count); + double* rfx_basis_ptr = extract_numeric_pointer(rfx_basis, "rfx_basis", protect_count); + + // Load the BARTData struct + // Consider reading directly from the R objects or at least checking for matches with the R object dimensions) + StochTree::BARTData data; + data.X_test = X_ptr; + data.basis_test = leaf_basis_ptr; + data.p = p; + data.n_test = n; + data.basis_dim = num_basis; + data.obs_weights_test = obs_weights_ptr; + data.rfx_group_ids_test = rfx_group_ids_ptr; + data.rfx_basis_test = rfx_basis_ptr; + data.rfx_num_groups = rfx_num_groups; + data.rfx_basis_dim = rfx_basis_dim; + + // Load the BCF model and config from the model list + StochTree::BARTPredictionMetadata pred_metadata; + pred_metadata.num_samples = Rf_asInteger(bart_model_metadata["num_samples"]); + pred_metadata.num_obs = n; + pred_metadata.num_basis = num_basis; + pred_metadata.y_bar = Rf_asReal(bart_model_metadata["y_bar"]); + pred_metadata.y_std = Rf_asReal(bart_model_metadata["y_std"]); + pred_metadata.has_variance_forest = (bool)Rf_asLogical(bart_model_metadata["include_variance_forest"]); + pred_metadata.has_rfx = (bool)Rf_asLogical(bart_model_metadata["has_rfx"]); + pred_metadata.cloglog_num_classes = Rf_asInteger(bart_model_metadata["cloglog_num_classes"]); + { + SEXP rfx_spec_sexp = bart_model_metadata["rfx_model_spec"]; + std::string rfx_model_spec_str = Rf_isNull(rfx_spec_sexp) ? "" : std::string(CHAR(STRING_ELT(rfx_spec_sexp, 0))); + if (rfx_model_spec_str == "intercept_only") { + pred_metadata.rfx_model_spec = StochTree::BARTRFXModelSpec::InterceptOnly; + } else { + pred_metadata.rfx_model_spec = StochTree::BARTRFXModelSpec::Custom; + } + } + pred_metadata.pred_type = posterior ? StochTree::PredType::kPosterior : StochTree::PredType::kMean; + if (scale == 0) { + pred_metadata.pred_scale = StochTree::PredScale::kLinear; + } else if (scale == 1) { + pred_metadata.pred_scale = StochTree::PredScale::kProbability; + } else { + pred_metadata.pred_scale = StochTree::PredScale::kClass; + } + pred_metadata.pred_terms.y_hat = predict_y_hat; + pred_metadata.pred_terms.mean_forest = predict_mean_forest; + pred_metadata.pred_terms.variance_forest = predict_variance_forest; + pred_metadata.pred_terms.random_effects = predict_random_effects; + { + SEXP link_function_sexp = bart_model_metadata["link_function"]; + std::string link_function_str = Rf_isNull(link_function_sexp) ? "" : std::string(CHAR(STRING_ELT(link_function_sexp, 0))); + if (link_function_str == "identity") { + pred_metadata.link_function = StochTree::LinkFunction::Identity; + } else if (link_function_str == "probit") { + pred_metadata.link_function = StochTree::LinkFunction::Probit; + } else if (link_function_str == "cloglog") { + pred_metadata.link_function = StochTree::LinkFunction::Cloglog; + } else { + StochTree::Log::Fatal("Unsupported link function specified in model list"); + } + } + { + SEXP outcome_type_sexp = bart_model_metadata["outcome_type"]; + std::string outcome_type_str = Rf_isNull(outcome_type_sexp) ? "" : std::string(CHAR(STRING_ELT(outcome_type_sexp, 0))); + if (outcome_type_str == "continuous") { + pred_metadata.outcome_type = StochTree::OutcomeType::Continuous; + } else if (outcome_type_str == "binary") { + pred_metadata.outcome_type = StochTree::OutcomeType::Binary; + } else if (outcome_type_str == "ordinal") { + pred_metadata.outcome_type = StochTree::OutcomeType::Ordinal; + } else { + StochTree::Log::Fatal("Unsupported outcome type specified in model list"); + } + } + + // Run the prediction function + StochTree::BARTPredictionResult pred_results = predict_bart_model(data, *bart_samples_ptr, pred_metadata); + + // Unprotect protected R objects + UNPROTECT(protect_count); + + // Unpack outputs + cpp11::writable::list output_list = convert_bart_preds_to_list(pred_results); + return output_list; +} diff --git a/src/R_bcf.cpp b/src/R_bcf.cpp new file mode 100644 index 00000000..89307c5e --- /dev/null +++ b/src/R_bcf.cpp @@ -0,0 +1,392 @@ +#include +#include +#include +#include +#include +#include +#include +#include "stochtree/prediction.h" +#include "stochtree_types.h" +#include +#include + +StochTree::BCFConfig convert_list_to_bcf_config(cpp11::list config) { + StochTree::BCFConfig output; + + // Global model parameters + output.standardize_outcome = get_config_scalar_default(config, "standardize_outcome", true); + output.num_threads = get_config_scalar_default(config, "num_threads", 1); + output.verbose = get_config_scalar_default(config, "verbose", false); + output.cutpoint_grid_size = get_config_scalar_default(config, "cutpoint_grid_size", 100); + output.link_function = static_cast(get_config_scalar_default(config, "link_function", 0)); + output.outcome_type = static_cast(get_config_scalar_default(config, "outcome_type", 0)); + output.random_seed = get_config_scalar_default(config, "random_seed", 1); + output.keep_gfr = get_config_scalar_default(config, "keep_gfr", true); + output.keep_burnin = get_config_scalar_default(config, "keep_burnin", false); + output.adaptive_coding = get_config_scalar_default(config, "adaptive_coding", false); + output.b_0_init = get_config_scalar_default(config, "b_0_init", 0.0); + output.b_1_init = get_config_scalar_default(config, "b_1_init", 1.0); + + // Global error variance parameters + output.a_sigma2_global = get_config_scalar_default(config, "a_sigma2_global", 0.0); + output.b_sigma2_global = get_config_scalar_default(config, "b_sigma2_global", 0.0); + output.sigma2_global_init = get_config_scalar_default(config, "sigma2_global_init", 1.0); + output.sample_sigma2_global = get_config_scalar_default(config, "sample_sigma2_global", true); + + // Prognostic forest parameters + output.num_trees_mu = get_config_scalar_default(config, "num_trees_mu", 200); + output.alpha_mu = get_config_scalar_default(config, "alpha_mu", 0.95); + output.beta_mu = get_config_scalar_default(config, "beta_mu", 2.0); + output.min_samples_leaf_mu = get_config_scalar_default(config, "min_samples_leaf_mu", 5); + output.max_depth_mu = get_config_scalar_default(config, "max_depth_mu", -1); + output.leaf_constant_mu = get_config_scalar_default(config, "leaf_constant_mu", true); + output.leaf_dim_mu = get_config_scalar_default(config, "leaf_dim_mu", 1); + output.exponentiated_leaf_mu = get_config_scalar_default(config, "exponentiated_leaf_mu", false); + output.num_features_subsample_mu = get_config_scalar_default(config, "num_features_subsample_mu", 0); + output.a_sigma2_mu = get_config_scalar_default(config, "a_sigma2_mu", 3.0); + output.b_sigma2_mu = get_config_scalar_default(config, "b_sigma2_mu", -1.0); + output.sigma2_mu_init = get_config_scalar_default(config, "sigma2_mu_init", -1.0); + output.sample_sigma2_leaf_mu = get_config_scalar_default(config, "sample_sigma2_leaf_mu", false); + + // Treatment effect forest parameters + output.num_trees_tau = get_config_scalar_default(config, "num_trees_tau", 50); + output.alpha_tau = get_config_scalar_default(config, "alpha_tau", 0.95); + output.beta_tau = get_config_scalar_default(config, "beta_tau", 2.0); + output.min_samples_leaf_tau = get_config_scalar_default(config, "min_samples_leaf_tau", 5); + output.max_depth_tau = get_config_scalar_default(config, "max_depth_tau", -1); + output.leaf_constant_tau = get_config_scalar_default(config, "leaf_constant_tau", true); + output.leaf_dim_tau = get_config_scalar_default(config, "leaf_dim_tau", 1); + output.exponentiated_leaf_tau = get_config_scalar_default(config, "exponentiated_leaf_tau", false); + output.num_features_subsample_tau = get_config_scalar_default(config, "num_features_subsample_tau", 0); + output.a_sigma2_tau = get_config_scalar_default(config, "a_sigma2_tau", 3.0); + output.b_sigma2_tau = get_config_scalar_default(config, "b_sigma2_tau", -1.0); + output.sigma2_tau_init = get_config_scalar_default(config, "sigma2_tau_init", -1.0); + output.sample_sigma2_leaf_tau = get_config_scalar_default(config, "sample_sigma2_leaf_tau", false); + output.tau_leaf_model_type = static_cast(get_config_scalar_default(config, "tau_leaf_model_type", 1)); + output.sample_tau_0 = get_config_scalar_default(config, "sample_tau_0", true); + output.tau_0_prior_var_scalar = get_config_scalar_default(config, "tau_0_prior_var_scalar", -1.0); + + // Variance forest parameters + output.num_trees_variance = get_config_scalar_default(config, "num_trees_variance", 0); + output.leaf_prior_calibration_param = get_config_scalar_default(config, "leaf_prior_calibration_param", 1.5); + output.shape_variance_forest = get_config_scalar_default(config, "shape_variance_forest", -1.0); + output.scale_variance_forest = get_config_scalar_default(config, "scale_variance_forest", -1.0); + output.variance_forest_leaf_init = get_config_scalar_default(config, "variance_forest_leaf_init", -1.0); + output.alpha_variance = get_config_scalar_default(config, "alpha_variance", 0.5); + output.beta_variance = get_config_scalar_default(config, "beta_variance", 2.0); + output.min_samples_leaf_variance = get_config_scalar_default(config, "min_samples_leaf_variance", 5); + output.max_depth_variance = get_config_scalar_default(config, "max_depth_variance", -1); + output.leaf_constant_variance = get_config_scalar_default(config, "leaf_constant_variance", true); + output.leaf_dim_variance = get_config_scalar_default(config, "leaf_dim_variance", 1); + output.exponentiated_leaf_variance = get_config_scalar_default(config, "exponentiated_leaf_variance", true); + output.num_features_subsample_variance = get_config_scalar_default(config, "num_features_subsample_variance", 0); + + // Random effect parameters + output.has_random_effects = get_config_scalar_default(config, "has_random_effects", false); + output.rfx_model_spec = static_cast(get_config_scalar_default(config, "rfx_model_spec", 0)); + output.rfx_variance_prior_shape = get_config_scalar_default(config, "rfx_variance_prior_shape", 1.0); + output.rfx_variance_prior_scale = get_config_scalar_default(config, "rfx_variance_prior_scale", 1.0); + + // Handle vector conversions separately + SEXP feature_type_raw = static_cast(config["feature_types"]); + if (!Rf_isNull(feature_type_raw)) { + cpp11::integers feature_types_r_vec(feature_type_raw); + for (auto i : feature_types_r_vec) { + output.feature_types.push_back(static_cast(i)); + } + } + SEXP sweep_update_indices_mu_raw = static_cast(config["sweep_update_indices_mu"]); + if (!Rf_isNull(sweep_update_indices_mu_raw)) { + cpp11::integers sweep_update_indices_mu_r_vec(sweep_update_indices_mu_raw); + output.sweep_update_indices_mu.assign(sweep_update_indices_mu_r_vec.begin(), sweep_update_indices_mu_r_vec.end()); + } + SEXP sweep_update_indices_tau_raw = static_cast(config["sweep_update_indices_tau"]); + if (!Rf_isNull(sweep_update_indices_tau_raw)) { + cpp11::integers sweep_update_indices_tau_r_vec(sweep_update_indices_tau_raw); + output.sweep_update_indices_tau.assign(sweep_update_indices_tau_r_vec.begin(), sweep_update_indices_tau_r_vec.end()); + } + SEXP sweep_update_indices_variance_raw = static_cast(config["sweep_update_indices_variance"]); + if (!Rf_isNull(sweep_update_indices_variance_raw)) { + cpp11::integers sweep_update_indices_variance_r_vec(sweep_update_indices_variance_raw); + output.sweep_update_indices_variance.assign(sweep_update_indices_variance_r_vec.begin(), sweep_update_indices_variance_r_vec.end()); + } + SEXP var_weights_mu_raw = static_cast(config["var_weights_mu"]); + if (!Rf_isNull(var_weights_mu_raw)) { + cpp11::doubles var_weights_mu_r_vec(var_weights_mu_raw); + output.var_weights_mu.assign(var_weights_mu_r_vec.begin(), var_weights_mu_r_vec.end()); + } + SEXP var_weights_tau_raw = static_cast(config["var_weights_tau"]); + if (!Rf_isNull(var_weights_tau_raw)) { + cpp11::doubles var_weights_tau_r_vec(var_weights_tau_raw); + output.var_weights_tau.assign(var_weights_tau_r_vec.begin(), var_weights_tau_r_vec.end()); + } + SEXP sigma2_leaf_tau_matrix_raw = static_cast(config["sigma2_leaf_tau_matrix"]); + if (!Rf_isNull(sigma2_leaf_tau_matrix_raw)) { + cpp11::doubles sigma2_leaf_tau_matrix_r_vec(sigma2_leaf_tau_matrix_raw); + output.sigma2_leaf_tau_matrix.assign(sigma2_leaf_tau_matrix_r_vec.begin(), sigma2_leaf_tau_matrix_r_vec.end()); + } + SEXP tau_0_prior_var_raw = static_cast(config["tau_0_prior_var_multivariate"]); + if (!Rf_isNull(tau_0_prior_var_raw)) { + cpp11::doubles tau_0_prior_var_r_vec(tau_0_prior_var_raw); + output.tau_0_prior_var_multivariate.assign(tau_0_prior_var_r_vec.begin(), tau_0_prior_var_r_vec.end()); + } + SEXP var_weights_variance_raw = static_cast(config["var_weights_variance"]); + if (!Rf_isNull(var_weights_variance_raw)) { + cpp11::doubles var_weights_variance_r_vec(var_weights_variance_raw); + output.var_weights_variance.assign(var_weights_variance_r_vec.begin(), var_weights_variance_r_vec.end()); + } + SEXP rfx_working_parameter_mean_prior_raw = static_cast(config["rfx_working_parameter_mean_prior"]); + if (!Rf_isNull(rfx_working_parameter_mean_prior_raw)) { + cpp11::doubles rfx_working_parameter_mean_prior_r_vec(rfx_working_parameter_mean_prior_raw); + output.rfx_working_parameter_mean_prior.assign(rfx_working_parameter_mean_prior_r_vec.begin(), rfx_working_parameter_mean_prior_r_vec.end()); + } + SEXP rfx_group_parameter_mean_prior_raw = static_cast(config["rfx_group_parameter_mean_prior"]); + if (!Rf_isNull(rfx_group_parameter_mean_prior_raw)) { + cpp11::doubles rfx_group_parameter_mean_prior_r_vec(rfx_group_parameter_mean_prior_raw); + output.rfx_group_parameter_mean_prior.assign(rfx_group_parameter_mean_prior_r_vec.begin(), rfx_group_parameter_mean_prior_r_vec.end()); + } + SEXP rfx_working_parameter_cov_prior_raw = static_cast(config["rfx_working_parameter_cov_prior"]); + if (!Rf_isNull(rfx_working_parameter_cov_prior_raw)) { + cpp11::doubles rfx_working_parameter_cov_prior_r_vec(rfx_working_parameter_cov_prior_raw); + output.rfx_working_parameter_cov_prior.assign(rfx_working_parameter_cov_prior_r_vec.begin(), rfx_working_parameter_cov_prior_r_vec.end()); + } + SEXP rfx_group_parameter_cov_prior_raw = static_cast(config["rfx_group_parameter_cov_prior"]); + if (!Rf_isNull(rfx_group_parameter_cov_prior_raw)) { + cpp11::doubles rfx_group_parameter_cov_prior_r_vec(rfx_group_parameter_cov_prior_raw); + output.rfx_group_parameter_cov_prior.assign(rfx_group_parameter_cov_prior_r_vec.begin(), rfx_group_parameter_cov_prior_r_vec.end()); + } + return output; +} + +cpp11::writable::list create_bcf_metadata(StochTree::BCFConfig& config) { + cpp11::writable::list result; + result.push_back(cpp11::named_arg("sigma2_global_init") = config.sigma2_global_init); + result.push_back(cpp11::named_arg("sigma2_mu_init") = config.sigma2_mu_init); + result.push_back(cpp11::named_arg("sigma2_tau_init") = config.sigma2_tau_init); + result.push_back(cpp11::named_arg("b_sigma2_mu") = config.b_sigma2_mu); + result.push_back(cpp11::named_arg("b_sigma2_tau") = config.b_sigma2_tau); + result.push_back(cpp11::named_arg("shape_variance_forest") = config.shape_variance_forest); + result.push_back(cpp11::named_arg("scale_variance_forest") = config.scale_variance_forest); + return result; +} + +[[cpp11::register]] +cpp11::writable::list bcf_sample_cpp( + cpp11::external_pointer samples, + cpp11::sexp X_train, + cpp11::sexp Z_train, + cpp11::sexp y_train, + cpp11::sexp X_test, + cpp11::sexp Z_test, + int n_train, + int n_test, + int p, + int treatment_dim, + cpp11::sexp obs_weights_train, + cpp11::sexp obs_weights_test, + cpp11::sexp rfx_group_ids_train, + cpp11::sexp rfx_group_ids_test, + cpp11::sexp rfx_basis_train, + cpp11::sexp rfx_basis_test, + int rfx_num_groups, + int rfx_basis_dim, + int num_gfr, + int num_burnin, + int keep_every, + int num_mcmc, + int num_chains, + bool adaptive_coding, + cpp11::list config_input) { + // Extract pointers to raw data + int protect_count = 0; + double* X_train_ptr = extract_numeric_pointer(X_train, "X_train", protect_count); + double* Z_train_ptr = extract_numeric_pointer(Z_train, "Z_train", protect_count); + double* y_train_ptr = extract_numeric_pointer(y_train, "y_train", protect_count); + double* X_test_ptr = extract_numeric_pointer(X_test, "X_test", protect_count); + double* Z_test_ptr = extract_numeric_pointer(Z_test, "Z_test", protect_count); + double* obs_weights_train_ptr = extract_numeric_pointer(obs_weights_train, "obs_weights_train", protect_count); + double* obs_weights_test_ptr = extract_numeric_pointer(obs_weights_test, "obs_weights_test", protect_count); + int* rfx_group_ids_train_ptr = extract_integer_pointer(rfx_group_ids_train, "rfx_group_ids_train", protect_count); + int* rfx_group_ids_test_ptr = extract_integer_pointer(rfx_group_ids_test, "rfx_group_ids_test", protect_count); + double* rfx_basis_train_ptr = extract_numeric_pointer(rfx_basis_train, "rfx_basis_train", protect_count); + double* rfx_basis_test_ptr = extract_numeric_pointer(rfx_basis_test, "rfx_basis_test", protect_count); + + // Load the BCFData struct + // Consider reading directly from the R objects or at least checking for matches with the R object dimensions) + StochTree::BCFData data; + data.X_train = X_train_ptr; + data.treatment_train = Z_train_ptr; + data.y_train = y_train_ptr; + data.X_test = X_test_ptr; + data.treatment_test = Z_test_ptr; + data.n_train = n_train; + data.p = p; + data.n_test = n_test; + data.treatment_dim = treatment_dim; + data.obs_weights_train = obs_weights_train_ptr; + data.obs_weights_test = obs_weights_test_ptr; + data.rfx_group_ids_train = rfx_group_ids_train_ptr; + data.rfx_group_ids_test = rfx_group_ids_test_ptr; + data.rfx_basis_train = rfx_basis_train_ptr; + data.rfx_basis_test = rfx_basis_test_ptr; + data.rfx_num_groups = rfx_num_groups; + data.rfx_basis_dim = rfx_basis_dim; + + // Create the BCFConfig object + StochTree::BCFConfig config = convert_list_to_bcf_config(config_input); + + // Initialize a BCF sampler + StochTree::BCFSampler bcf_sampler(*samples, config, data); + + // Run the sampler + bcf_sampler.run_gfr(*samples, num_gfr, config.keep_gfr, num_chains); + if (num_chains > 1) { + bcf_sampler.run_mcmc_chains(*samples, num_chains, num_burnin, keep_every, num_mcmc); + } else { + bcf_sampler.run_mcmc(*samples, num_burnin, keep_every, num_mcmc); + } + bcf_sampler.postprocess_samples(*samples); + + // Unprotect protected R objects + UNPROTECT(protect_count); + + // Unpack outputs + cpp11::writable::list output_list = create_bcf_metadata(config); + return output_list; +} + +cpp11::writable::list convert_bcf_preds_to_list(StochTree::BCFPredictionResult& bcf_preds) { + cpp11::writable::list output; + + // Predictions + SEXP y_hat_sexp = !bcf_preds.y_hat.empty() + ? static_cast(cpp11::writable::doubles(bcf_preds.y_hat.begin(), bcf_preds.y_hat.end())) + : R_NilValue; + output.push_back(cpp11::named_arg("y_hat") = y_hat_sexp); + + SEXP mu_x_sexp = !bcf_preds.mu_x.empty() + ? static_cast(cpp11::writable::doubles(bcf_preds.mu_x.begin(), bcf_preds.mu_x.end())) + : R_NilValue; + output.push_back(cpp11::named_arg("mu_x") = mu_x_sexp); + + SEXP tau_x_sexp = !bcf_preds.tau_x.empty() + ? static_cast(cpp11::writable::doubles(bcf_preds.tau_x.begin(), bcf_preds.tau_x.end())) + : R_NilValue; + output.push_back(cpp11::named_arg("tau_x") = tau_x_sexp); + + SEXP prognostic_function_sexp = !bcf_preds.prognostic_function.empty() + ? static_cast(cpp11::writable::doubles(bcf_preds.prognostic_function.begin(), bcf_preds.prognostic_function.end())) + : R_NilValue; + output.push_back(cpp11::named_arg("prognostic_function") = prognostic_function_sexp); + + SEXP cate_sexp = !bcf_preds.cate.empty() + ? static_cast(cpp11::writable::doubles(bcf_preds.cate.begin(), bcf_preds.cate.end())) + : R_NilValue; + output.push_back(cpp11::named_arg("cate") = cate_sexp); + + SEXP conditional_variance_sexp = !bcf_preds.conditional_variance.empty() + ? static_cast(cpp11::writable::doubles(bcf_preds.conditional_variance.begin(), bcf_preds.conditional_variance.end())) + : R_NilValue; + output.push_back(cpp11::named_arg("conditional_variance") = conditional_variance_sexp); + + SEXP random_effects_sexp = !bcf_preds.random_effects.empty() + ? static_cast(cpp11::writable::doubles(bcf_preds.random_effects.begin(), bcf_preds.random_effects.end())) + : R_NilValue; + output.push_back(cpp11::named_arg("random_effects") = random_effects_sexp); + + return output; +} + +[[cpp11::register]] +cpp11::writable::list bcf_predict_cpp( + cpp11::external_pointer bcf_samples_ptr, + cpp11::list bcf_model_metadata, + cpp11::sexp X, + cpp11::sexp Z, + int n, + int p, + int treatment_dim, + cpp11::sexp obs_weights, + cpp11::sexp rfx_group_ids, + cpp11::sexp rfx_basis, + int rfx_num_groups, + int rfx_basis_dim, + bool posterior, + int scale, + bool predict_y_hat, + bool predict_mu_x, + bool predict_tau_x, + bool predict_prognostic_function, + bool predict_cate, + bool predict_conditional_variance, + bool predict_random_effects) { + // Extract pointers to raw data + int protect_count = 0; + double* X_ptr = extract_numeric_pointer(X, "X", protect_count); + double* Z_ptr = extract_numeric_pointer(Z, "Z", protect_count); + double* obs_weights_ptr = extract_numeric_pointer(obs_weights, "obs_weights", protect_count); + int* rfx_group_ids_ptr = extract_integer_pointer(rfx_group_ids, "rfx_group_ids", protect_count); + double* rfx_basis_ptr = extract_numeric_pointer(rfx_basis, "rfx_basis", protect_count); + + // Load the BCFData struct + // Consider reading directly from the R objects or at least checking for matches with the R object dimensions) + StochTree::BCFData data; + data.X_test = X_ptr; + data.treatment_test = Z_ptr; + data.p = p; + data.n_test = n; + data.treatment_dim = treatment_dim; + data.obs_weights_test = obs_weights_ptr; + data.rfx_group_ids_test = rfx_group_ids_ptr; + data.rfx_basis_test = rfx_basis_ptr; + data.rfx_num_groups = rfx_num_groups; + data.rfx_basis_dim = rfx_basis_dim; + + // Load the BCF model and config from the model list + StochTree::BCFPredictionMetadata pred_metadata; + pred_metadata.num_samples = Rf_asInteger(bcf_model_metadata["num_samples"]); + pred_metadata.num_obs = n; + pred_metadata.treatment_dim = treatment_dim; + pred_metadata.y_bar = Rf_asReal(bcf_model_metadata["y_bar"]); + pred_metadata.y_std = Rf_asReal(bcf_model_metadata["y_std"]); + pred_metadata.has_variance_forest = (bool)Rf_asLogical(bcf_model_metadata["include_variance_forest"]); + pred_metadata.has_rfx = (bool)Rf_asLogical(bcf_model_metadata["has_rfx"]); + { + SEXP rfx_spec_sexp = bcf_model_metadata["rfx_model_spec"]; + std::string rfx_model_spec_str = Rf_isNull(rfx_spec_sexp) ? "" : std::string(CHAR(STRING_ELT(rfx_spec_sexp, 0))); + if (rfx_model_spec_str == "intercept_only") { + pred_metadata.rfx_model_spec = StochTree::BCFRFXModelSpec::InterceptOnly; + } else if (rfx_model_spec_str == "intercept_plus_treatment") { + pred_metadata.rfx_model_spec = StochTree::BCFRFXModelSpec::InterceptPlusTreatment; + } else { + pred_metadata.rfx_model_spec = StochTree::BCFRFXModelSpec::Custom; + } + } + pred_metadata.adaptive_coding = (bool)Rf_asLogical(bcf_model_metadata["adaptive_coding"]); + pred_metadata.sample_tau_0 = (bool)Rf_asLogical(bcf_model_metadata["sample_tau_0"]); + pred_metadata.pred_type = posterior ? StochTree::PredType::kPosterior : StochTree::PredType::kMean; + if (scale == 0) { + pred_metadata.pred_scale = StochTree::PredScale::kLinear; + } else if (scale == 1) { + pred_metadata.pred_scale = StochTree::PredScale::kProbability; + } else { + pred_metadata.pred_scale = StochTree::PredScale::kClass; + } + pred_metadata.pred_terms.y_hat = predict_y_hat; + pred_metadata.pred_terms.mu_x = predict_mu_x; + pred_metadata.pred_terms.tau_x = predict_tau_x; + pred_metadata.pred_terms.prognostic_function = predict_prognostic_function; + pred_metadata.pred_terms.cate = predict_cate; + pred_metadata.pred_terms.conditional_variance = predict_conditional_variance; + pred_metadata.pred_terms.random_effects = predict_random_effects; + + // Run the prediction function + StochTree::BCFPredictionResult pred_results = predict_bcf_model(data, *bcf_samples_ptr, pred_metadata); + + // Unprotect protected R objects + UNPROTECT(protect_count); + + // Unpack outputs + cpp11::writable::list output_list = convert_bcf_preds_to_list(pred_results); + return output_list; +} diff --git a/src/R_data.cpp b/src/R_data.cpp index 3e96e0fc..681ca622 100644 --- a/src/R_data.cpp +++ b/src/R_data.cpp @@ -8,383 +8,383 @@ [[cpp11::register]] cpp11::external_pointer create_forest_dataset_cpp() { - // Create smart pointer to newly allocated object - std::unique_ptr dataset_ptr_ = std::make_unique(); - - // Release management of the pointer to R session - return cpp11::external_pointer(dataset_ptr_.release()); + // Create smart pointer to newly allocated object + std::unique_ptr dataset_ptr_ = std::make_unique(); + + // Release management of the pointer to R session + return cpp11::external_pointer(dataset_ptr_.release()); } [[cpp11::register]] int dataset_num_rows_cpp(cpp11::external_pointer dataset) { - return dataset->NumObservations(); + return dataset->NumObservations(); } [[cpp11::register]] int dataset_num_covariates_cpp(cpp11::external_pointer dataset) { - return dataset->NumCovariates(); + return dataset->NumCovariates(); } [[cpp11::register]] int dataset_num_basis_cpp(cpp11::external_pointer dataset) { - return dataset->NumBasis(); + return dataset->NumBasis(); } [[cpp11::register]] bool dataset_has_basis_cpp(cpp11::external_pointer dataset) { - return dataset->HasBasis(); + return dataset->HasBasis(); } [[cpp11::register]] bool dataset_has_variance_weights_cpp(cpp11::external_pointer dataset) { - return dataset->HasVarWeights(); + return dataset->HasVarWeights(); } [[cpp11::register]] void forest_dataset_add_covariates_cpp(cpp11::external_pointer dataset_ptr, cpp11::doubles_matrix<> covariates) { - // TODO: add handling code on the R side to ensure matrices are column-major - bool row_major{false}; + // TODO: add handling code on the R side to ensure matrices are column-major + bool row_major{false}; + + // Add covariates + StochTree::data_size_t n = covariates.nrow(); + int num_covariates = covariates.ncol(); + double* covariate_data_ptr = REAL(PROTECT(covariates)); + dataset_ptr->AddCovariates(covariate_data_ptr, n, num_covariates, row_major); - // Add covariates - StochTree::data_size_t n = covariates.nrow(); - int num_covariates = covariates.ncol(); - double* covariate_data_ptr = REAL(PROTECT(covariates)); - dataset_ptr->AddCovariates(covariate_data_ptr, n, num_covariates, row_major); - - // Unprotect pointers to R data - UNPROTECT(1); + // Unprotect pointers to R data + UNPROTECT(1); } [[cpp11::register]] void forest_dataset_add_basis_cpp(cpp11::external_pointer dataset_ptr, cpp11::doubles_matrix<> basis) { - // TODO: add handling code on the R side to ensure matrices are column-major - bool row_major{false}; + // TODO: add handling code on the R side to ensure matrices are column-major + bool row_major{false}; - // Add basis - StochTree::data_size_t n = basis.nrow(); - int num_basis = basis.ncol(); - double* basis_data_ptr = REAL(PROTECT(basis)); - dataset_ptr->AddBasis(basis_data_ptr, n, num_basis, row_major); - - // Unprotect pointers to R data - UNPROTECT(1); + // Add basis + StochTree::data_size_t n = basis.nrow(); + int num_basis = basis.ncol(); + double* basis_data_ptr = REAL(PROTECT(basis)); + dataset_ptr->AddBasis(basis_data_ptr, n, num_basis, row_major); + + // Unprotect pointers to R data + UNPROTECT(1); } [[cpp11::register]] void forest_dataset_update_basis_cpp(cpp11::external_pointer dataset_ptr, cpp11::doubles_matrix<> basis) { - // TODO: add handling code on the R side to ensure matrices are column-major - bool row_major{false}; - - // Add basis - StochTree::data_size_t n = basis.nrow(); - int num_basis = basis.ncol(); - double* basis_data_ptr = REAL(PROTECT(basis)); - dataset_ptr->UpdateBasis(basis_data_ptr, n, num_basis, row_major); - - // Unprotect pointers to R data - UNPROTECT(1); + // TODO: add handling code on the R side to ensure matrices are column-major + bool row_major{false}; + + // Add basis + StochTree::data_size_t n = basis.nrow(); + int num_basis = basis.ncol(); + double* basis_data_ptr = REAL(PROTECT(basis)); + dataset_ptr->UpdateBasis(basis_data_ptr, n, num_basis, row_major); + + // Unprotect pointers to R data + UNPROTECT(1); } [[cpp11::register]] void forest_dataset_update_var_weights_cpp(cpp11::external_pointer dataset_ptr, cpp11::doubles weights, bool exponentiate) { - // Add weights - StochTree::data_size_t n = weights.size(); - double* weight_data_ptr = REAL(PROTECT(weights)); - dataset_ptr->UpdateVarWeights(weight_data_ptr, n, exponentiate); - - // Unprotect pointers to R data - UNPROTECT(1); + // Add weights + StochTree::data_size_t n = weights.size(); + double* weight_data_ptr = REAL(PROTECT(weights)); + dataset_ptr->UpdateVarWeights(weight_data_ptr, n, exponentiate); + + // Unprotect pointers to R data + UNPROTECT(1); } [[cpp11::register]] void forest_dataset_add_weights_cpp(cpp11::external_pointer dataset_ptr, cpp11::doubles weights) { - // Add weights - StochTree::data_size_t n = weights.size(); - double* weight_data_ptr = REAL(PROTECT(weights)); - dataset_ptr->AddVarianceWeights(weight_data_ptr, n); + // Add weights + StochTree::data_size_t n = weights.size(); + double* weight_data_ptr = REAL(PROTECT(weights)); + dataset_ptr->AddVarianceWeights(weight_data_ptr, n); - // Unprotect pointers to R data - UNPROTECT(1); + // Unprotect pointers to R data + UNPROTECT(1); } [[cpp11::register]] cpp11::writable::doubles_matrix<> forest_dataset_get_covariates_cpp(cpp11::external_pointer dataset_ptr) { - // Initialize output matrix - int num_row = dataset_ptr->NumObservations(); - int num_col = dataset_ptr->NumCovariates(); - cpp11::writable::doubles_matrix<> output(num_row, num_col); - - for (int i = 0; i < num_row; i++) { - for (int j = 0; j < num_col; j++) { - output(i, j) = dataset_ptr->CovariateValue(i, j); - } + // Initialize output matrix + int num_row = dataset_ptr->NumObservations(); + int num_col = dataset_ptr->NumCovariates(); + cpp11::writable::doubles_matrix<> output(num_row, num_col); + + for (int i = 0; i < num_row; i++) { + for (int j = 0; j < num_col; j++) { + output(i, j) = dataset_ptr->CovariateValue(i, j); } + } - return output; + return output; } [[cpp11::register]] cpp11::writable::doubles_matrix<> forest_dataset_get_basis_cpp(cpp11::external_pointer dataset_ptr) { - // Initialize output matrix - int num_row = dataset_ptr->NumObservations(); - int num_col = dataset_ptr->NumBasis(); - cpp11::writable::doubles_matrix<> output(num_row, num_col); - for (int i = 0; i < num_row; i++) { - for (int j = 0; j < num_col; j++) { - output(i, j) = dataset_ptr->BasisValue(i, j); - } - } - return output; + // Initialize output matrix + int num_row = dataset_ptr->NumObservations(); + int num_col = dataset_ptr->NumBasis(); + cpp11::writable::doubles_matrix<> output(num_row, num_col); + for (int i = 0; i < num_row; i++) { + for (int j = 0; j < num_col; j++) { + output(i, j) = dataset_ptr->BasisValue(i, j); + } + } + return output; } [[cpp11::register]] cpp11::writable::doubles forest_dataset_get_variance_weights_cpp(cpp11::external_pointer dataset_ptr) { - // Initialize output vector - int num_row = dataset_ptr->NumObservations(); - cpp11::writable::doubles output(num_row); - for (int i = 0; i < num_row; i++) { - output.at(i) = dataset_ptr->VarWeightValue(i); - } - return output; + // Initialize output vector + int num_row = dataset_ptr->NumObservations(); + cpp11::writable::doubles output(num_row); + for (int i = 0; i < num_row; i++) { + output.at(i) = dataset_ptr->VarWeightValue(i); + } + return output; } [[cpp11::register]] bool forest_dataset_has_auxiliary_dimension_cpp(cpp11::external_pointer dataset_ptr, int dim_idx) { - return dataset_ptr->HasAuxiliaryDimension(dim_idx); + return dataset_ptr->HasAuxiliaryDimension(dim_idx); } [[cpp11::register]] void forest_dataset_add_auxiliary_dimension_cpp(cpp11::external_pointer dataset_ptr, int dim_size) { - dataset_ptr->AddAuxiliaryDimension(dim_size); + dataset_ptr->AddAuxiliaryDimension(dim_size); } [[cpp11::register]] double forest_dataset_get_auxiliary_data_value_cpp(cpp11::external_pointer dataset_ptr, int dim_idx, int element_idx) { - return dataset_ptr->GetAuxiliaryDataValue(dim_idx, element_idx); + return dataset_ptr->GetAuxiliaryDataValue(dim_idx, element_idx); } [[cpp11::register]] void forest_dataset_set_auxiliary_data_value_cpp(cpp11::external_pointer dataset_ptr, int dim_idx, int element_idx, double value) { - dataset_ptr->SetAuxiliaryDataValue(dim_idx, element_idx, value); + dataset_ptr->SetAuxiliaryDataValue(dim_idx, element_idx, value); } [[cpp11::register]] cpp11::writable::doubles forest_dataset_get_auxiliary_data_vector_cpp(cpp11::external_pointer dataset_ptr, int dim_idx) { - const std::vector output_raw = dataset_ptr->GetAuxiliaryDataVector(dim_idx); - int n = output_raw.size(); - cpp11::writable::doubles output(n); - for (int i = 0; i < n; i++) { - output[i] = output_raw[i]; - } - return output; + const std::vector output_raw = dataset_ptr->GetAuxiliaryDataVector(dim_idx); + int n = output_raw.size(); + cpp11::writable::doubles output(n); + for (int i = 0; i < n; i++) { + output[i] = output_raw[i]; + } + return output; } [[cpp11::register]] cpp11::external_pointer create_column_vector_cpp(cpp11::doubles outcome) { - // Unpack pointers to data and dimensions - StochTree::data_size_t n = outcome.size(); - double* outcome_data_ptr = REAL(PROTECT(outcome)); + // Unpack pointers to data and dimensions + StochTree::data_size_t n = outcome.size(); + double* outcome_data_ptr = REAL(PROTECT(outcome)); + + // Create smart pointer + std::unique_ptr vector_ptr_ = std::make_unique(outcome_data_ptr, n); - // Create smart pointer - std::unique_ptr vector_ptr_ = std::make_unique(outcome_data_ptr, n); - - // Unprotect pointers to R data - UNPROTECT(1); - - // Release management of the pointer to R session - return cpp11::external_pointer(vector_ptr_.release()); + // Unprotect pointers to R data + UNPROTECT(1); + + // Release management of the pointer to R session + return cpp11::external_pointer(vector_ptr_.release()); } [[cpp11::register]] void add_to_column_vector_cpp(cpp11::external_pointer outcome, cpp11::doubles update_vector) { - // Unpack pointers to data and dimensions - StochTree::data_size_t n = update_vector.size(); - double* update_data_ptr = REAL(PROTECT(update_vector)); - - // Add to the outcome data using the C++ API - outcome->AddToData(update_data_ptr, n); - - // Unprotect pointers to R data - UNPROTECT(1); + // Unpack pointers to data and dimensions + StochTree::data_size_t n = update_vector.size(); + double* update_data_ptr = REAL(PROTECT(update_vector)); + + // Add to the outcome data using the C++ API + outcome->AddToData(update_data_ptr, n); + + // Unprotect pointers to R data + UNPROTECT(1); } [[cpp11::register]] void subtract_from_column_vector_cpp(cpp11::external_pointer outcome, cpp11::doubles update_vector) { - // Unpack pointers to data and dimensions - StochTree::data_size_t n = update_vector.size(); - double* update_data_ptr = REAL(PROTECT(update_vector)); - - // Add to the outcome data using the C++ API - outcome->SubtractFromData(update_data_ptr, n); - - // Unprotect pointers to R data - UNPROTECT(1); + // Unpack pointers to data and dimensions + StochTree::data_size_t n = update_vector.size(); + double* update_data_ptr = REAL(PROTECT(update_vector)); + + // Add to the outcome data using the C++ API + outcome->SubtractFromData(update_data_ptr, n); + + // Unprotect pointers to R data + UNPROTECT(1); } [[cpp11::register]] void overwrite_column_vector_cpp(cpp11::external_pointer outcome, cpp11::doubles new_vector) { - // Unpack pointers to data and dimensions - StochTree::data_size_t n = new_vector.size(); - double* update_data_ptr = REAL(PROTECT(new_vector)); - - // Add to the outcome data using the C++ API - outcome->OverwriteData(update_data_ptr, n); - - // Unprotect pointers to R data - UNPROTECT(1); + // Unpack pointers to data and dimensions + StochTree::data_size_t n = new_vector.size(); + double* update_data_ptr = REAL(PROTECT(new_vector)); + + // Add to the outcome data using the C++ API + outcome->OverwriteData(update_data_ptr, n); + + // Unprotect pointers to R data + UNPROTECT(1); } [[cpp11::register]] -void propagate_trees_column_vector_cpp(cpp11::external_pointer tracker, +void propagate_trees_column_vector_cpp(cpp11::external_pointer tracker, cpp11::external_pointer residual) { - StochTree::UpdateResidualNewOutcome(*tracker, *residual); + StochTree::UpdateResidualNewOutcome(*tracker, *residual); } [[cpp11::register]] cpp11::writable::doubles get_residual_cpp(cpp11::external_pointer vector_ptr) { - // Initialize output vector - StochTree::data_size_t n = vector_ptr->NumRows(); - cpp11::writable::doubles output(n); - - // Unpack data - for (StochTree::data_size_t i = 0; i < n; i++) { - output.at(i) = vector_ptr->GetElement(i); - } - - // Release management of the pointer to R session - return output; + // Initialize output vector + StochTree::data_size_t n = vector_ptr->NumRows(); + cpp11::writable::doubles output(n); + + // Unpack data + for (StochTree::data_size_t i = 0; i < n; i++) { + output.at(i) = vector_ptr->GetElement(i); + } + + // Release management of the pointer to R session + return output; } [[cpp11::register]] cpp11::external_pointer create_rfx_dataset_cpp() { - // Create smart pointer to newly allocated object - std::unique_ptr dataset_ptr_ = std::make_unique(); - - // Release management of the pointer to R session - return cpp11::external_pointer(dataset_ptr_.release()); + // Create smart pointer to newly allocated object + std::unique_ptr dataset_ptr_ = std::make_unique(); + + // Release management of the pointer to R session + return cpp11::external_pointer(dataset_ptr_.release()); } [[cpp11::register]] void rfx_dataset_update_basis_cpp(cpp11::external_pointer dataset_ptr, cpp11::doubles_matrix<> basis) { - // TODO: add handling code on the R side to ensure matrices are column-major - bool row_major{false}; - - // Add basis - StochTree::data_size_t n = basis.nrow(); - int num_basis = basis.ncol(); - double* basis_data_ptr = REAL(PROTECT(basis)); - dataset_ptr->UpdateBasis(basis_data_ptr, n, num_basis, row_major); - - // Unprotect pointers to R data - UNPROTECT(1); + // TODO: add handling code on the R side to ensure matrices are column-major + bool row_major{false}; + + // Add basis + StochTree::data_size_t n = basis.nrow(); + int num_basis = basis.ncol(); + double* basis_data_ptr = REAL(PROTECT(basis)); + dataset_ptr->UpdateBasis(basis_data_ptr, n, num_basis, row_major); + + // Unprotect pointers to R data + UNPROTECT(1); } [[cpp11::register]] void rfx_dataset_update_var_weights_cpp(cpp11::external_pointer dataset_ptr, cpp11::doubles weights, bool exponentiate) { - // Add weights - StochTree::data_size_t n = weights.size(); - double* weight_data_ptr = REAL(PROTECT(weights)); - dataset_ptr->UpdateVarWeights(weight_data_ptr, n, exponentiate); - - // Unprotect pointers to R data - UNPROTECT(1); + // Add weights + StochTree::data_size_t n = weights.size(); + double* weight_data_ptr = REAL(PROTECT(weights)); + dataset_ptr->UpdateVarWeights(weight_data_ptr, n, exponentiate); + + // Unprotect pointers to R data + UNPROTECT(1); } [[cpp11::register]] void rfx_dataset_update_group_labels_cpp(cpp11::external_pointer dataset_ptr, cpp11::integers group_labels) { - // Update group labels - int n = group_labels.size(); - std::vector group_labels_vec(group_labels.begin(), group_labels.end()); - dataset_ptr->UpdateGroupLabels(group_labels_vec, n); + // Update group labels + int n = group_labels.size(); + std::vector group_labels_vec(group_labels.begin(), group_labels.end()); + dataset_ptr->UpdateGroupLabels(group_labels_vec, n); } [[cpp11::register]] int rfx_dataset_num_basis_cpp(cpp11::external_pointer dataset) { - return dataset->NumBases(); + return dataset->NumBases(); } [[cpp11::register]] int rfx_dataset_num_rows_cpp(cpp11::external_pointer dataset) { - return dataset->NumObservations(); + return dataset->NumObservations(); } [[cpp11::register]] bool rfx_dataset_has_group_labels_cpp(cpp11::external_pointer dataset) { - return dataset->HasGroupLabels(); + return dataset->HasGroupLabels(); } [[cpp11::register]] bool rfx_dataset_has_basis_cpp(cpp11::external_pointer dataset) { - return dataset->HasBasis(); + return dataset->HasBasis(); } [[cpp11::register]] bool rfx_dataset_has_variance_weights_cpp(cpp11::external_pointer dataset) { - return dataset->HasVarWeights(); + return dataset->HasVarWeights(); } [[cpp11::register]] void rfx_dataset_add_group_labels_cpp(cpp11::external_pointer dataset_ptr, cpp11::integers group_labels) { - // Add group labels - std::vector group_labels_vec(group_labels.begin(), group_labels.end()); - dataset_ptr->AddGroupLabels(group_labels_vec); + // Add group labels + std::vector group_labels_vec(group_labels.begin(), group_labels.end()); + dataset_ptr->AddGroupLabels(group_labels_vec); } [[cpp11::register]] void rfx_dataset_add_basis_cpp(cpp11::external_pointer dataset_ptr, cpp11::doubles_matrix<> basis) { - // TODO: add handling code on the R side to ensure matrices are column-major - bool row_major{false}; - - // Add basis - StochTree::data_size_t n = basis.nrow(); - int num_basis = basis.ncol(); - double* basis_data_ptr = REAL(PROTECT(basis)); - dataset_ptr->AddBasis(basis_data_ptr, n, num_basis, row_major); - - // Unprotect pointers to R data - UNPROTECT(1); + // TODO: add handling code on the R side to ensure matrices are column-major + bool row_major{false}; + + // Add basis + StochTree::data_size_t n = basis.nrow(); + int num_basis = basis.ncol(); + double* basis_data_ptr = REAL(PROTECT(basis)); + dataset_ptr->AddBasis(basis_data_ptr, n, num_basis, row_major); + + // Unprotect pointers to R data + UNPROTECT(1); } [[cpp11::register]] void rfx_dataset_add_weights_cpp(cpp11::external_pointer dataset_ptr, cpp11::doubles weights) { - // Add weights - StochTree::data_size_t n = weights.size(); - double* weight_data_ptr = REAL(PROTECT(weights)); - dataset_ptr->AddVarianceWeights(weight_data_ptr, n); - - // Unprotect pointers to R data - UNPROTECT(1); + // Add weights + StochTree::data_size_t n = weights.size(); + double* weight_data_ptr = REAL(PROTECT(weights)); + dataset_ptr->AddVarianceWeights(weight_data_ptr, n); + + // Unprotect pointers to R data + UNPROTECT(1); } [[cpp11::register]] cpp11::writable::integers rfx_dataset_get_group_labels_cpp(cpp11::external_pointer dataset_ptr) { - int num_row = dataset_ptr->NumObservations(); - cpp11::writable::integers output(num_row); - for (int i = 0; i < num_row; i++) { - output.at(i) = dataset_ptr->GroupId(i); - } - return output; + int num_row = dataset_ptr->NumObservations(); + cpp11::writable::integers output(num_row); + for (int i = 0; i < num_row; i++) { + output.at(i) = dataset_ptr->GroupId(i); + } + return output; } [[cpp11::register]] cpp11::writable::doubles_matrix<> rfx_dataset_get_basis_cpp(cpp11::external_pointer dataset_ptr) { - int num_row = dataset_ptr->NumObservations(); - int num_col = dataset_ptr->NumBases(); - cpp11::writable::doubles_matrix<> output(num_row, num_col); - for (int i = 0; i < num_row; i++) { - for (int j = 0; j < num_col; j++) { - output(i, j) = dataset_ptr->BasisValue(i, j); - } - } - return output; + int num_row = dataset_ptr->NumObservations(); + int num_col = dataset_ptr->NumBases(); + cpp11::writable::doubles_matrix<> output(num_row, num_col); + for (int i = 0; i < num_row; i++) { + for (int j = 0; j < num_col; j++) { + output(i, j) = dataset_ptr->BasisValue(i, j); + } + } + return output; } [[cpp11::register]] cpp11::writable::doubles rfx_dataset_get_variance_weights_cpp(cpp11::external_pointer dataset_ptr) { - int num_row = dataset_ptr->NumObservations(); - cpp11::writable::doubles output(num_row); - for (int i = 0; i < num_row; i++) { - output.at(i) = dataset_ptr->VarWeightValue(i); - } - return output; + int num_row = dataset_ptr->NumObservations(); + cpp11::writable::doubles output(num_row); + for (int i = 0; i < num_row; i++) { + output.at(i) = dataset_ptr->VarWeightValue(i); + } + return output; } diff --git a/src/R_random_effects.cpp b/src/R_random_effects.cpp index fffba538..cddd0f9e 100644 --- a/src/R_random_effects.cpp +++ b/src/R_random_effects.cpp @@ -11,334 +11,334 @@ [[cpp11::register]] cpp11::external_pointer rfx_container_cpp(int num_components, int num_groups) { - // Create smart pointer to newly allocated object - std::unique_ptr rfx_container_ptr_ = std::make_unique(num_components, num_groups); - - // Release management of the pointer to R session - return cpp11::external_pointer(rfx_container_ptr_.release()); + // Create smart pointer to newly allocated object + std::unique_ptr rfx_container_ptr_ = std::make_unique(num_components, num_groups); + + // Release management of the pointer to R session + return cpp11::external_pointer(rfx_container_ptr_.release()); } [[cpp11::register]] cpp11::external_pointer rfx_container_from_json_cpp(cpp11::external_pointer json_ptr, std::string rfx_label) { - // Create smart pointer to newly allocated object - std::unique_ptr rfx_container_ptr_ = std::make_unique(); - - // Extract the random effect container's json - nlohmann::json rfx_json = json_ptr->at("random_effects").at(rfx_label); - - // Reset the forest sample container using the json - rfx_container_ptr_->Reset(); - rfx_container_ptr_->from_json(rfx_json); - - // Release management of the pointer to R session - return cpp11::external_pointer(rfx_container_ptr_.release()); + // Create smart pointer to newly allocated object + std::unique_ptr rfx_container_ptr_ = std::make_unique(); + + // Extract the random effect container's json + nlohmann::json rfx_json = json_ptr->at("random_effects").at(rfx_label); + + // Reset the forest sample container using the json + rfx_container_ptr_->Reset(); + rfx_container_ptr_->from_json(rfx_json); + + // Release management of the pointer to R session + return cpp11::external_pointer(rfx_container_ptr_.release()); } [[cpp11::register]] cpp11::external_pointer rfx_label_mapper_from_json_cpp(cpp11::external_pointer json_ptr, std::string rfx_label) { - // Create smart pointer to newly allocated object - std::unique_ptr label_mapper_ptr_ = std::make_unique(); - - // Extract the label mapper's json - nlohmann::json rfx_json = json_ptr->at("random_effects").at(rfx_label); - - // Reset the label mapper using the json - label_mapper_ptr_->Reset(); - label_mapper_ptr_->from_json(rfx_json); - - // Release management of the pointer to R session - return cpp11::external_pointer(label_mapper_ptr_.release()); + // Create smart pointer to newly allocated object + std::unique_ptr label_mapper_ptr_ = std::make_unique(); + + // Extract the label mapper's json + nlohmann::json rfx_json = json_ptr->at("random_effects").at(rfx_label); + + // Reset the label mapper using the json + label_mapper_ptr_->Reset(); + label_mapper_ptr_->from_json(rfx_json); + + // Release management of the pointer to R session + return cpp11::external_pointer(label_mapper_ptr_.release()); } [[cpp11::register]] cpp11::writable::integers rfx_group_ids_from_json_cpp(cpp11::external_pointer json_ptr, std::string rfx_label) { - // Create smart pointer to newly allocated object - cpp11::writable::integers output; - - // Extract the groupids' json - nlohmann::json rfx_json = json_ptr->at("random_effects").at(rfx_label); - - // Reset the forest sample container using the json - int num_groups = rfx_json.size(); - for (int i = 0; i < num_groups; i++) { - output.push_back(rfx_json.at(i)); - } - - return output; + // Create smart pointer to newly allocated object + cpp11::writable::integers output; + + // Extract the groupids' json + nlohmann::json rfx_json = json_ptr->at("random_effects").at(rfx_label); + + // Reset the forest sample container using the json + int num_groups = rfx_json.size(); + for (int i = 0; i < num_groups; i++) { + output.push_back(rfx_json.at(i)); + } + + return output; } [[cpp11::register]] void rfx_container_append_from_json_cpp(cpp11::external_pointer rfx_container_ptr, cpp11::external_pointer json_ptr, std::string rfx_label) { - // Extract the random effect container's json - nlohmann::json rfx_json = json_ptr->at("random_effects").at(rfx_label); - - // Reset the forest sample container using the json - rfx_container_ptr->append_from_json(rfx_json); + // Extract the random effect container's json + nlohmann::json rfx_json = json_ptr->at("random_effects").at(rfx_label); + + // Reset the forest sample container using the json + rfx_container_ptr->append_from_json(rfx_json); } [[cpp11::register]] cpp11::external_pointer rfx_container_from_json_string_cpp(std::string json_string, std::string rfx_label) { - // Create smart pointer to newly allocated object - std::unique_ptr rfx_container_ptr_ = std::make_unique(); - - // Create a nlohmann::json object from the string - nlohmann::json json_object = nlohmann::json::parse(json_string); - - // Extract the random effect container's json - nlohmann::json rfx_json = json_object.at("random_effects").at(rfx_label); - - // Reset the forest sample container using the json - rfx_container_ptr_->Reset(); - rfx_container_ptr_->from_json(rfx_json); - - // Release management of the pointer to R session - return cpp11::external_pointer(rfx_container_ptr_.release()); + // Create smart pointer to newly allocated object + std::unique_ptr rfx_container_ptr_ = std::make_unique(); + + // Create a nlohmann::json object from the string + nlohmann::json json_object = nlohmann::json::parse(json_string); + + // Extract the random effect container's json + nlohmann::json rfx_json = json_object.at("random_effects").at(rfx_label); + + // Reset the forest sample container using the json + rfx_container_ptr_->Reset(); + rfx_container_ptr_->from_json(rfx_json); + + // Release management of the pointer to R session + return cpp11::external_pointer(rfx_container_ptr_.release()); } [[cpp11::register]] cpp11::external_pointer rfx_label_mapper_from_json_string_cpp(std::string json_string, std::string rfx_label) { - // Create smart pointer to newly allocated object - std::unique_ptr label_mapper_ptr_ = std::make_unique(); - - // Create a nlohmann::json object from the string - nlohmann::json json_object = nlohmann::json::parse(json_string); - - // Extract the label mapper's json - nlohmann::json rfx_json = json_object.at("random_effects").at(rfx_label); - - // Reset the label mapper using the json - label_mapper_ptr_->Reset(); - label_mapper_ptr_->from_json(rfx_json); - - // Release management of the pointer to R session - return cpp11::external_pointer(label_mapper_ptr_.release()); + // Create smart pointer to newly allocated object + std::unique_ptr label_mapper_ptr_ = std::make_unique(); + + // Create a nlohmann::json object from the string + nlohmann::json json_object = nlohmann::json::parse(json_string); + + // Extract the label mapper's json + nlohmann::json rfx_json = json_object.at("random_effects").at(rfx_label); + + // Reset the label mapper using the json + label_mapper_ptr_->Reset(); + label_mapper_ptr_->from_json(rfx_json); + + // Release management of the pointer to R session + return cpp11::external_pointer(label_mapper_ptr_.release()); } [[cpp11::register]] cpp11::writable::integers rfx_group_ids_from_json_string_cpp(std::string json_string, std::string rfx_label) { - // Create smart pointer to newly allocated object - cpp11::writable::integers output; - - // Create a nlohmann::json object from the string - nlohmann::json json_object = nlohmann::json::parse(json_string); - - // Extract the groupids' json - nlohmann::json rfx_json = json_object.at("random_effects").at(rfx_label); - - // Reset the forest sample container using the json - int num_groups = rfx_json.size(); - for (int i = 0; i < num_groups; i++) { - output.push_back(rfx_json.at(i)); - } - - return output; + // Create smart pointer to newly allocated object + cpp11::writable::integers output; + + // Create a nlohmann::json object from the string + nlohmann::json json_object = nlohmann::json::parse(json_string); + + // Extract the groupids' json + nlohmann::json rfx_json = json_object.at("random_effects").at(rfx_label); + + // Reset the forest sample container using the json + int num_groups = rfx_json.size(); + for (int i = 0; i < num_groups; i++) { + output.push_back(rfx_json.at(i)); + } + + return output; } [[cpp11::register]] void rfx_container_append_from_json_string_cpp(cpp11::external_pointer rfx_container_ptr, std::string json_string, std::string rfx_label) { - // Create a nlohmann::json object from the string - nlohmann::json json_object = nlohmann::json::parse(json_string); - - // Extract the random effect container's json - nlohmann::json rfx_json = json_object.at("random_effects").at(rfx_label); - - // Reset the forest sample container using the json - rfx_container_ptr->append_from_json(rfx_json); + // Create a nlohmann::json object from the string + nlohmann::json json_object = nlohmann::json::parse(json_string); + + // Extract the random effect container's json + nlohmann::json rfx_json = json_object.at("random_effects").at(rfx_label); + + // Reset the forest sample container using the json + rfx_container_ptr->append_from_json(rfx_json); } [[cpp11::register]] cpp11::external_pointer rfx_model_cpp(int num_components, int num_groups) { - // Create smart pointer to newly allocated object - std::unique_ptr rfx_model_ptr_ = std::make_unique(num_components, num_groups); - - // Release management of the pointer to R session - return cpp11::external_pointer(rfx_model_ptr_.release()); + // Create smart pointer to newly allocated object + std::unique_ptr rfx_model_ptr_ = std::make_unique(num_components, num_groups); + + // Release management of the pointer to R session + return cpp11::external_pointer(rfx_model_ptr_.release()); } [[cpp11::register]] cpp11::external_pointer rfx_tracker_cpp(cpp11::integers group_labels) { - // Convert group_labels to a std::vector - std::vector group_labels_vec(group_labels.begin(), group_labels.end()); - - // Create smart pointer to newly allocated object - std::unique_ptr rfx_tracker_ptr_ = std::make_unique(group_labels_vec); - - // Release management of the pointer to R session - return cpp11::external_pointer(rfx_tracker_ptr_.release()); + // Convert group_labels to a std::vector + std::vector group_labels_vec(group_labels.begin(), group_labels.end()); + + // Create smart pointer to newly allocated object + std::unique_ptr rfx_tracker_ptr_ = std::make_unique(group_labels_vec); + + // Release management of the pointer to R session + return cpp11::external_pointer(rfx_tracker_ptr_.release()); } [[cpp11::register]] cpp11::external_pointer rfx_label_mapper_cpp(cpp11::external_pointer rfx_tracker) { - // Create smart pointer to newly allocated object - std::unique_ptr rfx_label_mapper_ptr_ = std::make_unique(rfx_tracker->GetLabelMap()); - - // Release management of the pointer to R session - return cpp11::external_pointer(rfx_label_mapper_ptr_.release()); + // Create smart pointer to newly allocated object + std::unique_ptr rfx_label_mapper_ptr_ = std::make_unique(rfx_tracker->GetLabelMap()); + + // Release management of the pointer to R session + return cpp11::external_pointer(rfx_label_mapper_ptr_.release()); } [[cpp11::register]] -void rfx_model_sample_random_effects_cpp(cpp11::external_pointer rfx_model, cpp11::external_pointer rfx_dataset, - cpp11::external_pointer residual, cpp11::external_pointer rfx_tracker, +void rfx_model_sample_random_effects_cpp(cpp11::external_pointer rfx_model, cpp11::external_pointer rfx_dataset, + cpp11::external_pointer residual, cpp11::external_pointer rfx_tracker, cpp11::external_pointer rfx_container, bool keep_sample, double global_variance, cpp11::external_pointer rng) { - rfx_model->SampleRandomEffects(*rfx_dataset, *residual, *rfx_tracker, global_variance, *rng); - if (keep_sample) rfx_container->AddSample(*rfx_model); + rfx_model->SampleRandomEffects(*rfx_dataset, *residual, *rfx_tracker, global_variance, *rng); + if (keep_sample) rfx_container->AddSample(*rfx_model); } [[cpp11::register]] -cpp11::writable::doubles rfx_model_predict_cpp(cpp11::external_pointer rfx_model, - cpp11::external_pointer rfx_dataset, +cpp11::writable::doubles rfx_model_predict_cpp(cpp11::external_pointer rfx_model, + cpp11::external_pointer rfx_dataset, cpp11::external_pointer rfx_tracker) { - std::vector output = rfx_model->Predict(*rfx_dataset, *rfx_tracker); - return output; + std::vector output = rfx_model->Predict(*rfx_dataset, *rfx_tracker); + return output; } [[cpp11::register]] -cpp11::writable::doubles rfx_container_predict_cpp(cpp11::external_pointer rfx_container, - cpp11::external_pointer rfx_dataset, +cpp11::writable::doubles rfx_container_predict_cpp(cpp11::external_pointer rfx_container, + cpp11::external_pointer rfx_dataset, cpp11::external_pointer label_mapper) { - int num_observations = rfx_dataset->NumObservations(); - int num_samples = rfx_container->NumSamples(); - std::vector output(num_observations*num_samples); - rfx_container->Predict(*rfx_dataset, *label_mapper, output); - return output; + int num_observations = rfx_dataset->NumObservations(); + int num_samples = rfx_container->NumSamples(); + std::vector output(num_observations * num_samples); + rfx_container->Predict(*rfx_dataset, *label_mapper, output); + return output; } [[cpp11::register]] int rfx_container_num_samples_cpp(cpp11::external_pointer rfx_container) { - return rfx_container->NumSamples(); + return rfx_container->NumSamples(); } [[cpp11::register]] int rfx_container_num_components_cpp(cpp11::external_pointer rfx_container) { - return rfx_container->NumComponents(); + return rfx_container->NumComponents(); } [[cpp11::register]] int rfx_container_num_groups_cpp(cpp11::external_pointer rfx_container) { - return rfx_container->NumGroups(); + return rfx_container->NumGroups(); } [[cpp11::register]] void rfx_container_delete_sample_cpp(cpp11::external_pointer rfx_container, int sample_num) { - rfx_container->DeleteSample(sample_num); + rfx_container->DeleteSample(sample_num); } [[cpp11::register]] void rfx_model_set_working_parameter_cpp(cpp11::external_pointer rfx_model, cpp11::doubles working_param_init) { - Eigen::VectorXd working_param_eigen(working_param_init.size()); - for (int i = 0; i < working_param_init.size(); i++) { - working_param_eigen(i) = working_param_init.at(i); - } - rfx_model->SetWorkingParameter(working_param_eigen); + Eigen::VectorXd working_param_eigen(working_param_init.size()); + for (int i = 0; i < working_param_init.size(); i++) { + working_param_eigen(i) = working_param_init.at(i); + } + rfx_model->SetWorkingParameter(working_param_eigen); } [[cpp11::register]] void rfx_model_set_group_parameters_cpp(cpp11::external_pointer rfx_model, cpp11::doubles_matrix<> group_params_init) { - Eigen::MatrixXd group_params_eigen(group_params_init.nrow(), group_params_init.ncol()); - for (int i = 0; i < group_params_init.nrow(); i++) { - for (int j = 0; j < group_params_init.ncol(); j++) { - group_params_eigen(i,j) = group_params_init(i,j); - } + Eigen::MatrixXd group_params_eigen(group_params_init.nrow(), group_params_init.ncol()); + for (int i = 0; i < group_params_init.nrow(); i++) { + for (int j = 0; j < group_params_init.ncol(); j++) { + group_params_eigen(i, j) = group_params_init(i, j); } - rfx_model->SetGroupParameters(group_params_eigen); + } + rfx_model->SetGroupParameters(group_params_eigen); } [[cpp11::register]] void rfx_model_set_working_parameter_covariance_cpp(cpp11::external_pointer rfx_model, cpp11::doubles_matrix<> working_param_cov_init) { - Eigen::MatrixXd working_param_cov_eigen(working_param_cov_init.nrow(), working_param_cov_init.ncol()); - for (int i = 0; i < working_param_cov_init.nrow(); i++) { - for (int j = 0; j < working_param_cov_init.ncol(); j++) { - working_param_cov_eigen(i,j) = working_param_cov_init(i,j); - } + Eigen::MatrixXd working_param_cov_eigen(working_param_cov_init.nrow(), working_param_cov_init.ncol()); + for (int i = 0; i < working_param_cov_init.nrow(); i++) { + for (int j = 0; j < working_param_cov_init.ncol(); j++) { + working_param_cov_eigen(i, j) = working_param_cov_init(i, j); } - rfx_model->SetWorkingParameterCovariance(working_param_cov_eigen); + } + rfx_model->SetWorkingParameterCovariance(working_param_cov_eigen); } [[cpp11::register]] void rfx_model_set_group_parameter_covariance_cpp(cpp11::external_pointer rfx_model, cpp11::doubles_matrix<> group_param_cov_init) { - Eigen::MatrixXd group_param_cov_eigen(group_param_cov_init.nrow(), group_param_cov_init.ncol()); - for (int i = 0; i < group_param_cov_init.nrow(); i++) { - for (int j = 0; j < group_param_cov_init.ncol(); j++) { - group_param_cov_eigen(i,j) = group_param_cov_init(i,j); - } + Eigen::MatrixXd group_param_cov_eigen(group_param_cov_init.nrow(), group_param_cov_init.ncol()); + for (int i = 0; i < group_param_cov_init.nrow(); i++) { + for (int j = 0; j < group_param_cov_init.ncol(); j++) { + group_param_cov_eigen(i, j) = group_param_cov_init(i, j); } - rfx_model->SetGroupParameterCovariance(group_param_cov_eigen); + } + rfx_model->SetGroupParameterCovariance(group_param_cov_eigen); } [[cpp11::register]] void rfx_model_set_variance_prior_shape_cpp(cpp11::external_pointer rfx_model, double shape) { - rfx_model->SetVariancePriorShape(shape); + rfx_model->SetVariancePriorShape(shape); } [[cpp11::register]] void rfx_model_set_variance_prior_scale_cpp(cpp11::external_pointer rfx_model, double scale) { - rfx_model->SetVariancePriorScale(scale); + rfx_model->SetVariancePriorScale(scale); } [[cpp11::register]] cpp11::writable::integers rfx_tracker_get_unique_group_ids_cpp(cpp11::external_pointer rfx_tracker) { - std::vector output = rfx_tracker->GetUniqueGroupIds(); - return output; + std::vector output = rfx_tracker->GetUniqueGroupIds(); + return output; } [[cpp11::register]] cpp11::writable::doubles rfx_container_get_beta_cpp(cpp11::external_pointer rfx_container_ptr) { - return rfx_container_ptr->GetBeta(); + return rfx_container_ptr->GetBeta(); } [[cpp11::register]] cpp11::writable::doubles rfx_container_get_alpha_cpp(cpp11::external_pointer rfx_container_ptr) { - return rfx_container_ptr->GetAlpha(); + return rfx_container_ptr->GetAlpha(); } [[cpp11::register]] cpp11::writable::doubles rfx_container_get_xi_cpp(cpp11::external_pointer rfx_container_ptr) { - return rfx_container_ptr->GetXi(); + return rfx_container_ptr->GetXi(); } [[cpp11::register]] cpp11::writable::doubles rfx_container_get_sigma_cpp(cpp11::external_pointer rfx_container_ptr) { - return rfx_container_ptr->GetSigma(); + return rfx_container_ptr->GetSigma(); } [[cpp11::register]] cpp11::list rfx_label_mapper_to_list_cpp(cpp11::external_pointer label_mapper_ptr) { - cpp11::writable::integers keys; - cpp11::writable::integers values; - std::map label_map = label_mapper_ptr->Map(); - for (const auto& [key, value] : label_map) { - keys.push_back(key); - values.push_back(value); - } - - cpp11::writable::list output; - output.push_back(keys); - output.push_back(values); - return output; + cpp11::writable::integers keys; + cpp11::writable::integers values; + std::map label_map = label_mapper_ptr->Map(); + for (const auto& [key, value] : label_map) { + keys.push_back(key); + values.push_back(value); + } + + cpp11::writable::list output; + output.push_back(keys); + output.push_back(values); + return output; } [[cpp11::register]] -void reset_rfx_model_cpp(cpp11::external_pointer rfx_model, - cpp11::external_pointer rfx_container, +void reset_rfx_model_cpp(cpp11::external_pointer rfx_model, + cpp11::external_pointer rfx_container, int sample_num) { - // Reset the RFX model from a previous sample - rfx_model->ResetFromSample(*rfx_container, sample_num); + // Reset the RFX model from a previous sample + rfx_model->ResetFromSample(*rfx_container, sample_num); } [[cpp11::register]] -void reset_rfx_tracker_cpp(cpp11::external_pointer tracker, - cpp11::external_pointer dataset, - cpp11::external_pointer residual, +void reset_rfx_tracker_cpp(cpp11::external_pointer tracker, + cpp11::external_pointer dataset, + cpp11::external_pointer residual, cpp11::external_pointer rfx_model) { - // Reset the RFX tracker from a previous sample - tracker->ResetFromSample(*rfx_model, *dataset, *residual); + // Reset the RFX tracker from a previous sample + tracker->ResetFromSample(*rfx_model, *dataset, *residual); } [[cpp11::register]] -void root_reset_rfx_tracker_cpp(cpp11::external_pointer tracker, - cpp11::external_pointer dataset, - cpp11::external_pointer residual, +void root_reset_rfx_tracker_cpp(cpp11::external_pointer tracker, + cpp11::external_pointer dataset, + cpp11::external_pointer residual, cpp11::external_pointer rfx_model) { - // Reset the RFX tracker from root - tracker->RootReset(*rfx_model, *dataset, *residual); + // Reset the RFX tracker from root + tracker->RootReset(*rfx_model, *dataset, *residual); } diff --git a/src/R_samples.cpp b/src/R_samples.cpp new file mode 100644 index 00000000..37883bd5 --- /dev/null +++ b/src/R_samples.cpp @@ -0,0 +1,468 @@ +#include +#include "stochtree/random_effects.h" +#include "stochtree_types.h" +#include +#include +#include +#include +#include + +// Deep-copy a forest container sample-by-sample (so the caller's EXTPTR keeps its own copy). +// File-local (static) helper, matching the convention used elsewhere (e.g. src/stochtree_types.h). +static std::unique_ptr clone_forest_container(StochTree::ForestContainer* src) { + auto copy_ptr = std::make_unique( + src->NumTrees(), src->OutputDimension(), src->IsLeafConstant(), src->IsExponentiated()); + for (int i = 0; i < src->NumSamples(); i++) copy_ptr->AddSample(*src->GetEnsemble(i)); + return copy_ptr; +} + +// Deep-copy a random effects container +static std::unique_ptr clone_rfx_container(StochTree::RandomEffectsContainer* src) { + auto copy_ptr = std::make_unique(); + copy_ptr->CopyFromOther(*src); + return copy_ptr; +} + +// Deep-copy a random effects container +static std::unique_ptr clone_label_mapper(StochTree::LabelMapper* src) { + auto copy_ptr = std::make_unique(); + copy_ptr->CopyFromOther(*src); + return copy_ptr; +} + +// Convert std::vector to cpp11::writable::doubles (for returning samples as an R-native vector). +static cpp11::writable::doubles vec_to_doubles(const std::vector& v) { + cpp11::writable::doubles out(static_cast(v.size())); + std::copy(v.begin(), v.end(), out.begin()); + return out; +} + +// Convert std::vector to cpp11::writable::doubles (for returning samples as an R-native vector). +static cpp11::writable::doubles vec_to_doubles_reshape(const std::vector& v, std::initializer_list dims) { + auto out = vec_to_doubles(v); + if (!v.empty()) { + out.attr("dim") = cpp11::writable::integers(dims); + } + return out; +} + +// -------------------------------- BARTSamples -------------------------------- + +[[cpp11::register]] +cpp11::external_pointer bart_samples_cpp() { + auto samples = std::make_unique(); + return cpp11::external_pointer(samples.release()); +} + +[[cpp11::register]] +cpp11::external_pointer bart_samples_from_json_cpp(cpp11::external_pointer json) { + auto samples = std::make_unique(); + samples->FromJson(*json); + return cpp11::external_pointer(samples.release()); +} + +[[cpp11::register]] +void append_bart_samples_to_json_cpp(cpp11::external_pointer samples, cpp11::external_pointer json) { + samples->AppendToJson(*json); +} + +[[cpp11::register]] +int bart_samples_num_samples_cpp(cpp11::external_pointer samples) { + return samples->num_samples; +} + +[[cpp11::register]] +double bart_samples_y_bar_cpp(cpp11::external_pointer samples) { + return samples->y_bar; +} + +[[cpp11::register]] +double bart_samples_y_std_cpp(cpp11::external_pointer samples) { + return samples->y_std; +} + +[[cpp11::register]] +bool bart_samples_has_yhat_train_cpp(cpp11::external_pointer samples) { + return samples->rfx_predictions_train.empty() == false || samples->mean_forest_predictions_train.empty() == false; +} + +[[cpp11::register]] +bool bart_samples_has_mean_forest_predictions_train_cpp(cpp11::external_pointer samples) { + return samples->mean_forest_predictions_train.empty() == false; +} + +[[cpp11::register]] +bool bart_samples_has_variance_forest_predictions_train_cpp(cpp11::external_pointer samples) { + return samples->variance_forest_predictions_train.empty() == false; +} + +[[cpp11::register]] +bool bart_samples_has_rfx_predictions_train_cpp(cpp11::external_pointer samples) { + return samples->rfx_predictions_train.empty() == false; +} + +[[cpp11::register]] +bool bart_samples_has_yhat_test_cpp(cpp11::external_pointer samples) { + return samples->rfx_predictions_test.empty() == false || samples->mean_forest_predictions_test.empty() == false; +} + +[[cpp11::register]] +bool bart_samples_has_mean_forest_predictions_test_cpp(cpp11::external_pointer samples) { + return samples->mean_forest_predictions_test.empty() == false; +} + +[[cpp11::register]] +bool bart_samples_has_variance_forest_predictions_test_cpp(cpp11::external_pointer samples) { + return samples->variance_forest_predictions_test.empty() == false; +} + +[[cpp11::register]] +bool bart_samples_has_rfx_predictions_test_cpp(cpp11::external_pointer samples) { + return samples->rfx_predictions_test.empty() == false; +} + +[[cpp11::register]] +bool bart_samples_has_cloglog_cutpoint_samples_cpp(cpp11::external_pointer samples) { + return samples->cloglog_cutpoint_samples.empty() == false; +} + +[[cpp11::register]] +bool bart_samples_has_global_var_samples_cpp(cpp11::external_pointer samples) { + return samples->global_error_variance_samples.empty() == false; +} + +[[cpp11::register]] +bool bart_samples_has_leaf_scale_samples_cpp(cpp11::external_pointer samples) { + return samples->leaf_scale_samples.empty() == false; +} + +[[cpp11::register]] +bool bart_samples_has_mean_forest_cpp(cpp11::external_pointer samples) { + return samples->mean_forests != nullptr; +} + +[[cpp11::register]] +bool bart_samples_has_rfx_cpp(cpp11::external_pointer samples) { + return samples->rfx_container != nullptr; +} + +[[cpp11::register]] +bool bart_samples_has_variance_forest_cpp(cpp11::external_pointer samples) { + return samples->variance_forests != nullptr; +} + +[[cpp11::register]] +cpp11::writable::doubles bart_samples_mean_forest_predictions_train_cpp(cpp11::external_pointer samples) { + return vec_to_doubles_reshape(samples->mean_forest_predictions_train, {samples->num_train, samples->num_samples}); +} + +[[cpp11::register]] +cpp11::writable::doubles bart_samples_variance_forest_predictions_train_cpp(cpp11::external_pointer samples) { + return vec_to_doubles_reshape(samples->variance_forest_predictions_train, {samples->num_train, samples->num_samples}); +} + +[[cpp11::register]] +cpp11::writable::doubles bart_samples_rfx_predictions_train_cpp(cpp11::external_pointer samples) { + return vec_to_doubles_reshape(samples->rfx_predictions_train, {samples->num_train, samples->num_samples}); +} + +[[cpp11::register]] +cpp11::writable::doubles bart_samples_yhat_train_cpp(cpp11::external_pointer samples) { + return vec_to_doubles_reshape(samples->OutcomePredictionsTrain(), {samples->num_train, samples->num_samples}); +} + +[[cpp11::register]] +cpp11::writable::doubles bart_samples_mean_forest_predictions_test_cpp(cpp11::external_pointer samples) { + return vec_to_doubles_reshape(samples->mean_forest_predictions_test, {samples->num_test, samples->num_samples}); +} + +[[cpp11::register]] +cpp11::writable::doubles bart_samples_variance_forest_predictions_test_cpp(cpp11::external_pointer samples) { + return vec_to_doubles_reshape(samples->variance_forest_predictions_test, {samples->num_test, samples->num_samples}); +} + +[[cpp11::register]] +cpp11::writable::doubles bart_samples_rfx_predictions_test_cpp(cpp11::external_pointer samples) { + return vec_to_doubles_reshape(samples->rfx_predictions_test, {samples->num_test, samples->num_samples}); +} + +[[cpp11::register]] +cpp11::writable::doubles bart_samples_yhat_test_cpp(cpp11::external_pointer samples) { + return vec_to_doubles_reshape(samples->OutcomePredictionsTest(), {samples->num_test, samples->num_samples}); +} + +[[cpp11::register]] +cpp11::writable::doubles bart_samples_global_var_samples_cpp(cpp11::external_pointer samples) { + return vec_to_doubles(samples->global_error_variance_samples); +} + +[[cpp11::register]] +cpp11::writable::doubles bart_samples_leaf_scale_samples_cpp(cpp11::external_pointer samples) { + return vec_to_doubles(samples->leaf_scale_samples); +} + +[[cpp11::register]] +cpp11::writable::doubles bart_samples_cloglog_cutpoint_samples_cpp(cpp11::external_pointer samples) { + int ns = samples->num_samples; + int len = static_cast(samples->cloglog_cutpoint_samples.size()); + int num_cutpoints = (ns > 0) ? len / ns : 0; + return vec_to_doubles_reshape(samples->cloglog_cutpoint_samples, {num_cutpoints, ns}); +} + +// Materialize a standalone deep copy of the mean forest container. +[[cpp11::register]] +cpp11::external_pointer bart_samples_materialize_mean_forest_cpp(cpp11::external_pointer samples) { + auto copy = clone_forest_container(samples->mean_forests.get()); + return cpp11::external_pointer(copy.release()); +} + +// Materialize a standalone deep copy of the variance forest container. +[[cpp11::register]] +cpp11::external_pointer bart_samples_materialize_variance_forest_cpp(cpp11::external_pointer samples) { + auto copy = clone_forest_container(samples->variance_forests.get()); + return cpp11::external_pointer(copy.release()); +} + +// Materialize a standalone deep copy of the random effects container. +[[cpp11::register]] +cpp11::external_pointer bart_samples_materialize_rfx_container_cpp(cpp11::external_pointer samples) { + auto copy = clone_rfx_container(samples->rfx_container.get()); + return cpp11::external_pointer(copy.release()); +} + +// Materialize a standalone deep copy of the random effects label mapper. +[[cpp11::register]] +cpp11::external_pointer bart_samples_materialize_rfx_label_mapper_cpp(cpp11::external_pointer samples) { + auto copy = clone_label_mapper(samples->rfx_label_mapper.get()); + return cpp11::external_pointer(copy.release()); +} + +// Borrowed (non-owning) pointers to the samples-owned forest containers, for read-through predict. +// The returned external_pointer does NOT own or finalize the container -- it aliases the one owned +// by `samples`, so it must not outlive it (predict uses it transiently within a single call). +[[cpp11::register]] +cpp11::external_pointer bart_samples_mean_forest_ptr_cpp(cpp11::external_pointer samples) { + return cpp11::external_pointer( + samples->mean_forests.get(), /*use_deleter=*/false, /*finalize_on_exit=*/false); +} + +[[cpp11::register]] +cpp11::external_pointer bart_samples_variance_forest_ptr_cpp(cpp11::external_pointer samples) { + return cpp11::external_pointer( + samples->variance_forests.get(), /*use_deleter=*/false, /*finalize_on_exit=*/false); +} + +[[cpp11::register]] +void bart_samples_merge_cpp(cpp11::external_pointer samples, + cpp11::external_pointer other) { + samples->Merge(*other); +} + +// -------------------------------- BCFSamples -------------------------------- + +[[cpp11::register]] +cpp11::external_pointer bcf_samples_cpp() { + auto samples = std::make_unique(); + return cpp11::external_pointer(samples.release()); +} + +[[cpp11::register]] +cpp11::external_pointer bcf_samples_from_json_cpp(cpp11::external_pointer json) { + auto samples = std::make_unique(); + samples->FromJson(*json); + return cpp11::external_pointer(samples.release()); +} + +[[cpp11::register]] +void append_bcf_samples_to_json_cpp(cpp11::external_pointer samples, cpp11::external_pointer json) { + samples->AppendToJson(*json); +} + +[[cpp11::register]] +int bcf_samples_num_samples_cpp(cpp11::external_pointer samples) { + return samples->num_samples; +} + +[[cpp11::register]] +int bcf_samples_treatment_dim_cpp(cpp11::external_pointer samples) { + return samples->treatment_dim; +} + +[[cpp11::register]] +double bcf_samples_y_bar_cpp(cpp11::external_pointer samples) { + return samples->y_bar; +} + +[[cpp11::register]] +double bcf_samples_y_std_cpp(cpp11::external_pointer samples) { + return samples->y_std; +} + +[[cpp11::register]] +bool bcf_samples_has_mu_forest_cpp(cpp11::external_pointer samples) { + return samples->mu_forests != nullptr; +} + +[[cpp11::register]] +bool bcf_samples_has_tau_forest_cpp(cpp11::external_pointer samples) { + return samples->tau_forests != nullptr; +} + +[[cpp11::register]] +bool bcf_samples_has_variance_forest_cpp(cpp11::external_pointer samples) { + return samples->variance_forests != nullptr; +} + +[[cpp11::register]] +bool bcf_samples_has_rfx_cpp(cpp11::external_pointer samples) { + return samples->rfx_container != nullptr; +} + +[[cpp11::register]] +cpp11::writable::doubles bcf_samples_global_var_samples_cpp(cpp11::external_pointer samples) { + return vec_to_doubles(samples->global_error_variance_samples); +} + +[[cpp11::register]] +cpp11::writable::doubles bcf_samples_leaf_scale_mu_samples_cpp(cpp11::external_pointer samples) { + return vec_to_doubles(samples->leaf_scale_mu_samples); +} + +[[cpp11::register]] +cpp11::writable::doubles bcf_samples_leaf_scale_tau_samples_cpp(cpp11::external_pointer samples) { + return vec_to_doubles(samples->leaf_scale_tau_samples); +} + +[[cpp11::register]] +cpp11::writable::doubles bcf_samples_tau_0_samples_cpp(cpp11::external_pointer samples) { + if (samples->treatment_dim <= 1) { + return vec_to_doubles(samples->tau_0_samples); + } else { + return vec_to_doubles_reshape(samples->tau_0_samples, {samples->treatment_dim, samples->num_samples}); + } +} + +[[cpp11::register]] +cpp11::writable::doubles bcf_samples_b0_samples_cpp(cpp11::external_pointer samples) { + return vec_to_doubles(samples->b0_samples); +} + +[[cpp11::register]] +cpp11::writable::doubles bcf_samples_b1_samples_cpp(cpp11::external_pointer samples) { + return vec_to_doubles(samples->b1_samples); +} + +[[cpp11::register]] +cpp11::writable::doubles bcf_samples_yhat_train_cpp(cpp11::external_pointer samples) { + return vec_to_doubles_reshape(samples->y_hat_train, {samples->num_train, samples->num_samples}); +} + +[[cpp11::register]] +cpp11::writable::doubles bcf_samples_mu_forest_predictions_train_cpp(cpp11::external_pointer samples) { + return vec_to_doubles_reshape(samples->mu_forest_predictions_train, {samples->num_train, samples->num_samples}); +} + +[[cpp11::register]] +cpp11::writable::doubles bcf_samples_tau_forest_predictions_train_cpp(cpp11::external_pointer samples) { + if (samples->treatment_dim <= 1) { + return vec_to_doubles_reshape(samples->tau_forest_predictions_train, {samples->num_train, samples->num_samples}); + } else { + return vec_to_doubles_reshape(samples->tau_forest_predictions_train, {samples->num_train, samples->treatment_dim, samples->num_samples}); + } +} + +[[cpp11::register]] +cpp11::writable::doubles bcf_samples_variance_forest_predictions_train_cpp(cpp11::external_pointer samples) { + return vec_to_doubles_reshape(samples->variance_forest_predictions_train, {samples->num_train, samples->num_samples}); +} + +[[cpp11::register]] +cpp11::writable::doubles bcf_samples_rfx_predictions_train_cpp(cpp11::external_pointer samples) { + return vec_to_doubles_reshape(samples->rfx_predictions_train, {samples->num_train, samples->num_samples}); +} + +[[cpp11::register]] +cpp11::writable::doubles bcf_samples_yhat_test_cpp(cpp11::external_pointer samples) { + return vec_to_doubles_reshape(samples->y_hat_test, {samples->num_test, samples->num_samples}); +} + +[[cpp11::register]] +cpp11::writable::doubles bcf_samples_mu_forest_predictions_test_cpp(cpp11::external_pointer samples) { + return vec_to_doubles_reshape(samples->mu_forest_predictions_test, {samples->num_test, samples->num_samples}); +} + +[[cpp11::register]] +cpp11::writable::doubles bcf_samples_tau_forest_predictions_test_cpp(cpp11::external_pointer samples) { + if (samples->treatment_dim <= 1) { + return vec_to_doubles_reshape(samples->tau_forest_predictions_test, {samples->num_test, samples->num_samples}); + } else { + return vec_to_doubles_reshape(samples->tau_forest_predictions_test, {samples->num_test, samples->treatment_dim, samples->num_samples}); + } +} + +[[cpp11::register]] +cpp11::writable::doubles bcf_samples_variance_forest_predictions_test_cpp(cpp11::external_pointer samples) { + return vec_to_doubles_reshape(samples->variance_forest_predictions_test, {samples->num_test, samples->num_samples}); +} + +[[cpp11::register]] +cpp11::writable::doubles bcf_samples_rfx_predictions_test_cpp(cpp11::external_pointer samples) { + return vec_to_doubles_reshape(samples->rfx_predictions_test, {samples->num_test, samples->num_samples}); +} + +[[cpp11::register]] +cpp11::external_pointer bcf_samples_materialize_mu_forest_cpp(cpp11::external_pointer samples) { + auto copy = clone_forest_container(samples->mu_forests.get()); + return cpp11::external_pointer(copy.release()); +} + +[[cpp11::register]] +cpp11::external_pointer bcf_samples_materialize_tau_forest_cpp(cpp11::external_pointer samples) { + auto copy = clone_forest_container(samples->tau_forests.get()); + return cpp11::external_pointer(copy.release()); +} + +[[cpp11::register]] +cpp11::external_pointer bcf_samples_materialize_variance_forest_cpp(cpp11::external_pointer samples) { + auto copy = clone_forest_container(samples->variance_forests.get()); + return cpp11::external_pointer(copy.release()); +} + +// Materialize a standalone deep copy of the random effects container. +[[cpp11::register]] +cpp11::external_pointer bcf_samples_materialize_rfx_container_cpp(cpp11::external_pointer samples) { + auto copy = clone_rfx_container(samples->rfx_container.get()); + return cpp11::external_pointer(copy.release()); +} + +// Materialize a standalone deep copy of the random effects label mapper. +[[cpp11::register]] +cpp11::external_pointer bcf_samples_materialize_rfx_label_mapper_cpp(cpp11::external_pointer samples) { + auto copy = clone_label_mapper(samples->rfx_label_mapper.get()); + return cpp11::external_pointer(copy.release()); +} + +// Borrowed (non-owning) pointers to the samples-owned forest containers, for read-through predict. +[[cpp11::register]] +cpp11::external_pointer bcf_samples_mu_forest_ptr_cpp(cpp11::external_pointer samples) { + return cpp11::external_pointer( + samples->mu_forests.get(), /*use_deleter=*/false, /*finalize_on_exit=*/false); +} + +[[cpp11::register]] +cpp11::external_pointer bcf_samples_tau_forest_ptr_cpp(cpp11::external_pointer samples) { + return cpp11::external_pointer( + samples->tau_forests.get(), /*use_deleter=*/false, /*finalize_on_exit=*/false); +} + +[[cpp11::register]] +cpp11::external_pointer bcf_samples_variance_forest_ptr_cpp(cpp11::external_pointer samples) { + return cpp11::external_pointer( + samples->variance_forests.get(), /*use_deleter=*/false, /*finalize_on_exit=*/false); +} + +[[cpp11::register]] +void bcf_samples_merge_cpp(cpp11::external_pointer samples, + cpp11::external_pointer other) { + samples->Merge(*other); +} diff --git a/src/R_utils.cpp b/src/R_utils.cpp index 038023cb..8df37da6 100644 --- a/src/R_utils.cpp +++ b/src/R_utils.cpp @@ -3,33 +3,33 @@ [[cpp11::register]] double sum_cpp(cpp11::doubles x) { - double output = 0.0; - for (int i = 0; i < x.size(); i++) { - output += x[i]; - } - return output; + double output = 0.0; + for (int i = 0; i < x.size(); i++) { + output += x[i]; + } + return output; } [[cpp11::register]] double mean_cpp(cpp11::doubles x) { - double output = 0.0; - for (int i = 0; i < x.size(); i++) { - output += x[i]; - } - return output / x.size(); + double output = 0.0; + for (int i = 0; i < x.size(); i++) { + output += x[i]; + } + return output / x.size(); } [[cpp11::register]] double var_cpp(cpp11::doubles x) { - double mean = mean_cpp(x); - double output = 0.0; - for (int i = 0; i < x.size(); i++) { - output += (x[i] - mean) * (x[i] - mean); - } - return output / (x.size() - 1); + double mean = mean_cpp(x); + double output = 0.0; + for (int i = 0; i < x.size(); i++) { + output += (x[i] - mean) * (x[i] - mean); + } + return output / (x.size() - 1); } [[cpp11::register]] double sd_cpp(cpp11::doubles x) { - return std::sqrt(var_cpp(x)); + return std::sqrt(var_cpp(x)); } diff --git a/src/bart_sampler.cpp b/src/bart_sampler.cpp new file mode 100644 index 00000000..2dfa5ecc --- /dev/null +++ b/src/bart_sampler.cpp @@ -0,0 +1,847 @@ +/*! Copyright (c) 2026 by stochtree authors */ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "stochtree/data.h" +#include "stochtree/random_effects.h" + +namespace StochTree { + +BARTSampler::BARTSampler(BARTSamples& samples, BARTConfig& config, BARTData& data, bool continuation) : config_{config}, data_{data}, mean_leaf_model_(GaussianConstantLeafModel(0.0)), variance_leaf_model_(0.0, 0.0) { + InitializeState(samples, continuation); +} + +void BARTSampler::InitializeState(BARTSamples& samples, bool continuation) { + // Validate y_train values match the expected support for discrete link functions + if (config_.link_function == LinkFunction::Probit) { + for (int i = 0; i < data_.n_train; i++) { + if (std::floor(data_.y_train[i]) != data_.y_train[i]) { + Log::Fatal("Outcomes must be integers for probit link model"); + } + if (data_.y_train[i] != 0.0 && data_.y_train[i] != 1.0) { + Log::Fatal("Outcomes must be 0 or 1 for probit link model"); + } + } + } else if (config_.link_function == LinkFunction::Cloglog) { + for (int i = 0; i < data_.n_train; i++) { + if (std::floor(data_.y_train[i]) != data_.y_train[i]) { + Log::Fatal("Outcomes must be integers for cloglog link model"); + } + if (config_.outcome_type == OutcomeType::Binary && (data_.y_train[i] < 0.0 || data_.y_train[i] > 1.0)) { + Log::Fatal("Outcomes must be between 0 and 1 for binary cloglog link model"); + } else if (config_.outcome_type == OutcomeType::Ordinal && (data_.y_train[i] < 0.0 || data_.y_train[i] >= config_.num_classes_cloglog)) { + Log::Fatal("Outcomes must be integers between 0 and num_classes_cloglog - 1 for ordinal cloglog link model"); + } + } + if (config_.outcome_type != OutcomeType::Binary && config_.outcome_type != OutcomeType::Ordinal) { + Log::Fatal("Cloglog link function is only supported for binary and ordinal outcomes"); + } + if (config_.num_classes_cloglog <= 1) { + Log::Fatal("num_classes_cloglog must be greater than 1 for cloglog link function"); + } + if (config_.has_random_effects) { + Log::Fatal("Random effects are not currently supported with the cloglog link function"); + } + } + + // Continued sampling currently supports only the basic Gaussian mean-forest case; + // reject configurations whose warm-start init is not yet implemented so that + // unsupported continuations fail loudly rather than silently re-initializing from root. + if (continuation) { + if (config_.num_trees_mean <= 0) { + Log::Fatal("Continued sampling requires an existing mean forest"); + } + if (config_.num_trees_variance > 0) { + Log::Fatal("Continued sampling is not yet supported for models with a variance forest"); + } + if (config_.has_random_effects) { + Log::Fatal("Continued sampling is not yet supported for models with random effects"); + } + if (config_.link_function != LinkFunction::Identity) { + Log::Fatal("Continued sampling is not yet supported for probit or cloglog link functions"); + } + } + + // Load data from BARTData object into ForestDataset object + forest_dataset_ = std::make_unique(); + forest_dataset_->AddCovariates(data_.X_train, data_.n_train, data_.p, /*row_major=*/false); + if (data_.basis_train != nullptr) { + forest_dataset_->AddBasis(data_.basis_train, data_.n_train, data_.basis_dim, /*row_major=*/false); + } + if (data_.obs_weights_train != nullptr) { + forest_dataset_->AddVarianceWeights(data_.obs_weights_train, data_.n_train); + } + samples.num_train = data_.n_train; + samples.num_test = data_.n_test; + residual_ = std::make_unique(data_.y_train, data_.n_train); + outcome_raw_ = std::make_unique(data_.y_train, data_.n_train); + if (data_.X_test != nullptr) { + forest_dataset_test_ = std::make_unique(); + forest_dataset_test_->AddCovariates(data_.X_test, data_.n_test, data_.p, /*row_major=*/false); + if (data_.basis_test != nullptr) { + forest_dataset_test_->AddBasis(data_.basis_test, data_.n_test, data_.basis_dim, /*row_major=*/false); + } + if (data_.obs_weights_test != nullptr) { + forest_dataset_test_->AddVarianceWeights(data_.obs_weights_test, data_.n_test); + } + has_test_ = true; + } + + // Precompute outcome mean and variance for standardization and calibration + double y_mean = 0.0, M2 = 0.0, y_mean_prev = 0.0; + for (int i = 0; i < data_.n_train; i++) { + y_mean_prev = y_mean; + y_mean = y_mean_prev + (data_.y_train[i] - y_mean_prev) / (i + 1); + M2 = M2 + (data_.y_train[i] - y_mean_prev) * (data_.y_train[i] - y_mean); + } + double y_var = M2 / data_.n_train; + + // Standardization, calibration, and initialization for mean forests + if (config_.num_trees_mean > 0) { + if (config_.link_function == LinkFunction::Probit) { + // Initialize forests to 0, no scaling, but offset by the probit transform of the mean outcome to improve mixing + samples.y_std = 1.0; + samples.y_bar = norm_inv_cdf(y_mean); + init_val_mean_ = 0.0; + if (config_.mean_leaf_model_type == MeanLeafModelType::GaussianMultivariateRegression) { + init_val_mean_vec_.assign(config_.leaf_dim_mean, 0.0); + } + } else if (config_.link_function == LinkFunction::Cloglog) { + // Initialize forests to 0, no scaling or location shifting of the outcome + // Outcomes are expected to already be 0-indexed by the caller + samples.y_std = 1.0; + samples.y_bar = 0.0; + init_val_mean_ = 0.0; + } else { + if (config_.mean_leaf_model_type == MeanLeafModelType::GaussianConstant) { + // Case 1: Constant leaf + if (config_.standardize_outcome) { + samples.y_bar = y_mean; + samples.y_std = std::sqrt(y_var); + init_val_mean_ = 0.0; + } else { + samples.y_bar = 0.0; + samples.y_std = 1.0; + init_val_mean_ = y_mean; + } + } else if (config_.mean_leaf_model_type == MeanLeafModelType::GaussianUnivariateRegression) { + // Case 2: Univariate leaf regression + if (config_.standardize_outcome) { + samples.y_bar = y_mean; + samples.y_std = std::sqrt(y_var); + } else { + samples.y_bar = 0.0; + samples.y_std = 1.0; + } + // Always map initial leaf value to zero + // Users fitting a univariate leaf regression (with a non-centered basis) should standardize their outcomes + // TODO: consider adding warning in R / Python if univariate regression leaf model is specified without standardization + init_val_mean_ = 0.0; + } else { + // Case 3: Multivariate leaf regression + if (config_.standardize_outcome) { + samples.y_bar = y_mean; + samples.y_std = std::sqrt(y_var); + } else { + samples.y_bar = 0.0; + samples.y_std = 1.0; + } + init_val_mean_ = 0.0; + init_val_mean_vec_.assign(config_.leaf_dim_mean, 0.0); + } + } + + // Calibrate leaf scale and variance model priors + if (config_.sigma2_mean_init < 0.0) { + if (config_.link_function == LinkFunction::Probit) { + config_.sigma2_mean_init = 1.0 / config_.num_trees_mean; + } else { + config_.sigma2_mean_init = y_var / config_.num_trees_mean; + } + } + if (config_.sample_sigma2_leaf_mean) { + if (config_.b_sigma2_mean <= 0.0) { + if (config_.link_function == LinkFunction::Probit) { + config_.b_sigma2_mean = 1.0 / (2 * config_.num_trees_mean); + } else { + config_.b_sigma2_mean = y_var / (2 * config_.num_trees_mean); + } + } + } + + // Initialize leaf model + if (config_.mean_leaf_model_type == MeanLeafModelType::GaussianConstant) { + mean_leaf_model_ = GaussianConstantLeafModel(config_.sigma2_mean_init); + } else if (config_.mean_leaf_model_type == MeanLeafModelType::GaussianUnivariateRegression) { + mean_leaf_model_ = GaussianUnivariateRegressionLeafModel(config_.sigma2_mean_init); + } else if (config_.mean_leaf_model_type == MeanLeafModelType::GaussianMultivariateRegression) { + Eigen::MatrixXd Sigma_0; + if (!config_.sigma2_leaf_mean_matrix.empty()) { + if ((int)config_.sigma2_leaf_mean_matrix.size() != config_.leaf_dim_mean * config_.leaf_dim_mean) { + Log::Fatal("sigma2_leaf_mean_matrix must have leaf_dim_mean * leaf_dim_mean = %d elements, but has %zu", + config_.leaf_dim_mean * config_.leaf_dim_mean, config_.sigma2_leaf_mean_matrix.size()); + } + // Column-major interpretation matches both R and Eigen (python must be reordered before passing to C++) + Sigma_0 = Eigen::Map(config_.sigma2_leaf_mean_matrix.data(), config_.leaf_dim_mean, config_.leaf_dim_mean); + } else { + Sigma_0 = config_.sigma2_mean_init * Eigen::MatrixXd::Identity(config_.leaf_dim_mean, config_.leaf_dim_mean); + } + mean_leaf_model_ = GaussianMultivariateRegressionLeafModel(Sigma_0); + } else if (config_.mean_leaf_model_type == MeanLeafModelType::CloglogOrdinal) { + mean_leaf_model_ = CloglogOrdinalLeafModel(config_.cloglog_leaf_prior_shape, config_.cloglog_leaf_prior_scale); + } else { + Log::Fatal("Unsupported leaf model type for mean forest"); + } + } else { + // Variance-only model (num_trees_mean == 0): no mean forest, but y_bar/y_std must + // still be valid so the residual initialisation below doesn't divide by zero. + if (config_.standardize_outcome) { + samples.y_bar = y_mean; + samples.y_std = std::sqrt(y_var); + } else { + samples.y_bar = 0.0; + samples.y_std = 1.0; + } + } + + // Calibration for variance forests + if (config_.num_trees_variance > 0) { + // NOTE: calibration only works for standardized outcomes + if (config_.shape_variance_forest <= 0.0 || config_.scale_variance_forest <= 0.0) { + if (config_.leaf_prior_calibration_param <= 0.0) { + config_.leaf_prior_calibration_param = 1.5; + } + if (config_.shape_variance_forest <= 0.0) { + config_.shape_variance_forest = config_.num_trees_variance / (config_.leaf_prior_calibration_param * config_.leaf_prior_calibration_param) + 0.5; + } + if (config_.scale_variance_forest <= 0.0) { + config_.scale_variance_forest = config_.num_trees_variance / (config_.leaf_prior_calibration_param * config_.leaf_prior_calibration_param); + } + } + if (config_.variance_forest_leaf_init > 0.0) { + init_val_variance_ = config_.variance_forest_leaf_init; + } else if (config_.standardize_outcome) { + init_val_variance_ = 1.0; + } else { + init_val_variance_ = y_var; + } + } + + // Standardize partial residuals in place; these are updated in each iteration but initialized to standardized outcomes + // Works for: + // 1. Standardized outcomes (since y_bar = mean(y) and y_std = sd(y)) + // 2. Non-standardized outcomes (since y_bar = 0 and y_std = 1, so this just transfers y_train as-is) + // 3. Probit link (since y_bar = norm_inv_cdf(mean(y)) and y_std = 1) + for (int i = 0; i < data_.n_train; i++) residual_->GetData()[i] = (data_.y_train[i] - samples.y_bar) / samples.y_std; + + // Initialize mean forest state (if present) + if (config_.num_trees_mean > 0) { + if (continuation) { + std::visit(MeanForestContinuationInitVisitor{*this, samples}, mean_leaf_model_); + } else { + std::visit(MeanForestInitVisitor{*this, samples}, mean_leaf_model_); + } + } + + // Initialize variance forest state (if present) + if (config_.num_trees_variance > 0) { + variance_leaf_model_ = LogLinearVarianceLeafModel(config_.shape_variance_forest, config_.scale_variance_forest); + variance_forest_ = std::make_unique(config_.num_trees_variance, config_.leaf_dim_variance, config_.leaf_constant_variance, config_.exponentiated_leaf_variance); + samples.variance_forests = std::make_unique(config_.num_trees_variance, config_.leaf_dim_variance, config_.leaf_constant_variance, config_.exponentiated_leaf_variance); + variance_forest_tracker_ = std::make_unique(forest_dataset_->GetCovariates(), config_.feature_types, config_.num_trees_variance, data_.n_train); + tree_prior_variance_ = std::make_unique(config_.alpha_variance, config_.beta_variance, config_.min_samples_leaf_variance, config_.max_depth_variance); + // Leaf values for the log-linear variance model are on the log scale; the ensemble sums + // log(sigma^2_i) contributions, so each tree starts at log(init_val) / num_trees. + variance_forest_->SetLeafValue(std::log(init_val_variance_) / config_.num_trees_variance); + variance_forest_tracker_->UpdatePredictions(variance_forest_.get(), *forest_dataset_.get()); + // UpdateVarModelTree (called inside GFRSampleOneIter / MCMCSampleOneIter) unconditionally + // reads and writes the dataset variance weight slot via VarWeightValue / SetVarWeightValue. + // This slot tracks the cumulative per-observation variance prediction + // (sigma^2_i = exp(sum of tree leaf values)) and is incompatible with case weights, which + // would need to be reapplied after every per-tree update. The R/Python APIs enforce this + // as a hard error; guard here for callers that use BARTSampler directly. + if (forest_dataset_->HasVarWeights()) { + Log::Fatal("observation_weights and a variance forest cannot be used together."); + } + std::vector initial_variance_preds(data_.n_train, init_val_variance_); + forest_dataset_->AddVarianceWeights(initial_variance_preds.data(), data_.n_train); + has_variance_forest_ = true; + } + + // Global error variance model + if (config_.sample_sigma2_global) { + var_model_ = std::make_unique(); + sample_sigma2_global_ = true; + } + + // Leaf scale model + if (config_.sample_sigma2_leaf_mean && config_.num_trees_mean > 0) { + leaf_scale_model_ = std::make_unique(); + sample_sigma2_leaf_ = true; + } + + // Random effects model + if (config_.has_random_effects) { + random_effects_dataset_ = std::make_unique(); + random_effects_dataset_->AddGroupLabels(data_.rfx_group_ids_train, data_.n_train); + if (data_.rfx_basis_train != nullptr) { + random_effects_dataset_->AddBasis(data_.rfx_basis_train, data_.n_train, data_.rfx_basis_dim, /*row_major=*/false); + } else { + if (config_.rfx_model_spec == BARTRFXModelSpec::InterceptOnly) { + // If no basis is provided, add an intercept basis (column of 1s) + // TODO: do we need to do this before we determine rfx_basis_dim and initialize the RFX data structures? + std::vector intercept_basis(data_.n_train, 1.0); + random_effects_dataset_->AddBasis(intercept_basis.data(), data_.n_train, 1, /*row_major=*/false); + // Override rfx_basis_dim to 1 for intercept-only model the basis is a 1-dimensional vector of ones + data_.rfx_basis_dim = 1; + } else { + Log::Fatal("Random effects basis data must be provided for non-intercept-only random effects model"); + } + } + // Tracking data structure for random effects groups + random_effects_tracker_ = std::make_unique(data_.rfx_group_ids_train, data_.n_train); + // Container of random effects samples + samples.rfx_container = std::make_unique(data_.rfx_basis_dim, data_.rfx_num_groups); + // Mapping from RFX labels to 0-indexed group IDs for efficient lookup in the sampler; populated from the RFX dataset group labels + samples.rfx_label_mapper = std::make_unique(random_effects_tracker_->GetLabelMap()); + + // Initialize random effects model object + random_effects_model_ = std::make_unique(data_.rfx_basis_dim, data_.rfx_num_groups); + + // Handle "working" parameter prior mean + Eigen::VectorXd working_parameter_prior_mean; + if (!config_.rfx_working_parameter_mean_prior.empty()) { + if ((int)config_.rfx_working_parameter_mean_prior.size() != data_.rfx_basis_dim) { + Log::Fatal("rfx_working_parameter_mean_prior must have rfx_basis_dim = %d elements, but has %zu", + data_.rfx_basis_dim, config_.rfx_working_parameter_mean_prior.size()); + } + // Column-major interpretation matches both R and Eigen (python must be reordered before passing to C++) + working_parameter_prior_mean = Eigen::Map(config_.rfx_working_parameter_mean_prior.data(), data_.rfx_basis_dim); + } else { + working_parameter_prior_mean = Eigen::VectorXd::Zero(data_.rfx_basis_dim); + } + random_effects_model_->SetWorkingParameter(working_parameter_prior_mean); + + // Handle "group" parameter prior mean + Eigen::MatrixXd group_parameter_prior_mean; + if (!config_.rfx_group_parameter_mean_prior.empty()) { + if ((int)config_.rfx_group_parameter_mean_prior.size() != data_.rfx_basis_dim * data_.rfx_num_groups) { + Log::Fatal("rfx_group_parameter_mean_prior must have rfx_basis_dim * rfx_num_groups = %d elements, but has %zu", + data_.rfx_basis_dim * data_.rfx_num_groups, config_.rfx_group_parameter_mean_prior.size()); + } + // Column-major interpretation matches both R and Eigen (python must be reordered before passing to C++) + group_parameter_prior_mean = Eigen::Map(config_.rfx_group_parameter_mean_prior.data(), data_.rfx_basis_dim, data_.rfx_num_groups); + } else { + group_parameter_prior_mean = Eigen::MatrixXd::Zero(data_.rfx_basis_dim, data_.rfx_num_groups); + } + random_effects_model_->SetGroupParameters(group_parameter_prior_mean); + + // Handle "working" parameter prior covariance + Eigen::MatrixXd working_parameter_prior_cov; + if (!config_.rfx_working_parameter_cov_prior.empty()) { + if ((int)config_.rfx_working_parameter_cov_prior.size() != data_.rfx_basis_dim * data_.rfx_basis_dim) { + Log::Fatal("rfx_working_parameter_cov_prior must have rfx_basis_dim * rfx_basis_dim = %d elements, but has %zu", + data_.rfx_basis_dim * data_.rfx_basis_dim, config_.rfx_working_parameter_cov_prior.size()); + } + // Column-major interpretation matches both R and Eigen (python must be reordered before passing to C++) + working_parameter_prior_cov = Eigen::Map(config_.rfx_working_parameter_cov_prior.data(), data_.rfx_basis_dim, data_.rfx_basis_dim); + } else { + working_parameter_prior_cov = Eigen::MatrixXd::Identity(data_.rfx_basis_dim, data_.rfx_basis_dim); + } + random_effects_model_->SetWorkingParameterCovariance(working_parameter_prior_cov); + + // Handle "group" parameter prior covariance + Eigen::MatrixXd group_parameter_prior_cov; + if (!config_.rfx_group_parameter_cov_prior.empty()) { + if ((int)config_.rfx_group_parameter_cov_prior.size() != data_.rfx_basis_dim * data_.rfx_basis_dim) { + Log::Fatal("rfx_group_parameter_cov_prior must have rfx_basis_dim * rfx_basis_dim = %d elements, but has %zu", + data_.rfx_basis_dim * data_.rfx_basis_dim, config_.rfx_group_parameter_cov_prior.size()); + } + // Column-major interpretation matches both R and Eigen (python must be reordered before passing to C++) + group_parameter_prior_cov = Eigen::Map(config_.rfx_group_parameter_cov_prior.data(), data_.rfx_basis_dim, data_.rfx_basis_dim); + } else { + group_parameter_prior_cov = Eigen::MatrixXd::Identity(data_.rfx_basis_dim, data_.rfx_basis_dim); + } + random_effects_model_->SetGroupParameterCovariance(group_parameter_prior_cov); + + // Handle variance model priors + if (config_.rfx_variance_prior_shape <= 0.0) { + config_.rfx_variance_prior_shape = 1.0; + } + if (config_.rfx_variance_prior_scale <= 0.0) { + config_.rfx_variance_prior_scale = 1.0; + } + random_effects_model_->SetVariancePriorShape(config_.rfx_variance_prior_shape); + random_effects_model_->SetVariancePriorScale(config_.rfx_variance_prior_scale); + + // Set has_random_effects_ flag to true so that the sampler will perform random effects updates at each iteration + has_random_effects_ = true; + } + + // RNG + rng_ = std::mt19937(config_.random_seed >= 0 ? config_.random_seed : std::random_device{}()); + + // Cloglog state + if (config_.link_function == LinkFunction::Cloglog) { + // Initialize the ordinal sampler + ordinal_sampler_ = std::make_unique(); + // Latent variable (Z in Alam et al (2025) notation) + forest_dataset_->AddAuxiliaryDimension(data_.n_train); + // Forest predictions (eta in Alam et al (2025) notation) + forest_dataset_->AddAuxiliaryDimension(data_.n_train); + // Log-scale non-cumulative cutpoint (gamma in Alam et al (2025) notation) + forest_dataset_->AddAuxiliaryDimension(config_.num_classes_cloglog - 1); + // Exponentiated cumulative cutpoints (exp(c_k) in Alam et al (2025) notation) + // This auxiliary series is designed so that the element stored at position `i` + // corresponds to the sum of all exponentiated gamma_j values for j < i. + // It has cloglog_num_categories elements instead of cloglog_num_categories - 1 because + // even the largest categorical index has a valid value of sum_{j < i} exp(gamma_j) + forest_dataset_->AddAuxiliaryDimension(config_.num_classes_cloglog); + + // Set initial values for auxiliary data + // Initialize latent variables to zero (slot 0) + for (int i = 0; i < data_.n_train; i++) { + forest_dataset_->SetAuxiliaryDataValue(0, i, 0.0); + } + // Initialize forest predictions to zero (slot 1) + for (int i = 0; i < data_.n_train; i++) { + forest_dataset_->SetAuxiliaryDataValue(1, i, 0.0); + } + // Initialize log-scale cutpoints to 0 + for (int i = 0; i < config_.num_classes_cloglog - 1; i++) { + forest_dataset_->SetAuxiliaryDataValue(2, i, 0.0); + } + // Convert to cumulative exponentiated cutpoints directly in C++ + ordinal_sampler_->UpdateCumulativeExpSums(*forest_dataset_); + } + + // Other internal model state + if (continuation) { + // Warm-start the scalar variance state from the last retained sample. + // Within the sampler lifecycle both arrays hold STANDARDIZED values + // (postprocess_samples applies the y_std^2 factor to global variance only at + // the very end); the continuation binding supplies standardized history, so + // read the last value directly here. + if (sample_sigma2_global_ && !samples.global_error_variance_samples.empty()) { + global_variance_ = samples.global_error_variance_samples.back(); + } else { + global_variance_ = config_.sigma2_global_init; + } + if (sample_sigma2_leaf_ && !samples.leaf_scale_samples.empty()) { + leaf_scale_ = samples.leaf_scale_samples.back(); + } else { + leaf_scale_ = config_.sigma2_mean_init; + } + // Sync the leaf model's scale (tau) with the warm-started leaf_scale_. The leaf model + // was constructed with sigma2_mean_init above; without this the first continued + // iteration would sample the mean forest using the initial tau rather than the last + // retained leaf scale (which the one-shot sampler carries forward from the prior iteration). + if (has_mean_forest_) { + std::visit(ScaleUpdateVisitor{*this, leaf_scale_}, mean_leaf_model_); + } + } else { + global_variance_ = config_.sigma2_global_init; + leaf_scale_ = config_.sigma2_mean_init; + } + // leaf_scale_multivariate_ = config_.sigma2_leaf_multivariate_init; + + initialized_ = true; +} + +void BARTSampler::run_gfr(BARTSamples& samples, int num_gfr, bool keep_gfr, int num_chains) { + // Reserve space for GFR predictions if they are to be retained + if (keep_gfr) { + if (has_mean_forest_) { + samples.mean_forest_predictions_train.reserve(data_.n_train * num_gfr); + } + if (has_variance_forest_) { + samples.variance_forest_predictions_train.reserve(data_.n_train * num_gfr); + } + } + + // NOTE: for serial sampling (which is all we currently support), chain 1 uses the live sampler state after + // the GFR loop (i.e. the state after GFR iteration num_gfr-1). Chains 2..N each need their own earlier + // GFR starting point so that the chains are initialized from distinct states. + // We save exactly num_chains-1 snapshots, one per "extra" chain: + // gfr_snapshots_[k] = state after GFR iteration (num_gfr - num_chains + k), for k = 0..num_chains-2. + // The last GFR iteration (i = num_gfr-1) is NOT snapshotted because chain 1 uses the live state. + // Chain j (1-indexed, j >= 2) uses gfr_snapshots_[num_chains-j] = state after GFR[num_gfr-j]. + // If num_chains is 1, we keep no snapshots. + int snapshot_start = (num_chains > 1) ? std::max(0, num_gfr - num_chains) : num_gfr; + if (num_chains > 1) { + gfr_snapshots_.clear(); + gfr_snapshots_.reserve(num_chains - 1); + } + + if (config_.verbose && num_gfr > 0) { + Log::Info("Running GFR sampler (%d iterations)", num_gfr); + } + const int gfr_report_every = std::max(1, num_gfr / 10); + bool write_snapshot = false; + for (int i = 0; i < num_gfr; i++) { + // Do not snapshot the final GFR iteration: chain 1 uses the live sampler state directly. + write_snapshot = (i >= snapshot_start) && (i < num_gfr - 1); + RunOneIteration(samples, /*gfr=*/true, /*keep_sample=*/keep_gfr, /*write_snapshot=*/write_snapshot); + if (config_.verbose && ((i + 1) % gfr_report_every == 0 || i + 1 == num_gfr)) { + Log::Info("GFR: %d%% (%d/%d)", (100 * (i + 1)) / num_gfr, i + 1, num_gfr); + } + } +} + +void BARTSampler::run_mcmc(BARTSamples& samples, int num_burnin, int keep_every, int num_mcmc) { + // Reserve space for MCMC predictions if they are to be retained + if (has_mean_forest_) { + samples.mean_forest_predictions_train.reserve(data_.n_train * num_mcmc); + if (has_test_) { + samples.mean_forest_predictions_test.reserve(data_.n_test * num_mcmc); + } + } + if (has_variance_forest_) { + samples.variance_forest_predictions_train.reserve(data_.n_train * num_mcmc); + if (has_test_) { + samples.variance_forest_predictions_test.reserve(data_.n_test * num_mcmc); + } + } + + // Create leaf models and pass them to the RunOneIteration function; these are updated in place and will reflect the current state of the leaf scale parameters (if they are being sampled) + bool keep_forest = false; + const int mcmc_total = num_burnin + keep_every * num_mcmc; + const int mcmc_report_every = std::max(1, mcmc_total / 10); + for (int i = 0; i < mcmc_total; i++) { + if (i >= num_burnin && (i - num_burnin) % keep_every == 0) + keep_forest = true; + else + keep_forest = false; + RunOneIteration(samples, /*gfr=*/false, /*keep_sample=*/keep_forest, /*write_snapshot=*/false); + if (config_.verbose && ((i + 1) % mcmc_report_every == 0 || i + 1 == mcmc_total)) { + Log::Info("MCMC: %d%% (%d/%d)", (100 * (i + 1)) / mcmc_total, i + 1, mcmc_total); + } + } +} + +void BARTSampler::run_mcmc_chains(BARTSamples& samples, int num_chains, int num_burnin, int keep_every, int num_mcmc) { + for (int chain_idx = 0; chain_idx < num_chains; chain_idx++) { + if (chain_idx > 0 && !gfr_snapshots_.empty()) { + // Re-initialize the sampler state for each new chain. + // gfr_snapshots_ holds num_chains-1 states (oldest-first): index 0 = GFR[num_gfr-num_chains], + // index num_chains-2 = GFR[num_gfr-2]. Chain j (1-indexed, j>=2) uses index num_chains-j. + // When num_gfr < num_chains we may not have enough distinct snapshots; in that case + // fall back to running chains from whatever state is available (same behavior as R). + int snapshot_idx = num_chains - 1 - chain_idx; + if (snapshot_idx >= 0 && snapshot_idx < static_cast(gfr_snapshots_.size())) { + RestoreStateFromGFRSnapshot(samples, snapshot_idx); + } + } + if (config_.verbose) { + Log::Info("Running MCMC chain %d/%d (%d samples)", chain_idx + 1, num_chains, num_mcmc); + } + run_mcmc(samples, num_burnin, keep_every, num_mcmc); + } +} + +void BARTSampler::postprocess_samples(BARTSamples& samples) { + // Unpack test set predictions for mean and variance forest + if (has_test_) { + if (has_mean_forest_) { + std::vector predictions = samples.mean_forests->Predict(*forest_dataset_test_); + samples.mean_forest_predictions_test.insert(samples.mean_forest_predictions_test.end(), + predictions.data(), predictions.data() + predictions.size()); + } + if (has_variance_forest_) { + std::vector predictions = samples.variance_forests->Predict(*forest_dataset_test_); + samples.variance_forest_predictions_test.insert(samples.variance_forest_predictions_test.end(), + predictions.data(), predictions.data() + predictions.size()); + } + if (has_random_effects_) { + RandomEffectsDataset rfx_dataset_test; + rfx_dataset_test.AddGroupLabels(data_.rfx_group_ids_test, data_.n_test); + if (data_.rfx_basis_test != nullptr) { + rfx_dataset_test.AddBasis(data_.rfx_basis_test, data_.n_test, data_.rfx_basis_dim, /*row_major=*/false); + } else { + std::vector ones(data_.n_test, 1.0); + rfx_dataset_test.AddBasis(ones.data(), data_.n_test, 1, /*row_major=*/false); + } + samples.rfx_predictions_test.resize(data_.n_test * samples.num_samples); + samples.rfx_container->Predict(rfx_dataset_test, *samples.rfx_label_mapper, samples.rfx_predictions_test); + } + } + + // Convert variance forest predictions and global error variance from + // standardized space to original outcome scale. + // - Train predictions come from ForestTracker::GetSumPredictions() (log-scale leaf sums), + // so apply exp() then multiply by y_std^2. + // - Test predictions come from ForestContainer::Predict() with is_exponentiated_=true, + // which already applies exp() internally, so just multiply by y_std^2. + // - Global error variance samples are in standardized space; multiply by y_std^2. + if (has_variance_forest_) { + double y_std2 = samples.y_std * samples.y_std; + for (double& v : samples.variance_forest_predictions_train) v = std::exp(v) * y_std2; + for (double& v : samples.variance_forest_predictions_test) v *= y_std2; + } + if (sample_sigma2_global_) { + double y_std2 = samples.y_std * samples.y_std; + for (double& v : samples.global_error_variance_samples) v *= y_std2; + } + + // Convert mean forest and random effects predictions from standardized space to the + // original outcome scale, matching predict(): the mean forest carries the location + // shift (y_bar), random effects carry only the scale factor. For probit/cloglog + // outcomes y_bar=0 and y_std=1, so these adjustments are the identity. + if (has_mean_forest_) { + for (double& v : samples.mean_forest_predictions_train) v = v * samples.y_std + samples.y_bar; + for (double& v : samples.mean_forest_predictions_test) v = v * samples.y_std + samples.y_bar; + } + if (has_random_effects_) { + for (double& v : samples.rfx_predictions_train) v *= samples.y_std; + for (double& v : samples.rfx_predictions_test) v *= samples.y_std; + } +} + +void BARTSampler::RunOneIteration(BARTSamples& samples, bool gfr, bool keep_sample, bool write_snapshot) { + if (has_mean_forest_) { + if (gfr) { + std::visit(GFROneIterationVisitor{*this, samples, keep_sample}, mean_leaf_model_); + } else { + std::visit(MCMCOneIterationVisitor{*this, samples, keep_sample}, mean_leaf_model_); + } + } + + if (has_variance_forest_) { + if (gfr) { + GFRSampleOneIter( + *variance_forest_, *variance_forest_tracker_, *samples.variance_forests, variance_leaf_model_, + *forest_dataset_, *residual_, *tree_prior_variance_, rng_, + config_.var_weights_variance, config_.sweep_update_indices_variance, global_variance_, config_.feature_types, + config_.cutpoint_grid_size, /*keep_forest=*/keep_sample, + /*pre_initialized=*/true, /*backfitting=*/false, + /*num_features_subsample=*/config_.num_features_subsample_variance, config_.num_threads); + } else { + MCMCSampleOneIter( + *variance_forest_, *variance_forest_tracker_, *samples.variance_forests, variance_leaf_model_, + *forest_dataset_, *residual_, *tree_prior_variance_, rng_, + config_.var_weights_variance, config_.sweep_update_indices_variance, global_variance_, /*keep_forest=*/keep_sample, + /*pre_initialized=*/true, /*backfitting=*/false, + /*num_threads=*/config_.num_threads); + } + } + + if (config_.link_function == LinkFunction::Probit) { + sample_probit_latent_outcome(rng_, outcome_raw_->GetData().data(), mean_forest_tracker_->GetSumPredictions(), + residual_->GetData().data(), samples.y_bar, data_.n_train); + } + + if (sample_sigma2_global_) { + global_variance_ = var_model_->SampleVarianceParameter( + residual_->GetData(), config_.a_sigma2_global, config_.b_sigma2_global, rng_); + } + + if (sample_sigma2_leaf_) { + leaf_scale_ = leaf_scale_model_->SampleVarianceParameter( + mean_forest_.get(), config_.a_sigma2_mean, config_.b_sigma2_mean, rng_); + std::visit(ScaleUpdateVisitor{*this, leaf_scale_}, mean_leaf_model_); + } + + // Gibbs updates for the cloglog model + if (config_.link_function == LinkFunction::Cloglog) { + // Update auxiliary data to current forest predictions + for (int i = 0; i < data_.n_train; i++) { + forest_dataset_->SetAuxiliaryDataValue(1, i, mean_forest_tracker_->GetSamplePrediction(i)); + } + + // Sample latent z_i's using truncated exponential + ordinal_sampler_->UpdateLatentVariables(*forest_dataset_, residual_->GetData(), rng_); + + // Sample gamma parameters (cutpoints) + ordinal_sampler_->UpdateGammaParams(*forest_dataset_, residual_->GetData(), config_.cloglog_leaf_prior_shape, config_.cloglog_leaf_prior_scale, config_.cloglog_cutpoint_0, rng_); + + // Update cumulative sum of exp(gamma) values + ordinal_sampler_->UpdateCumulativeExpSums(*forest_dataset_); + } + + // Gibbs updates for random effects model + if (has_random_effects_) { + random_effects_model_->SampleRandomEffects(*random_effects_dataset_, *residual_, *random_effects_tracker_, global_variance_, rng_); + if (keep_sample) { + samples.rfx_container->AddSample(*random_effects_model_); + for (int i = 0; i < data_.n_train; i++) { + samples.rfx_predictions_train.push_back(random_effects_tracker_->GetPrediction(i)); + } + } + } + + if (keep_sample) { + // Add parameter and prediction samples + samples.num_samples++; + if (sample_sigma2_global_) samples.global_error_variance_samples.push_back(global_variance_); + if (sample_sigma2_leaf_) samples.leaf_scale_samples.push_back(leaf_scale_); + if (has_mean_forest_) { + double* mean_forest_preds_train = mean_forest_tracker_->GetSumPredictions(); + samples.mean_forest_predictions_train.insert(samples.mean_forest_predictions_train.end(), + mean_forest_preds_train, + mean_forest_preds_train + samples.num_train); + } + if (has_variance_forest_) { + double* variance_forest_preds_train = variance_forest_tracker_->GetSumPredictions(); + samples.variance_forest_predictions_train.insert(samples.variance_forest_predictions_train.end(), + variance_forest_preds_train, + variance_forest_preds_train + samples.num_train); + } + if (config_.outcome_type == OutcomeType::Ordinal) { + // Store cutpoint samples (ordinal only; binary cloglog has a single fixed cutpoint) + std::vector cloglog_cutpoints(config_.num_classes_cloglog - 1); + for (int i = 0; i < config_.num_classes_cloglog - 1; i++) { + cloglog_cutpoints[i] = forest_dataset_->GetAuxiliaryDataValue(2, i); + } + samples.cloglog_cutpoint_samples.insert(samples.cloglog_cutpoint_samples.end(), + cloglog_cutpoints.data(), cloglog_cutpoints.data() + cloglog_cutpoints.size()); + } + } + + if (write_snapshot) { + GFRSnapshot snap; + if (has_mean_forest_) snap.mean_forest = std::make_unique(*mean_forest_); + if (has_variance_forest_) snap.variance_forest = std::make_unique(*variance_forest_); + snap.sigma2 = global_variance_; + if (has_mean_forest_) { + if (config_.mean_leaf_model_type == MeanLeafModelType::GaussianMultivariateRegression) { + snap.leaf_scale_multivariate = leaf_scale_multivariate_; + } else if (config_.mean_leaf_model_type == MeanLeafModelType::GaussianUnivariateRegression || config_.mean_leaf_model_type == MeanLeafModelType::GaussianConstant) { + snap.leaf_scale = leaf_scale_; + } else if (config_.mean_leaf_model_type == MeanLeafModelType::CloglogOrdinal) { + snap.cloglog_forest_preds.clear(); + snap.cloglog_forest_preds.resize(data_.n_train); + snap.cloglog_forest_preds.assign(mean_forest_tracker_->GetSumPredictions(), mean_forest_tracker_->GetSumPredictions() + data_.n_train); + snap.cloglog_latent_outcome.clear(); + snap.cloglog_latent_outcome.resize(data_.n_train); + for (int i = 0; i < data_.n_train; i++) { + snap.cloglog_latent_outcome[i] = forest_dataset_->GetAuxiliaryDataValue(0, i); + } + snap.cloglog_logscale_cutpoints.clear(); + snap.cloglog_logscale_cutpoints.resize(config_.num_classes_cloglog - 1); + for (int i = 0; i < config_.num_classes_cloglog - 1; i++) { + snap.cloglog_logscale_cutpoints[i] = forest_dataset_->GetAuxiliaryDataValue(2, i); + } + } + } + snap.residual.clear(); + snap.residual.resize(data_.n_train); + snap.residual.assign(residual_->GetData().data(), residual_->GetData().data() + data_.n_train); + if (has_variance_forest_) { + snap.variance_weights.clear(); + snap.variance_weights.resize(data_.n_train); + snap.variance_weights.assign(forest_dataset_->GetVarWeights().data(), forest_dataset_->GetVarWeights().data() + data_.n_train); + } + if (config_.has_random_effects) { + snap.rfx_working_parameter = random_effects_model_->GetWorkingParameter(); + snap.rfx_group_parameters = random_effects_model_->GetGroupParameters(); + snap.rfx_group_parameter_covariance = random_effects_model_->GetGroupParameterCovariance(); + snap.rfx_working_parameter_covariance = random_effects_model_->GetWorkingParameterCovariance(); + snap.rfx_variance_prior_shape = random_effects_model_->GetVariancePriorShape(); + snap.rfx_variance_prior_scale = random_effects_model_->GetVariancePriorScale(); + } + gfr_snapshots_.push_back(std::move(snap)); + } +} + +void BARTSampler::RestoreStateFromGFRSnapshot(BARTSamples& samples, int snapshot_index) { + GFRSnapshot& snap = gfr_snapshots_[snapshot_index]; + + // Restore mean forest state (if present). + // ReconstituteFromForest increments the residual by (prev_tree_pred - new_tree_pred) for + // every tree, swapping the chain-N forest contribution out and the GFR-snapshot contribution + // in. The residual must still hold the chain-N state here so that this swap is correct. + if (config_.num_trees_mean > 0) { + std::visit(MeanForestResetVisitor{*this, samples, *snap.mean_forest}, mean_leaf_model_); + } + + // Initialize variance forest state (if present) + if (config_.num_trees_variance > 0) { + variance_leaf_model_ = LogLinearVarianceLeafModel(config_.shape_variance_forest, config_.scale_variance_forest); + variance_forest_ = std::make_unique(config_.num_trees_variance, config_.leaf_dim_variance, config_.leaf_constant_variance, config_.exponentiated_leaf_variance); + samples.variance_forests = std::make_unique(config_.num_trees_variance, config_.leaf_dim_variance, config_.leaf_constant_variance, config_.exponentiated_leaf_variance); + variance_forest_tracker_ = std::make_unique(forest_dataset_->GetCovariates(), config_.feature_types, config_.num_trees_variance, data_.n_train); + tree_prior_variance_ = std::make_unique(config_.alpha_variance, config_.beta_variance, config_.min_samples_leaf_variance, config_.max_depth_variance); + // Leaf values for the log-linear variance model are on the log scale; the ensemble sums + // log(sigma^2_i) contributions, so each tree starts at log(init_val) / num_trees. + variance_forest_->SetLeafValue(std::log(init_val_variance_) / config_.num_trees_variance); + variance_forest_tracker_->UpdatePredictions(variance_forest_.get(), *forest_dataset_.get()); + // UpdateVarModelTree (called inside GFRSampleOneIter / MCMCSampleOneIter) unconditionally + // reads and writes the dataset variance weight slot via VarWeightValue / SetVarWeightValue. + // This slot tracks the cumulative per-observation variance prediction + // (sigma^2_i = exp(sum of tree leaf values)) and is incompatible with case weights, which + // would need to be reapplied after every per-tree update. The R/Python APIs enforce this + // as a hard error; guard here for callers that use BARTSampler directly. + if (forest_dataset_->HasVarWeights()) { + Log::Fatal("observation_weights and a variance forest cannot be used together."); + } + std::vector initial_variance_preds(data_.n_train, init_val_variance_); + forest_dataset_->AddVarianceWeights(initial_variance_preds.data(), data_.n_train); + has_variance_forest_ = true; + } + + // Random effects model + if (config_.has_random_effects) { + // Restore "working" parameter prior mean + random_effects_model_->SetWorkingParameter(snap.rfx_working_parameter); + + // Restore "group" parameter prior mean + random_effects_model_->SetGroupParameters(snap.rfx_group_parameters); + + // Restore "working" parameter prior covariance + random_effects_model_->SetWorkingParameterCovariance(snap.rfx_working_parameter_covariance); + + // Restore "group" parameter prior covariance + random_effects_model_->SetGroupParameterCovariance(snap.rfx_group_parameter_covariance); + + // Restore variance model priors + random_effects_model_->SetVariancePriorShape(snap.rfx_variance_prior_shape); + random_effects_model_->SetVariancePriorScale(snap.rfx_variance_prior_scale); + + // Swap the chain-N RFX contribution out of the residual and the GFR-snapshot + // contribution in, exactly as the R sampler does via resetRandomEffectsTracker. + // At this point residual_ = y - f_gfr - rfx_chain_N (forest already swapped above), + // so ResetFromSample produces: residual_ += rfx_chain_N - rfx_gfr = y - f_gfr - rfx_gfr. + random_effects_tracker_->ResetFromSample(*random_effects_model_, *random_effects_dataset_, *residual_); + } + + // Cloglog state + if (config_.link_function == LinkFunction::Cloglog) { + // Reset auxiliary data values + // Restore latent outcome values (slot 0) + for (int i = 0; i < data_.n_train; i++) { + forest_dataset_->SetAuxiliaryDataValue(0, i, snap.cloglog_latent_outcome[i]); + } + // Restore forest predictions (slot 1) + for (int i = 0; i < data_.n_train; i++) { + forest_dataset_->SetAuxiliaryDataValue(1, i, snap.cloglog_forest_preds[i]); + } + // Restore log-scale cutpoints + for (int i = 0; i < config_.num_classes_cloglog - 1; i++) { + forest_dataset_->SetAuxiliaryDataValue(2, i, snap.cloglog_logscale_cutpoints[i]); + } + // Convert to cumulative exponentiated cutpoints directly in C++ + ordinal_sampler_->UpdateCumulativeExpSums(*forest_dataset_); + // For cloglog, residual_ holds raw y values (not y - f(X)), so the incremental + // tree-prediction swap in ReconstituteFromForest is wrong. Restore from snapshot. + residual_->OverwriteData(snap.residual.data(), data_.n_train); + } + + // Other internal model state + global_variance_ = snap.sigma2; + if (has_mean_forest_) { + if (config_.mean_leaf_model_type == MeanLeafModelType::GaussianMultivariateRegression) { + leaf_scale_multivariate_ = snap.leaf_scale_multivariate; + } else if (config_.mean_leaf_model_type == MeanLeafModelType::GaussianUnivariateRegression || config_.mean_leaf_model_type == MeanLeafModelType::GaussianConstant) { + leaf_scale_ = snap.leaf_scale; + } + } +} + +} // namespace StochTree diff --git a/src/bcf_sampler.cpp b/src/bcf_sampler.cpp new file mode 100644 index 00000000..8731ee4c --- /dev/null +++ b/src/bcf_sampler.cpp @@ -0,0 +1,1326 @@ +/*! Copyright (c) 2026 by stochtree authors */ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace StochTree { + +void AddModelTermsForProbit(double* outcome_preds, ForestTracker* mu_forest_tracker, ForestTracker* tau_forest_tracker, RandomEffectsTracker* random_effects_tracker, int n) { + double* mu_preds = mu_forest_tracker->GetSumPredictions(); + double* tau_preds = tau_forest_tracker->GetSumPredictions(); + if (random_effects_tracker != nullptr) { + double* rfx_preds = random_effects_tracker->GetPredictions(); + for (int i = 0; i < n; i++) { + outcome_preds[i] = mu_preds[i] + tau_preds[i] + rfx_preds[i]; + } + } else { + for (int i = 0; i < n; i++) { + outcome_preds[i] = mu_preds[i] + tau_preds[i]; + } + } +} + +BCFSampler::BCFSampler(BCFSamples& samples, BCFConfig& config, BCFData& data, bool continuation) : config_{config}, data_{data}, mu_leaf_model_(GaussianConstantLeafModel(0.0)), tau_leaf_model_(GaussianUnivariateRegressionLeafModel(0.0)), variance_leaf_model_(0.0, 0.0) { + InitializeState(samples, continuation); +} + +void BCFSampler::InitializeState(BCFSamples& samples, bool continuation) { + // Validate y_train values match the expected support for discrete link functions + if (config_.link_function == LinkFunction::Probit) { + for (int i = 0; i < data_.n_train; i++) { + if (data_.y_train[i] != 0.0 && data_.y_train[i] != 1.0) { + Log::Fatal("Outcomes must be 0 or 1 for probit link function"); + } + } + // Initialize model_preds_ vector for probit latent outcome sampling + model_preds_.resize(data_.n_train, 0.0); + } else if (config_.link_function == LinkFunction::Cloglog) { + Log::Fatal("Cloglog link function is not currently supported in BCF"); + } + + // Validate that both num_trees_mu and num_trees_tau are positive + if (config_.num_trees_mu <= 0) { + Log::Fatal("num_trees_mu must be >0"); + } + if (config_.num_trees_tau <= 0) { + Log::Fatal("num_trees_tau must be >0"); + } + + // Validate outcome type + if (config_.outcome_type == OutcomeType::Ordinal) { + Log::Fatal("Ordinal outcome type is not currently supported in BCF"); + } + + // Continued sampling currently supports only the basic Gaussian, identity-link, + // univariate-treatment case; reject configurations whose warm-start init is not yet + // implemented so that unsupported continuations fail loudly rather than silently + // re-initializing from root. + if (continuation) { + if (config_.num_trees_variance > 0) { + Log::Fatal("Continued sampling is not yet supported for models with a variance forest"); + } + if (config_.has_random_effects) { + Log::Fatal("Continued sampling is not yet supported for models with random effects"); + } + if (config_.link_function != LinkFunction::Identity) { + Log::Fatal("Continued sampling is not yet supported for probit or cloglog link functions"); + } + if (config_.adaptive_coding) { + Log::Fatal("Continued sampling is not yet supported for adaptive coding"); + } + if (data_.treatment_dim != 1 || config_.tau_leaf_model_type != MeanLeafModelType::GaussianUnivariateRegression) { + Log::Fatal("Continued sampling is not yet supported for multivariate treatment effects"); + } + } + + // Switch off treatment forest leaf scale sampling if treatment is multivariate + if (config_.sample_sigma2_leaf_tau && config_.tau_leaf_model_type != MeanLeafModelType::GaussianUnivariateRegression) { + Log::Info("sample_sigma2_leaf_tau can only be true when tau_leaf_model_type is GaussianUnivariateRegression, setting sample_sigma2_leaf_tau to false"); + config_.sample_sigma2_leaf_tau = false; + } + + // Adaptive coding model + if (config_.adaptive_coding) { + if (data_.treatment_dim != 1) { + Log::Fatal("Adaptive coding is currently only supported for binary treatments (treatment_dim=1)"); + } + adaptive_coding_ = true; + b_0_ = config_.b_0_init; + b_1_ = config_.b_1_init; + } + + // Load data from BARTData object into ForestDataset object + forest_dataset_ = std::make_unique(); + forest_dataset_->AddCovariates(data_.X_train, data_.n_train, data_.p, /*row_major=*/false); + if (data_.treatment_train != nullptr) { + if (adaptive_coding_) { + // Basis becomes b_0 * (1-Z) + b_1 * Z + tau_basis_vector_train_.resize(data_.n_train); + for (int i = 0; i < data_.n_train; i++) { + double z = data_.treatment_train[i]; + tau_basis_vector_train_[i] = b_0_ * (1.0 - z) + b_1_ * z; + } + forest_dataset_->AddBasis(tau_basis_vector_train_.data(), data_.n_train, data_.treatment_dim, /*row_major=*/false); + } else { + forest_dataset_->AddBasis(data_.treatment_train, data_.n_train, data_.treatment_dim, /*row_major=*/false); + } + } + if (data_.obs_weights_train != nullptr) { + forest_dataset_->AddVarianceWeights(data_.obs_weights_train, data_.n_train); + } + samples.num_train = data_.n_train; + samples.num_test = data_.n_test; + samples.treatment_dim = data_.treatment_dim; + residual_ = std::make_unique(data_.y_train, data_.n_train); + outcome_raw_ = std::make_unique(data_.y_train, data_.n_train); + if (data_.X_test != nullptr) { + forest_dataset_test_ = std::make_unique(); + forest_dataset_test_->AddCovariates(data_.X_test, data_.n_test, data_.p, /*row_major=*/false); + if (data_.treatment_test != nullptr) { + if (adaptive_coding_) { + // Basis becomes b_0 * (1-Z) + b_1 * Z + tau_basis_vector_test_.resize(data_.n_test); + for (int i = 0; i < data_.n_test; i++) { + double z = data_.treatment_test[i]; + tau_basis_vector_test_[i] = b_0_ * (1.0 - z) + b_1_ * z; + } + forest_dataset_test_->AddBasis(tau_basis_vector_test_.data(), data_.n_test, data_.treatment_dim, /*row_major=*/false); + } else { + forest_dataset_test_->AddBasis(data_.treatment_test, data_.n_test, data_.treatment_dim, /*row_major=*/false); + } + } + if (data_.obs_weights_test != nullptr) { + forest_dataset_test_->AddVarianceWeights(data_.obs_weights_test, data_.n_test); + } + has_test_ = true; + } + + // Precompute outcome mean and variance for standardization and calibration + double y_mean = 0.0, M2 = 0.0, y_mean_prev = 0.0; + for (int i = 0; i < data_.n_train; i++) { + y_mean_prev = y_mean; + y_mean = y_mean_prev + (data_.y_train[i] - y_mean_prev) / (i + 1); + M2 = M2 + (data_.y_train[i] - y_mean_prev) * (data_.y_train[i] - y_mean); + } + double y_var = M2 / data_.n_train; + + // Outcome standardization and forest initial value setup + if (config_.link_function == LinkFunction::Probit) { + // Initialize forests to 0, no scaling, but offset by the probit transform of the mean outcome to improve mixing + samples.y_std = 1.0; + samples.y_bar = norm_inv_cdf(y_mean); + init_val_mu_ = 0.0; + init_val_tau_ = 0.0; + if (config_.tau_leaf_model_type == MeanLeafModelType::GaussianMultivariateRegression) { + init_val_tau_vec_.assign(config_.leaf_dim_tau, 0.0); + } + } else { + if (config_.standardize_outcome) { + samples.y_bar = y_mean; + samples.y_std = std::sqrt(y_var); + init_val_mu_ = 0.0; + init_val_tau_ = 0.0; + if (config_.tau_leaf_model_type == MeanLeafModelType::GaussianMultivariateRegression) { + init_val_tau_vec_.assign(config_.leaf_dim_tau, 0.0); + } + } else { + samples.y_bar = 0.0; + samples.y_std = 1.0; + init_val_mu_ = y_mean; + init_val_tau_ = 0.0; + if (config_.tau_leaf_model_type == MeanLeafModelType::GaussianMultivariateRegression) { + init_val_tau_vec_.assign(config_.leaf_dim_tau, 0.0); + } + } + } + + // Calibration for mu forest + if (config_.sigma2_mu_init < 0.0) { + if (config_.link_function == LinkFunction::Probit) { + config_.sigma2_mu_init = 1.0 / config_.num_trees_mu; + } else { + if (config_.standardize_outcome) + config_.sigma2_mu_init = 1.0 / config_.num_trees_mu; + else + config_.sigma2_mu_init = y_var / config_.num_trees_mu; + } + } + if (config_.sample_sigma2_leaf_mu) { + if (config_.b_sigma2_mu <= 0.0) { + if (config_.link_function == LinkFunction::Probit) { + config_.b_sigma2_mu = 1.0 / (2 * config_.num_trees_mu); + } else { + if (config_.standardize_outcome) + config_.b_sigma2_mu = 1.0 / (2 * config_.num_trees_mu); + else + config_.b_sigma2_mu = y_var / (2 * config_.num_trees_mu); + } + } + } + + // Calibration for tau forest + if (config_.sigma2_tau_init < 0.0) { + if (config_.link_function == LinkFunction::Probit) { + config_.sigma2_tau_init = 1.0 / config_.num_trees_tau; + } else { + if (config_.standardize_outcome) + config_.sigma2_tau_init = 1.0 / config_.num_trees_tau; + else + config_.sigma2_tau_init = y_var / config_.num_trees_tau; + } + } + if (config_.sample_sigma2_leaf_tau) { + if (config_.b_sigma2_tau <= 0.0) { + if (config_.link_function == LinkFunction::Probit) { + config_.b_sigma2_tau = 1.0 / (2 * config_.num_trees_tau); + } else { + if (config_.standardize_outcome) + config_.b_sigma2_tau = 1.0 / (2 * config_.num_trees_tau); + else + config_.b_sigma2_tau = y_var / (2 * config_.num_trees_tau); + } + } + } + + // Initialize mu leaf model + mu_leaf_model_ = GaussianConstantLeafModel(config_.sigma2_mu_init); + + // Initialize tau leaf model + if (config_.tau_leaf_model_type == MeanLeafModelType::GaussianUnivariateRegression) { + tau_leaf_model_ = GaussianUnivariateRegressionLeafModel(config_.sigma2_tau_init); + } else if (config_.tau_leaf_model_type == MeanLeafModelType::GaussianMultivariateRegression) { + Eigen::MatrixXd Sigma_0; + if (!config_.sigma2_leaf_tau_matrix.empty()) { + if ((int)config_.sigma2_leaf_tau_matrix.size() != config_.leaf_dim_tau * config_.leaf_dim_tau) { + Log::Fatal("sigma2_leaf_tau_matrix must have leaf_dim_tau * leaf_dim_tau = %d elements, but has %zu", + config_.leaf_dim_tau * config_.leaf_dim_tau, config_.sigma2_leaf_tau_matrix.size()); + } + // Column-major interpretation matches both R and Eigen (python must be reordered before passing to C++) + Sigma_0 = Eigen::Map(config_.sigma2_leaf_tau_matrix.data(), config_.leaf_dim_tau, config_.leaf_dim_tau); + } else { + Sigma_0 = config_.sigma2_tau_init * Eigen::MatrixXd::Identity(config_.leaf_dim_tau, config_.leaf_dim_tau); + } + tau_leaf_model_ = GaussianMultivariateRegressionLeafModel(Sigma_0); + } else { + Log::Fatal("Unsupported leaf model type for treatment forest"); + } + + // Calibration for variance forests + if (config_.num_trees_variance > 0) { + // NOTE: calibration only works for standardized outcomes + if (config_.shape_variance_forest <= 0.0 || config_.scale_variance_forest <= 0.0) { + if (config_.leaf_prior_calibration_param <= 0.0) { + config_.leaf_prior_calibration_param = 1.5; + } + if (config_.shape_variance_forest <= 0.0) { + config_.shape_variance_forest = config_.num_trees_variance / (config_.leaf_prior_calibration_param * config_.leaf_prior_calibration_param) + 0.5; + } + if (config_.scale_variance_forest <= 0.0) { + config_.scale_variance_forest = config_.num_trees_variance / (config_.leaf_prior_calibration_param * config_.leaf_prior_calibration_param); + } + } + if (config_.variance_forest_leaf_init > 0.0) { + init_val_variance_ = config_.variance_forest_leaf_init; + } else if (config_.standardize_outcome) { + init_val_variance_ = 1.0; + } else { + init_val_variance_ = y_var; + } + } + + // Standardize partial residuals in place; these are updated in each iteration but initialized to standardized outcomes + // Works for: + // 1. Standardized outcomes (since y_bar = mean(y) and y_std = sd(y)) + // 2. Non-standardized outcomes (since y_bar = 0 and y_std = 1, so this just transfers y_train as-is) + // 3. Probit link (since y_bar = norm_inv_cdf(mean(y)) and y_std = 1) + for (int i = 0; i < data_.n_train; i++) residual_->GetData()[i] = (data_.y_train[i] - samples.y_bar) / samples.y_std; + + // Initialize mean forest state. + // On continuation, do NOT re-create samples.mu_forests so that new draws are appended to the + // existing container; otherwise allocate a fresh container. + mu_forest_ = std::make_unique(config_.num_trees_mu, config_.leaf_dim_mu, config_.leaf_constant_mu, config_.exponentiated_leaf_mu); + if (!continuation) { + samples.mu_forests = std::make_unique(config_.num_trees_mu, config_.leaf_dim_mu, config_.leaf_constant_mu, config_.exponentiated_leaf_mu); + } + mu_forest_tracker_ = std::make_unique(forest_dataset_->GetCovariates(), config_.feature_types, config_.num_trees_mu, data_.n_train); + tree_prior_mu_ = std::make_unique(config_.alpha_mu, config_.beta_mu, config_.min_samples_leaf_mu, config_.max_depth_mu); + mu_forest_->SetLeafValue(init_val_mu_ / config_.num_trees_mu); + UpdateResidualEntireForest(*mu_forest_tracker_, *forest_dataset_, *residual_, mu_forest_.get(), !config_.leaf_constant_mu, std::minus()); + mu_forest_tracker_->UpdatePredictions(mu_forest_.get(), *forest_dataset_.get()); + if (continuation) { + // Warm-start from the last retained mu sample. The reset requires the already-populated + // tracker established by the root init above (mirrors RestoreStateFromGFRSnapshot). + int last_mu_idx = samples.mu_forests->NumSamples() - 1; + TreeEnsemble& last_mu = *samples.mu_forests->GetEnsemble(last_mu_idx); + mu_forest_->ReconstituteFromForest(last_mu); + mu_forest_tracker_->ReconstituteFromForest(last_mu, *forest_dataset_, *residual_, true); + mu_forest_tracker_->UpdatePredictions(mu_forest_.get(), *forest_dataset_.get()); + } + + // Initialize treatment effect forest state. + // On continuation, do NOT re-create samples.tau_forests (append to the existing container). + if (config_.tau_leaf_model_type == MeanLeafModelType::GaussianUnivariateRegression) { + tau_forest_ = std::make_unique(config_.num_trees_tau, config_.leaf_dim_tau, config_.leaf_constant_tau, config_.exponentiated_leaf_tau); + if (!continuation) { + samples.tau_forests = std::make_unique(config_.num_trees_tau, config_.leaf_dim_tau, config_.leaf_constant_tau, config_.exponentiated_leaf_tau); + } + tau_forest_tracker_ = std::make_unique(forest_dataset_->GetCovariates(), config_.feature_types, config_.num_trees_tau, data_.n_train); + tree_prior_tau_ = std::make_unique(config_.alpha_tau, config_.beta_tau, config_.min_samples_leaf_tau, config_.max_depth_tau); + tau_forest_->SetLeafValue(init_val_tau_ / config_.num_trees_tau); + UpdateResidualEntireForest(*tau_forest_tracker_, *forest_dataset_, *residual_, tau_forest_.get(), !config_.leaf_constant_tau, std::minus()); + tau_forest_tracker_->UpdatePredictions(tau_forest_.get(), *forest_dataset_.get()); + if (continuation) { + // Warm-start from the last retained tau sample. At this point the residual already has the + // warm-started mu contribution removed; reconstitution swaps the root tau contribution out + // and the last-sample tau contribution in. The treatment intercept tau_0 is restored + // separately (below), after both forests are warm-started. + int last_tau_idx = samples.tau_forests->NumSamples() - 1; + TreeEnsemble& last_tau = *samples.tau_forests->GetEnsemble(last_tau_idx); + tau_forest_->ReconstituteFromForest(last_tau); + tau_forest_tracker_->ReconstituteFromForest(last_tau, *forest_dataset_, *residual_, true); + tau_forest_tracker_->UpdatePredictions(tau_forest_.get(), *forest_dataset_.get()); + } + } else if (config_.tau_leaf_model_type == MeanLeafModelType::GaussianMultivariateRegression) { + tau_forest_ = std::make_unique(config_.num_trees_tau, config_.leaf_dim_tau, config_.leaf_constant_tau, config_.exponentiated_leaf_tau); + samples.tau_forests = std::make_unique(config_.num_trees_tau, config_.leaf_dim_tau, config_.leaf_constant_tau, config_.exponentiated_leaf_tau); + tau_forest_tracker_ = std::make_unique(forest_dataset_->GetCovariates(), config_.feature_types, config_.num_trees_tau, data_.n_train); + tree_prior_tau_ = std::make_unique(config_.alpha_tau, config_.beta_tau, config_.min_samples_leaf_tau, config_.max_depth_tau); + tau_forest_->SetLeafVector(init_val_tau_vec_); + UpdateResidualEntireForest(*tau_forest_tracker_, *forest_dataset_, *residual_, tau_forest_.get(), true, std::minus()); + tau_forest_tracker_->UpdatePredictions(tau_forest_.get(), *forest_dataset_.get()); + } else { + Log::Fatal("Unsupported leaf model type for treatment forest"); + } + + // Initialize variance forest state (if present) + if (config_.num_trees_variance > 0) { + variance_leaf_model_ = LogLinearVarianceLeafModel(config_.shape_variance_forest, config_.scale_variance_forest); + variance_forest_ = std::make_unique(config_.num_trees_variance, config_.leaf_dim_variance, config_.leaf_constant_variance, config_.exponentiated_leaf_variance); + samples.variance_forests = std::make_unique(config_.num_trees_variance, config_.leaf_dim_variance, config_.leaf_constant_variance, config_.exponentiated_leaf_variance); + variance_forest_tracker_ = std::make_unique(forest_dataset_->GetCovariates(), config_.feature_types, config_.num_trees_variance, data_.n_train); + tree_prior_variance_ = std::make_unique(config_.alpha_variance, config_.beta_variance, config_.min_samples_leaf_variance, config_.max_depth_variance); + // Leaf values for the log-linear variance model are on the log scale; the ensemble sums + // log(sigma^2_i) contributions, so each tree starts at log(init_val) / num_trees. + variance_forest_->SetLeafValue(std::log(init_val_variance_) / config_.num_trees_variance); + variance_forest_tracker_->UpdatePredictions(variance_forest_.get(), *forest_dataset_.get()); + // UpdateVarModelTree (called inside GFRSampleOneIter / MCMCSampleOneIter) unconditionally + // reads and writes the dataset variance weight slot via VarWeightValue / SetVarWeightValue. + // This slot tracks the cumulative per-observation variance prediction + // (sigma^2_i = exp(sum of tree leaf values)) and is incompatible with case weights, which + // would need to be reapplied after every per-tree update. The R/Python APIs enforce this + // as a hard error; guard here for callers that use BARTSampler directly. + if (forest_dataset_->HasVarWeights()) { + Log::Fatal("observation_weights and a variance forest cannot be used together."); + } + std::vector initial_variance_preds(data_.n_train, init_val_variance_); + forest_dataset_->AddVarianceWeights(initial_variance_preds.data(), data_.n_train); + has_variance_forest_ = true; + } + + // Global error variance model + if (config_.sample_sigma2_global) { + var_model_ = std::make_unique(); + sample_sigma2_global_ = true; + } + + // Leaf scale models + if (config_.sample_sigma2_leaf_mu) { + leaf_scale_model_mu_ = std::make_unique(); + sample_sigma2_leaf_mu_ = true; + } + if (config_.sample_sigma2_leaf_tau) { + leaf_scale_model_tau_ = std::make_unique(); + sample_sigma2_leaf_tau_ = true; + } + + // Treatment intercept model + if (config_.sample_tau_0) { + sample_tau_0_ = true; + if (data_.treatment_dim > 1) { + tau_0_vector_.assign(data_.treatment_dim, 0.0); + if (config_.tau_0_prior_var_multivariate.empty()) { + config_.tau_0_prior_var_multivariate.assign(data_.treatment_dim, config_.sigma2_tau_init * config_.num_trees_tau); + } else { + if ((int)config_.tau_0_prior_var_multivariate.size() != data_.treatment_dim) { + Log::Fatal("tau_0_prior_var_multivariate must have treatment_dim = %d elements, but has %zu", + data_.treatment_dim, config_.tau_0_prior_var_multivariate.size()); + } + } + } else { + tau_0_scalar_ = 0.0; + if (config_.tau_0_prior_var_scalar <= 0.0) { + config_.tau_0_prior_var_scalar = config_.sigma2_tau_init * config_.num_trees_tau; + } + } + } + + // Random effects model + if (config_.has_random_effects) { + random_effects_dataset_ = std::make_unique(); + random_effects_dataset_->AddGroupLabels(data_.rfx_group_ids_train, data_.n_train); + if (data_.rfx_basis_train != nullptr) { + random_effects_dataset_->AddBasis(data_.rfx_basis_train, data_.n_train, data_.rfx_basis_dim, /*row_major=*/false); + } else { + if (config_.rfx_model_spec == BCFRFXModelSpec::InterceptOnly) { + // If no basis is provided, add an intercept basis (column of 1s) + // TODO: do we need to do this before we determine rfx_basis_dim and initialize the RFX data structures? + std::vector rfx_basis(data_.n_train, 1.0); + random_effects_dataset_->AddBasis(rfx_basis.data(), data_.n_train, 1, /*row_major=*/false); + // Override rfx_basis_dim to 1 for intercept-only model the basis is a 1-dimensional vector of ones + data_.rfx_basis_dim = 1; + } else if (config_.rfx_model_spec == BCFRFXModelSpec::InterceptPlusTreatment) { + // If no basis is provided, add an intercept basis (column of 1s) and the treatment variable(s) as the basis + // TODO: do we need to do this before we determine rfx_basis_dim and initialize the RFX data structures? + std::vector rfx_basis(data_.n_train * (1 + data_.treatment_dim), 1.0); + for (int i = 0; i < data_.n_train; i++) { + for (int j = 0; j < data_.treatment_dim; j++) { + rfx_basis[(j + 1) * data_.n_train + i] = data_.treatment_train[j * data_.n_train + i]; + } + } + random_effects_dataset_->AddBasis(rfx_basis.data(), data_.n_train, 1 + data_.treatment_dim, /*row_major=*/false); + // Override rfx_basis_dim to 1 for intercept-only model the basis is a 1-dimensional vector of ones + data_.rfx_basis_dim = 1 + data_.treatment_dim; + } else { + Log::Fatal("Random effects basis data must be provided for non-intercept-only random effects model"); + } + } + // Tracking data structure for random effects groups + random_effects_tracker_ = std::make_unique(data_.rfx_group_ids_train, data_.n_train); + // Container of random effects samples + samples.rfx_container = std::make_unique(data_.rfx_basis_dim, data_.rfx_num_groups); + // Mapping from RFX labels to 0-indexed group IDs for efficient lookup in the sampler; populated from the RFX dataset group labels + samples.rfx_label_mapper = std::make_unique(random_effects_tracker_->GetLabelMap()); + + // Initialize random effects model object + random_effects_model_ = std::make_unique(data_.rfx_basis_dim, data_.rfx_num_groups); + + // Handle "working" parameter prior mean + Eigen::VectorXd working_parameter_prior_mean; + if (!config_.rfx_working_parameter_mean_prior.empty()) { + if ((int)config_.rfx_working_parameter_mean_prior.size() != data_.rfx_basis_dim) { + Log::Fatal("rfx_working_parameter_mean_prior must have rfx_basis_dim = %d elements, but has %zu", + data_.rfx_basis_dim, config_.rfx_working_parameter_mean_prior.size()); + } + // Column-major interpretation matches both R and Eigen (python must be reordered before passing to C++) + working_parameter_prior_mean = Eigen::Map(config_.rfx_working_parameter_mean_prior.data(), data_.rfx_basis_dim); + } else { + working_parameter_prior_mean = Eigen::VectorXd::Zero(data_.rfx_basis_dim); + } + random_effects_model_->SetWorkingParameter(working_parameter_prior_mean); + + // Handle "group" parameter prior mean + Eigen::MatrixXd group_parameter_prior_mean; + if (!config_.rfx_group_parameter_mean_prior.empty()) { + if ((int)config_.rfx_group_parameter_mean_prior.size() != data_.rfx_basis_dim * data_.rfx_num_groups) { + Log::Fatal("rfx_group_parameter_mean_prior must have rfx_basis_dim * rfx_num_groups = %d elements, but has %zu", + data_.rfx_basis_dim * data_.rfx_num_groups, config_.rfx_group_parameter_mean_prior.size()); + } + // Column-major interpretation matches both R and Eigen (python must be reordered before passing to C++) + group_parameter_prior_mean = Eigen::Map(config_.rfx_group_parameter_mean_prior.data(), data_.rfx_basis_dim, data_.rfx_num_groups); + } else { + group_parameter_prior_mean = Eigen::MatrixXd::Zero(data_.rfx_basis_dim, data_.rfx_num_groups); + } + random_effects_model_->SetGroupParameters(group_parameter_prior_mean); + + // Handle "working" parameter prior covariance + Eigen::MatrixXd working_parameter_prior_cov; + if (!config_.rfx_working_parameter_cov_prior.empty()) { + if ((int)config_.rfx_working_parameter_cov_prior.size() != data_.rfx_basis_dim * data_.rfx_basis_dim) { + Log::Fatal("rfx_working_parameter_cov_prior must have rfx_basis_dim * rfx_basis_dim = %d elements, but has %zu", + data_.rfx_basis_dim * data_.rfx_basis_dim, config_.rfx_working_parameter_cov_prior.size()); + } + // Column-major interpretation matches both R and Eigen (python must be reordered before passing to C++) + working_parameter_prior_cov = Eigen::Map(config_.rfx_working_parameter_cov_prior.data(), data_.rfx_basis_dim, data_.rfx_basis_dim); + } else { + working_parameter_prior_cov = Eigen::MatrixXd::Identity(data_.rfx_basis_dim, data_.rfx_basis_dim); + } + random_effects_model_->SetWorkingParameterCovariance(working_parameter_prior_cov); + + // Handle "group" parameter prior covariance + Eigen::MatrixXd group_parameter_prior_cov; + if (!config_.rfx_group_parameter_cov_prior.empty()) { + if ((int)config_.rfx_group_parameter_cov_prior.size() != data_.rfx_basis_dim * data_.rfx_basis_dim) { + Log::Fatal("rfx_group_parameter_cov_prior must have rfx_basis_dim * rfx_basis_dim = %d elements, but has %zu", + data_.rfx_basis_dim * data_.rfx_basis_dim, config_.rfx_group_parameter_cov_prior.size()); + } + // Column-major interpretation matches both R and Eigen (python must be reordered before passing to C++) + group_parameter_prior_cov = Eigen::Map(config_.rfx_group_parameter_cov_prior.data(), data_.rfx_basis_dim, data_.rfx_basis_dim); + } else { + group_parameter_prior_cov = Eigen::MatrixXd::Identity(data_.rfx_basis_dim, data_.rfx_basis_dim); + } + random_effects_model_->SetGroupParameterCovariance(group_parameter_prior_cov); + + // Handle variance model priors + if (config_.rfx_variance_prior_shape <= 0.0) { + config_.rfx_variance_prior_shape = 1.0; + } + if (config_.rfx_variance_prior_scale <= 0.0) { + config_.rfx_variance_prior_scale = 1.0; + } + random_effects_model_->SetVariancePriorShape(config_.rfx_variance_prior_shape); + random_effects_model_->SetVariancePriorScale(config_.rfx_variance_prior_scale); + + // Set has_random_effects_ flag to true so that the sampler will perform random effects updates at each iteration + has_random_effects_ = true; + } + + // RNG + rng_ = std::mt19937(config_.random_seed >= 0 ? config_.random_seed : std::random_device{}()); + + // Other internal model state + if (continuation) { + // Warm-start the scalar state from the last retained sample. Within the sampler lifecycle + // all of these arrays hold STANDARDIZED values (postprocess_samples applies the y_std^2 + // factor to global variance only at the very end; the continuation binding supplies + // standardized history), so read the last value directly here. + if (sample_sigma2_global_ && !samples.global_error_variance_samples.empty()) { + global_variance_ = samples.global_error_variance_samples.back(); + } else { + global_variance_ = config_.sigma2_global_init; + } + if (sample_sigma2_leaf_mu_ && !samples.leaf_scale_mu_samples.empty()) { + leaf_scale_mu_ = samples.leaf_scale_mu_samples.back(); + } else { + leaf_scale_mu_ = config_.sigma2_mu_init; + } + if (sample_sigma2_leaf_tau_ && !samples.leaf_scale_tau_samples.empty()) { + leaf_scale_tau_ = samples.leaf_scale_tau_samples.back(); + } else { + leaf_scale_tau_ = config_.sigma2_tau_init; + } + leaf_scale_tau_multivariate_ = config_.sigma2_leaf_tau_matrix; + // Sync the leaf models' scales with the warm-started leaf scales so the first continued + // iteration samples each forest using the last retained scale (matching the one-shot run). + mu_leaf_model_.SetScale(leaf_scale_mu_); + std::visit(ScaleUpdateVisitor{*this, leaf_scale_tau_}, tau_leaf_model_); + // Restore the treatment intercept and remove its contribution from the residual. + // After the forest warm-starts above, residual_ = y_std - mu_last - Z*tau_last; the one-shot + // sampler additionally carries -Z*tau_0_last at the start of the next iteration. + if (sample_tau_0_ && !samples.tau_0_samples.empty()) { + tau_0_scalar_ = samples.tau_0_samples.back(); + double* resid_ptr = residual_->GetData().data(); + const double* basis = data_.treatment_train; // adaptive coding is guarded off for continuation + for (int i = 0; i < data_.n_train; i++) { + resid_ptr[i] -= tau_0_scalar_ * basis[i]; + } + } + } else { + global_variance_ = config_.sigma2_global_init; + leaf_scale_mu_ = config_.sigma2_mu_init; + leaf_scale_tau_ = config_.sigma2_tau_init; + leaf_scale_tau_multivariate_ = config_.sigma2_leaf_tau_matrix; + } + + tau_raw_sum_preds_.assign(data_.n_train * data_.treatment_dim, 0.0); + + initialized_ = true; +} + +void BCFSampler::run_gfr(BCFSamples& samples, int num_gfr, bool keep_gfr, int num_chains) { + // Reserve space for GFR predictions if they are to be retained + if (keep_gfr) { + samples.mu_forest_predictions_train.reserve(data_.n_train * num_gfr); + samples.tau_forest_predictions_train.reserve(data_.n_train * data_.treatment_dim * num_gfr); + if (has_variance_forest_) { + samples.variance_forest_predictions_train.reserve(data_.n_train * num_gfr); + } + } + + // NOTE: for serial sampling (which is all we currently support), chain 1 uses the live sampler state after + // the GFR loop (i.e. the state after GFR iteration num_gfr-1). Chains 2..N each need their own earlier + // GFR starting point so that the chains are initialized from distinct states. + // We save exactly num_chains-1 snapshots, one per "extra" chain: + // gfr_snapshots_[k] = state after GFR iteration (num_gfr - num_chains + k), for k = 0..num_chains-2. + // The last GFR iteration (i = num_gfr-1) is NOT snapshotted because chain 1 uses the live state. + // Chain j (1-indexed, j >= 2) uses gfr_snapshots_[num_chains-j] = state after GFR[num_gfr-j]. + // If num_chains is 1, we keep no snapshots. + int snapshot_start = (num_chains > 1) ? std::max(0, num_gfr - num_chains) : num_gfr; + if (num_chains > 1) { + gfr_snapshots_.clear(); + gfr_snapshots_.reserve(num_chains - 1); + } + + if (config_.verbose && num_gfr > 0) { + Log::Info("Running GFR sampler (%d iterations)", num_gfr); + } + const int gfr_report_every = std::max(1, num_gfr / 10); + bool write_snapshot = false; + for (int i = 0; i < num_gfr; i++) { + // Do not snapshot the final GFR iteration: chain 1 uses the live sampler state directly. + write_snapshot = (i >= snapshot_start) && (i < num_gfr - 1); + RunOneIteration(samples, /*gfr=*/true, /*keep_sample=*/keep_gfr, /*write_snapshot=*/write_snapshot); + if (config_.verbose && ((i + 1) % gfr_report_every == 0 || i + 1 == num_gfr)) { + Log::Info("GFR: %d%% (%d/%d)", (100 * (i + 1)) / num_gfr, i + 1, num_gfr); + } + } +} + +void BCFSampler::run_mcmc(BCFSamples& samples, int num_burnin, int keep_every, int num_mcmc) { + // Reserve space for MCMC predictions if they are to be retained + samples.mu_forest_predictions_train.reserve(data_.n_train * num_mcmc); + samples.tau_forest_predictions_train.reserve(data_.n_train * data_.treatment_dim * num_mcmc); + if (has_test_) { + samples.mu_forest_predictions_test.reserve(data_.n_test * num_mcmc); + samples.tau_forest_predictions_test.reserve(data_.n_test * data_.treatment_dim * num_mcmc); + } + if (has_variance_forest_) { + samples.variance_forest_predictions_train.reserve(data_.n_train * num_mcmc); + if (has_test_) { + samples.variance_forest_predictions_test.reserve(data_.n_test * num_mcmc); + } + } + + // Create leaf models and pass them to the RunOneIteration function; these are updated in place and will reflect the current state of the leaf scale parameters (if they are being sampled) + bool keep_forest = false; + const int mcmc_total = num_burnin + keep_every * num_mcmc; + const int mcmc_report_every = std::max(1, mcmc_total / 10); + for (int i = 0; i < mcmc_total; i++) { + if (i >= num_burnin && (i - num_burnin) % keep_every == 0) + keep_forest = true; + else + keep_forest = false; + RunOneIteration(samples, /*gfr=*/false, /*keep_sample=*/keep_forest, /*write_snapshot=*/false); + if (config_.verbose && ((i + 1) % mcmc_report_every == 0 || i + 1 == mcmc_total)) { + Log::Info("MCMC: %d%% (%d/%d)", (100 * (i + 1)) / mcmc_total, i + 1, mcmc_total); + } + } +} + +void BCFSampler::run_mcmc_chains(BCFSamples& samples, int num_chains, int num_burnin, int keep_every, int num_mcmc) { + for (int chain_idx = 0; chain_idx < num_chains; chain_idx++) { + if (chain_idx > 0 && !gfr_snapshots_.empty()) { + // Re-initialize the sampler state for each new chain. + // gfr_snapshots_ holds num_chains-1 states (oldest-first): index 0 = GFR[num_gfr-num_chains], + // index num_chains-2 = GFR[num_gfr-2]. Chain j (1-indexed, j>=2) uses index num_chains-j. + // When num_gfr < num_chains we may not have enough distinct snapshots; in that case + // fall back to running chains from whatever state is available (same behavior as R). + int snapshot_idx = num_chains - 1 - chain_idx; + if (snapshot_idx >= 0 && snapshot_idx < static_cast(gfr_snapshots_.size())) { + RestoreStateFromGFRSnapshot(samples, snapshot_idx); + } + } + if (config_.verbose) { + Log::Info("Running MCMC chain %d/%d (%d samples)", chain_idx + 1, num_chains, num_mcmc); + } + run_mcmc(samples, num_burnin, keep_every, num_mcmc); + } +} + +void BCFSampler::postprocess_samples(BCFSamples& samples) { + // Compute outcome predictions on the linear (link) scale: E[eta|X,Z] = mu(X) + Z*tau(X) + rfx + // tau_forest_predictions stores raw tau(x) (no z multiplication), so we multiply by z here. + // Callers that need probability-scale predictions (probit, cloglog) apply the inverse link themselves. + samples.y_hat_train.resize(data_.n_train * samples.num_samples); + double mu_term, tau_term, y_term; + const int treatment_dim = data_.treatment_dim; + for (int j = 0; j < samples.num_samples; j++) { + for (int i = 0; i < data_.n_train; i++) { + // Data index for the two terms that are guaranteed to be univariate - mu(x) and y_hat + const int k = j * data_.n_train + i; + mu_term = samples.mu_forest_predictions_train[k]; + if (treatment_dim > 1) { + tau_term = 0; + for (int treatment_idx = 0; treatment_idx < treatment_dim; treatment_idx++) { + // Starting data index for multivariate treatment case, where tau(x) is col-major with dimensions (n_train, treatment_dim, num_samples) + const int k_tau = j * data_.n_train * treatment_dim + data_.n_train * treatment_idx + i; + tau_term += samples.tau_forest_predictions_train[k_tau] * data_.treatment_train[data_.n_train * treatment_idx + i]; + } + } else { + tau_term = samples.tau_forest_predictions_train[k] * data_.treatment_train[i]; + } + y_term = mu_term + tau_term; + if (has_random_effects_) y_term += samples.rfx_predictions_train[k]; + samples.y_hat_train[k] = y_term * samples.y_std + samples.y_bar; + } + } + + // Unpack test set predictions for mean and variance forest + if (has_test_) { + std::vector mu_predictions = samples.mu_forests->Predict(*forest_dataset_test_); + std::vector tau_predictions = samples.tau_forests->PredictRaw(*forest_dataset_test_, /*row_major=*/false); + // Add tau_0 to the treatment effect function predictions if it was sampled. + // tau_0_samples layout: col-major (treatment dim k, sample j) -> j * treatment_dim + k. + // For treatment_dim==1 this collapses to samples.tau_0_samples[j]. + if (sample_tau_0_) { + const int treatment_dim = data_.treatment_dim; + for (int j = 0; j < samples.num_samples; j++) { + for (int k = 0; k < treatment_dim; k++) { + for (int i = 0; i < data_.n_test; i++) { + const int idx = j * data_.n_test * treatment_dim + data_.n_test * k + i; + tau_predictions[idx] += samples.tau_0_samples[j * treatment_dim + k]; + } + } + } + } + // Handle adaptive coding correctly: + // When treatment is b_0 (1-Z) + b_1 Z, the conditional mean model: + // mu(x) + [tau_0 + tau(x)] * (b_0 * (1-Z) + b_1 * Z) + // turns into + // [mu(x) + b_0 * (tau_0 + tau(x))] + (tau_0 + tau(x)) * (b_1 - b_0) * Z + // So the treatment effect function that gets multiplied by Z is actually (b_1 - b_0) * (tau_0 + tau(x)) + // and the prognostic function has an added contribution of b_0 * (tau_0 + tau(x)) + if (adaptive_coding_) { + for (int i = 0; i < samples.num_samples; i++) { + double b_0 = samples.b0_samples[i]; + double b_1 = samples.b1_samples[i]; + for (int j = 0; j < data_.n_test; j++) { + const int idx = i * data_.n_test + j; + // Add b_0 * (tau_0 + tau(x)) to the prognostic function predictions + mu_predictions[idx] += b_0 * tau_predictions[idx]; + // Scale tau_predictions by (b_1 - b_0) + tau_predictions[idx] *= (b_1 - b_0); + } + } + } + samples.mu_forest_predictions_test.insert(samples.mu_forest_predictions_test.end(), + mu_predictions.data(), mu_predictions.data() + mu_predictions.size()); + samples.tau_forest_predictions_test.insert(samples.tau_forest_predictions_test.end(), + tau_predictions.data(), tau_predictions.data() + tau_predictions.size()); + if (has_variance_forest_) { + std::vector predictions = samples.variance_forests->Predict(*forest_dataset_test_); + samples.variance_forest_predictions_test.insert(samples.variance_forest_predictions_test.end(), + predictions.data(), predictions.data() + predictions.size()); + } + if (has_random_effects_) { + RandomEffectsDataset rfx_dataset_test; + rfx_dataset_test.AddGroupLabels(data_.rfx_group_ids_test, data_.n_test); + if (data_.rfx_basis_test != nullptr) { + rfx_dataset_test.AddBasis(data_.rfx_basis_test, data_.n_test, data_.rfx_basis_dim, /*row_major=*/false); + } else if (config_.rfx_model_spec == BCFRFXModelSpec::InterceptOnly) { + std::vector rfx_basis(data_.n_test, 1.0); + rfx_dataset_test.AddBasis(rfx_basis.data(), data_.n_test, 1, /*row_major=*/false); + } else if (config_.rfx_model_spec == BCFRFXModelSpec::InterceptPlusTreatment) { + // Column-major rfx basis + std::vector rfx_basis(data_.n_test * (1 + data_.treatment_dim)); + for (int i = 0; i < data_.n_test; i++) { + rfx_basis[i] = 1.0; + } + for (int j = 0; j < data_.treatment_dim; j++) { + for (int i = 0; i < data_.n_test; i++) { + rfx_basis[(j + 1) * data_.n_test + i] = data_.treatment_test[j * data_.n_test + i]; + } + } + rfx_dataset_test.AddBasis(rfx_basis.data(), data_.n_test, 1 + data_.treatment_dim, /*row_major=*/false); + } else { + Log::Fatal("BCF model random effects term was not sampled with intercept_only or intercept_plus_treatment specification, but not random effect basis was provided for prediction"); + } + samples.rfx_predictions_test.resize(data_.n_test * samples.num_samples); + samples.rfx_container->Predict(rfx_dataset_test, *samples.rfx_label_mapper, samples.rfx_predictions_test); + } + + samples.y_hat_test.resize(data_.n_test * samples.num_samples); + for (int j = 0; j < samples.num_samples; j++) { + for (int i = 0; i < data_.n_test; i++) { + // Data index for the two terms that are guaranteed to be univariate - mu(x) and y_hat + const int k = j * data_.n_test + i; + mu_term = samples.mu_forest_predictions_test[k]; + if (treatment_dim > 1) { + tau_term = 0; + for (int treatment_idx = 0; treatment_idx < treatment_dim; treatment_idx++) { + // Starting data index for multivariate treatment case, where tau(x) is col-major with dimensions (n_test, treatment_dim, num_samples) + const int k_tau = j * data_.n_test * treatment_dim + data_.n_test * treatment_idx + i; + tau_term += samples.tau_forest_predictions_test[k_tau] * data_.treatment_test[data_.n_test * treatment_idx + i]; + } + } else { + tau_term = samples.tau_forest_predictions_test[k] * data_.treatment_test[i]; + } + y_term = mu_term + tau_term; + if (has_random_effects_) y_term += samples.rfx_predictions_test[k]; + samples.y_hat_test[k] = y_term * samples.y_std + samples.y_bar; + } + } + } + + // Convert variance forest predictions and global error variance from + // standardized space to original outcome scale. + // - Train predictions come from ForestTracker::GetSumPredictions() (log-scale leaf sums), + // so apply exp() then multiply by y_std^2. + // - Test predictions come from ForestContainer::Predict() with is_exponentiated_=true, + // which already applies exp() internally, so just multiply by y_std^2. + // - Global error variance samples are in standardized space; multiply by y_std^2. + if (has_variance_forest_) { + double y_std2 = samples.y_std * samples.y_std; + for (double& v : samples.variance_forest_predictions_train) v = std::exp(v) * y_std2; + for (double& v : samples.variance_forest_predictions_test) v *= y_std2; + } + if (sample_sigma2_global_) { + double y_std2 = samples.y_std * samples.y_std; + for (double& v : samples.global_error_variance_samples) v *= y_std2; + } + + // Convert the cached prognostic / treatment-effect / random-effects predictions from standardized + // space to the original outcome scale, matching the R main-branch cached attributes and predict(): + // the prognostic function mu(x) carries the location shift (y_bar); the treatment effect tau(x) + // (already the full CATE tau_0 + tau(x), with any adaptive-coding (b1 - b0) factor folded in during + // sampling) and the random effects carry only the scale factor (y_std). y_hat_train / y_hat_test + // were already placed on the original scale above, computed from these caches while still standardized. + for (double& v : samples.mu_forest_predictions_train) v = v * samples.y_std + samples.y_bar; + for (double& v : samples.mu_forest_predictions_test) v = v * samples.y_std + samples.y_bar; + for (double& v : samples.tau_forest_predictions_train) v *= samples.y_std; + for (double& v : samples.tau_forest_predictions_test) v *= samples.y_std; + if (has_random_effects_) { + for (double& v : samples.rfx_predictions_train) v *= samples.y_std; + for (double& v : samples.rfx_predictions_test) v *= samples.y_std; + } +} + +void BCFSampler::RunOneIteration(BCFSamples& samples, bool gfr, bool keep_sample, bool write_snapshot) { + // Prognostic forest + if (gfr) { + GFRSampleOneIter( + *mu_forest_, *mu_forest_tracker_, *samples.mu_forests, mu_leaf_model_, + *forest_dataset_, *residual_, *tree_prior_mu_, rng_, + config_.var_weights_mu, config_.sweep_update_indices_mu, global_variance_, config_.feature_types, + config_.cutpoint_grid_size, /*keep_forest=*/keep_sample, + /*pre_initialized=*/true, /*backfitting=*/true, + /*num_features_subsample=*/config_.num_features_subsample_mu, config_.num_threads); + } else { + MCMCSampleOneIter( + *mu_forest_, *mu_forest_tracker_, *samples.mu_forests, mu_leaf_model_, + *forest_dataset_, *residual_, *tree_prior_mu_, rng_, + config_.var_weights_mu, config_.sweep_update_indices_mu, global_variance_, /*keep_forest=*/keep_sample, + /*pre_initialized=*/true, /*backfitting=*/true, + /*num_threads=*/config_.num_threads); + } + + // Parametric treatment intercept term + if (sample_tau_0_) { + SampleParametricTreatmentEffect(); + } + + // Treatment effect forest + if (gfr) { + std::visit(GFROneIterationVisitorTau{*this, samples, keep_sample}, tau_leaf_model_); + } else { + std::visit(MCMCOneIterationVisitorTau{*this, samples, keep_sample}, tau_leaf_model_); + } + + // Update raw tau(x): sum leaf values across trees for each dimension of the tau leaf. + // Uses node IDs already cached in the tracker — no tree traversal needed. + // Stored col-major: tau_raw_sum_preds_[k * n_train + i] matches postprocess_samples indexing. + const int tau_dim = data_.treatment_dim; + const int data_dim = data_.n_train; + double tau_0 = 0.0; + for (int k = 0; k < tau_dim; k++) { + if (sample_tau_0_) { + if (data_.treatment_dim > 1) { + tau_0 = tau_0_vector_[k]; + } else { + tau_0 = tau_0_scalar_; + } + } + for (int i = 0; i < data_dim; i++) { + tau_raw_sum_preds_[k * data_dim + i] = tau_0; + for (int j = 0; j < config_.num_trees_tau; j++) { + data_size_t leaf = tau_forest_tracker_->GetNodeId(i, j); + tau_raw_sum_preds_[k * data_dim + i] += tau_forest_->GetTree(j)->LeafValue(leaf, k); + } + } + } + + // Adaptive coding parameters + if (adaptive_coding_) { + SampleAdaptiveCodingParameters(); + } + + // Variance forest term + if (has_variance_forest_) { + if (gfr) { + GFRSampleOneIter( + *variance_forest_, *variance_forest_tracker_, *samples.variance_forests, variance_leaf_model_, + *forest_dataset_, *residual_, *tree_prior_variance_, rng_, + config_.var_weights_variance, config_.sweep_update_indices_variance, global_variance_, config_.feature_types, + config_.cutpoint_grid_size, /*keep_forest=*/keep_sample, + /*pre_initialized=*/true, /*backfitting=*/false, + /*num_features_subsample=*/config_.num_features_subsample_variance, config_.num_threads); + } else { + MCMCSampleOneIter( + *variance_forest_, *variance_forest_tracker_, *samples.variance_forests, variance_leaf_model_, + *forest_dataset_, *residual_, *tree_prior_variance_, rng_, + config_.var_weights_variance, config_.sweep_update_indices_variance, global_variance_, /*keep_forest=*/keep_sample, + /*pre_initialized=*/true, /*backfitting=*/false, + /*num_threads=*/config_.num_threads); + } + } + + // Latent continuous outcome for probit link + if (config_.link_function == LinkFunction::Probit) { + // Add mu(x) + Z*tau(x) + rfx to model_preds_ for each training observation, then sample the latent outcome given the observed binary outcome and the model prediction + AddModelTermsForProbit(model_preds_.data(), mu_forest_tracker_.get(), tau_forest_tracker_.get(), random_effects_tracker_.get(), data_.n_train); + // If tau_0 is sampled, then add it to model_preds_ as well + if (sample_tau_0_) { + if (data_.treatment_dim > 1) { + for (int i = 0; i < data_.n_train; i++) { + for (int k = 0; k < data_.treatment_dim; k++) { + model_preds_[i] += data_.treatment_train[data_.n_train * k + i] * tau_0_vector_[k]; + } + } + } else { + const double* treatment_ptr = adaptive_coding_ ? tau_basis_vector_train_.data() : data_.treatment_train; + for (int i = 0; i < data_.n_train; i++) { + model_preds_[i] += tau_0_scalar_ * treatment_ptr[i]; + } + } + } + // Sample latent outcome into outcome_raw_ (overwriting the previous iteration's raw predictions, which are not needed for the probit likelihood) + sample_probit_latent_outcome(rng_, outcome_raw_->GetData().data(), model_preds_.data(), + residual_->GetData().data(), samples.y_bar, data_.n_train); + } + + // Global error scale + if (sample_sigma2_global_) { + global_variance_ = var_model_->SampleVarianceParameter( + residual_->GetData(), config_.a_sigma2_global, config_.b_sigma2_global, rng_); + } + + // Prognostic forest leaf scale + if (sample_sigma2_leaf_mu_) { + leaf_scale_mu_ = leaf_scale_model_mu_->SampleVarianceParameter( + mu_forest_.get(), config_.a_sigma2_mu, config_.b_sigma2_mu, rng_); + mu_leaf_model_.SetScale(leaf_scale_mu_); + } + + // Treatment effect forest leaf scale + if (sample_sigma2_leaf_tau_) { + leaf_scale_tau_ = leaf_scale_model_tau_->SampleVarianceParameter( + tau_forest_.get(), config_.a_sigma2_tau, config_.b_sigma2_tau, rng_); + std::visit(ScaleUpdateVisitor{*this, leaf_scale_tau_}, tau_leaf_model_); + } + + // Gibbs updates for random effects model + if (has_random_effects_) { + random_effects_model_->SampleRandomEffects(*random_effects_dataset_, *residual_, *random_effects_tracker_, global_variance_, rng_); + // NOTE: we keep this code in the random effects sampling block (as opposed to the parameter / prediction storage block below) to mirror the way that forests are retained within a sampling step + if (keep_sample) { + samples.rfx_container->AddSample(*random_effects_model_); + } + } + + if (keep_sample) { + // Add parameter and prediction samples + samples.num_samples++; + // Global error variance + if (sample_sigma2_global_) samples.global_error_variance_samples.push_back(global_variance_); + // Prognostic forest leaf scale + if (sample_sigma2_leaf_mu_) samples.leaf_scale_mu_samples.push_back(leaf_scale_mu_); + // Treatment effect forest leaf scale + if (sample_sigma2_leaf_tau_) samples.leaf_scale_tau_samples.push_back(leaf_scale_tau_); + // Treatment intercept + if (sample_tau_0_) { + if (data_.treatment_dim > 1) { + samples.tau_0_samples.insert(samples.tau_0_samples.end(), tau_0_vector_.begin(), tau_0_vector_.end()); + } else { + samples.tau_0_samples.push_back(tau_0_scalar_); + } + } + // Adaptive coding parameters + if (adaptive_coding_) { + samples.b0_samples.push_back(b_0_); + samples.b1_samples.push_back(b_1_); + } + // Prognostic and treatment forest predictions + double* mu_forest_preds_train = mu_forest_tracker_->GetSumPredictions(); + if (adaptive_coding_) { + // TODO: refactor this or at least cache to avoid unnecessary malloc + std::vector mu_adj(data_.n_train); + std::vector tau_adj(data_.n_train); + for (int i = 0; i < data_.n_train; i++) { + mu_adj[i] = mu_forest_preds_train[i] + b_0_ * tau_raw_sum_preds_[i]; + tau_adj[i] = tau_raw_sum_preds_[i] * (b_1_ - b_0_); + } + samples.mu_forest_predictions_train.insert(samples.mu_forest_predictions_train.end(), + mu_adj.begin(), mu_adj.end()); + samples.tau_forest_predictions_train.insert(samples.tau_forest_predictions_train.end(), + tau_adj.begin(), tau_adj.end()); + } else { + samples.mu_forest_predictions_train.insert(samples.mu_forest_predictions_train.end(), + mu_forest_preds_train, + mu_forest_preds_train + samples.num_train); + samples.tau_forest_predictions_train.insert(samples.tau_forest_predictions_train.end(), + tau_raw_sum_preds_.begin(), tau_raw_sum_preds_.end()); + } + // Variance forest predictions + if (has_variance_forest_) { + double* variance_forest_preds_train = variance_forest_tracker_->GetSumPredictions(); + samples.variance_forest_predictions_train.insert(samples.variance_forest_predictions_train.end(), + variance_forest_preds_train, + variance_forest_preds_train + samples.num_train); + } + // Random effects predictions + if (has_random_effects_) { + for (int i = 0; i < data_.n_train; i++) { + samples.rfx_predictions_train.push_back(random_effects_tracker_->GetPrediction(i)); + } + } + } + + if (write_snapshot) { + GFRSnapshot snap; + // Forests + snap.mu_forest = std::make_unique(*mu_forest_); + snap.tau_forest = std::make_unique(*tau_forest_); + if (has_variance_forest_) snap.variance_forest = std::make_unique(*variance_forest_); + // Scale parameters + snap.sigma2 = global_variance_; + snap.leaf_scale_mu = leaf_scale_mu_; + if (config_.tau_leaf_model_type == MeanLeafModelType::GaussianMultivariateRegression) { + snap.leaf_scale_tau_multivariate = leaf_scale_tau_multivariate_; + } else if (config_.tau_leaf_model_type == MeanLeafModelType::GaussianUnivariateRegression) { + snap.leaf_scale_tau = leaf_scale_tau_; + } + // Treatment intercept + if (sample_tau_0_) { + if (data_.treatment_dim > 1) { + snap.tau_0_vector = tau_0_vector_; + } else { + snap.tau_0_scalar = tau_0_scalar_; + } + } + // Adaptive coding + if (adaptive_coding_) { + snap.b_0 = b_0_; + snap.b_1 = b_1_; + } + // Residual + snap.residual.clear(); + snap.residual.resize(data_.n_train); + snap.residual.assign(residual_->GetData().data(), residual_->GetData().data() + data_.n_train); + // Variance weights (from variance forest) + if (has_variance_forest_) { + snap.variance_weights.clear(); + snap.variance_weights.resize(data_.n_train); + snap.variance_weights.assign(forest_dataset_->GetVarWeights().data(), forest_dataset_->GetVarWeights().data() + data_.n_train); + } + // Random effects terms + if (config_.has_random_effects) { + snap.rfx_working_parameter = random_effects_model_->GetWorkingParameter(); + snap.rfx_group_parameters = random_effects_model_->GetGroupParameters(); + snap.rfx_group_parameter_covariance = random_effects_model_->GetGroupParameterCovariance(); + snap.rfx_working_parameter_covariance = random_effects_model_->GetWorkingParameterCovariance(); + snap.rfx_variance_prior_shape = random_effects_model_->GetVariancePriorShape(); + snap.rfx_variance_prior_scale = random_effects_model_->GetVariancePriorScale(); + } + gfr_snapshots_.push_back(std::move(snap)); + } +} + +void BCFSampler::RestoreStateFromGFRSnapshot(BCFSamples& samples, int snapshot_index) { + GFRSnapshot& snap = gfr_snapshots_[snapshot_index]; + + // Remove the contribution of tau_0 from the residual before all forest-based state restoration, which modifies the residual in-place + if (sample_tau_0_) { + double* resid_ptr = residual_->GetData().data(); + if (data_.treatment_dim == 1) { + const double* previous_basis = adaptive_coding_ ? tau_basis_vector_train_.data() : data_.treatment_train; + for (int i = 0; i < data_.n_train; i++) { + resid_ptr[i] += tau_0_scalar_ * previous_basis[i]; + } + } else { + for (int i = 0; i < data_.n_train; i++) { + for (int k = 0; k < data_.treatment_dim; k++) { + resid_ptr[i] += data_.treatment_train[k * data_.n_train + i] * tau_0_vector_[k]; + } + } + } + } + + // Restore mu and tau forest state + // Prognostic forest + mu_forest_->ReconstituteFromForest(*snap.mu_forest); + mu_forest_tracker_->ReconstituteFromForest(*snap.mu_forest, *forest_dataset_, *residual_, true); + mu_forest_tracker_->UpdatePredictions(mu_forest_.get(), *forest_dataset_.get()); + + // Adaptive coding parameters and their implied basis + if (adaptive_coding_) { + b_0_ = snap.b_0; + b_1_ = snap.b_1; + for (int i = 0; i < data_.n_train; i++) { + double z = data_.treatment_train[i]; + tau_basis_vector_train_[i] = b_0_ * (1 - z) + b_1_ * z; + } + forest_dataset_->UpdateBasis(tau_basis_vector_train_.data(), /*num_row=*/data_.n_train, /*num_col=*/1, /*row_major=*/false); + if (has_test_ && data_.treatment_test != nullptr) { + for (int i = 0; i < data_.n_test; i++) { + double z = data_.treatment_test[i]; + tau_basis_vector_test_[i] = b_0_ * (1 - z) + b_1_ * z; + } + forest_dataset_test_->UpdateBasis(tau_basis_vector_test_.data(), /*num_row=*/data_.n_test, /*num_col=*/1, /*row_major=*/false); + } + } + + // Treatment intercept + if (sample_tau_0_) { + if (data_.treatment_dim > 1) { + tau_0_vector_ = snap.tau_0_vector; + } else { + tau_0_scalar_ = snap.tau_0_scalar; + } + } + + // Remove tau_0 from residual + if (sample_tau_0_) { + double* resid_ptr = residual_->GetData().data(); + if (data_.treatment_dim == 1) { + const double* current_basis = adaptive_coding_ ? tau_basis_vector_train_.data() : data_.treatment_train; + for (int i = 0; i < data_.n_train; i++) { + resid_ptr[i] -= tau_0_scalar_ * current_basis[i]; + } + } else { + for (int i = 0; i < data_.n_train; i++) { + for (int k = 0; k < data_.treatment_dim; k++) { + resid_ptr[i] -= data_.treatment_train[k * data_.n_train + i] * tau_0_vector_[k]; + } + } + } + } + + // Treatment effect forest + std::visit(TauForestResetVisitor{*this, samples, *snap.tau_forest}, tau_leaf_model_); + + // Initialize variance forest state (if present) + if (config_.num_trees_variance > 0) { + variance_leaf_model_ = LogLinearVarianceLeafModel(config_.shape_variance_forest, config_.scale_variance_forest); + variance_forest_ = std::make_unique(config_.num_trees_variance, config_.leaf_dim_variance, config_.leaf_constant_variance, config_.exponentiated_leaf_variance); + samples.variance_forests = std::make_unique(config_.num_trees_variance, config_.leaf_dim_variance, config_.leaf_constant_variance, config_.exponentiated_leaf_variance); + variance_forest_tracker_ = std::make_unique(forest_dataset_->GetCovariates(), config_.feature_types, config_.num_trees_variance, data_.n_train); + tree_prior_variance_ = std::make_unique(config_.alpha_variance, config_.beta_variance, config_.min_samples_leaf_variance, config_.max_depth_variance); + // Leaf values for the log-linear variance model are on the log scale; the ensemble sums + // log(sigma^2_i) contributions, so each tree starts at log(init_val) / num_trees. + variance_forest_->SetLeafValue(std::log(init_val_variance_) / config_.num_trees_variance); + variance_forest_tracker_->UpdatePredictions(variance_forest_.get(), *forest_dataset_.get()); + // UpdateVarModelTree (called inside GFRSampleOneIter / MCMCSampleOneIter) unconditionally + // reads and writes the dataset variance weight slot via VarWeightValue / SetVarWeightValue. + // This slot tracks the cumulative per-observation variance prediction + // (sigma^2_i = exp(sum of tree leaf values)) and is incompatible with case weights, which + // would need to be reapplied after every per-tree update. The R/Python APIs enforce this + // as a hard error; guard here for callers that use BARTSampler directly. + if (forest_dataset_->HasVarWeights()) { + Log::Fatal("observation_weights and a variance forest cannot be used together."); + } + std::vector initial_variance_preds(data_.n_train, init_val_variance_); + forest_dataset_->AddVarianceWeights(initial_variance_preds.data(), data_.n_train); + has_variance_forest_ = true; + } + + // Random effects model + if (config_.has_random_effects) { + // Restore "working" parameter prior mean + random_effects_model_->SetWorkingParameter(snap.rfx_working_parameter); + + // Restore "group" parameter prior mean + random_effects_model_->SetGroupParameters(snap.rfx_group_parameters); + + // Restore "working" parameter prior covariance + random_effects_model_->SetWorkingParameterCovariance(snap.rfx_working_parameter_covariance); + + // Restore "group" parameter prior covariance + random_effects_model_->SetGroupParameterCovariance(snap.rfx_group_parameter_covariance); + + // Restore variance model priors + random_effects_model_->SetVariancePriorShape(snap.rfx_variance_prior_shape); + random_effects_model_->SetVariancePriorScale(snap.rfx_variance_prior_scale); + + // Swap the chain-N RFX contribution out of the residual and the GFR-snapshot + // contribution in, exactly as the R sampler does via resetRandomEffectsTracker. + // At this point residual_ = y - f_gfr - rfx_chain_N (forest already swapped above), + // so ResetFromSample produces: residual_ += rfx_chain_N - rfx_gfr = y - f_gfr - rfx_gfr. + random_effects_tracker_->ResetFromSample(*random_effects_model_, *random_effects_dataset_, *residual_); + } + + // Other internal model state + global_variance_ = snap.sigma2; + leaf_scale_mu_ = snap.leaf_scale_mu; + if (config_.tau_leaf_model_type == MeanLeafModelType::GaussianMultivariateRegression) { + leaf_scale_tau_multivariate_ = snap.leaf_scale_tau_multivariate; + } else if (config_.tau_leaf_model_type == MeanLeafModelType::GaussianUnivariateRegression || config_.tau_leaf_model_type == MeanLeafModelType::GaussianConstant) { + leaf_scale_tau_ = snap.leaf_scale_tau; + } +} + +void BCFSampler::SampleParametricTreatmentEffect() { + // Determine whether treatment is univariate, bivariate, or multivariate + const int tau_dim = data_.treatment_dim; + const int n_train = data_.n_train; + if (tau_dim == 1) { + // Dispatch univariate specialization of regression sampler + // Add tau_0 * Z to residual to get partial residual + double* partial_resid_ptr = residual_->GetData().data(); + double* treatment_ptr = adaptive_coding_ ? tau_basis_vector_train_.data() : data_.treatment_train; + for (int i = 0; i < n_train; i++) { + partial_resid_ptr[i] += treatment_ptr[i] * tau_0_scalar_; + } + // Sample tau_0 via sample_univariate_gaussian_regression_coefficient + double tau_0_update = sample_univariate_gaussian_regression_coefficient(partial_resid_ptr, treatment_ptr, global_variance_, config_.tau_0_prior_var_scalar, n_train, rng_); + tau_0_scalar_ = tau_0_update; + // Subtract tau_0 * Z from partial residual + for (int i = 0; i < n_train; i++) { + partial_resid_ptr[i] -= treatment_ptr[i] * tau_0_scalar_; + } + } else if (tau_dim == 2) { + // Add tau_0 * Z to residual to get partial residual + double* partial_resid_ptr = residual_->GetData().data(); + for (int i = 0; i < n_train; i++) { + for (int k = 0; k < 2; k++) { + partial_resid_ptr[i] += data_.treatment_train[k * n_train + i] * tau_0_vector_[k]; + } + } + // Dispatch bivariate specialization of regression sampler (with diagonal covariance) + std::vector tau_0_update(2, 0.0); + sample_diagonal_bivariate_gaussian_regression_coefficients(tau_0_update.data(), partial_resid_ptr, data_.treatment_train, data_.treatment_train + n_train, global_variance_, config_.tau_0_prior_var_multivariate[0], config_.tau_0_prior_var_multivariate[1], n_train, rng_); + // Push results back to tau_0_vector_ + tau_0_vector_[0] = tau_0_update[0]; + tau_0_vector_[1] = tau_0_update[1]; + // Subtract tau_0 * Z from partial residual + for (int i = 0; i < n_train; i++) { + for (int k = 0; k < 2; k++) { + partial_resid_ptr[i] -= data_.treatment_train[k * n_train + i] * tau_0_vector_[k]; + } + } + } else { + // Dispatch general-purpose multivariate regression sampler, which returns parameters as an Eigen::VectorXd + // Add tau_0 * Z to residual to get partial residual + double* partial_resid_ptr = residual_->GetData().data(); + for (int i = 0; i < n_train; i++) { + for (int k = 0; k < tau_dim; k++) { + partial_resid_ptr[i] += data_.treatment_train[k * n_train + i] * tau_0_vector_[k]; + } + } + // Wrap an Eigen map around the partial residual, treatment, and prior covariance for efficient vectorized operations + Eigen::Map partial_resid(partial_resid_ptr, n_train); + Eigen::Map treatment(data_.treatment_train, n_train, tau_dim); + // Construct diagonal prior covariance matrix from config_.tau_0_prior_var_multivariate + Eigen::Map tau_0_prior_var_vec(config_.tau_0_prior_var_multivariate.data(), tau_dim); + const Eigen::MatrixXd tau_0_prior_cov = tau_0_prior_var_vec.asDiagonal(); + // Sample tau_0 via sample_general_gaussian_regression_coefficients + Eigen::VectorXd tau_0_update = sample_general_gaussian_regression_coefficients(partial_resid, treatment, global_variance_, tau_0_prior_cov, n_train, rng_); + // Push results back to tau_0_vector_ + for (int k = 0; k < tau_dim; k++) { + tau_0_vector_[k] = tau_0_update[k]; + } + // Subtract tau_0 * Z from partial residual + for (int i = 0; i < n_train; i++) { + for (int k = 0; k < tau_dim; k++) { + partial_resid_ptr[i] -= data_.treatment_train[k * n_train + i] * tau_0_vector_[k]; + } + } + } +} + +void BCFSampler::SampleAdaptiveCodingParameters() { + // Extract data dimensions and pointers + const int n = data_.n_train; + double* resid_ptr = residual_->GetData().data(); + double* treatment_ptr = data_.treatment_train; + + // Add [b_0 * (1-Z) + b_1 * Z] * tau(x) to residual to get partial residual + std::vector partial_resid(n, 0.0); + for (int i = 0; i < n; i++) { + partial_resid[i] = resid_ptr[i] + tau_raw_sum_preds_[i] * tau_basis_vector_train_[i]; + } + + // Compute sufficient statistics for b_0 and b_1 + double xtx_control = 0.0; // sum of squared regression basis for control group: sum of (1-Z)^2 * tau(x)^2 + double xtx_treatment = 0.0; // sum of squared regression basis for treatment group: sum of Z^2 * tau(x)^2 + double xty_control = 0.0; // sum of (1-Z) * tau(x) * y + double xty_treatment = 0.0; // sum of Z * tau(x) * y + for (int i = 0; i < n; i++) { + double x_i = tau_raw_sum_preds_[i]; + double y_i = partial_resid[i]; + double z_i = treatment_ptr[i]; + if (z_i == 0.0) { + xtx_control += x_i * x_i; + xty_control += x_i * y_i; + } else if (z_i == 1.0) { + xtx_treatment += x_i * x_i; + xty_treatment += x_i * y_i; + } + } + + // Perform regression Gibbs update for b_0 and b_1 + // We use a fixed prior of b_0, b_1 ~ N(0, 1/2) (independent across b_0 and b_1) + const double prior_var = 0.5; + double posterior_var_control = global_variance_ / (xtx_control + (global_variance_ / prior_var)); + double posterior_var_treatment = global_variance_ / (xtx_treatment + (global_variance_ / prior_var)); + double posterior_mean_control = xty_control / (xtx_control + (global_variance_ / prior_var)); + double posterior_mean_treatment = xty_treatment / (xtx_treatment + (global_variance_ / prior_var)); + b_0_ = sample_standard_normal(posterior_mean_control, std::sqrt(posterior_var_control), rng_); + b_1_ = sample_standard_normal(posterior_mean_treatment, std::sqrt(posterior_var_treatment), rng_); + + // Update basis + std::vector prev_tau_basis = tau_basis_vector_train_; + for (int i = 0; i < n; i++) { + double z = treatment_ptr[i]; + tau_basis_vector_train_[i] = (b_0_ * (1.0 - z) + b_1_ * z); + } + forest_dataset_->UpdateBasis(tau_basis_vector_train_.data(), n, 1, false); + if (has_test_ && data_.treatment_test != nullptr) { + for (int i = 0; i < data_.n_test; i++) { + double z = data_.treatment_test[i]; + tau_basis_vector_test_[i] = (b_0_ * (1.0 - z) + b_1_ * z); + } + forest_dataset_test_->UpdateBasis(tau_basis_vector_test_.data(), data_.n_test, 1, false); + } + + // Propagate basis changes through to the trackers + UpdateResidualNewBasis(*tau_forest_tracker_, *forest_dataset_, *residual_, tau_forest_.get()); + + // If a tau_0 treatment intercept term is sampled, we must also subtract tau_0 * (new_basis - old_basis) from the residual + if (sample_tau_0_) { + for (int i = 0; i < n; i++) { + resid_ptr[i] -= tau_0_scalar_ * (tau_basis_vector_train_[i] - prev_tau_basis[i]); + } + } +} + +} // namespace StochTree diff --git a/src/container.cpp b/src/container.cpp index 0d7d3548..b106fd80 100644 --- a/src/container.cpp +++ b/src/container.cpp @@ -48,7 +48,7 @@ void ForestContainer::InitializeRoot(double leaf_value) { CHECK_EQ(forests_.size(), 0); forests_.resize(1); forests_[0].reset(new TreeEnsemble(num_trees_, output_dimension_, is_leaf_constant_, is_exponentiated_)); - // NOTE: not setting num_samples = 1, since we are just initializing constant root + // NOTE: not setting num_samples = 1, since we are just initializing constant root // nodes and the forest still needs to be sampled by either MCMC or GFR num_samples_ = 0; SetLeafValue(0, leaf_value); @@ -60,7 +60,7 @@ void ForestContainer::InitializeRoot(std::vector& leaf_vector) { CHECK_EQ(forests_.size(), 0); forests_.resize(1); forests_[0].reset(new TreeEnsemble(num_trees_, output_dimension_, is_leaf_constant_, is_exponentiated_)); - // NOTE: not setting num_samples = 1, since we are just initializing constant root + // NOTE: not setting num_samples = 1, since we are just initializing constant root // nodes and the forest still needs to be sampled by either MCMC or GFR num_samples_ = 0; SetLeafVector(0, leaf_vector); @@ -78,17 +78,17 @@ void ForestContainer::AddSamples(int num_samples) { std::vector ForestContainer::Predict(ForestDataset& dataset) { data_size_t n = dataset.NumObservations(); - data_size_t total_output_size = n*num_samples_; + data_size_t total_output_size = n * num_samples_; std::vector output(total_output_size); PredictInPlace(dataset, output); return output; } -std::vector ForestContainer::PredictRaw(ForestDataset& dataset) { +std::vector ForestContainer::PredictRaw(ForestDataset& dataset, bool row_major) { data_size_t n = dataset.NumObservations(); data_size_t total_output_size = n * output_dimension_ * num_samples_; std::vector output(total_output_size); - PredictRawInPlace(dataset, output); + PredictRawInPlace(dataset, output, row_major); return output; } @@ -110,7 +110,7 @@ std::vector ForestContainer::PredictRawSingleTree(ForestDataset& dataset void ForestContainer::PredictInPlace(ForestDataset& dataset, std::vector& output) { data_size_t n = dataset.NumObservations(); - data_size_t total_output_size = n*num_samples_; + data_size_t total_output_size = n * num_samples_; CHECK_EQ(total_output_size, output.size()); data_size_t offset = 0; for (int i = 0; i < num_samples_; i++) { @@ -120,14 +120,14 @@ void ForestContainer::PredictInPlace(ForestDataset& dataset, std::vector } } -void ForestContainer::PredictRawInPlace(ForestDataset& dataset, std::vector& output) { +void ForestContainer::PredictRawInPlace(ForestDataset& dataset, std::vector& output, bool row_major) { data_size_t n = dataset.NumObservations(); data_size_t total_output_size = n * output_dimension_ * num_samples_; CHECK_EQ(total_output_size, output.size()); data_size_t offset = 0; for (int i = 0; i < num_samples_; i++) { auto num_trees = forests_[i]->NumTrees(); - forests_[i]->PredictRawInplace(dataset, output, 0, num_trees, offset); + forests_[i]->PredictRawInplace(dataset, output, 0, num_trees, offset, row_major); offset += n * output_dimension_; } } @@ -146,14 +146,13 @@ void ForestContainer::PredictRawSingleTreeInPlace(ForestDataset& dataset, int fo data_size_t total_output_size = n * output_dimension_; CHECK_EQ(total_output_size, output.size()); data_size_t offset = 0; - forests_[forest_num]->PredictRawInplace(dataset, output, tree_num, tree_num+1, offset); + forests_[forest_num]->PredictRawInplace(dataset, output, tree_num, tree_num + 1, offset); } void ForestContainer::PredictLeafIndicesInplace( - Eigen::Map>& covariates, - Eigen::Map>& output, - std::vector& forest_indices, int num_trees, data_size_t n -) { + Eigen::Map>& covariates, + Eigen::Map>& output, + std::vector& forest_indices, int num_trees, data_size_t n) { int num_forests = forest_indices.size(); int forest_id; for (int i = 0; i < num_forests; i++) { @@ -177,7 +176,7 @@ json ForestContainer::to_json() { forest_label = "forest_" + std::to_string(i); result_obj.emplace(forest_label, forests_[i]->to_json()); } - + return result_obj; } @@ -222,4 +221,4 @@ void ForestContainer::append_from_json(const json& forest_container_json) { this->num_samples_ += new_num_samples; } -} // namespace StochTree +} // namespace StochTree diff --git a/src/cpp11.cpp b/src/cpp11.cpp index 3626eb0b..d753ba70 100644 --- a/src/cpp11.cpp +++ b/src/cpp11.cpp @@ -5,6 +5,34 @@ #include "cpp11/declarations.hpp" #include +// R_bart.cpp +cpp11::writable::list bart_sample_cpp(cpp11::external_pointer samples, cpp11::sexp X_train, cpp11::sexp y_train, cpp11::sexp X_test, int n_train, int n_test, int p, cpp11::sexp basis_train, cpp11::sexp basis_test, int basis_dim, cpp11::sexp obs_weights_train, cpp11::sexp obs_weights_test, cpp11::sexp rfx_group_ids_train, cpp11::sexp rfx_group_ids_test, cpp11::sexp rfx_basis_train, cpp11::sexp rfx_basis_test, int rfx_num_groups, int rfx_basis_dim, int num_gfr, int num_burnin, int keep_every, int num_mcmc, int num_chains, cpp11::list config_input); +extern "C" SEXP _stochtree_bart_sample_cpp(SEXP samples, SEXP X_train, SEXP y_train, SEXP X_test, SEXP n_train, SEXP n_test, SEXP p, SEXP basis_train, SEXP basis_test, SEXP basis_dim, SEXP obs_weights_train, SEXP obs_weights_test, SEXP rfx_group_ids_train, SEXP rfx_group_ids_test, SEXP rfx_basis_train, SEXP rfx_basis_test, SEXP rfx_num_groups, SEXP rfx_basis_dim, SEXP num_gfr, SEXP num_burnin, SEXP keep_every, SEXP num_mcmc, SEXP num_chains, SEXP config_input) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_sample_cpp(cpp11::as_cpp>>(samples), cpp11::as_cpp>(X_train), cpp11::as_cpp>(y_train), cpp11::as_cpp>(X_test), cpp11::as_cpp>(n_train), cpp11::as_cpp>(n_test), cpp11::as_cpp>(p), cpp11::as_cpp>(basis_train), cpp11::as_cpp>(basis_test), cpp11::as_cpp>(basis_dim), cpp11::as_cpp>(obs_weights_train), cpp11::as_cpp>(obs_weights_test), cpp11::as_cpp>(rfx_group_ids_train), cpp11::as_cpp>(rfx_group_ids_test), cpp11::as_cpp>(rfx_basis_train), cpp11::as_cpp>(rfx_basis_test), cpp11::as_cpp>(rfx_num_groups), cpp11::as_cpp>(rfx_basis_dim), cpp11::as_cpp>(num_gfr), cpp11::as_cpp>(num_burnin), cpp11::as_cpp>(keep_every), cpp11::as_cpp>(num_mcmc), cpp11::as_cpp>(num_chains), cpp11::as_cpp>(config_input))); + END_CPP11 +} +// R_bart.cpp +cpp11::writable::list bart_predict_cpp(cpp11::external_pointer bart_samples_ptr, cpp11::list bart_model_metadata, cpp11::sexp X, cpp11::sexp leaf_basis, int n, int p, int num_basis, cpp11::sexp obs_weights, cpp11::sexp rfx_group_ids, cpp11::sexp rfx_basis, int rfx_num_groups, int rfx_basis_dim, bool posterior, int scale, bool predict_y_hat, bool predict_mean_forest, bool predict_variance_forest, bool predict_random_effects); +extern "C" SEXP _stochtree_bart_predict_cpp(SEXP bart_samples_ptr, SEXP bart_model_metadata, SEXP X, SEXP leaf_basis, SEXP n, SEXP p, SEXP num_basis, SEXP obs_weights, SEXP rfx_group_ids, SEXP rfx_basis, SEXP rfx_num_groups, SEXP rfx_basis_dim, SEXP posterior, SEXP scale, SEXP predict_y_hat, SEXP predict_mean_forest, SEXP predict_variance_forest, SEXP predict_random_effects) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_predict_cpp(cpp11::as_cpp>>(bart_samples_ptr), cpp11::as_cpp>(bart_model_metadata), cpp11::as_cpp>(X), cpp11::as_cpp>(leaf_basis), cpp11::as_cpp>(n), cpp11::as_cpp>(p), cpp11::as_cpp>(num_basis), cpp11::as_cpp>(obs_weights), cpp11::as_cpp>(rfx_group_ids), cpp11::as_cpp>(rfx_basis), cpp11::as_cpp>(rfx_num_groups), cpp11::as_cpp>(rfx_basis_dim), cpp11::as_cpp>(posterior), cpp11::as_cpp>(scale), cpp11::as_cpp>(predict_y_hat), cpp11::as_cpp>(predict_mean_forest), cpp11::as_cpp>(predict_variance_forest), cpp11::as_cpp>(predict_random_effects))); + END_CPP11 +} +// R_bcf.cpp +cpp11::writable::list bcf_sample_cpp(cpp11::external_pointer samples, cpp11::sexp X_train, cpp11::sexp Z_train, cpp11::sexp y_train, cpp11::sexp X_test, cpp11::sexp Z_test, int n_train, int n_test, int p, int treatment_dim, cpp11::sexp obs_weights_train, cpp11::sexp obs_weights_test, cpp11::sexp rfx_group_ids_train, cpp11::sexp rfx_group_ids_test, cpp11::sexp rfx_basis_train, cpp11::sexp rfx_basis_test, int rfx_num_groups, int rfx_basis_dim, int num_gfr, int num_burnin, int keep_every, int num_mcmc, int num_chains, bool adaptive_coding, cpp11::list config_input); +extern "C" SEXP _stochtree_bcf_sample_cpp(SEXP samples, SEXP X_train, SEXP Z_train, SEXP y_train, SEXP X_test, SEXP Z_test, SEXP n_train, SEXP n_test, SEXP p, SEXP treatment_dim, SEXP obs_weights_train, SEXP obs_weights_test, SEXP rfx_group_ids_train, SEXP rfx_group_ids_test, SEXP rfx_basis_train, SEXP rfx_basis_test, SEXP rfx_num_groups, SEXP rfx_basis_dim, SEXP num_gfr, SEXP num_burnin, SEXP keep_every, SEXP num_mcmc, SEXP num_chains, SEXP adaptive_coding, SEXP config_input) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_sample_cpp(cpp11::as_cpp>>(samples), cpp11::as_cpp>(X_train), cpp11::as_cpp>(Z_train), cpp11::as_cpp>(y_train), cpp11::as_cpp>(X_test), cpp11::as_cpp>(Z_test), cpp11::as_cpp>(n_train), cpp11::as_cpp>(n_test), cpp11::as_cpp>(p), cpp11::as_cpp>(treatment_dim), cpp11::as_cpp>(obs_weights_train), cpp11::as_cpp>(obs_weights_test), cpp11::as_cpp>(rfx_group_ids_train), cpp11::as_cpp>(rfx_group_ids_test), cpp11::as_cpp>(rfx_basis_train), cpp11::as_cpp>(rfx_basis_test), cpp11::as_cpp>(rfx_num_groups), cpp11::as_cpp>(rfx_basis_dim), cpp11::as_cpp>(num_gfr), cpp11::as_cpp>(num_burnin), cpp11::as_cpp>(keep_every), cpp11::as_cpp>(num_mcmc), cpp11::as_cpp>(num_chains), cpp11::as_cpp>(adaptive_coding), cpp11::as_cpp>(config_input))); + END_CPP11 +} +// R_bcf.cpp +cpp11::writable::list bcf_predict_cpp(cpp11::external_pointer bcf_samples_ptr, cpp11::list bcf_model_metadata, cpp11::sexp X, cpp11::sexp Z, int n, int p, int treatment_dim, cpp11::sexp obs_weights, cpp11::sexp rfx_group_ids, cpp11::sexp rfx_basis, int rfx_num_groups, int rfx_basis_dim, bool posterior, int scale, bool predict_y_hat, bool predict_mu_x, bool predict_tau_x, bool predict_prognostic_function, bool predict_cate, bool predict_conditional_variance, bool predict_random_effects); +extern "C" SEXP _stochtree_bcf_predict_cpp(SEXP bcf_samples_ptr, SEXP bcf_model_metadata, SEXP X, SEXP Z, SEXP n, SEXP p, SEXP treatment_dim, SEXP obs_weights, SEXP rfx_group_ids, SEXP rfx_basis, SEXP rfx_num_groups, SEXP rfx_basis_dim, SEXP posterior, SEXP scale, SEXP predict_y_hat, SEXP predict_mu_x, SEXP predict_tau_x, SEXP predict_prognostic_function, SEXP predict_cate, SEXP predict_conditional_variance, SEXP predict_random_effects) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_predict_cpp(cpp11::as_cpp>>(bcf_samples_ptr), cpp11::as_cpp>(bcf_model_metadata), cpp11::as_cpp>(X), cpp11::as_cpp>(Z), cpp11::as_cpp>(n), cpp11::as_cpp>(p), cpp11::as_cpp>(treatment_dim), cpp11::as_cpp>(obs_weights), cpp11::as_cpp>(rfx_group_ids), cpp11::as_cpp>(rfx_basis), cpp11::as_cpp>(rfx_num_groups), cpp11::as_cpp>(rfx_basis_dim), cpp11::as_cpp>(posterior), cpp11::as_cpp>(scale), cpp11::as_cpp>(predict_y_hat), cpp11::as_cpp>(predict_mu_x), cpp11::as_cpp>(predict_tau_x), cpp11::as_cpp>(predict_prognostic_function), cpp11::as_cpp>(predict_cate), cpp11::as_cpp>(predict_conditional_variance), cpp11::as_cpp>(predict_random_effects))); + END_CPP11 +} // R_data.cpp cpp11::external_pointer create_forest_dataset_cpp(); extern "C" SEXP _stochtree_create_forest_dataset_cpp() { @@ -553,6 +581,528 @@ extern "C" SEXP _stochtree_root_reset_rfx_tracker_cpp(SEXP tracker, SEXP dataset return R_NilValue; END_CPP11 } +// R_samples.cpp +cpp11::external_pointer bart_samples_cpp(); +extern "C" SEXP _stochtree_bart_samples_cpp() { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_cpp()); + END_CPP11 +} +// R_samples.cpp +cpp11::external_pointer bart_samples_from_json_cpp(cpp11::external_pointer json); +extern "C" SEXP _stochtree_bart_samples_from_json_cpp(SEXP json) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_from_json_cpp(cpp11::as_cpp>>(json))); + END_CPP11 +} +// R_samples.cpp +void append_bart_samples_to_json_cpp(cpp11::external_pointer samples, cpp11::external_pointer json); +extern "C" SEXP _stochtree_append_bart_samples_to_json_cpp(SEXP samples, SEXP json) { + BEGIN_CPP11 + append_bart_samples_to_json_cpp(cpp11::as_cpp>>(samples), cpp11::as_cpp>>(json)); + return R_NilValue; + END_CPP11 +} +// R_samples.cpp +int bart_samples_num_samples_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_num_samples_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_num_samples_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +double bart_samples_y_bar_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_y_bar_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_y_bar_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +double bart_samples_y_std_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_y_std_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_y_std_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +bool bart_samples_has_yhat_train_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_has_yhat_train_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_has_yhat_train_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +bool bart_samples_has_mean_forest_predictions_train_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_has_mean_forest_predictions_train_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_has_mean_forest_predictions_train_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +bool bart_samples_has_variance_forest_predictions_train_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_has_variance_forest_predictions_train_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_has_variance_forest_predictions_train_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +bool bart_samples_has_rfx_predictions_train_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_has_rfx_predictions_train_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_has_rfx_predictions_train_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +bool bart_samples_has_yhat_test_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_has_yhat_test_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_has_yhat_test_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +bool bart_samples_has_mean_forest_predictions_test_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_has_mean_forest_predictions_test_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_has_mean_forest_predictions_test_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +bool bart_samples_has_variance_forest_predictions_test_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_has_variance_forest_predictions_test_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_has_variance_forest_predictions_test_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +bool bart_samples_has_rfx_predictions_test_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_has_rfx_predictions_test_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_has_rfx_predictions_test_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +bool bart_samples_has_cloglog_cutpoint_samples_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_has_cloglog_cutpoint_samples_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_has_cloglog_cutpoint_samples_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +bool bart_samples_has_global_var_samples_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_has_global_var_samples_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_has_global_var_samples_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +bool bart_samples_has_leaf_scale_samples_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_has_leaf_scale_samples_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_has_leaf_scale_samples_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +bool bart_samples_has_mean_forest_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_has_mean_forest_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_has_mean_forest_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +bool bart_samples_has_rfx_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_has_rfx_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_has_rfx_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +bool bart_samples_has_variance_forest_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_has_variance_forest_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_has_variance_forest_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bart_samples_mean_forest_predictions_train_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_mean_forest_predictions_train_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_mean_forest_predictions_train_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bart_samples_variance_forest_predictions_train_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_variance_forest_predictions_train_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_variance_forest_predictions_train_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bart_samples_rfx_predictions_train_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_rfx_predictions_train_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_rfx_predictions_train_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bart_samples_yhat_train_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_yhat_train_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_yhat_train_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bart_samples_mean_forest_predictions_test_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_mean_forest_predictions_test_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_mean_forest_predictions_test_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bart_samples_variance_forest_predictions_test_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_variance_forest_predictions_test_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_variance_forest_predictions_test_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bart_samples_rfx_predictions_test_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_rfx_predictions_test_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_rfx_predictions_test_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bart_samples_yhat_test_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_yhat_test_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_yhat_test_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bart_samples_global_var_samples_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_global_var_samples_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_global_var_samples_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bart_samples_leaf_scale_samples_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_leaf_scale_samples_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_leaf_scale_samples_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bart_samples_cloglog_cutpoint_samples_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_cloglog_cutpoint_samples_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_cloglog_cutpoint_samples_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::external_pointer bart_samples_materialize_mean_forest_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_materialize_mean_forest_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_materialize_mean_forest_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::external_pointer bart_samples_materialize_variance_forest_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_materialize_variance_forest_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_materialize_variance_forest_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::external_pointer bart_samples_materialize_rfx_container_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_materialize_rfx_container_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_materialize_rfx_container_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::external_pointer bart_samples_materialize_rfx_label_mapper_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_materialize_rfx_label_mapper_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_materialize_rfx_label_mapper_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::external_pointer bart_samples_mean_forest_ptr_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_mean_forest_ptr_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_mean_forest_ptr_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::external_pointer bart_samples_variance_forest_ptr_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bart_samples_variance_forest_ptr_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bart_samples_variance_forest_ptr_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +void bart_samples_merge_cpp(cpp11::external_pointer samples, cpp11::external_pointer other); +extern "C" SEXP _stochtree_bart_samples_merge_cpp(SEXP samples, SEXP other) { + BEGIN_CPP11 + bart_samples_merge_cpp(cpp11::as_cpp>>(samples), cpp11::as_cpp>>(other)); + return R_NilValue; + END_CPP11 +} +// R_samples.cpp +cpp11::external_pointer bcf_samples_cpp(); +extern "C" SEXP _stochtree_bcf_samples_cpp() { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_cpp()); + END_CPP11 +} +// R_samples.cpp +cpp11::external_pointer bcf_samples_from_json_cpp(cpp11::external_pointer json); +extern "C" SEXP _stochtree_bcf_samples_from_json_cpp(SEXP json) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_from_json_cpp(cpp11::as_cpp>>(json))); + END_CPP11 +} +// R_samples.cpp +void append_bcf_samples_to_json_cpp(cpp11::external_pointer samples, cpp11::external_pointer json); +extern "C" SEXP _stochtree_append_bcf_samples_to_json_cpp(SEXP samples, SEXP json) { + BEGIN_CPP11 + append_bcf_samples_to_json_cpp(cpp11::as_cpp>>(samples), cpp11::as_cpp>>(json)); + return R_NilValue; + END_CPP11 +} +// R_samples.cpp +int bcf_samples_num_samples_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_num_samples_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_num_samples_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +int bcf_samples_treatment_dim_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_treatment_dim_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_treatment_dim_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +double bcf_samples_y_bar_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_y_bar_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_y_bar_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +double bcf_samples_y_std_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_y_std_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_y_std_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +bool bcf_samples_has_mu_forest_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_has_mu_forest_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_has_mu_forest_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +bool bcf_samples_has_tau_forest_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_has_tau_forest_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_has_tau_forest_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +bool bcf_samples_has_variance_forest_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_has_variance_forest_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_has_variance_forest_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +bool bcf_samples_has_rfx_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_has_rfx_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_has_rfx_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bcf_samples_global_var_samples_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_global_var_samples_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_global_var_samples_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bcf_samples_leaf_scale_mu_samples_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_leaf_scale_mu_samples_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_leaf_scale_mu_samples_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bcf_samples_leaf_scale_tau_samples_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_leaf_scale_tau_samples_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_leaf_scale_tau_samples_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bcf_samples_tau_0_samples_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_tau_0_samples_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_tau_0_samples_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bcf_samples_b0_samples_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_b0_samples_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_b0_samples_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bcf_samples_b1_samples_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_b1_samples_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_b1_samples_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bcf_samples_yhat_train_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_yhat_train_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_yhat_train_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bcf_samples_mu_forest_predictions_train_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_mu_forest_predictions_train_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_mu_forest_predictions_train_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bcf_samples_tau_forest_predictions_train_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_tau_forest_predictions_train_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_tau_forest_predictions_train_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bcf_samples_variance_forest_predictions_train_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_variance_forest_predictions_train_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_variance_forest_predictions_train_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bcf_samples_rfx_predictions_train_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_rfx_predictions_train_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_rfx_predictions_train_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bcf_samples_yhat_test_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_yhat_test_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_yhat_test_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bcf_samples_mu_forest_predictions_test_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_mu_forest_predictions_test_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_mu_forest_predictions_test_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bcf_samples_tau_forest_predictions_test_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_tau_forest_predictions_test_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_tau_forest_predictions_test_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bcf_samples_variance_forest_predictions_test_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_variance_forest_predictions_test_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_variance_forest_predictions_test_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::writable::doubles bcf_samples_rfx_predictions_test_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_rfx_predictions_test_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_rfx_predictions_test_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::external_pointer bcf_samples_materialize_mu_forest_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_materialize_mu_forest_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_materialize_mu_forest_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::external_pointer bcf_samples_materialize_tau_forest_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_materialize_tau_forest_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_materialize_tau_forest_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::external_pointer bcf_samples_materialize_variance_forest_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_materialize_variance_forest_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_materialize_variance_forest_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::external_pointer bcf_samples_materialize_rfx_container_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_materialize_rfx_container_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_materialize_rfx_container_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::external_pointer bcf_samples_materialize_rfx_label_mapper_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_materialize_rfx_label_mapper_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_materialize_rfx_label_mapper_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::external_pointer bcf_samples_mu_forest_ptr_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_mu_forest_ptr_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_mu_forest_ptr_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::external_pointer bcf_samples_tau_forest_ptr_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_tau_forest_ptr_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_tau_forest_ptr_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +cpp11::external_pointer bcf_samples_variance_forest_ptr_cpp(cpp11::external_pointer samples); +extern "C" SEXP _stochtree_bcf_samples_variance_forest_ptr_cpp(SEXP samples) { + BEGIN_CPP11 + return cpp11::as_sexp(bcf_samples_variance_forest_ptr_cpp(cpp11::as_cpp>>(samples))); + END_CPP11 +} +// R_samples.cpp +void bcf_samples_merge_cpp(cpp11::external_pointer samples, cpp11::external_pointer other); +extern "C" SEXP _stochtree_bcf_samples_merge_cpp(SEXP samples, SEXP other) { + BEGIN_CPP11 + bcf_samples_merge_cpp(cpp11::as_cpp>>(samples), cpp11::as_cpp>>(other)); + return R_NilValue; + END_CPP11 +} // R_utils.cpp double sum_cpp(cpp11::doubles x); extern "C" SEXP _stochtree_sum_cpp(SEXP x) { @@ -1511,6 +2061,38 @@ extern "C" SEXP _stochtree_json_contains_field_cpp(SEXP json_ptr, SEXP field_nam END_CPP11 } // serialization.cpp +void json_erase_field_cpp(cpp11::external_pointer json_ptr, std::string field_name); +extern "C" SEXP _stochtree_json_erase_field_cpp(SEXP json_ptr, SEXP field_name) { + BEGIN_CPP11 + json_erase_field_cpp(cpp11::as_cpp>>(json_ptr), cpp11::as_cpp>(field_name)); + return R_NilValue; + END_CPP11 +} +// serialization.cpp +void json_erase_field_subfolder_cpp(cpp11::external_pointer json_ptr, std::string subfolder_name, std::string field_name); +extern "C" SEXP _stochtree_json_erase_field_subfolder_cpp(SEXP json_ptr, SEXP subfolder_name, SEXP field_name) { + BEGIN_CPP11 + json_erase_field_subfolder_cpp(cpp11::as_cpp>>(json_ptr), cpp11::as_cpp>(subfolder_name), cpp11::as_cpp>(field_name)); + return R_NilValue; + END_CPP11 +} +// serialization.cpp +void json_rename_field_cpp(cpp11::external_pointer json_ptr, std::string old_name, std::string new_name); +extern "C" SEXP _stochtree_json_rename_field_cpp(SEXP json_ptr, SEXP old_name, SEXP new_name) { + BEGIN_CPP11 + json_rename_field_cpp(cpp11::as_cpp>>(json_ptr), cpp11::as_cpp>(old_name), cpp11::as_cpp>(new_name)); + return R_NilValue; + END_CPP11 +} +// serialization.cpp +void json_rename_field_subfolder_cpp(cpp11::external_pointer json_ptr, std::string subfolder_name, std::string old_name, std::string new_name); +extern "C" SEXP _stochtree_json_rename_field_subfolder_cpp(SEXP json_ptr, SEXP subfolder_name, SEXP old_name, SEXP new_name) { + BEGIN_CPP11 + json_rename_field_subfolder_cpp(cpp11::as_cpp>>(json_ptr), cpp11::as_cpp>(subfolder_name), cpp11::as_cpp>(old_name), cpp11::as_cpp>(new_name)); + return R_NilValue; + END_CPP11 +} +// serialization.cpp double json_extract_double_subfolder_cpp(cpp11::external_pointer json_ptr, std::string subfolder_name, std::string field_name); extern "C" SEXP _stochtree_json_extract_double_subfolder_cpp(SEXP json_ptr, SEXP subfolder_name, SEXP field_name) { BEGIN_CPP11 @@ -1609,10 +2191,10 @@ extern "C" SEXP _stochtree_json_extract_string_vector_cpp(SEXP json_ptr, SEXP fi END_CPP11 } // serialization.cpp -std::string json_add_forest_cpp(cpp11::external_pointer json_ptr, cpp11::external_pointer forest_samples); -extern "C" SEXP _stochtree_json_add_forest_cpp(SEXP json_ptr, SEXP forest_samples) { +std::string json_add_forest_cpp(cpp11::external_pointer json_ptr, cpp11::external_pointer forest_samples, std::string forest_label); +extern "C" SEXP _stochtree_json_add_forest_cpp(SEXP json_ptr, SEXP forest_samples, SEXP forest_label) { BEGIN_CPP11 - return cpp11::as_sexp(json_add_forest_cpp(cpp11::as_cpp>>(json_ptr), cpp11::as_cpp>>(forest_samples))); + return cpp11::as_sexp(json_add_forest_cpp(cpp11::as_cpp>>(json_ptr), cpp11::as_cpp>>(forest_samples), cpp11::as_cpp>(forest_label))); END_CPP11 } // serialization.cpp @@ -1638,6 +2220,13 @@ extern "C" SEXP _stochtree_json_add_rfx_label_mapper_cpp(SEXP json_ptr, SEXP lab END_CPP11 } // serialization.cpp +cpp11::writable::integers rfx_label_mapper_unique_group_ids_cpp(cpp11::external_pointer label_mapper); +extern "C" SEXP _stochtree_rfx_label_mapper_unique_group_ids_cpp(SEXP label_mapper) { + BEGIN_CPP11 + return cpp11::as_sexp(rfx_label_mapper_unique_group_ids_cpp(cpp11::as_cpp>>(label_mapper))); + END_CPP11 +} +// serialization.cpp std::string json_add_rfx_groupids_cpp(cpp11::external_pointer json_ptr, cpp11::integers groupids); extern "C" SEXP _stochtree_json_add_rfx_groupids_cpp(SEXP json_ptr, SEXP groupids) { BEGIN_CPP11 @@ -1678,232 +2267,315 @@ extern "C" SEXP _stochtree_json_load_string_cpp(SEXP json_ptr, SEXP json_string) extern "C" { static const R_CallMethodDef CallEntries[] = { - {"_stochtree_active_forest_cpp", (DL_FUNC) &_stochtree_active_forest_cpp, 4}, - {"_stochtree_add_numeric_split_tree_value_active_forest_cpp", (DL_FUNC) &_stochtree_add_numeric_split_tree_value_active_forest_cpp, 7}, - {"_stochtree_add_numeric_split_tree_value_forest_container_cpp", (DL_FUNC) &_stochtree_add_numeric_split_tree_value_forest_container_cpp, 8}, - {"_stochtree_add_numeric_split_tree_vector_active_forest_cpp", (DL_FUNC) &_stochtree_add_numeric_split_tree_vector_active_forest_cpp, 7}, - {"_stochtree_add_numeric_split_tree_vector_forest_container_cpp", (DL_FUNC) &_stochtree_add_numeric_split_tree_vector_forest_container_cpp, 8}, - {"_stochtree_add_sample_forest_container_cpp", (DL_FUNC) &_stochtree_add_sample_forest_container_cpp, 1}, - {"_stochtree_add_sample_value_forest_container_cpp", (DL_FUNC) &_stochtree_add_sample_value_forest_container_cpp, 2}, - {"_stochtree_add_sample_vector_forest_container_cpp", (DL_FUNC) &_stochtree_add_sample_vector_forest_container_cpp, 2}, - {"_stochtree_add_to_column_vector_cpp", (DL_FUNC) &_stochtree_add_to_column_vector_cpp, 2}, - {"_stochtree_add_to_forest_forest_container_cpp", (DL_FUNC) &_stochtree_add_to_forest_forest_container_cpp, 3}, - {"_stochtree_adjust_residual_active_forest_cpp", (DL_FUNC) &_stochtree_adjust_residual_active_forest_cpp, 6}, - {"_stochtree_adjust_residual_forest_container_cpp", (DL_FUNC) &_stochtree_adjust_residual_forest_container_cpp, 7}, - {"_stochtree_all_roots_active_forest_cpp", (DL_FUNC) &_stochtree_all_roots_active_forest_cpp, 1}, - {"_stochtree_all_roots_forest_container_cpp", (DL_FUNC) &_stochtree_all_roots_forest_container_cpp, 2}, - {"_stochtree_average_max_depth_active_forest_cpp", (DL_FUNC) &_stochtree_average_max_depth_active_forest_cpp, 1}, - {"_stochtree_average_max_depth_forest_container_cpp", (DL_FUNC) &_stochtree_average_max_depth_forest_container_cpp, 1}, - {"_stochtree_combine_forests_forest_container_cpp", (DL_FUNC) &_stochtree_combine_forests_forest_container_cpp, 2}, - {"_stochtree_compute_leaf_indices_cpp", (DL_FUNC) &_stochtree_compute_leaf_indices_cpp, 3}, - {"_stochtree_create_column_vector_cpp", (DL_FUNC) &_stochtree_create_column_vector_cpp, 1}, - {"_stochtree_create_forest_dataset_cpp", (DL_FUNC) &_stochtree_create_forest_dataset_cpp, 0}, - {"_stochtree_create_rfx_dataset_cpp", (DL_FUNC) &_stochtree_create_rfx_dataset_cpp, 0}, - {"_stochtree_dataset_has_basis_cpp", (DL_FUNC) &_stochtree_dataset_has_basis_cpp, 1}, - {"_stochtree_dataset_has_variance_weights_cpp", (DL_FUNC) &_stochtree_dataset_has_variance_weights_cpp, 1}, - {"_stochtree_dataset_num_basis_cpp", (DL_FUNC) &_stochtree_dataset_num_basis_cpp, 1}, - {"_stochtree_dataset_num_covariates_cpp", (DL_FUNC) &_stochtree_dataset_num_covariates_cpp, 1}, - {"_stochtree_dataset_num_rows_cpp", (DL_FUNC) &_stochtree_dataset_num_rows_cpp, 1}, - {"_stochtree_ensemble_average_max_depth_forest_container_cpp", (DL_FUNC) &_stochtree_ensemble_average_max_depth_forest_container_cpp, 2}, - {"_stochtree_ensemble_tree_max_depth_active_forest_cpp", (DL_FUNC) &_stochtree_ensemble_tree_max_depth_active_forest_cpp, 2}, - {"_stochtree_ensemble_tree_max_depth_forest_container_cpp", (DL_FUNC) &_stochtree_ensemble_tree_max_depth_forest_container_cpp, 3}, - {"_stochtree_forest_add_constant_cpp", (DL_FUNC) &_stochtree_forest_add_constant_cpp, 2}, - {"_stochtree_forest_container_append_from_json_cpp", (DL_FUNC) &_stochtree_forest_container_append_from_json_cpp, 3}, - {"_stochtree_forest_container_append_from_json_string_cpp", (DL_FUNC) &_stochtree_forest_container_append_from_json_string_cpp, 3}, - {"_stochtree_forest_container_cpp", (DL_FUNC) &_stochtree_forest_container_cpp, 4}, - {"_stochtree_forest_container_from_json_cpp", (DL_FUNC) &_stochtree_forest_container_from_json_cpp, 2}, - {"_stochtree_forest_container_from_json_string_cpp", (DL_FUNC) &_stochtree_forest_container_from_json_string_cpp, 2}, - {"_stochtree_forest_container_get_max_leaf_index_cpp", (DL_FUNC) &_stochtree_forest_container_get_max_leaf_index_cpp, 2}, - {"_stochtree_forest_dataset_add_auxiliary_dimension_cpp", (DL_FUNC) &_stochtree_forest_dataset_add_auxiliary_dimension_cpp, 2}, - {"_stochtree_forest_dataset_add_basis_cpp", (DL_FUNC) &_stochtree_forest_dataset_add_basis_cpp, 2}, - {"_stochtree_forest_dataset_add_covariates_cpp", (DL_FUNC) &_stochtree_forest_dataset_add_covariates_cpp, 2}, - {"_stochtree_forest_dataset_add_weights_cpp", (DL_FUNC) &_stochtree_forest_dataset_add_weights_cpp, 2}, - {"_stochtree_forest_dataset_get_auxiliary_data_value_cpp", (DL_FUNC) &_stochtree_forest_dataset_get_auxiliary_data_value_cpp, 3}, - {"_stochtree_forest_dataset_get_auxiliary_data_vector_cpp", (DL_FUNC) &_stochtree_forest_dataset_get_auxiliary_data_vector_cpp, 2}, - {"_stochtree_forest_dataset_get_basis_cpp", (DL_FUNC) &_stochtree_forest_dataset_get_basis_cpp, 1}, - {"_stochtree_forest_dataset_get_covariates_cpp", (DL_FUNC) &_stochtree_forest_dataset_get_covariates_cpp, 1}, - {"_stochtree_forest_dataset_get_variance_weights_cpp", (DL_FUNC) &_stochtree_forest_dataset_get_variance_weights_cpp, 1}, - {"_stochtree_forest_dataset_has_auxiliary_dimension_cpp", (DL_FUNC) &_stochtree_forest_dataset_has_auxiliary_dimension_cpp, 2}, - {"_stochtree_forest_dataset_set_auxiliary_data_value_cpp", (DL_FUNC) &_stochtree_forest_dataset_set_auxiliary_data_value_cpp, 4}, - {"_stochtree_forest_dataset_update_basis_cpp", (DL_FUNC) &_stochtree_forest_dataset_update_basis_cpp, 2}, - {"_stochtree_forest_dataset_update_var_weights_cpp", (DL_FUNC) &_stochtree_forest_dataset_update_var_weights_cpp, 3}, - {"_stochtree_forest_merge_cpp", (DL_FUNC) &_stochtree_forest_merge_cpp, 2}, - {"_stochtree_forest_multiply_constant_cpp", (DL_FUNC) &_stochtree_forest_multiply_constant_cpp, 2}, - {"_stochtree_forest_tracker_cpp", (DL_FUNC) &_stochtree_forest_tracker_cpp, 4}, - {"_stochtree_get_alpha_tree_prior_cpp", (DL_FUNC) &_stochtree_get_alpha_tree_prior_cpp, 1}, - {"_stochtree_get_beta_tree_prior_cpp", (DL_FUNC) &_stochtree_get_beta_tree_prior_cpp, 1}, - {"_stochtree_get_cached_forest_predictions_cpp", (DL_FUNC) &_stochtree_get_cached_forest_predictions_cpp, 1}, - {"_stochtree_get_forest_split_counts_forest_container_cpp", (DL_FUNC) &_stochtree_get_forest_split_counts_forest_container_cpp, 3}, - {"_stochtree_get_granular_split_count_array_active_forest_cpp", (DL_FUNC) &_stochtree_get_granular_split_count_array_active_forest_cpp, 2}, - {"_stochtree_get_granular_split_count_array_forest_container_cpp", (DL_FUNC) &_stochtree_get_granular_split_count_array_forest_container_cpp, 2}, - {"_stochtree_get_json_string_cpp", (DL_FUNC) &_stochtree_get_json_string_cpp, 1}, - {"_stochtree_get_max_depth_tree_prior_cpp", (DL_FUNC) &_stochtree_get_max_depth_tree_prior_cpp, 1}, - {"_stochtree_get_min_samples_leaf_tree_prior_cpp", (DL_FUNC) &_stochtree_get_min_samples_leaf_tree_prior_cpp, 1}, - {"_stochtree_get_overall_split_counts_active_forest_cpp", (DL_FUNC) &_stochtree_get_overall_split_counts_active_forest_cpp, 2}, - {"_stochtree_get_overall_split_counts_forest_container_cpp", (DL_FUNC) &_stochtree_get_overall_split_counts_forest_container_cpp, 2}, - {"_stochtree_get_residual_cpp", (DL_FUNC) &_stochtree_get_residual_cpp, 1}, - {"_stochtree_get_tree_leaves_active_forest_cpp", (DL_FUNC) &_stochtree_get_tree_leaves_active_forest_cpp, 2}, - {"_stochtree_get_tree_leaves_forest_container_cpp", (DL_FUNC) &_stochtree_get_tree_leaves_forest_container_cpp, 3}, - {"_stochtree_get_tree_split_counts_active_forest_cpp", (DL_FUNC) &_stochtree_get_tree_split_counts_active_forest_cpp, 3}, - {"_stochtree_get_tree_split_counts_forest_container_cpp", (DL_FUNC) &_stochtree_get_tree_split_counts_forest_container_cpp, 4}, - {"_stochtree_init_json_cpp", (DL_FUNC) &_stochtree_init_json_cpp, 0}, - {"_stochtree_initialize_forest_model_active_forest_cpp", (DL_FUNC) &_stochtree_initialize_forest_model_active_forest_cpp, 6}, - {"_stochtree_initialize_forest_model_cpp", (DL_FUNC) &_stochtree_initialize_forest_model_cpp, 6}, - {"_stochtree_is_categorical_split_node_forest_container_cpp", (DL_FUNC) &_stochtree_is_categorical_split_node_forest_container_cpp, 4}, - {"_stochtree_is_exponentiated_active_forest_cpp", (DL_FUNC) &_stochtree_is_exponentiated_active_forest_cpp, 1}, - {"_stochtree_is_exponentiated_forest_container_cpp", (DL_FUNC) &_stochtree_is_exponentiated_forest_container_cpp, 1}, - {"_stochtree_is_leaf_constant_active_forest_cpp", (DL_FUNC) &_stochtree_is_leaf_constant_active_forest_cpp, 1}, - {"_stochtree_is_leaf_constant_forest_container_cpp", (DL_FUNC) &_stochtree_is_leaf_constant_forest_container_cpp, 1}, - {"_stochtree_is_leaf_node_forest_container_cpp", (DL_FUNC) &_stochtree_is_leaf_node_forest_container_cpp, 4}, - {"_stochtree_is_numeric_split_node_forest_container_cpp", (DL_FUNC) &_stochtree_is_numeric_split_node_forest_container_cpp, 4}, - {"_stochtree_json_add_bool_cpp", (DL_FUNC) &_stochtree_json_add_bool_cpp, 3}, - {"_stochtree_json_add_bool_subfolder_cpp", (DL_FUNC) &_stochtree_json_add_bool_subfolder_cpp, 4}, - {"_stochtree_json_add_double_cpp", (DL_FUNC) &_stochtree_json_add_double_cpp, 3}, - {"_stochtree_json_add_double_subfolder_cpp", (DL_FUNC) &_stochtree_json_add_double_subfolder_cpp, 4}, - {"_stochtree_json_add_forest_cpp", (DL_FUNC) &_stochtree_json_add_forest_cpp, 2}, - {"_stochtree_json_add_integer_cpp", (DL_FUNC) &_stochtree_json_add_integer_cpp, 3}, - {"_stochtree_json_add_integer_subfolder_cpp", (DL_FUNC) &_stochtree_json_add_integer_subfolder_cpp, 4}, - {"_stochtree_json_add_integer_vector_cpp", (DL_FUNC) &_stochtree_json_add_integer_vector_cpp, 3}, - {"_stochtree_json_add_integer_vector_subfolder_cpp", (DL_FUNC) &_stochtree_json_add_integer_vector_subfolder_cpp, 4}, - {"_stochtree_json_add_rfx_container_cpp", (DL_FUNC) &_stochtree_json_add_rfx_container_cpp, 2}, - {"_stochtree_json_add_rfx_groupids_cpp", (DL_FUNC) &_stochtree_json_add_rfx_groupids_cpp, 2}, - {"_stochtree_json_add_rfx_label_mapper_cpp", (DL_FUNC) &_stochtree_json_add_rfx_label_mapper_cpp, 2}, - {"_stochtree_json_add_string_cpp", (DL_FUNC) &_stochtree_json_add_string_cpp, 3}, - {"_stochtree_json_add_string_subfolder_cpp", (DL_FUNC) &_stochtree_json_add_string_subfolder_cpp, 4}, - {"_stochtree_json_add_string_vector_cpp", (DL_FUNC) &_stochtree_json_add_string_vector_cpp, 3}, - {"_stochtree_json_add_string_vector_subfolder_cpp", (DL_FUNC) &_stochtree_json_add_string_vector_subfolder_cpp, 4}, - {"_stochtree_json_add_vector_cpp", (DL_FUNC) &_stochtree_json_add_vector_cpp, 3}, - {"_stochtree_json_add_vector_subfolder_cpp", (DL_FUNC) &_stochtree_json_add_vector_subfolder_cpp, 4}, - {"_stochtree_json_contains_field_cpp", (DL_FUNC) &_stochtree_json_contains_field_cpp, 2}, - {"_stochtree_json_contains_field_subfolder_cpp", (DL_FUNC) &_stochtree_json_contains_field_subfolder_cpp, 3}, - {"_stochtree_json_extract_bool_cpp", (DL_FUNC) &_stochtree_json_extract_bool_cpp, 2}, - {"_stochtree_json_extract_bool_subfolder_cpp", (DL_FUNC) &_stochtree_json_extract_bool_subfolder_cpp, 3}, - {"_stochtree_json_extract_double_cpp", (DL_FUNC) &_stochtree_json_extract_double_cpp, 2}, - {"_stochtree_json_extract_double_subfolder_cpp", (DL_FUNC) &_stochtree_json_extract_double_subfolder_cpp, 3}, - {"_stochtree_json_extract_integer_cpp", (DL_FUNC) &_stochtree_json_extract_integer_cpp, 2}, - {"_stochtree_json_extract_integer_subfolder_cpp", (DL_FUNC) &_stochtree_json_extract_integer_subfolder_cpp, 3}, - {"_stochtree_json_extract_integer_vector_cpp", (DL_FUNC) &_stochtree_json_extract_integer_vector_cpp, 2}, - {"_stochtree_json_extract_integer_vector_subfolder_cpp", (DL_FUNC) &_stochtree_json_extract_integer_vector_subfolder_cpp, 3}, - {"_stochtree_json_extract_string_cpp", (DL_FUNC) &_stochtree_json_extract_string_cpp, 2}, - {"_stochtree_json_extract_string_subfolder_cpp", (DL_FUNC) &_stochtree_json_extract_string_subfolder_cpp, 3}, - {"_stochtree_json_extract_string_vector_cpp", (DL_FUNC) &_stochtree_json_extract_string_vector_cpp, 2}, - {"_stochtree_json_extract_string_vector_subfolder_cpp", (DL_FUNC) &_stochtree_json_extract_string_vector_subfolder_cpp, 3}, - {"_stochtree_json_extract_vector_cpp", (DL_FUNC) &_stochtree_json_extract_vector_cpp, 2}, - {"_stochtree_json_extract_vector_subfolder_cpp", (DL_FUNC) &_stochtree_json_extract_vector_subfolder_cpp, 3}, - {"_stochtree_json_increment_rfx_count_cpp", (DL_FUNC) &_stochtree_json_increment_rfx_count_cpp, 1}, - {"_stochtree_json_load_file_cpp", (DL_FUNC) &_stochtree_json_load_file_cpp, 2}, - {"_stochtree_json_load_forest_container_cpp", (DL_FUNC) &_stochtree_json_load_forest_container_cpp, 2}, - {"_stochtree_json_load_string_cpp", (DL_FUNC) &_stochtree_json_load_string_cpp, 2}, - {"_stochtree_json_save_file_cpp", (DL_FUNC) &_stochtree_json_save_file_cpp, 2}, - {"_stochtree_json_save_forest_container_cpp", (DL_FUNC) &_stochtree_json_save_forest_container_cpp, 2}, - {"_stochtree_leaf_dimension_active_forest_cpp", (DL_FUNC) &_stochtree_leaf_dimension_active_forest_cpp, 1}, - {"_stochtree_leaf_dimension_forest_container_cpp", (DL_FUNC) &_stochtree_leaf_dimension_forest_container_cpp, 1}, - {"_stochtree_leaf_values_forest_container_cpp", (DL_FUNC) &_stochtree_leaf_values_forest_container_cpp, 4}, - {"_stochtree_leaves_forest_container_cpp", (DL_FUNC) &_stochtree_leaves_forest_container_cpp, 3}, - {"_stochtree_left_child_node_forest_container_cpp", (DL_FUNC) &_stochtree_left_child_node_forest_container_cpp, 4}, - {"_stochtree_mean_cpp", (DL_FUNC) &_stochtree_mean_cpp, 1}, - {"_stochtree_multiply_forest_forest_container_cpp", (DL_FUNC) &_stochtree_multiply_forest_forest_container_cpp, 3}, - {"_stochtree_node_depth_forest_container_cpp", (DL_FUNC) &_stochtree_node_depth_forest_container_cpp, 4}, - {"_stochtree_nodes_forest_container_cpp", (DL_FUNC) &_stochtree_nodes_forest_container_cpp, 3}, - {"_stochtree_num_leaf_parents_forest_container_cpp", (DL_FUNC) &_stochtree_num_leaf_parents_forest_container_cpp, 3}, - {"_stochtree_num_leaves_ensemble_forest_container_cpp", (DL_FUNC) &_stochtree_num_leaves_ensemble_forest_container_cpp, 2}, - {"_stochtree_num_leaves_forest_container_cpp", (DL_FUNC) &_stochtree_num_leaves_forest_container_cpp, 3}, - {"_stochtree_num_nodes_forest_container_cpp", (DL_FUNC) &_stochtree_num_nodes_forest_container_cpp, 3}, - {"_stochtree_num_samples_forest_container_cpp", (DL_FUNC) &_stochtree_num_samples_forest_container_cpp, 1}, - {"_stochtree_num_split_nodes_forest_container_cpp", (DL_FUNC) &_stochtree_num_split_nodes_forest_container_cpp, 3}, - {"_stochtree_num_trees_active_forest_cpp", (DL_FUNC) &_stochtree_num_trees_active_forest_cpp, 1}, - {"_stochtree_num_trees_forest_container_cpp", (DL_FUNC) &_stochtree_num_trees_forest_container_cpp, 1}, - {"_stochtree_ordinal_sampler_cpp", (DL_FUNC) &_stochtree_ordinal_sampler_cpp, 0}, - {"_stochtree_ordinal_sampler_update_cumsum_exp_cpp", (DL_FUNC) &_stochtree_ordinal_sampler_update_cumsum_exp_cpp, 2}, - {"_stochtree_ordinal_sampler_update_gamma_params_cpp", (DL_FUNC) &_stochtree_ordinal_sampler_update_gamma_params_cpp, 7}, - {"_stochtree_ordinal_sampler_update_latent_variables_cpp", (DL_FUNC) &_stochtree_ordinal_sampler_update_latent_variables_cpp, 4}, - {"_stochtree_overwrite_column_vector_cpp", (DL_FUNC) &_stochtree_overwrite_column_vector_cpp, 2}, - {"_stochtree_parent_node_forest_container_cpp", (DL_FUNC) &_stochtree_parent_node_forest_container_cpp, 4}, - {"_stochtree_predict_active_forest_cpp", (DL_FUNC) &_stochtree_predict_active_forest_cpp, 2}, - {"_stochtree_predict_forest_cpp", (DL_FUNC) &_stochtree_predict_forest_cpp, 2}, - {"_stochtree_predict_forest_raw_cpp", (DL_FUNC) &_stochtree_predict_forest_raw_cpp, 2}, - {"_stochtree_predict_forest_raw_single_forest_cpp", (DL_FUNC) &_stochtree_predict_forest_raw_single_forest_cpp, 3}, - {"_stochtree_predict_forest_raw_single_tree_cpp", (DL_FUNC) &_stochtree_predict_forest_raw_single_tree_cpp, 4}, - {"_stochtree_predict_raw_active_forest_cpp", (DL_FUNC) &_stochtree_predict_raw_active_forest_cpp, 2}, - {"_stochtree_propagate_basis_update_active_forest_cpp", (DL_FUNC) &_stochtree_propagate_basis_update_active_forest_cpp, 4}, - {"_stochtree_propagate_basis_update_forest_container_cpp", (DL_FUNC) &_stochtree_propagate_basis_update_forest_container_cpp, 5}, - {"_stochtree_propagate_trees_column_vector_cpp", (DL_FUNC) &_stochtree_propagate_trees_column_vector_cpp, 2}, - {"_stochtree_remove_sample_forest_container_cpp", (DL_FUNC) &_stochtree_remove_sample_forest_container_cpp, 2}, - {"_stochtree_reset_active_forest_cpp", (DL_FUNC) &_stochtree_reset_active_forest_cpp, 3}, - {"_stochtree_reset_forest_model_cpp", (DL_FUNC) &_stochtree_reset_forest_model_cpp, 5}, - {"_stochtree_reset_rfx_model_cpp", (DL_FUNC) &_stochtree_reset_rfx_model_cpp, 3}, - {"_stochtree_reset_rfx_tracker_cpp", (DL_FUNC) &_stochtree_reset_rfx_tracker_cpp, 4}, - {"_stochtree_rfx_container_append_from_json_cpp", (DL_FUNC) &_stochtree_rfx_container_append_from_json_cpp, 3}, - {"_stochtree_rfx_container_append_from_json_string_cpp", (DL_FUNC) &_stochtree_rfx_container_append_from_json_string_cpp, 3}, - {"_stochtree_rfx_container_cpp", (DL_FUNC) &_stochtree_rfx_container_cpp, 2}, - {"_stochtree_rfx_container_delete_sample_cpp", (DL_FUNC) &_stochtree_rfx_container_delete_sample_cpp, 2}, - {"_stochtree_rfx_container_from_json_cpp", (DL_FUNC) &_stochtree_rfx_container_from_json_cpp, 2}, - {"_stochtree_rfx_container_from_json_string_cpp", (DL_FUNC) &_stochtree_rfx_container_from_json_string_cpp, 2}, - {"_stochtree_rfx_container_get_alpha_cpp", (DL_FUNC) &_stochtree_rfx_container_get_alpha_cpp, 1}, - {"_stochtree_rfx_container_get_beta_cpp", (DL_FUNC) &_stochtree_rfx_container_get_beta_cpp, 1}, - {"_stochtree_rfx_container_get_sigma_cpp", (DL_FUNC) &_stochtree_rfx_container_get_sigma_cpp, 1}, - {"_stochtree_rfx_container_get_xi_cpp", (DL_FUNC) &_stochtree_rfx_container_get_xi_cpp, 1}, - {"_stochtree_rfx_container_num_components_cpp", (DL_FUNC) &_stochtree_rfx_container_num_components_cpp, 1}, - {"_stochtree_rfx_container_num_groups_cpp", (DL_FUNC) &_stochtree_rfx_container_num_groups_cpp, 1}, - {"_stochtree_rfx_container_num_samples_cpp", (DL_FUNC) &_stochtree_rfx_container_num_samples_cpp, 1}, - {"_stochtree_rfx_container_predict_cpp", (DL_FUNC) &_stochtree_rfx_container_predict_cpp, 3}, - {"_stochtree_rfx_dataset_add_basis_cpp", (DL_FUNC) &_stochtree_rfx_dataset_add_basis_cpp, 2}, - {"_stochtree_rfx_dataset_add_group_labels_cpp", (DL_FUNC) &_stochtree_rfx_dataset_add_group_labels_cpp, 2}, - {"_stochtree_rfx_dataset_add_weights_cpp", (DL_FUNC) &_stochtree_rfx_dataset_add_weights_cpp, 2}, - {"_stochtree_rfx_dataset_get_basis_cpp", (DL_FUNC) &_stochtree_rfx_dataset_get_basis_cpp, 1}, - {"_stochtree_rfx_dataset_get_group_labels_cpp", (DL_FUNC) &_stochtree_rfx_dataset_get_group_labels_cpp, 1}, - {"_stochtree_rfx_dataset_get_variance_weights_cpp", (DL_FUNC) &_stochtree_rfx_dataset_get_variance_weights_cpp, 1}, - {"_stochtree_rfx_dataset_has_basis_cpp", (DL_FUNC) &_stochtree_rfx_dataset_has_basis_cpp, 1}, - {"_stochtree_rfx_dataset_has_group_labels_cpp", (DL_FUNC) &_stochtree_rfx_dataset_has_group_labels_cpp, 1}, - {"_stochtree_rfx_dataset_has_variance_weights_cpp", (DL_FUNC) &_stochtree_rfx_dataset_has_variance_weights_cpp, 1}, - {"_stochtree_rfx_dataset_num_basis_cpp", (DL_FUNC) &_stochtree_rfx_dataset_num_basis_cpp, 1}, - {"_stochtree_rfx_dataset_num_rows_cpp", (DL_FUNC) &_stochtree_rfx_dataset_num_rows_cpp, 1}, - {"_stochtree_rfx_dataset_update_basis_cpp", (DL_FUNC) &_stochtree_rfx_dataset_update_basis_cpp, 2}, - {"_stochtree_rfx_dataset_update_group_labels_cpp", (DL_FUNC) &_stochtree_rfx_dataset_update_group_labels_cpp, 2}, - {"_stochtree_rfx_dataset_update_var_weights_cpp", (DL_FUNC) &_stochtree_rfx_dataset_update_var_weights_cpp, 3}, - {"_stochtree_rfx_group_ids_from_json_cpp", (DL_FUNC) &_stochtree_rfx_group_ids_from_json_cpp, 2}, - {"_stochtree_rfx_group_ids_from_json_string_cpp", (DL_FUNC) &_stochtree_rfx_group_ids_from_json_string_cpp, 2}, - {"_stochtree_rfx_label_mapper_cpp", (DL_FUNC) &_stochtree_rfx_label_mapper_cpp, 1}, - {"_stochtree_rfx_label_mapper_from_json_cpp", (DL_FUNC) &_stochtree_rfx_label_mapper_from_json_cpp, 2}, - {"_stochtree_rfx_label_mapper_from_json_string_cpp", (DL_FUNC) &_stochtree_rfx_label_mapper_from_json_string_cpp, 2}, - {"_stochtree_rfx_label_mapper_to_list_cpp", (DL_FUNC) &_stochtree_rfx_label_mapper_to_list_cpp, 1}, - {"_stochtree_rfx_model_cpp", (DL_FUNC) &_stochtree_rfx_model_cpp, 2}, - {"_stochtree_rfx_model_predict_cpp", (DL_FUNC) &_stochtree_rfx_model_predict_cpp, 3}, - {"_stochtree_rfx_model_sample_random_effects_cpp", (DL_FUNC) &_stochtree_rfx_model_sample_random_effects_cpp, 8}, - {"_stochtree_rfx_model_set_group_parameter_covariance_cpp", (DL_FUNC) &_stochtree_rfx_model_set_group_parameter_covariance_cpp, 2}, - {"_stochtree_rfx_model_set_group_parameters_cpp", (DL_FUNC) &_stochtree_rfx_model_set_group_parameters_cpp, 2}, - {"_stochtree_rfx_model_set_variance_prior_scale_cpp", (DL_FUNC) &_stochtree_rfx_model_set_variance_prior_scale_cpp, 2}, - {"_stochtree_rfx_model_set_variance_prior_shape_cpp", (DL_FUNC) &_stochtree_rfx_model_set_variance_prior_shape_cpp, 2}, - {"_stochtree_rfx_model_set_working_parameter_covariance_cpp", (DL_FUNC) &_stochtree_rfx_model_set_working_parameter_covariance_cpp, 2}, - {"_stochtree_rfx_model_set_working_parameter_cpp", (DL_FUNC) &_stochtree_rfx_model_set_working_parameter_cpp, 2}, - {"_stochtree_rfx_tracker_cpp", (DL_FUNC) &_stochtree_rfx_tracker_cpp, 1}, - {"_stochtree_rfx_tracker_get_unique_group_ids_cpp", (DL_FUNC) &_stochtree_rfx_tracker_get_unique_group_ids_cpp, 1}, - {"_stochtree_right_child_node_forest_container_cpp", (DL_FUNC) &_stochtree_right_child_node_forest_container_cpp, 4}, - {"_stochtree_rng_cpp", (DL_FUNC) &_stochtree_rng_cpp, 1}, - {"_stochtree_root_reset_active_forest_cpp", (DL_FUNC) &_stochtree_root_reset_active_forest_cpp, 1}, - {"_stochtree_root_reset_rfx_tracker_cpp", (DL_FUNC) &_stochtree_root_reset_rfx_tracker_cpp, 4}, - {"_stochtree_sample_gfr_one_iteration_cpp", (DL_FUNC) &_stochtree_sample_gfr_one_iteration_cpp, 19}, - {"_stochtree_sample_mcmc_one_iteration_cpp", (DL_FUNC) &_stochtree_sample_mcmc_one_iteration_cpp, 18}, - {"_stochtree_sample_sigma2_one_iteration_cpp", (DL_FUNC) &_stochtree_sample_sigma2_one_iteration_cpp, 5}, - {"_stochtree_sample_tau_one_iteration_cpp", (DL_FUNC) &_stochtree_sample_tau_one_iteration_cpp, 4}, - {"_stochtree_sample_without_replacement_integer_cpp", (DL_FUNC) &_stochtree_sample_without_replacement_integer_cpp, 3}, - {"_stochtree_sd_cpp", (DL_FUNC) &_stochtree_sd_cpp, 1}, - {"_stochtree_set_leaf_value_active_forest_cpp", (DL_FUNC) &_stochtree_set_leaf_value_active_forest_cpp, 2}, - {"_stochtree_set_leaf_value_forest_container_cpp", (DL_FUNC) &_stochtree_set_leaf_value_forest_container_cpp, 2}, - {"_stochtree_set_leaf_vector_active_forest_cpp", (DL_FUNC) &_stochtree_set_leaf_vector_active_forest_cpp, 2}, - {"_stochtree_set_leaf_vector_forest_container_cpp", (DL_FUNC) &_stochtree_set_leaf_vector_forest_container_cpp, 2}, - {"_stochtree_split_categories_forest_container_cpp", (DL_FUNC) &_stochtree_split_categories_forest_container_cpp, 4}, - {"_stochtree_split_index_forest_container_cpp", (DL_FUNC) &_stochtree_split_index_forest_container_cpp, 4}, - {"_stochtree_split_theshold_forest_container_cpp", (DL_FUNC) &_stochtree_split_theshold_forest_container_cpp, 4}, - {"_stochtree_subtract_from_column_vector_cpp", (DL_FUNC) &_stochtree_subtract_from_column_vector_cpp, 2}, - {"_stochtree_sum_cpp", (DL_FUNC) &_stochtree_sum_cpp, 1}, - {"_stochtree_sum_leaves_squared_ensemble_forest_container_cpp", (DL_FUNC) &_stochtree_sum_leaves_squared_ensemble_forest_container_cpp, 2}, - {"_stochtree_tree_prior_cpp", (DL_FUNC) &_stochtree_tree_prior_cpp, 4}, - {"_stochtree_update_alpha_tree_prior_cpp", (DL_FUNC) &_stochtree_update_alpha_tree_prior_cpp, 2}, - {"_stochtree_update_beta_tree_prior_cpp", (DL_FUNC) &_stochtree_update_beta_tree_prior_cpp, 2}, - {"_stochtree_update_max_depth_tree_prior_cpp", (DL_FUNC) &_stochtree_update_max_depth_tree_prior_cpp, 2}, - {"_stochtree_update_min_samples_leaf_tree_prior_cpp", (DL_FUNC) &_stochtree_update_min_samples_leaf_tree_prior_cpp, 2}, - {"_stochtree_var_cpp", (DL_FUNC) &_stochtree_var_cpp, 1}, + {"_stochtree_active_forest_cpp", (DL_FUNC) &_stochtree_active_forest_cpp, 4}, + {"_stochtree_add_numeric_split_tree_value_active_forest_cpp", (DL_FUNC) &_stochtree_add_numeric_split_tree_value_active_forest_cpp, 7}, + {"_stochtree_add_numeric_split_tree_value_forest_container_cpp", (DL_FUNC) &_stochtree_add_numeric_split_tree_value_forest_container_cpp, 8}, + {"_stochtree_add_numeric_split_tree_vector_active_forest_cpp", (DL_FUNC) &_stochtree_add_numeric_split_tree_vector_active_forest_cpp, 7}, + {"_stochtree_add_numeric_split_tree_vector_forest_container_cpp", (DL_FUNC) &_stochtree_add_numeric_split_tree_vector_forest_container_cpp, 8}, + {"_stochtree_add_sample_forest_container_cpp", (DL_FUNC) &_stochtree_add_sample_forest_container_cpp, 1}, + {"_stochtree_add_sample_value_forest_container_cpp", (DL_FUNC) &_stochtree_add_sample_value_forest_container_cpp, 2}, + {"_stochtree_add_sample_vector_forest_container_cpp", (DL_FUNC) &_stochtree_add_sample_vector_forest_container_cpp, 2}, + {"_stochtree_add_to_column_vector_cpp", (DL_FUNC) &_stochtree_add_to_column_vector_cpp, 2}, + {"_stochtree_add_to_forest_forest_container_cpp", (DL_FUNC) &_stochtree_add_to_forest_forest_container_cpp, 3}, + {"_stochtree_adjust_residual_active_forest_cpp", (DL_FUNC) &_stochtree_adjust_residual_active_forest_cpp, 6}, + {"_stochtree_adjust_residual_forest_container_cpp", (DL_FUNC) &_stochtree_adjust_residual_forest_container_cpp, 7}, + {"_stochtree_all_roots_active_forest_cpp", (DL_FUNC) &_stochtree_all_roots_active_forest_cpp, 1}, + {"_stochtree_all_roots_forest_container_cpp", (DL_FUNC) &_stochtree_all_roots_forest_container_cpp, 2}, + {"_stochtree_append_bart_samples_to_json_cpp", (DL_FUNC) &_stochtree_append_bart_samples_to_json_cpp, 2}, + {"_stochtree_append_bcf_samples_to_json_cpp", (DL_FUNC) &_stochtree_append_bcf_samples_to_json_cpp, 2}, + {"_stochtree_average_max_depth_active_forest_cpp", (DL_FUNC) &_stochtree_average_max_depth_active_forest_cpp, 1}, + {"_stochtree_average_max_depth_forest_container_cpp", (DL_FUNC) &_stochtree_average_max_depth_forest_container_cpp, 1}, + {"_stochtree_bart_predict_cpp", (DL_FUNC) &_stochtree_bart_predict_cpp, 18}, + {"_stochtree_bart_sample_cpp", (DL_FUNC) &_stochtree_bart_sample_cpp, 24}, + {"_stochtree_bart_samples_cloglog_cutpoint_samples_cpp", (DL_FUNC) &_stochtree_bart_samples_cloglog_cutpoint_samples_cpp, 1}, + {"_stochtree_bart_samples_cpp", (DL_FUNC) &_stochtree_bart_samples_cpp, 0}, + {"_stochtree_bart_samples_from_json_cpp", (DL_FUNC) &_stochtree_bart_samples_from_json_cpp, 1}, + {"_stochtree_bart_samples_global_var_samples_cpp", (DL_FUNC) &_stochtree_bart_samples_global_var_samples_cpp, 1}, + {"_stochtree_bart_samples_has_cloglog_cutpoint_samples_cpp", (DL_FUNC) &_stochtree_bart_samples_has_cloglog_cutpoint_samples_cpp, 1}, + {"_stochtree_bart_samples_has_global_var_samples_cpp", (DL_FUNC) &_stochtree_bart_samples_has_global_var_samples_cpp, 1}, + {"_stochtree_bart_samples_has_leaf_scale_samples_cpp", (DL_FUNC) &_stochtree_bart_samples_has_leaf_scale_samples_cpp, 1}, + {"_stochtree_bart_samples_has_mean_forest_cpp", (DL_FUNC) &_stochtree_bart_samples_has_mean_forest_cpp, 1}, + {"_stochtree_bart_samples_has_mean_forest_predictions_test_cpp", (DL_FUNC) &_stochtree_bart_samples_has_mean_forest_predictions_test_cpp, 1}, + {"_stochtree_bart_samples_has_mean_forest_predictions_train_cpp", (DL_FUNC) &_stochtree_bart_samples_has_mean_forest_predictions_train_cpp, 1}, + {"_stochtree_bart_samples_has_rfx_cpp", (DL_FUNC) &_stochtree_bart_samples_has_rfx_cpp, 1}, + {"_stochtree_bart_samples_has_rfx_predictions_test_cpp", (DL_FUNC) &_stochtree_bart_samples_has_rfx_predictions_test_cpp, 1}, + {"_stochtree_bart_samples_has_rfx_predictions_train_cpp", (DL_FUNC) &_stochtree_bart_samples_has_rfx_predictions_train_cpp, 1}, + {"_stochtree_bart_samples_has_variance_forest_cpp", (DL_FUNC) &_stochtree_bart_samples_has_variance_forest_cpp, 1}, + {"_stochtree_bart_samples_has_variance_forest_predictions_test_cpp", (DL_FUNC) &_stochtree_bart_samples_has_variance_forest_predictions_test_cpp, 1}, + {"_stochtree_bart_samples_has_variance_forest_predictions_train_cpp", (DL_FUNC) &_stochtree_bart_samples_has_variance_forest_predictions_train_cpp, 1}, + {"_stochtree_bart_samples_has_yhat_test_cpp", (DL_FUNC) &_stochtree_bart_samples_has_yhat_test_cpp, 1}, + {"_stochtree_bart_samples_has_yhat_train_cpp", (DL_FUNC) &_stochtree_bart_samples_has_yhat_train_cpp, 1}, + {"_stochtree_bart_samples_leaf_scale_samples_cpp", (DL_FUNC) &_stochtree_bart_samples_leaf_scale_samples_cpp, 1}, + {"_stochtree_bart_samples_materialize_mean_forest_cpp", (DL_FUNC) &_stochtree_bart_samples_materialize_mean_forest_cpp, 1}, + {"_stochtree_bart_samples_materialize_rfx_container_cpp", (DL_FUNC) &_stochtree_bart_samples_materialize_rfx_container_cpp, 1}, + {"_stochtree_bart_samples_materialize_rfx_label_mapper_cpp", (DL_FUNC) &_stochtree_bart_samples_materialize_rfx_label_mapper_cpp, 1}, + {"_stochtree_bart_samples_materialize_variance_forest_cpp", (DL_FUNC) &_stochtree_bart_samples_materialize_variance_forest_cpp, 1}, + {"_stochtree_bart_samples_mean_forest_predictions_test_cpp", (DL_FUNC) &_stochtree_bart_samples_mean_forest_predictions_test_cpp, 1}, + {"_stochtree_bart_samples_mean_forest_predictions_train_cpp", (DL_FUNC) &_stochtree_bart_samples_mean_forest_predictions_train_cpp, 1}, + {"_stochtree_bart_samples_mean_forest_ptr_cpp", (DL_FUNC) &_stochtree_bart_samples_mean_forest_ptr_cpp, 1}, + {"_stochtree_bart_samples_merge_cpp", (DL_FUNC) &_stochtree_bart_samples_merge_cpp, 2}, + {"_stochtree_bart_samples_num_samples_cpp", (DL_FUNC) &_stochtree_bart_samples_num_samples_cpp, 1}, + {"_stochtree_bart_samples_rfx_predictions_test_cpp", (DL_FUNC) &_stochtree_bart_samples_rfx_predictions_test_cpp, 1}, + {"_stochtree_bart_samples_rfx_predictions_train_cpp", (DL_FUNC) &_stochtree_bart_samples_rfx_predictions_train_cpp, 1}, + {"_stochtree_bart_samples_variance_forest_predictions_test_cpp", (DL_FUNC) &_stochtree_bart_samples_variance_forest_predictions_test_cpp, 1}, + {"_stochtree_bart_samples_variance_forest_predictions_train_cpp", (DL_FUNC) &_stochtree_bart_samples_variance_forest_predictions_train_cpp, 1}, + {"_stochtree_bart_samples_variance_forest_ptr_cpp", (DL_FUNC) &_stochtree_bart_samples_variance_forest_ptr_cpp, 1}, + {"_stochtree_bart_samples_y_bar_cpp", (DL_FUNC) &_stochtree_bart_samples_y_bar_cpp, 1}, + {"_stochtree_bart_samples_y_std_cpp", (DL_FUNC) &_stochtree_bart_samples_y_std_cpp, 1}, + {"_stochtree_bart_samples_yhat_test_cpp", (DL_FUNC) &_stochtree_bart_samples_yhat_test_cpp, 1}, + {"_stochtree_bart_samples_yhat_train_cpp", (DL_FUNC) &_stochtree_bart_samples_yhat_train_cpp, 1}, + {"_stochtree_bcf_predict_cpp", (DL_FUNC) &_stochtree_bcf_predict_cpp, 21}, + {"_stochtree_bcf_sample_cpp", (DL_FUNC) &_stochtree_bcf_sample_cpp, 25}, + {"_stochtree_bcf_samples_b0_samples_cpp", (DL_FUNC) &_stochtree_bcf_samples_b0_samples_cpp, 1}, + {"_stochtree_bcf_samples_b1_samples_cpp", (DL_FUNC) &_stochtree_bcf_samples_b1_samples_cpp, 1}, + {"_stochtree_bcf_samples_cpp", (DL_FUNC) &_stochtree_bcf_samples_cpp, 0}, + {"_stochtree_bcf_samples_from_json_cpp", (DL_FUNC) &_stochtree_bcf_samples_from_json_cpp, 1}, + {"_stochtree_bcf_samples_global_var_samples_cpp", (DL_FUNC) &_stochtree_bcf_samples_global_var_samples_cpp, 1}, + {"_stochtree_bcf_samples_has_mu_forest_cpp", (DL_FUNC) &_stochtree_bcf_samples_has_mu_forest_cpp, 1}, + {"_stochtree_bcf_samples_has_rfx_cpp", (DL_FUNC) &_stochtree_bcf_samples_has_rfx_cpp, 1}, + {"_stochtree_bcf_samples_has_tau_forest_cpp", (DL_FUNC) &_stochtree_bcf_samples_has_tau_forest_cpp, 1}, + {"_stochtree_bcf_samples_has_variance_forest_cpp", (DL_FUNC) &_stochtree_bcf_samples_has_variance_forest_cpp, 1}, + {"_stochtree_bcf_samples_leaf_scale_mu_samples_cpp", (DL_FUNC) &_stochtree_bcf_samples_leaf_scale_mu_samples_cpp, 1}, + {"_stochtree_bcf_samples_leaf_scale_tau_samples_cpp", (DL_FUNC) &_stochtree_bcf_samples_leaf_scale_tau_samples_cpp, 1}, + {"_stochtree_bcf_samples_materialize_mu_forest_cpp", (DL_FUNC) &_stochtree_bcf_samples_materialize_mu_forest_cpp, 1}, + {"_stochtree_bcf_samples_materialize_rfx_container_cpp", (DL_FUNC) &_stochtree_bcf_samples_materialize_rfx_container_cpp, 1}, + {"_stochtree_bcf_samples_materialize_rfx_label_mapper_cpp", (DL_FUNC) &_stochtree_bcf_samples_materialize_rfx_label_mapper_cpp, 1}, + {"_stochtree_bcf_samples_materialize_tau_forest_cpp", (DL_FUNC) &_stochtree_bcf_samples_materialize_tau_forest_cpp, 1}, + {"_stochtree_bcf_samples_materialize_variance_forest_cpp", (DL_FUNC) &_stochtree_bcf_samples_materialize_variance_forest_cpp, 1}, + {"_stochtree_bcf_samples_merge_cpp", (DL_FUNC) &_stochtree_bcf_samples_merge_cpp, 2}, + {"_stochtree_bcf_samples_mu_forest_predictions_test_cpp", (DL_FUNC) &_stochtree_bcf_samples_mu_forest_predictions_test_cpp, 1}, + {"_stochtree_bcf_samples_mu_forest_predictions_train_cpp", (DL_FUNC) &_stochtree_bcf_samples_mu_forest_predictions_train_cpp, 1}, + {"_stochtree_bcf_samples_mu_forest_ptr_cpp", (DL_FUNC) &_stochtree_bcf_samples_mu_forest_ptr_cpp, 1}, + {"_stochtree_bcf_samples_num_samples_cpp", (DL_FUNC) &_stochtree_bcf_samples_num_samples_cpp, 1}, + {"_stochtree_bcf_samples_rfx_predictions_test_cpp", (DL_FUNC) &_stochtree_bcf_samples_rfx_predictions_test_cpp, 1}, + {"_stochtree_bcf_samples_rfx_predictions_train_cpp", (DL_FUNC) &_stochtree_bcf_samples_rfx_predictions_train_cpp, 1}, + {"_stochtree_bcf_samples_tau_0_samples_cpp", (DL_FUNC) &_stochtree_bcf_samples_tau_0_samples_cpp, 1}, + {"_stochtree_bcf_samples_tau_forest_predictions_test_cpp", (DL_FUNC) &_stochtree_bcf_samples_tau_forest_predictions_test_cpp, 1}, + {"_stochtree_bcf_samples_tau_forest_predictions_train_cpp", (DL_FUNC) &_stochtree_bcf_samples_tau_forest_predictions_train_cpp, 1}, + {"_stochtree_bcf_samples_tau_forest_ptr_cpp", (DL_FUNC) &_stochtree_bcf_samples_tau_forest_ptr_cpp, 1}, + {"_stochtree_bcf_samples_treatment_dim_cpp", (DL_FUNC) &_stochtree_bcf_samples_treatment_dim_cpp, 1}, + {"_stochtree_bcf_samples_variance_forest_predictions_test_cpp", (DL_FUNC) &_stochtree_bcf_samples_variance_forest_predictions_test_cpp, 1}, + {"_stochtree_bcf_samples_variance_forest_predictions_train_cpp", (DL_FUNC) &_stochtree_bcf_samples_variance_forest_predictions_train_cpp, 1}, + {"_stochtree_bcf_samples_variance_forest_ptr_cpp", (DL_FUNC) &_stochtree_bcf_samples_variance_forest_ptr_cpp, 1}, + {"_stochtree_bcf_samples_y_bar_cpp", (DL_FUNC) &_stochtree_bcf_samples_y_bar_cpp, 1}, + {"_stochtree_bcf_samples_y_std_cpp", (DL_FUNC) &_stochtree_bcf_samples_y_std_cpp, 1}, + {"_stochtree_bcf_samples_yhat_test_cpp", (DL_FUNC) &_stochtree_bcf_samples_yhat_test_cpp, 1}, + {"_stochtree_bcf_samples_yhat_train_cpp", (DL_FUNC) &_stochtree_bcf_samples_yhat_train_cpp, 1}, + {"_stochtree_combine_forests_forest_container_cpp", (DL_FUNC) &_stochtree_combine_forests_forest_container_cpp, 2}, + {"_stochtree_compute_leaf_indices_cpp", (DL_FUNC) &_stochtree_compute_leaf_indices_cpp, 3}, + {"_stochtree_create_column_vector_cpp", (DL_FUNC) &_stochtree_create_column_vector_cpp, 1}, + {"_stochtree_create_forest_dataset_cpp", (DL_FUNC) &_stochtree_create_forest_dataset_cpp, 0}, + {"_stochtree_create_rfx_dataset_cpp", (DL_FUNC) &_stochtree_create_rfx_dataset_cpp, 0}, + {"_stochtree_dataset_has_basis_cpp", (DL_FUNC) &_stochtree_dataset_has_basis_cpp, 1}, + {"_stochtree_dataset_has_variance_weights_cpp", (DL_FUNC) &_stochtree_dataset_has_variance_weights_cpp, 1}, + {"_stochtree_dataset_num_basis_cpp", (DL_FUNC) &_stochtree_dataset_num_basis_cpp, 1}, + {"_stochtree_dataset_num_covariates_cpp", (DL_FUNC) &_stochtree_dataset_num_covariates_cpp, 1}, + {"_stochtree_dataset_num_rows_cpp", (DL_FUNC) &_stochtree_dataset_num_rows_cpp, 1}, + {"_stochtree_ensemble_average_max_depth_forest_container_cpp", (DL_FUNC) &_stochtree_ensemble_average_max_depth_forest_container_cpp, 2}, + {"_stochtree_ensemble_tree_max_depth_active_forest_cpp", (DL_FUNC) &_stochtree_ensemble_tree_max_depth_active_forest_cpp, 2}, + {"_stochtree_ensemble_tree_max_depth_forest_container_cpp", (DL_FUNC) &_stochtree_ensemble_tree_max_depth_forest_container_cpp, 3}, + {"_stochtree_forest_add_constant_cpp", (DL_FUNC) &_stochtree_forest_add_constant_cpp, 2}, + {"_stochtree_forest_container_append_from_json_cpp", (DL_FUNC) &_stochtree_forest_container_append_from_json_cpp, 3}, + {"_stochtree_forest_container_append_from_json_string_cpp", (DL_FUNC) &_stochtree_forest_container_append_from_json_string_cpp, 3}, + {"_stochtree_forest_container_cpp", (DL_FUNC) &_stochtree_forest_container_cpp, 4}, + {"_stochtree_forest_container_from_json_cpp", (DL_FUNC) &_stochtree_forest_container_from_json_cpp, 2}, + {"_stochtree_forest_container_from_json_string_cpp", (DL_FUNC) &_stochtree_forest_container_from_json_string_cpp, 2}, + {"_stochtree_forest_container_get_max_leaf_index_cpp", (DL_FUNC) &_stochtree_forest_container_get_max_leaf_index_cpp, 2}, + {"_stochtree_forest_dataset_add_auxiliary_dimension_cpp", (DL_FUNC) &_stochtree_forest_dataset_add_auxiliary_dimension_cpp, 2}, + {"_stochtree_forest_dataset_add_basis_cpp", (DL_FUNC) &_stochtree_forest_dataset_add_basis_cpp, 2}, + {"_stochtree_forest_dataset_add_covariates_cpp", (DL_FUNC) &_stochtree_forest_dataset_add_covariates_cpp, 2}, + {"_stochtree_forest_dataset_add_weights_cpp", (DL_FUNC) &_stochtree_forest_dataset_add_weights_cpp, 2}, + {"_stochtree_forest_dataset_get_auxiliary_data_value_cpp", (DL_FUNC) &_stochtree_forest_dataset_get_auxiliary_data_value_cpp, 3}, + {"_stochtree_forest_dataset_get_auxiliary_data_vector_cpp", (DL_FUNC) &_stochtree_forest_dataset_get_auxiliary_data_vector_cpp, 2}, + {"_stochtree_forest_dataset_get_basis_cpp", (DL_FUNC) &_stochtree_forest_dataset_get_basis_cpp, 1}, + {"_stochtree_forest_dataset_get_covariates_cpp", (DL_FUNC) &_stochtree_forest_dataset_get_covariates_cpp, 1}, + {"_stochtree_forest_dataset_get_variance_weights_cpp", (DL_FUNC) &_stochtree_forest_dataset_get_variance_weights_cpp, 1}, + {"_stochtree_forest_dataset_has_auxiliary_dimension_cpp", (DL_FUNC) &_stochtree_forest_dataset_has_auxiliary_dimension_cpp, 2}, + {"_stochtree_forest_dataset_set_auxiliary_data_value_cpp", (DL_FUNC) &_stochtree_forest_dataset_set_auxiliary_data_value_cpp, 4}, + {"_stochtree_forest_dataset_update_basis_cpp", (DL_FUNC) &_stochtree_forest_dataset_update_basis_cpp, 2}, + {"_stochtree_forest_dataset_update_var_weights_cpp", (DL_FUNC) &_stochtree_forest_dataset_update_var_weights_cpp, 3}, + {"_stochtree_forest_merge_cpp", (DL_FUNC) &_stochtree_forest_merge_cpp, 2}, + {"_stochtree_forest_multiply_constant_cpp", (DL_FUNC) &_stochtree_forest_multiply_constant_cpp, 2}, + {"_stochtree_forest_tracker_cpp", (DL_FUNC) &_stochtree_forest_tracker_cpp, 4}, + {"_stochtree_get_alpha_tree_prior_cpp", (DL_FUNC) &_stochtree_get_alpha_tree_prior_cpp, 1}, + {"_stochtree_get_beta_tree_prior_cpp", (DL_FUNC) &_stochtree_get_beta_tree_prior_cpp, 1}, + {"_stochtree_get_cached_forest_predictions_cpp", (DL_FUNC) &_stochtree_get_cached_forest_predictions_cpp, 1}, + {"_stochtree_get_forest_split_counts_forest_container_cpp", (DL_FUNC) &_stochtree_get_forest_split_counts_forest_container_cpp, 3}, + {"_stochtree_get_granular_split_count_array_active_forest_cpp", (DL_FUNC) &_stochtree_get_granular_split_count_array_active_forest_cpp, 2}, + {"_stochtree_get_granular_split_count_array_forest_container_cpp", (DL_FUNC) &_stochtree_get_granular_split_count_array_forest_container_cpp, 2}, + {"_stochtree_get_json_string_cpp", (DL_FUNC) &_stochtree_get_json_string_cpp, 1}, + {"_stochtree_get_max_depth_tree_prior_cpp", (DL_FUNC) &_stochtree_get_max_depth_tree_prior_cpp, 1}, + {"_stochtree_get_min_samples_leaf_tree_prior_cpp", (DL_FUNC) &_stochtree_get_min_samples_leaf_tree_prior_cpp, 1}, + {"_stochtree_get_overall_split_counts_active_forest_cpp", (DL_FUNC) &_stochtree_get_overall_split_counts_active_forest_cpp, 2}, + {"_stochtree_get_overall_split_counts_forest_container_cpp", (DL_FUNC) &_stochtree_get_overall_split_counts_forest_container_cpp, 2}, + {"_stochtree_get_residual_cpp", (DL_FUNC) &_stochtree_get_residual_cpp, 1}, + {"_stochtree_get_tree_leaves_active_forest_cpp", (DL_FUNC) &_stochtree_get_tree_leaves_active_forest_cpp, 2}, + {"_stochtree_get_tree_leaves_forest_container_cpp", (DL_FUNC) &_stochtree_get_tree_leaves_forest_container_cpp, 3}, + {"_stochtree_get_tree_split_counts_active_forest_cpp", (DL_FUNC) &_stochtree_get_tree_split_counts_active_forest_cpp, 3}, + {"_stochtree_get_tree_split_counts_forest_container_cpp", (DL_FUNC) &_stochtree_get_tree_split_counts_forest_container_cpp, 4}, + {"_stochtree_init_json_cpp", (DL_FUNC) &_stochtree_init_json_cpp, 0}, + {"_stochtree_initialize_forest_model_active_forest_cpp", (DL_FUNC) &_stochtree_initialize_forest_model_active_forest_cpp, 6}, + {"_stochtree_initialize_forest_model_cpp", (DL_FUNC) &_stochtree_initialize_forest_model_cpp, 6}, + {"_stochtree_is_categorical_split_node_forest_container_cpp", (DL_FUNC) &_stochtree_is_categorical_split_node_forest_container_cpp, 4}, + {"_stochtree_is_exponentiated_active_forest_cpp", (DL_FUNC) &_stochtree_is_exponentiated_active_forest_cpp, 1}, + {"_stochtree_is_exponentiated_forest_container_cpp", (DL_FUNC) &_stochtree_is_exponentiated_forest_container_cpp, 1}, + {"_stochtree_is_leaf_constant_active_forest_cpp", (DL_FUNC) &_stochtree_is_leaf_constant_active_forest_cpp, 1}, + {"_stochtree_is_leaf_constant_forest_container_cpp", (DL_FUNC) &_stochtree_is_leaf_constant_forest_container_cpp, 1}, + {"_stochtree_is_leaf_node_forest_container_cpp", (DL_FUNC) &_stochtree_is_leaf_node_forest_container_cpp, 4}, + {"_stochtree_is_numeric_split_node_forest_container_cpp", (DL_FUNC) &_stochtree_is_numeric_split_node_forest_container_cpp, 4}, + {"_stochtree_json_add_bool_cpp", (DL_FUNC) &_stochtree_json_add_bool_cpp, 3}, + {"_stochtree_json_add_bool_subfolder_cpp", (DL_FUNC) &_stochtree_json_add_bool_subfolder_cpp, 4}, + {"_stochtree_json_add_double_cpp", (DL_FUNC) &_stochtree_json_add_double_cpp, 3}, + {"_stochtree_json_add_double_subfolder_cpp", (DL_FUNC) &_stochtree_json_add_double_subfolder_cpp, 4}, + {"_stochtree_json_add_forest_cpp", (DL_FUNC) &_stochtree_json_add_forest_cpp, 3}, + {"_stochtree_json_add_integer_cpp", (DL_FUNC) &_stochtree_json_add_integer_cpp, 3}, + {"_stochtree_json_add_integer_subfolder_cpp", (DL_FUNC) &_stochtree_json_add_integer_subfolder_cpp, 4}, + {"_stochtree_json_add_integer_vector_cpp", (DL_FUNC) &_stochtree_json_add_integer_vector_cpp, 3}, + {"_stochtree_json_add_integer_vector_subfolder_cpp", (DL_FUNC) &_stochtree_json_add_integer_vector_subfolder_cpp, 4}, + {"_stochtree_json_add_rfx_container_cpp", (DL_FUNC) &_stochtree_json_add_rfx_container_cpp, 2}, + {"_stochtree_json_add_rfx_groupids_cpp", (DL_FUNC) &_stochtree_json_add_rfx_groupids_cpp, 2}, + {"_stochtree_json_add_rfx_label_mapper_cpp", (DL_FUNC) &_stochtree_json_add_rfx_label_mapper_cpp, 2}, + {"_stochtree_json_add_string_cpp", (DL_FUNC) &_stochtree_json_add_string_cpp, 3}, + {"_stochtree_json_add_string_subfolder_cpp", (DL_FUNC) &_stochtree_json_add_string_subfolder_cpp, 4}, + {"_stochtree_json_add_string_vector_cpp", (DL_FUNC) &_stochtree_json_add_string_vector_cpp, 3}, + {"_stochtree_json_add_string_vector_subfolder_cpp", (DL_FUNC) &_stochtree_json_add_string_vector_subfolder_cpp, 4}, + {"_stochtree_json_add_vector_cpp", (DL_FUNC) &_stochtree_json_add_vector_cpp, 3}, + {"_stochtree_json_add_vector_subfolder_cpp", (DL_FUNC) &_stochtree_json_add_vector_subfolder_cpp, 4}, + {"_stochtree_json_contains_field_cpp", (DL_FUNC) &_stochtree_json_contains_field_cpp, 2}, + {"_stochtree_json_contains_field_subfolder_cpp", (DL_FUNC) &_stochtree_json_contains_field_subfolder_cpp, 3}, + {"_stochtree_json_erase_field_cpp", (DL_FUNC) &_stochtree_json_erase_field_cpp, 2}, + {"_stochtree_json_erase_field_subfolder_cpp", (DL_FUNC) &_stochtree_json_erase_field_subfolder_cpp, 3}, + {"_stochtree_json_extract_bool_cpp", (DL_FUNC) &_stochtree_json_extract_bool_cpp, 2}, + {"_stochtree_json_extract_bool_subfolder_cpp", (DL_FUNC) &_stochtree_json_extract_bool_subfolder_cpp, 3}, + {"_stochtree_json_extract_double_cpp", (DL_FUNC) &_stochtree_json_extract_double_cpp, 2}, + {"_stochtree_json_extract_double_subfolder_cpp", (DL_FUNC) &_stochtree_json_extract_double_subfolder_cpp, 3}, + {"_stochtree_json_extract_integer_cpp", (DL_FUNC) &_stochtree_json_extract_integer_cpp, 2}, + {"_stochtree_json_extract_integer_subfolder_cpp", (DL_FUNC) &_stochtree_json_extract_integer_subfolder_cpp, 3}, + {"_stochtree_json_extract_integer_vector_cpp", (DL_FUNC) &_stochtree_json_extract_integer_vector_cpp, 2}, + {"_stochtree_json_extract_integer_vector_subfolder_cpp", (DL_FUNC) &_stochtree_json_extract_integer_vector_subfolder_cpp, 3}, + {"_stochtree_json_extract_string_cpp", (DL_FUNC) &_stochtree_json_extract_string_cpp, 2}, + {"_stochtree_json_extract_string_subfolder_cpp", (DL_FUNC) &_stochtree_json_extract_string_subfolder_cpp, 3}, + {"_stochtree_json_extract_string_vector_cpp", (DL_FUNC) &_stochtree_json_extract_string_vector_cpp, 2}, + {"_stochtree_json_extract_string_vector_subfolder_cpp", (DL_FUNC) &_stochtree_json_extract_string_vector_subfolder_cpp, 3}, + {"_stochtree_json_extract_vector_cpp", (DL_FUNC) &_stochtree_json_extract_vector_cpp, 2}, + {"_stochtree_json_extract_vector_subfolder_cpp", (DL_FUNC) &_stochtree_json_extract_vector_subfolder_cpp, 3}, + {"_stochtree_json_increment_rfx_count_cpp", (DL_FUNC) &_stochtree_json_increment_rfx_count_cpp, 1}, + {"_stochtree_json_load_file_cpp", (DL_FUNC) &_stochtree_json_load_file_cpp, 2}, + {"_stochtree_json_load_forest_container_cpp", (DL_FUNC) &_stochtree_json_load_forest_container_cpp, 2}, + {"_stochtree_json_load_string_cpp", (DL_FUNC) &_stochtree_json_load_string_cpp, 2}, + {"_stochtree_json_rename_field_cpp", (DL_FUNC) &_stochtree_json_rename_field_cpp, 3}, + {"_stochtree_json_rename_field_subfolder_cpp", (DL_FUNC) &_stochtree_json_rename_field_subfolder_cpp, 4}, + {"_stochtree_json_save_file_cpp", (DL_FUNC) &_stochtree_json_save_file_cpp, 2}, + {"_stochtree_json_save_forest_container_cpp", (DL_FUNC) &_stochtree_json_save_forest_container_cpp, 2}, + {"_stochtree_leaf_dimension_active_forest_cpp", (DL_FUNC) &_stochtree_leaf_dimension_active_forest_cpp, 1}, + {"_stochtree_leaf_dimension_forest_container_cpp", (DL_FUNC) &_stochtree_leaf_dimension_forest_container_cpp, 1}, + {"_stochtree_leaf_values_forest_container_cpp", (DL_FUNC) &_stochtree_leaf_values_forest_container_cpp, 4}, + {"_stochtree_leaves_forest_container_cpp", (DL_FUNC) &_stochtree_leaves_forest_container_cpp, 3}, + {"_stochtree_left_child_node_forest_container_cpp", (DL_FUNC) &_stochtree_left_child_node_forest_container_cpp, 4}, + {"_stochtree_mean_cpp", (DL_FUNC) &_stochtree_mean_cpp, 1}, + {"_stochtree_multiply_forest_forest_container_cpp", (DL_FUNC) &_stochtree_multiply_forest_forest_container_cpp, 3}, + {"_stochtree_node_depth_forest_container_cpp", (DL_FUNC) &_stochtree_node_depth_forest_container_cpp, 4}, + {"_stochtree_nodes_forest_container_cpp", (DL_FUNC) &_stochtree_nodes_forest_container_cpp, 3}, + {"_stochtree_num_leaf_parents_forest_container_cpp", (DL_FUNC) &_stochtree_num_leaf_parents_forest_container_cpp, 3}, + {"_stochtree_num_leaves_ensemble_forest_container_cpp", (DL_FUNC) &_stochtree_num_leaves_ensemble_forest_container_cpp, 2}, + {"_stochtree_num_leaves_forest_container_cpp", (DL_FUNC) &_stochtree_num_leaves_forest_container_cpp, 3}, + {"_stochtree_num_nodes_forest_container_cpp", (DL_FUNC) &_stochtree_num_nodes_forest_container_cpp, 3}, + {"_stochtree_num_samples_forest_container_cpp", (DL_FUNC) &_stochtree_num_samples_forest_container_cpp, 1}, + {"_stochtree_num_split_nodes_forest_container_cpp", (DL_FUNC) &_stochtree_num_split_nodes_forest_container_cpp, 3}, + {"_stochtree_num_trees_active_forest_cpp", (DL_FUNC) &_stochtree_num_trees_active_forest_cpp, 1}, + {"_stochtree_num_trees_forest_container_cpp", (DL_FUNC) &_stochtree_num_trees_forest_container_cpp, 1}, + {"_stochtree_ordinal_sampler_cpp", (DL_FUNC) &_stochtree_ordinal_sampler_cpp, 0}, + {"_stochtree_ordinal_sampler_update_cumsum_exp_cpp", (DL_FUNC) &_stochtree_ordinal_sampler_update_cumsum_exp_cpp, 2}, + {"_stochtree_ordinal_sampler_update_gamma_params_cpp", (DL_FUNC) &_stochtree_ordinal_sampler_update_gamma_params_cpp, 7}, + {"_stochtree_ordinal_sampler_update_latent_variables_cpp", (DL_FUNC) &_stochtree_ordinal_sampler_update_latent_variables_cpp, 4}, + {"_stochtree_overwrite_column_vector_cpp", (DL_FUNC) &_stochtree_overwrite_column_vector_cpp, 2}, + {"_stochtree_parent_node_forest_container_cpp", (DL_FUNC) &_stochtree_parent_node_forest_container_cpp, 4}, + {"_stochtree_predict_active_forest_cpp", (DL_FUNC) &_stochtree_predict_active_forest_cpp, 2}, + {"_stochtree_predict_forest_cpp", (DL_FUNC) &_stochtree_predict_forest_cpp, 2}, + {"_stochtree_predict_forest_raw_cpp", (DL_FUNC) &_stochtree_predict_forest_raw_cpp, 2}, + {"_stochtree_predict_forest_raw_single_forest_cpp", (DL_FUNC) &_stochtree_predict_forest_raw_single_forest_cpp, 3}, + {"_stochtree_predict_forest_raw_single_tree_cpp", (DL_FUNC) &_stochtree_predict_forest_raw_single_tree_cpp, 4}, + {"_stochtree_predict_raw_active_forest_cpp", (DL_FUNC) &_stochtree_predict_raw_active_forest_cpp, 2}, + {"_stochtree_propagate_basis_update_active_forest_cpp", (DL_FUNC) &_stochtree_propagate_basis_update_active_forest_cpp, 4}, + {"_stochtree_propagate_basis_update_forest_container_cpp", (DL_FUNC) &_stochtree_propagate_basis_update_forest_container_cpp, 5}, + {"_stochtree_propagate_trees_column_vector_cpp", (DL_FUNC) &_stochtree_propagate_trees_column_vector_cpp, 2}, + {"_stochtree_remove_sample_forest_container_cpp", (DL_FUNC) &_stochtree_remove_sample_forest_container_cpp, 2}, + {"_stochtree_reset_active_forest_cpp", (DL_FUNC) &_stochtree_reset_active_forest_cpp, 3}, + {"_stochtree_reset_forest_model_cpp", (DL_FUNC) &_stochtree_reset_forest_model_cpp, 5}, + {"_stochtree_reset_rfx_model_cpp", (DL_FUNC) &_stochtree_reset_rfx_model_cpp, 3}, + {"_stochtree_reset_rfx_tracker_cpp", (DL_FUNC) &_stochtree_reset_rfx_tracker_cpp, 4}, + {"_stochtree_rfx_container_append_from_json_cpp", (DL_FUNC) &_stochtree_rfx_container_append_from_json_cpp, 3}, + {"_stochtree_rfx_container_append_from_json_string_cpp", (DL_FUNC) &_stochtree_rfx_container_append_from_json_string_cpp, 3}, + {"_stochtree_rfx_container_cpp", (DL_FUNC) &_stochtree_rfx_container_cpp, 2}, + {"_stochtree_rfx_container_delete_sample_cpp", (DL_FUNC) &_stochtree_rfx_container_delete_sample_cpp, 2}, + {"_stochtree_rfx_container_from_json_cpp", (DL_FUNC) &_stochtree_rfx_container_from_json_cpp, 2}, + {"_stochtree_rfx_container_from_json_string_cpp", (DL_FUNC) &_stochtree_rfx_container_from_json_string_cpp, 2}, + {"_stochtree_rfx_container_get_alpha_cpp", (DL_FUNC) &_stochtree_rfx_container_get_alpha_cpp, 1}, + {"_stochtree_rfx_container_get_beta_cpp", (DL_FUNC) &_stochtree_rfx_container_get_beta_cpp, 1}, + {"_stochtree_rfx_container_get_sigma_cpp", (DL_FUNC) &_stochtree_rfx_container_get_sigma_cpp, 1}, + {"_stochtree_rfx_container_get_xi_cpp", (DL_FUNC) &_stochtree_rfx_container_get_xi_cpp, 1}, + {"_stochtree_rfx_container_num_components_cpp", (DL_FUNC) &_stochtree_rfx_container_num_components_cpp, 1}, + {"_stochtree_rfx_container_num_groups_cpp", (DL_FUNC) &_stochtree_rfx_container_num_groups_cpp, 1}, + {"_stochtree_rfx_container_num_samples_cpp", (DL_FUNC) &_stochtree_rfx_container_num_samples_cpp, 1}, + {"_stochtree_rfx_container_predict_cpp", (DL_FUNC) &_stochtree_rfx_container_predict_cpp, 3}, + {"_stochtree_rfx_dataset_add_basis_cpp", (DL_FUNC) &_stochtree_rfx_dataset_add_basis_cpp, 2}, + {"_stochtree_rfx_dataset_add_group_labels_cpp", (DL_FUNC) &_stochtree_rfx_dataset_add_group_labels_cpp, 2}, + {"_stochtree_rfx_dataset_add_weights_cpp", (DL_FUNC) &_stochtree_rfx_dataset_add_weights_cpp, 2}, + {"_stochtree_rfx_dataset_get_basis_cpp", (DL_FUNC) &_stochtree_rfx_dataset_get_basis_cpp, 1}, + {"_stochtree_rfx_dataset_get_group_labels_cpp", (DL_FUNC) &_stochtree_rfx_dataset_get_group_labels_cpp, 1}, + {"_stochtree_rfx_dataset_get_variance_weights_cpp", (DL_FUNC) &_stochtree_rfx_dataset_get_variance_weights_cpp, 1}, + {"_stochtree_rfx_dataset_has_basis_cpp", (DL_FUNC) &_stochtree_rfx_dataset_has_basis_cpp, 1}, + {"_stochtree_rfx_dataset_has_group_labels_cpp", (DL_FUNC) &_stochtree_rfx_dataset_has_group_labels_cpp, 1}, + {"_stochtree_rfx_dataset_has_variance_weights_cpp", (DL_FUNC) &_stochtree_rfx_dataset_has_variance_weights_cpp, 1}, + {"_stochtree_rfx_dataset_num_basis_cpp", (DL_FUNC) &_stochtree_rfx_dataset_num_basis_cpp, 1}, + {"_stochtree_rfx_dataset_num_rows_cpp", (DL_FUNC) &_stochtree_rfx_dataset_num_rows_cpp, 1}, + {"_stochtree_rfx_dataset_update_basis_cpp", (DL_FUNC) &_stochtree_rfx_dataset_update_basis_cpp, 2}, + {"_stochtree_rfx_dataset_update_group_labels_cpp", (DL_FUNC) &_stochtree_rfx_dataset_update_group_labels_cpp, 2}, + {"_stochtree_rfx_dataset_update_var_weights_cpp", (DL_FUNC) &_stochtree_rfx_dataset_update_var_weights_cpp, 3}, + {"_stochtree_rfx_group_ids_from_json_cpp", (DL_FUNC) &_stochtree_rfx_group_ids_from_json_cpp, 2}, + {"_stochtree_rfx_group_ids_from_json_string_cpp", (DL_FUNC) &_stochtree_rfx_group_ids_from_json_string_cpp, 2}, + {"_stochtree_rfx_label_mapper_cpp", (DL_FUNC) &_stochtree_rfx_label_mapper_cpp, 1}, + {"_stochtree_rfx_label_mapper_from_json_cpp", (DL_FUNC) &_stochtree_rfx_label_mapper_from_json_cpp, 2}, + {"_stochtree_rfx_label_mapper_from_json_string_cpp", (DL_FUNC) &_stochtree_rfx_label_mapper_from_json_string_cpp, 2}, + {"_stochtree_rfx_label_mapper_to_list_cpp", (DL_FUNC) &_stochtree_rfx_label_mapper_to_list_cpp, 1}, + {"_stochtree_rfx_label_mapper_unique_group_ids_cpp", (DL_FUNC) &_stochtree_rfx_label_mapper_unique_group_ids_cpp, 1}, + {"_stochtree_rfx_model_cpp", (DL_FUNC) &_stochtree_rfx_model_cpp, 2}, + {"_stochtree_rfx_model_predict_cpp", (DL_FUNC) &_stochtree_rfx_model_predict_cpp, 3}, + {"_stochtree_rfx_model_sample_random_effects_cpp", (DL_FUNC) &_stochtree_rfx_model_sample_random_effects_cpp, 8}, + {"_stochtree_rfx_model_set_group_parameter_covariance_cpp", (DL_FUNC) &_stochtree_rfx_model_set_group_parameter_covariance_cpp, 2}, + {"_stochtree_rfx_model_set_group_parameters_cpp", (DL_FUNC) &_stochtree_rfx_model_set_group_parameters_cpp, 2}, + {"_stochtree_rfx_model_set_variance_prior_scale_cpp", (DL_FUNC) &_stochtree_rfx_model_set_variance_prior_scale_cpp, 2}, + {"_stochtree_rfx_model_set_variance_prior_shape_cpp", (DL_FUNC) &_stochtree_rfx_model_set_variance_prior_shape_cpp, 2}, + {"_stochtree_rfx_model_set_working_parameter_covariance_cpp", (DL_FUNC) &_stochtree_rfx_model_set_working_parameter_covariance_cpp, 2}, + {"_stochtree_rfx_model_set_working_parameter_cpp", (DL_FUNC) &_stochtree_rfx_model_set_working_parameter_cpp, 2}, + {"_stochtree_rfx_tracker_cpp", (DL_FUNC) &_stochtree_rfx_tracker_cpp, 1}, + {"_stochtree_rfx_tracker_get_unique_group_ids_cpp", (DL_FUNC) &_stochtree_rfx_tracker_get_unique_group_ids_cpp, 1}, + {"_stochtree_right_child_node_forest_container_cpp", (DL_FUNC) &_stochtree_right_child_node_forest_container_cpp, 4}, + {"_stochtree_rng_cpp", (DL_FUNC) &_stochtree_rng_cpp, 1}, + {"_stochtree_root_reset_active_forest_cpp", (DL_FUNC) &_stochtree_root_reset_active_forest_cpp, 1}, + {"_stochtree_root_reset_rfx_tracker_cpp", (DL_FUNC) &_stochtree_root_reset_rfx_tracker_cpp, 4}, + {"_stochtree_sample_gfr_one_iteration_cpp", (DL_FUNC) &_stochtree_sample_gfr_one_iteration_cpp, 19}, + {"_stochtree_sample_mcmc_one_iteration_cpp", (DL_FUNC) &_stochtree_sample_mcmc_one_iteration_cpp, 18}, + {"_stochtree_sample_sigma2_one_iteration_cpp", (DL_FUNC) &_stochtree_sample_sigma2_one_iteration_cpp, 5}, + {"_stochtree_sample_tau_one_iteration_cpp", (DL_FUNC) &_stochtree_sample_tau_one_iteration_cpp, 4}, + {"_stochtree_sample_without_replacement_integer_cpp", (DL_FUNC) &_stochtree_sample_without_replacement_integer_cpp, 3}, + {"_stochtree_sd_cpp", (DL_FUNC) &_stochtree_sd_cpp, 1}, + {"_stochtree_set_leaf_value_active_forest_cpp", (DL_FUNC) &_stochtree_set_leaf_value_active_forest_cpp, 2}, + {"_stochtree_set_leaf_value_forest_container_cpp", (DL_FUNC) &_stochtree_set_leaf_value_forest_container_cpp, 2}, + {"_stochtree_set_leaf_vector_active_forest_cpp", (DL_FUNC) &_stochtree_set_leaf_vector_active_forest_cpp, 2}, + {"_stochtree_set_leaf_vector_forest_container_cpp", (DL_FUNC) &_stochtree_set_leaf_vector_forest_container_cpp, 2}, + {"_stochtree_split_categories_forest_container_cpp", (DL_FUNC) &_stochtree_split_categories_forest_container_cpp, 4}, + {"_stochtree_split_index_forest_container_cpp", (DL_FUNC) &_stochtree_split_index_forest_container_cpp, 4}, + {"_stochtree_split_theshold_forest_container_cpp", (DL_FUNC) &_stochtree_split_theshold_forest_container_cpp, 4}, + {"_stochtree_subtract_from_column_vector_cpp", (DL_FUNC) &_stochtree_subtract_from_column_vector_cpp, 2}, + {"_stochtree_sum_cpp", (DL_FUNC) &_stochtree_sum_cpp, 1}, + {"_stochtree_sum_leaves_squared_ensemble_forest_container_cpp", (DL_FUNC) &_stochtree_sum_leaves_squared_ensemble_forest_container_cpp, 2}, + {"_stochtree_tree_prior_cpp", (DL_FUNC) &_stochtree_tree_prior_cpp, 4}, + {"_stochtree_update_alpha_tree_prior_cpp", (DL_FUNC) &_stochtree_update_alpha_tree_prior_cpp, 2}, + {"_stochtree_update_beta_tree_prior_cpp", (DL_FUNC) &_stochtree_update_beta_tree_prior_cpp, 2}, + {"_stochtree_update_max_depth_tree_prior_cpp", (DL_FUNC) &_stochtree_update_max_depth_tree_prior_cpp, 2}, + {"_stochtree_update_min_samples_leaf_tree_prior_cpp", (DL_FUNC) &_stochtree_update_min_samples_leaf_tree_prior_cpp, 2}, + {"_stochtree_var_cpp", (DL_FUNC) &_stochtree_var_cpp, 1}, {NULL, NULL, 0} }; } diff --git a/src/cutpoint_candidates.cpp b/src/cutpoint_candidates.cpp index e43b8219..5c7848f0 100644 --- a/src/cutpoint_candidates.cpp +++ b/src/cutpoint_candidates.cpp @@ -26,15 +26,15 @@ void FeatureCutpointGrid::CalculateStridesNumeric(Eigen::MatrixXd& covariates, E data_size_t node_size = node_end - node_begin; // Check if node has fewer observations than cutpoint_grid_size if (node_size <= cutpoint_grid_size_) { - // In this case it is still possible to have "duplicates" if the values of - // a numeric feature are very close together which in practice will only + // In this case it is still possible to have "duplicates" if the values of + // a numeric feature are very close together which in practice will only // occur when a categorical was imported incorrectly as numeric. - // For this case, we run through the sorted data, determining the stride length + // For this case, we run through the sorted data, determining the stride length // of all unique values. EnumerateNumericCutpointsDeduplication(covariates, residuals, feature_node_sort_tracker, node_id, node_begin, node_end, node_size, feature_index); } else { // Here we must essentially "thin out" the possible cutpoints - // First, we determine a step size that ensures there will be as + // First, we determine a step size that ensures there will be as // many potential cutpoints as articulated in cutpoint_grid_size ScanNumericCutpoints(covariates, residuals, feature_node_sort_tracker, node_id, node_begin, node_end, node_size, feature_index); } @@ -42,7 +42,7 @@ void FeatureCutpointGrid::CalculateStridesNumeric(Eigen::MatrixXd& covariates, E void FeatureCutpointGrid::CalculateStridesOrderedCategorical(Eigen::MatrixXd& covariates, Eigen::VectorXd& residuals, SortedNodeSampleTracker* feature_node_sort_tracker, int32_t node_id, data_size_t node_begin, data_size_t node_end, int32_t feature_index) { data_size_t node_size = node_end - node_begin; - + // Edge case 1: single observation double single_value; if (node_end - node_begin == 1) { @@ -63,7 +63,7 @@ void FeatureCutpointGrid::CalculateStridesOrderedCategorical(Eigen::MatrixXd& co cutpoint_values_.push_back(static_cast(single_value)); return; } - + // Run the "regular" algorithm for computing categorical strides data_size_t stride_begin = node_begin; data_size_t stride_length = 0; @@ -71,14 +71,14 @@ void FeatureCutpointGrid::CalculateStridesOrderedCategorical(Eigen::MatrixXd& co bool last_element; bool stride_complete; double current_val, next_val; - for (data_size_t i = node_begin; i < node_end; i++){ + for (data_size_t i = node_begin; i < node_end; i++) { current_sort_ind = feature_node_sort_tracker->SortIndex(i, feature_index); current_val = covariates(current_sort_ind, feature_index); last_element = ((i == node_end - 1)); // Increment stride length and bin_sum stride_length += 1; - + if (last_element) { // Update bin vectors node_stride_begin_.push_back(stride_begin); @@ -106,7 +106,7 @@ void FeatureCutpointGrid::CalculateStridesUnorderedCategorical(Eigen::MatrixXd& // TODO: refactor so that this initial code is shared between ordered and unordered categorical cutpoint calculation data_size_t node_size = node_end - node_begin; std::vector bin_sums; - + // Edge case 1: single observation double single_value; if (node_end - node_begin == 1) { @@ -127,7 +127,7 @@ void FeatureCutpointGrid::CalculateStridesUnorderedCategorical(Eigen::MatrixXd& cutpoint_values_.push_back(static_cast(single_value)); return; } - + // Run the "regular" algorithm for computing categorical strides data_size_t stride_begin = node_begin; data_size_t stride_length = 0; @@ -137,11 +137,11 @@ void FeatureCutpointGrid::CalculateStridesUnorderedCategorical(Eigen::MatrixXd& double current_val, next_val; double current_outcome, next_outcome; double bin_sum = 0; - for (data_size_t i = node_begin; i < node_end; i++){ + for (data_size_t i = node_begin; i < node_end; i++) { current_sort_ind = feature_node_sort_tracker->SortIndex(i, feature_index); current_val = covariates(current_sort_ind, feature_index); last_element = ((i == node_end - 1)); - + // Increment stride length and bin_sum stride_length += 1; bin_sum += residuals(current_sort_ind); @@ -156,7 +156,7 @@ void FeatureCutpointGrid::CalculateStridesUnorderedCategorical(Eigen::MatrixXd& next_sort_ind = feature_node_sort_tracker->SortIndex(i + 1, feature_index); next_val = covariates(next_sort_ind, feature_index); stride_complete = (static_cast(next_val) != static_cast(current_val)); - + if (stride_complete) { // Update bin vectors node_stride_begin_.push_back(stride_begin); @@ -173,16 +173,16 @@ void FeatureCutpointGrid::CalculateStridesUnorderedCategorical(Eigen::MatrixXd& } // Now re-arrange the categories according to the average outcome as in Fisher (1958) -// CHECK_EQ(residuals.cols(), 1); + // CHECK_EQ(residuals.cols(), 1); std::vector bin_avgs(bin_sums.size()); for (int i = 0; i < bin_sums.size(); i++) { bin_avgs[i] = bin_sums[i] / node_stride_length_[i]; } std::vector bin_sort_inds(bin_avgs.size()); std::iota(bin_sort_inds.begin(), bin_sort_inds.end(), 0); - auto comp_op = [&](size_t const &l, size_t const &r) { return std::less{}(bin_avgs[l], bin_avgs[r]); }; + auto comp_op = [&](size_t const& l, size_t const& r) { return std::less{}(bin_avgs[l], bin_avgs[r]); }; std::stable_sort(bin_sort_inds.begin(), bin_sort_inds.end(), comp_op); - + std::vector temp_stride_begin_; std::vector temp_stride_length_; std::vector temp_cutpoint_value_; @@ -192,9 +192,9 @@ void FeatureCutpointGrid::CalculateStridesUnorderedCategorical(Eigen::MatrixXd& std::copy(cutpoint_values_.begin(), cutpoint_values_.end(), std::back_inserter(temp_cutpoint_value_)); for (int i = 0; i < node_stride_begin_.size(); i++) { - node_stride_begin_[i] = temp_stride_begin_[bin_sort_inds[i]]; - node_stride_length_[i] = temp_stride_length_[bin_sort_inds[i]]; - cutpoint_values_[i] = temp_cutpoint_value_[bin_sort_inds[i]]; + node_stride_begin_[i] = temp_stride_begin_[bin_sort_inds[i]]; + node_stride_length_[i] = temp_stride_length_[bin_sort_inds[i]]; + cutpoint_values_[i] = temp_cutpoint_value_[bin_sort_inds[i]]; } } @@ -218,7 +218,7 @@ void FeatureCutpointGrid::EnumerateNumericCutpointsDeduplication(Eigen::MatrixXd cutpoint_values_.push_back(first_val); return; } - + // Run the "regular" algorithm for computing categorical strides data_size_t stride_begin = node_begin; data_size_t stride_length = 0; @@ -226,14 +226,14 @@ void FeatureCutpointGrid::EnumerateNumericCutpointsDeduplication(Eigen::MatrixXd bool last_element; bool stride_complete; double current_val, next_val; - for (data_size_t i = node_begin; i < node_end; i++){ + for (data_size_t i = node_begin; i < node_end; i++) { current_sort_ind = feature_node_sort_tracker->SortIndex(i, feature_index); current_val = covariates(current_sort_ind, feature_index); last_element = ((i == node_end - 1)); // Increment stride length stride_length += 1; - + if (last_element) { // Update bin vectors node_stride_begin_.push_back(stride_begin); @@ -277,7 +277,7 @@ void FeatureCutpointGrid::ScanNumericCutpoints(Eigen::MatrixXd& covariates, Eige cutpoint_values_.push_back(first_val); return; } - + // Run the "regular" algorithm for computing categorical strides data_size_t stride_begin = node_begin; data_size_t stride_length = 0; @@ -287,14 +287,14 @@ void FeatureCutpointGrid::ScanNumericCutpoints(Eigen::MatrixXd& covariates, Eige bool bin_complete; double step_size = node_size / cutpoint_grid_size_; double current_val, next_val; - for (data_size_t i = node_begin; i < node_end; i++){ + for (data_size_t i = node_begin; i < node_end; i++) { current_sort_ind = feature_node_sort_tracker->SortIndex(i, feature_index); current_val = covariates(current_sort_ind, feature_index); last_element = ((i == node_end - 1)); // Increment stride length stride_length += 1; - + if (last_element) { // Update bin vectors node_stride_begin_.push_back(stride_begin); @@ -319,4 +319,4 @@ void FeatureCutpointGrid::ScanNumericCutpoints(Eigen::MatrixXd& covariates, Eige } } -} // namespace StochTree +} // namespace StochTree diff --git a/src/data.cpp b/src/data.cpp index e48e9255..6a3ed492 100644 --- a/src/data.cpp +++ b/src/data.cpp @@ -10,20 +10,20 @@ ColumnMatrix::ColumnMatrix(double* data_ptr, data_size_t num_row, int num_col, b ColumnMatrix::ColumnMatrix(std::string filename, std::string column_index_string, bool header, bool precise_float_parser) { // Convert string to vector of indices - std::vector column_indices = Str2FeatureVec(column_index_string.c_str()); - + std::vector column_indices = Str2FeatureVec(column_index_string.c_str()); + // Set up CSV parser data_size_t num_global_data = 0; auto parser = std::unique_ptr(Parser::CreateParser(filename.c_str(), header, 0, precise_float_parser)); if (parser == nullptr) { Log::Fatal("Could not recognize data format of %s", filename.c_str()); } - + // Determine number of columns in the data file int num_columns = parser->NumFeatures(); // Check compatibility between column_indices and num_columns - int32_t max_col = *std::max_element(column_indices.begin(), column_indices.end()); + int max_col = *std::max_element(column_indices.begin(), column_indices.end()); if (max_col >= num_columns) Log::Fatal("Some column indices requested do not exist in the CSV file"); // Read data to memory @@ -45,7 +45,7 @@ void ColumnMatrix::LoadData(double* data_ptr, data_size_t num_row, int num_col, double temp_value; for (data_size_t i = 0; i < num_row; ++i) { for (int j = 0; j < num_col; ++j) { - if (is_row_major){ + if (is_row_major) { // Numpy 2-d arrays are stored in "row major" order temp_value = static_cast(*(data_ptr + static_cast(num_col) * i + j)); } else { @@ -61,14 +61,14 @@ ColumnVector::ColumnVector(double* data_ptr, data_size_t num_row) { LoadData(data_ptr, num_row); } -ColumnVector::ColumnVector(std::string filename, int32_t column_index, bool header, bool precise_float_parser) { +ColumnVector::ColumnVector(std::string filename, int column_index, bool header, bool precise_float_parser) { // Set up CSV parser data_size_t num_global_data = 0; auto parser = std::unique_ptr(Parser::CreateParser(filename.c_str(), header, 0, precise_float_parser)); if (parser == nullptr) { Log::Fatal("Could not recognize data format of %s", filename.c_str()); } - + // Read data to memory auto text_data = LoadTextDataToMemory(filename.c_str(), &num_global_data, header); int num_observations = static_cast(text_data.size()); @@ -131,7 +131,7 @@ void LoadData(double* data_ptr, int num_row, int num_col, bool is_row_major, Eig double temp_value; for (data_size_t i = 0; i < num_row; ++i) { for (int j = 0; j < num_col; ++j) { - if (is_row_major){ + if (is_row_major) { // Numpy 2-d arrays are stored in "row major" order temp_value = static_cast(*(data_ptr + static_cast(num_col) * i + j)); } else { @@ -154,4 +154,4 @@ void LoadData(double* data_ptr, int num_row, Eigen::VectorXd& data_vector) { } } -} // namespace StochTree +} // namespace StochTree diff --git a/src/forest.cpp b/src/forest.cpp index 77c61534..ce258821 100644 --- a/src/forest.cpp +++ b/src/forest.cpp @@ -11,899 +11,914 @@ [[cpp11::register]] cpp11::external_pointer active_forest_cpp(int num_trees, int output_dimension = 1, bool is_leaf_constant = true, bool is_exponentiated = false) { - // Create smart pointer to newly allocated object - std::unique_ptr forest_ptr_ = std::make_unique(num_trees, output_dimension, is_leaf_constant, is_exponentiated); - - // Release management of the pointer to R session - return cpp11::external_pointer(forest_ptr_.release()); + // Create smart pointer to newly allocated object + std::unique_ptr forest_ptr_ = std::make_unique(num_trees, output_dimension, is_leaf_constant, is_exponentiated); + + // Release management of the pointer to R session + return cpp11::external_pointer(forest_ptr_.release()); } [[cpp11::register]] cpp11::external_pointer forest_container_cpp(int num_trees, int output_dimension = 1, bool is_leaf_constant = true, bool is_exponentiated = false) { - // Create smart pointer to newly allocated object - std::unique_ptr forest_sample_ptr_ = std::make_unique(num_trees, output_dimension, is_leaf_constant, is_exponentiated); - - // Release management of the pointer to R session - return cpp11::external_pointer(forest_sample_ptr_.release()); + // Create smart pointer to newly allocated object + std::unique_ptr forest_sample_ptr_ = std::make_unique(num_trees, output_dimension, is_leaf_constant, is_exponentiated); + + // Release management of the pointer to R session + return cpp11::external_pointer(forest_sample_ptr_.release()); } [[cpp11::register]] cpp11::external_pointer forest_container_from_json_cpp(cpp11::external_pointer json_ptr, std::string forest_label) { - // Create smart pointer to newly allocated object - std::unique_ptr forest_sample_ptr_ = std::make_unique(0, 1, true); - - // Extract the forest's json - nlohmann::json forest_json = json_ptr->at("forests").at(forest_label); - - // Reset the forest sample container using the json - forest_sample_ptr_->Reset(); - forest_sample_ptr_->from_json(forest_json); - - // Release management of the pointer to R session - return cpp11::external_pointer(forest_sample_ptr_.release()); + // Create smart pointer to newly allocated object + std::unique_ptr forest_sample_ptr_ = std::make_unique(0, 1, true); + + // Extract the forest's json + nlohmann::json forest_json = json_ptr->at("forests").at(forest_label); + + // Reset the forest sample container using the json + forest_sample_ptr_->Reset(); + forest_sample_ptr_->from_json(forest_json); + + // Release management of the pointer to R session + return cpp11::external_pointer(forest_sample_ptr_.release()); } [[cpp11::register]] void forest_container_append_from_json_cpp(cpp11::external_pointer forest_sample_ptr, cpp11::external_pointer json_ptr, std::string forest_label) { - // Extract the forest's json - nlohmann::json forest_json = json_ptr->at("forests").at(forest_label); - - // Append to the forest sample container using the json - forest_sample_ptr->append_from_json(forest_json); + // Extract the forest's json + nlohmann::json forest_json = json_ptr->at("forests").at(forest_label); + + // Append to the forest sample container using the json + forest_sample_ptr->append_from_json(forest_json); } [[cpp11::register]] cpp11::external_pointer forest_container_from_json_string_cpp(std::string json_string, std::string forest_label) { - // Create smart pointer to newly allocated object - std::unique_ptr forest_sample_ptr_ = std::make_unique(0, 1, true); - - // Create a nlohmann::json object from the string - nlohmann::json json_object = nlohmann::json::parse(json_string); - - // Extract the forest's json - nlohmann::json forest_json = json_object.at("forests").at(forest_label); - - // Reset the forest sample container using the json - forest_sample_ptr_->Reset(); - forest_sample_ptr_->from_json(forest_json); - - // Release management of the pointer to R session - return cpp11::external_pointer(forest_sample_ptr_.release()); + // Create smart pointer to newly allocated object + std::unique_ptr forest_sample_ptr_ = std::make_unique(0, 1, true); + + // Create a nlohmann::json object from the string + nlohmann::json json_object = nlohmann::json::parse(json_string); + + // Extract the forest's json + nlohmann::json forest_json = json_object.at("forests").at(forest_label); + + // Reset the forest sample container using the json + forest_sample_ptr_->Reset(); + forest_sample_ptr_->from_json(forest_json); + + // Release management of the pointer to R session + return cpp11::external_pointer(forest_sample_ptr_.release()); } [[cpp11::register]] void forest_merge_cpp(cpp11::external_pointer inbound_forest_ptr, cpp11::external_pointer outbound_forest_ptr) { - inbound_forest_ptr->MergeForest(*outbound_forest_ptr); + inbound_forest_ptr->MergeForest(*outbound_forest_ptr); } [[cpp11::register]] void forest_add_constant_cpp(cpp11::external_pointer forest_ptr, double constant_value) { - forest_ptr->AddValueToLeaves(constant_value); + forest_ptr->AddValueToLeaves(constant_value); } [[cpp11::register]] void forest_multiply_constant_cpp(cpp11::external_pointer forest_ptr, double constant_multiple) { - forest_ptr->MultiplyLeavesByValue(constant_multiple); + forest_ptr->MultiplyLeavesByValue(constant_multiple); } [[cpp11::register]] void forest_container_append_from_json_string_cpp(cpp11::external_pointer forest_sample_ptr, std::string json_string, std::string forest_label) { - // Create a nlohmann::json object from the string - nlohmann::json json_object = nlohmann::json::parse(json_string); - - // Extract the forest's json - nlohmann::json forest_json = json_object.at("forests").at(forest_label); - - // Append to the forest sample container using the json - forest_sample_ptr->append_from_json(forest_json); + // Create a nlohmann::json object from the string + nlohmann::json json_object = nlohmann::json::parse(json_string); + + // Extract the forest's json + nlohmann::json forest_json = json_object.at("forests").at(forest_label); + + // Append to the forest sample container using the json + forest_sample_ptr->append_from_json(forest_json); } [[cpp11::register]] void combine_forests_forest_container_cpp(cpp11::external_pointer forest_samples, cpp11::integers forest_inds) { - int num_forests = forest_inds.size(); - for (int j = 1; j < num_forests; j++) { - forest_samples->MergeForests(forest_inds[0], forest_inds[j]); - } - // double combined_forest_scale_factor = 1.0 / num_forests; - // forest_samples->MultiplyForest(forest_inds[0], combined_forest_scale_factor); + int num_forests = forest_inds.size(); + for (int j = 1; j < num_forests; j++) { + forest_samples->MergeForests(forest_inds[0], forest_inds[j]); + } + // double combined_forest_scale_factor = 1.0 / num_forests; + // forest_samples->MultiplyForest(forest_inds[0], combined_forest_scale_factor); } [[cpp11::register]] void add_to_forest_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_index, double constant_value) { - forest_samples->AddToForest(forest_index, constant_value); + forest_samples->AddToForest(forest_index, constant_value); } [[cpp11::register]] void multiply_forest_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_index, double constant_multiple) { - forest_samples->MultiplyForest(forest_index, constant_multiple); + forest_samples->MultiplyForest(forest_index, constant_multiple); } [[cpp11::register]] int num_samples_forest_container_cpp(cpp11::external_pointer forest_samples) { - return forest_samples->NumSamples(); + return forest_samples->NumSamples(); } [[cpp11::register]] int ensemble_tree_max_depth_forest_container_cpp(cpp11::external_pointer forest_samples, int ensemble_num, int tree_num) { - return forest_samples->EnsembleTreeMaxDepth(ensemble_num, tree_num); + return forest_samples->EnsembleTreeMaxDepth(ensemble_num, tree_num); } [[cpp11::register]] double ensemble_average_max_depth_forest_container_cpp(cpp11::external_pointer forest_samples, int ensemble_num) { - return forest_samples->EnsembleAverageMaxDepth(ensemble_num); + return forest_samples->EnsembleAverageMaxDepth(ensemble_num); } [[cpp11::register]] double average_max_depth_forest_container_cpp(cpp11::external_pointer forest_samples) { - return forest_samples->AverageMaxDepth(); + return forest_samples->AverageMaxDepth(); } [[cpp11::register]] int num_leaves_ensemble_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num) { - StochTree::TreeEnsemble* forest = forest_samples->GetEnsemble(forest_num); - return forest->NumLeaves(); + StochTree::TreeEnsemble* forest = forest_samples->GetEnsemble(forest_num); + return forest->NumLeaves(); } [[cpp11::register]] double sum_leaves_squared_ensemble_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num) { - StochTree::TreeEnsemble* forest = forest_samples->GetEnsemble(forest_num); - return forest->SumLeafSquared(); + StochTree::TreeEnsemble* forest = forest_samples->GetEnsemble(forest_num); + return forest->SumLeafSquared(); } [[cpp11::register]] int num_trees_forest_container_cpp(cpp11::external_pointer forest_samples) { - return forest_samples->NumTrees(); + return forest_samples->NumTrees(); } [[cpp11::register]] void json_save_forest_container_cpp(cpp11::external_pointer forest_samples, std::string json_filename) { - forest_samples->SaveToJsonFile(json_filename); + forest_samples->SaveToJsonFile(json_filename); } [[cpp11::register]] void json_load_forest_container_cpp(cpp11::external_pointer forest_samples, std::string json_filename) { - forest_samples->LoadFromJsonFile(json_filename); + forest_samples->LoadFromJsonFile(json_filename); } [[cpp11::register]] int leaf_dimension_forest_container_cpp(cpp11::external_pointer forest_samples) { - return forest_samples->OutputDimension(); + return forest_samples->OutputDimension(); } [[cpp11::register]] int is_leaf_constant_forest_container_cpp(cpp11::external_pointer forest_samples) { - return forest_samples->IsLeafConstant(); + return forest_samples->IsLeafConstant(); } [[cpp11::register]] int is_exponentiated_forest_container_cpp(cpp11::external_pointer forest_samples) { - return forest_samples->IsExponentiated(); + return forest_samples->IsExponentiated(); } [[cpp11::register]] bool all_roots_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num) { - return forest_samples->AllRoots(forest_num); + return forest_samples->AllRoots(forest_num); } [[cpp11::register]] void add_sample_forest_container_cpp(cpp11::external_pointer forest_samples) { - forest_samples->AddSamples(1); + forest_samples->AddSamples(1); } [[cpp11::register]] void set_leaf_value_forest_container_cpp(cpp11::external_pointer forest_samples, double leaf_value) { - forest_samples->InitializeRoot(leaf_value); + forest_samples->InitializeRoot(leaf_value); } [[cpp11::register]] void add_sample_value_forest_container_cpp(cpp11::external_pointer forest_samples, double leaf_value) { - if (forest_samples->OutputDimension() != 1) { - cpp11::stop("leaf_value must match forest leaf dimension"); - } - int num_samples = forest_samples->NumSamples(); - forest_samples->AddSamples(1); - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(num_samples); - int num_trees = ensemble->NumTrees(); - for (int i = 0; i < num_trees; i++) { - StochTree::Tree* tree = ensemble->GetTree(i); - tree->SetLeaf(0, leaf_value); - } + if (forest_samples->OutputDimension() != 1) { + cpp11::stop("leaf_value must match forest leaf dimension"); + } + int num_samples = forest_samples->NumSamples(); + forest_samples->AddSamples(1); + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(num_samples); + int num_trees = ensemble->NumTrees(); + for (int i = 0; i < num_trees; i++) { + StochTree::Tree* tree = ensemble->GetTree(i); + tree->SetLeaf(0, leaf_value); + } } [[cpp11::register]] void add_sample_vector_forest_container_cpp(cpp11::external_pointer forest_samples, cpp11::doubles leaf_vector) { - if (forest_samples->OutputDimension() != leaf_vector.size()) { - cpp11::stop("leaf_vector must match forest leaf dimension"); - } - int num_samples = forest_samples->NumSamples(); - forest_samples->AddSamples(1); - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(num_samples); - int num_trees = ensemble->NumTrees(); - std::vector leaf_vector_cast(leaf_vector.begin(), leaf_vector.end()); - for (int i = 0; i < num_trees; i++) { - StochTree::Tree* tree = ensemble->GetTree(i); - tree->SetLeafVector(0, leaf_vector_cast); - } + if (forest_samples->OutputDimension() != leaf_vector.size()) { + cpp11::stop("leaf_vector must match forest leaf dimension"); + } + int num_samples = forest_samples->NumSamples(); + forest_samples->AddSamples(1); + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(num_samples); + int num_trees = ensemble->NumTrees(); + std::vector leaf_vector_cast(leaf_vector.begin(), leaf_vector.end()); + for (int i = 0; i < num_trees; i++) { + StochTree::Tree* tree = ensemble->GetTree(i); + tree->SetLeafVector(0, leaf_vector_cast); + } } [[cpp11::register]] void add_numeric_split_tree_value_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num, int leaf_num, int feature_num, double split_threshold, double left_leaf_value, double right_leaf_value) { - if (forest_samples->OutputDimension() != 1) { - cpp11::stop("leaf_vector must match forest leaf dimension"); - } - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - if (!tree->IsLeaf(leaf_num)) { - cpp11::stop("leaf_num is not a leaf"); - } - tree->ExpandNode(leaf_num, feature_num, split_threshold, left_leaf_value, right_leaf_value); + if (forest_samples->OutputDimension() != 1) { + cpp11::stop("leaf_vector must match forest leaf dimension"); + } + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + if (!tree->IsLeaf(leaf_num)) { + cpp11::stop("leaf_num is not a leaf"); + } + tree->ExpandNode(leaf_num, feature_num, split_threshold, left_leaf_value, right_leaf_value); } [[cpp11::register]] void add_numeric_split_tree_vector_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num, int leaf_num, int feature_num, double split_threshold, cpp11::doubles left_leaf_vector, cpp11::doubles right_leaf_vector) { - if (forest_samples->OutputDimension() != left_leaf_vector.size()) { - cpp11::stop("left_leaf_vector must match forest leaf dimension"); - } - if (forest_samples->OutputDimension() != right_leaf_vector.size()) { - cpp11::stop("right_leaf_vector must match forest leaf dimension"); - } - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - std::vector left_leaf_vector_cast(left_leaf_vector.begin(), left_leaf_vector.end()); - std::vector right_leaf_vector_cast(right_leaf_vector.begin(), right_leaf_vector.end()); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - if (!tree->IsLeaf(leaf_num)) { - cpp11::stop("leaf_num is not a leaf"); - } - tree->ExpandNode(leaf_num, feature_num, split_threshold, left_leaf_vector_cast, right_leaf_vector_cast); + if (forest_samples->OutputDimension() != left_leaf_vector.size()) { + cpp11::stop("left_leaf_vector must match forest leaf dimension"); + } + if (forest_samples->OutputDimension() != right_leaf_vector.size()) { + cpp11::stop("right_leaf_vector must match forest leaf dimension"); + } + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + std::vector left_leaf_vector_cast(left_leaf_vector.begin(), left_leaf_vector.end()); + std::vector right_leaf_vector_cast(right_leaf_vector.begin(), right_leaf_vector.end()); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + if (!tree->IsLeaf(leaf_num)) { + cpp11::stop("leaf_num is not a leaf"); + } + tree->ExpandNode(leaf_num, feature_num, split_threshold, left_leaf_vector_cast, right_leaf_vector_cast); } [[cpp11::register]] cpp11::writable::integers get_tree_leaves_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num) { - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - std::vector leaves_raw = tree->GetLeaves(); - cpp11::writable::integers leaves(leaves_raw.begin(), leaves_raw.end()); - return leaves; + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + std::vector leaves_raw = tree->GetLeaves(); + cpp11::writable::integers leaves(leaves_raw.begin(), leaves_raw.end()); + return leaves; } [[cpp11::register]] cpp11::writable::integers get_tree_split_counts_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num, int num_features) { - cpp11::writable::integers output(num_features); - for (int i = 0; i < output.size(); i++) output.at(i) = 0; - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - std::vector split_nodes = tree->GetInternalNodes(); - for (int i = 0; i < split_nodes.size(); i++) { - auto node_id = split_nodes.at(i); - auto split_feature = tree->SplitIndex(node_id); - output.at(split_feature)++; - } - return output; + cpp11::writable::integers output(num_features); + for (int i = 0; i < output.size(); i++) output.at(i) = 0; + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + std::vector split_nodes = tree->GetInternalNodes(); + for (int i = 0; i < split_nodes.size(); i++) { + auto node_id = split_nodes.at(i); + auto split_feature = tree->SplitIndex(node_id); + output.at(split_feature)++; + } + return output; } [[cpp11::register]] cpp11::writable::integers get_forest_split_counts_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int num_features) { - cpp11::writable::integers output(num_features); - for (int i = 0; i < output.size(); i++) output.at(i) = 0; - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - int num_trees = ensemble->NumTrees(); - for (int i = 0; i < num_trees; i++) { - StochTree::Tree* tree = ensemble->GetTree(i); - std::vector split_nodes = tree->GetInternalNodes(); - for (int j = 0; j < split_nodes.size(); j++) { - auto node_id = split_nodes.at(j); - auto split_feature = tree->SplitIndex(node_id); - output.at(split_feature)++; - } + cpp11::writable::integers output(num_features); + for (int i = 0; i < output.size(); i++) output.at(i) = 0; + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + int num_trees = ensemble->NumTrees(); + for (int i = 0; i < num_trees; i++) { + StochTree::Tree* tree = ensemble->GetTree(i); + std::vector split_nodes = tree->GetInternalNodes(); + for (int j = 0; j < split_nodes.size(); j++) { + auto node_id = split_nodes.at(j); + auto split_feature = tree->SplitIndex(node_id); + output.at(split_feature)++; } - return output; + } + return output; } [[cpp11::register]] cpp11::writable::integers get_overall_split_counts_forest_container_cpp(cpp11::external_pointer forest_samples, int num_features) { - cpp11::writable::integers output(num_features); - for (int i = 0; i < output.size(); i++) output.at(i) = 0; - int num_samples = forest_samples->NumSamples(); - int num_trees = forest_samples->NumTrees(); - for (int i = 0; i < num_samples; i++) { - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(i); - for (int j = 0; j < num_trees; j++) { - StochTree::Tree* tree = ensemble->GetTree(j); - std::vector split_nodes = tree->GetInternalNodes(); - for (int k = 0; k < split_nodes.size(); k++) { - auto node_id = split_nodes.at(k); - auto split_feature = tree->SplitIndex(node_id); - output.at(split_feature)++; - } - } + cpp11::writable::integers output(num_features); + for (int i = 0; i < output.size(); i++) output.at(i) = 0; + int num_samples = forest_samples->NumSamples(); + int num_trees = forest_samples->NumTrees(); + for (int i = 0; i < num_samples; i++) { + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(i); + for (int j = 0; j < num_trees; j++) { + StochTree::Tree* tree = ensemble->GetTree(j); + std::vector split_nodes = tree->GetInternalNodes(); + for (int k = 0; k < split_nodes.size(); k++) { + auto node_id = split_nodes.at(k); + auto split_feature = tree->SplitIndex(node_id); + output.at(split_feature)++; + } } - return output; + } + return output; } [[cpp11::register]] cpp11::writable::integers get_granular_split_count_array_forest_container_cpp(cpp11::external_pointer forest_samples, int num_features) { - int num_samples = forest_samples->NumSamples(); - int num_trees = forest_samples->NumTrees(); - cpp11::writable::integers output(num_features*num_samples*num_trees); - for (int elem = 0; elem < output.size(); elem++) output.at(elem) = 0; - for (int i = 0; i < num_samples; i++) { - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(i); - for (int j = 0; j < num_trees; j++) { - StochTree::Tree* tree = ensemble->GetTree(j); - std::vector split_nodes = tree->GetInternalNodes(); - for (int k = 0; k < split_nodes.size(); k++) { - auto node_id = split_nodes.at(k); - auto split_feature = tree->SplitIndex(node_id); - output.at(split_feature*num_samples*num_trees + j*num_samples + i)++; - } - } + int num_samples = forest_samples->NumSamples(); + int num_trees = forest_samples->NumTrees(); + cpp11::writable::integers output(num_features * num_samples * num_trees); + for (int elem = 0; elem < output.size(); elem++) output.at(elem) = 0; + for (int i = 0; i < num_samples; i++) { + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(i); + for (int j = 0; j < num_trees; j++) { + StochTree::Tree* tree = ensemble->GetTree(j); + std::vector split_nodes = tree->GetInternalNodes(); + for (int k = 0; k < split_nodes.size(); k++) { + auto node_id = split_nodes.at(k); + auto split_feature = tree->SplitIndex(node_id); + output.at(split_feature * num_samples * num_trees + j * num_samples + i)++; + } } - return output; + } + return output; } [[cpp11::register]] void set_leaf_vector_forest_container_cpp(cpp11::external_pointer forest_samples, cpp11::doubles leaf_vector) { - std::vector leaf_vector_converted(leaf_vector.size()); - for (int i = 0; i < leaf_vector.size(); i++) { - leaf_vector_converted[i] = leaf_vector[i]; - } - forest_samples->InitializeRoot(leaf_vector_converted); + std::vector leaf_vector_converted(leaf_vector.size()); + for (int i = 0; i < leaf_vector.size(); i++) { + leaf_vector_converted[i] = leaf_vector[i]; + } + forest_samples->InitializeRoot(leaf_vector_converted); } [[cpp11::register]] bool is_leaf_node_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num, int node_id) { - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - return tree->IsLeaf(node_id); + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + return tree->IsLeaf(node_id); } [[cpp11::register]] bool is_numeric_split_node_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num, int node_id) { - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - return tree->IsNumericSplitNode(node_id); + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + return tree->IsNumericSplitNode(node_id); } [[cpp11::register]] bool is_categorical_split_node_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num, int node_id) { - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - return tree->IsCategoricalSplitNode(node_id); + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + return tree->IsCategoricalSplitNode(node_id); } [[cpp11::register]] int parent_node_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num, int node_id) { - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - return tree->Parent(node_id); + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + return tree->Parent(node_id); } [[cpp11::register]] int left_child_node_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num, int node_id) { - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - return tree->LeftChild(node_id); + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + return tree->LeftChild(node_id); } [[cpp11::register]] int right_child_node_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num, int node_id) { - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - return tree->RightChild(node_id); + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + return tree->RightChild(node_id); } [[cpp11::register]] int node_depth_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num, int node_id) { - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - return tree->GetDepth(node_id); + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + return tree->GetDepth(node_id); } [[cpp11::register]] int split_index_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num, int node_id) { - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - return tree->SplitIndex(node_id); + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + return tree->SplitIndex(node_id); } [[cpp11::register]] double split_theshold_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num, int node_id) { - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - return tree->Threshold(node_id); + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + return tree->Threshold(node_id); } [[cpp11::register]] cpp11::writable::integers split_categories_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num, int node_id) { - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - std::vector raw_categories = tree->CategoryList(node_id); - cpp11::writable::integers output(raw_categories.begin(), raw_categories.end()); - return output; + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + std::vector raw_categories = tree->CategoryList(node_id); + cpp11::writable::integers output(raw_categories.begin(), raw_categories.end()); + return output; } [[cpp11::register]] cpp11::writable::doubles leaf_values_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num, int node_id) { - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - int num_outputs = tree->OutputDimension(); - cpp11::writable::doubles output(num_outputs); - for (int i = 0; i < num_outputs; i++) { - output[i] = tree->LeafValue(node_id, i); - } - return output; + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + int num_outputs = tree->OutputDimension(); + cpp11::writable::doubles output(num_outputs); + for (int i = 0; i < num_outputs; i++) { + output[i] = tree->LeafValue(node_id, i); + } + return output; } [[cpp11::register]] int num_nodes_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num) { - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - return tree->NumValidNodes(); + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + return tree->NumValidNodes(); } [[cpp11::register]] int num_leaves_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num) { - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - return tree->NumLeaves(); + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + return tree->NumLeaves(); } [[cpp11::register]] int num_leaf_parents_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num) { - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - return tree->NumLeafParents(); + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + return tree->NumLeafParents(); } [[cpp11::register]] int num_split_nodes_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num) { - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - return tree->NumSplitNodes(); + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + return tree->NumSplitNodes(); } [[cpp11::register]] cpp11::writable::integers nodes_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num) { - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - std::vector leaves = tree->GetNodes(); - cpp11::writable::integers output(leaves.begin(), leaves.end()); - return output; + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + std::vector leaves = tree->GetNodes(); + cpp11::writable::integers output(leaves.begin(), leaves.end()); + return output; } [[cpp11::register]] cpp11::writable::integers leaves_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num, int tree_num) { - StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); - StochTree::Tree* tree = ensemble->GetTree(tree_num); - std::vector leaves = tree->GetLeaves(); - cpp11::writable::integers output(leaves.begin(), leaves.end()); - return output; -} - -[[cpp11::register]] -void initialize_forest_model_cpp(cpp11::external_pointer data, - cpp11::external_pointer residual, - cpp11::external_pointer forest_samples, - cpp11::external_pointer tracker, - cpp11::doubles init_values, int leaf_model_int){ - // Convert leaf model type to enum - StochTree::ModelType model_type; - if (leaf_model_int == 0) model_type = StochTree::ModelType::kConstantLeafGaussian; - else if (leaf_model_int == 1) model_type = StochTree::ModelType::kUnivariateRegressionLeafGaussian; - else if (leaf_model_int == 2) model_type = StochTree::ModelType::kMultivariateRegressionLeafGaussian; - else if (leaf_model_int == 3) model_type = StochTree::ModelType::kLogLinearVariance; - else StochTree::Log::Fatal("Invalid model type"); - - // Unpack initial value - int num_trees = forest_samples->NumTrees(); - double init_val; - std::vector init_value_vector; - if ((model_type == StochTree::ModelType::kConstantLeafGaussian) || - (model_type == StochTree::ModelType::kUnivariateRegressionLeafGaussian) || - (model_type == StochTree::ModelType::kLogLinearVariance)) { - init_val = init_values.at(0); - } else if (model_type == StochTree::ModelType::kMultivariateRegressionLeafGaussian) { - int leaf_dim = init_values.size(); - init_value_vector.resize(leaf_dim); - for (int i = 0; i < leaf_dim; i++) { - init_value_vector[i] = init_values[i] / static_cast(num_trees); - } - } - - // Initialize the models accordingly - if (model_type == StochTree::ModelType::kConstantLeafGaussian) { - forest_samples->InitializeRoot(init_val / static_cast(num_trees)); - UpdateResidualEntireForest(*tracker, *data, *residual, forest_samples->GetEnsemble(0), false, std::minus()); - tracker->UpdatePredictions(forest_samples->GetEnsemble(0), *data); - } else if (model_type == StochTree::ModelType::kUnivariateRegressionLeafGaussian) { - forest_samples->InitializeRoot(init_val / static_cast(num_trees)); - UpdateResidualEntireForest(*tracker, *data, *residual, forest_samples->GetEnsemble(0), true, std::minus()); - tracker->UpdatePredictions(forest_samples->GetEnsemble(0), *data); - } else if (model_type == StochTree::ModelType::kMultivariateRegressionLeafGaussian) { - forest_samples->InitializeRoot(init_value_vector); - UpdateResidualEntireForest(*tracker, *data, *residual, forest_samples->GetEnsemble(0), true, std::minus()); - tracker->UpdatePredictions(forest_samples->GetEnsemble(0), *data); - } else if (model_type == StochTree::ModelType::kLogLinearVariance) { - forest_samples->InitializeRoot(std::log(init_val) / static_cast(num_trees)); - tracker->UpdatePredictions(forest_samples->GetEnsemble(0), *data); - int n = data->NumObservations(); - std::vector initial_preds(n, init_val); - data->AddVarianceWeights(initial_preds.data(), n); - } -} - -[[cpp11::register]] -void adjust_residual_forest_container_cpp(cpp11::external_pointer data, - cpp11::external_pointer residual, - cpp11::external_pointer forest_samples, - cpp11::external_pointer tracker, + StochTree::TreeEnsemble* ensemble = forest_samples->GetEnsemble(forest_num); + StochTree::Tree* tree = ensemble->GetTree(tree_num); + std::vector leaves = tree->GetLeaves(); + cpp11::writable::integers output(leaves.begin(), leaves.end()); + return output; +} + +[[cpp11::register]] +void initialize_forest_model_cpp(cpp11::external_pointer data, + cpp11::external_pointer residual, + cpp11::external_pointer forest_samples, + cpp11::external_pointer tracker, + cpp11::doubles init_values, int leaf_model_int) { + // Convert leaf model type to enum + StochTree::ModelType model_type; + if (leaf_model_int == 0) + model_type = StochTree::ModelType::kConstantLeafGaussian; + else if (leaf_model_int == 1) + model_type = StochTree::ModelType::kUnivariateRegressionLeafGaussian; + else if (leaf_model_int == 2) + model_type = StochTree::ModelType::kMultivariateRegressionLeafGaussian; + else if (leaf_model_int == 3) + model_type = StochTree::ModelType::kLogLinearVariance; + else + StochTree::Log::Fatal("Invalid model type"); + + // Unpack initial value + int num_trees = forest_samples->NumTrees(); + double init_val; + std::vector init_value_vector; + if ((model_type == StochTree::ModelType::kConstantLeafGaussian) || + (model_type == StochTree::ModelType::kUnivariateRegressionLeafGaussian) || + (model_type == StochTree::ModelType::kLogLinearVariance)) { + init_val = init_values.at(0); + } else if (model_type == StochTree::ModelType::kMultivariateRegressionLeafGaussian) { + int leaf_dim = init_values.size(); + init_value_vector.resize(leaf_dim); + for (int i = 0; i < leaf_dim; i++) { + init_value_vector[i] = init_values[i] / static_cast(num_trees); + } + } + + // Initialize the models accordingly + if (model_type == StochTree::ModelType::kConstantLeafGaussian) { + forest_samples->InitializeRoot(init_val / static_cast(num_trees)); + UpdateResidualEntireForest(*tracker, *data, *residual, forest_samples->GetEnsemble(0), false, std::minus()); + tracker->UpdatePredictions(forest_samples->GetEnsemble(0), *data); + } else if (model_type == StochTree::ModelType::kUnivariateRegressionLeafGaussian) { + forest_samples->InitializeRoot(init_val / static_cast(num_trees)); + UpdateResidualEntireForest(*tracker, *data, *residual, forest_samples->GetEnsemble(0), true, std::minus()); + tracker->UpdatePredictions(forest_samples->GetEnsemble(0), *data); + } else if (model_type == StochTree::ModelType::kMultivariateRegressionLeafGaussian) { + forest_samples->InitializeRoot(init_value_vector); + UpdateResidualEntireForest(*tracker, *data, *residual, forest_samples->GetEnsemble(0), true, std::minus()); + tracker->UpdatePredictions(forest_samples->GetEnsemble(0), *data); + } else if (model_type == StochTree::ModelType::kLogLinearVariance) { + forest_samples->InitializeRoot(std::log(init_val) / static_cast(num_trees)); + tracker->UpdatePredictions(forest_samples->GetEnsemble(0), *data); + int n = data->NumObservations(); + std::vector initial_preds(n, init_val); + data->AddVarianceWeights(initial_preds.data(), n); + } +} + +[[cpp11::register]] +void adjust_residual_forest_container_cpp(cpp11::external_pointer data, + cpp11::external_pointer residual, + cpp11::external_pointer forest_samples, + cpp11::external_pointer tracker, bool requires_basis, int forest_num, bool add) { - // Determine whether or not we are adding forest_num to the residuals - std::function op; - if (add) op = std::plus(); - else op = std::minus(); - - // Perform the update (addition / subtraction) operation - StochTree::UpdateResidualEntireForest(*tracker, *data, *residual, forest_samples->GetEnsemble(forest_num), requires_basis, op); + // Determine whether or not we are adding forest_num to the residuals + std::function op; + if (add) + op = std::plus(); + else + op = std::minus(); + + // Perform the update (addition / subtraction) operation + StochTree::UpdateResidualEntireForest(*tracker, *data, *residual, forest_samples->GetEnsemble(forest_num), requires_basis, op); } [[cpp11::register]] -void propagate_basis_update_forest_container_cpp(cpp11::external_pointer data, - cpp11::external_pointer residual, - cpp11::external_pointer forest_samples, - cpp11::external_pointer tracker, +void propagate_basis_update_forest_container_cpp(cpp11::external_pointer data, + cpp11::external_pointer residual, + cpp11::external_pointer forest_samples, + cpp11::external_pointer tracker, int forest_num) { - // Perform the update (addition / subtraction) operation - StochTree::UpdateResidualNewBasis(*tracker, *data, *residual, forest_samples->GetEnsemble(forest_num)); + // Perform the update (addition / subtraction) operation + StochTree::UpdateResidualNewBasis(*tracker, *data, *residual, forest_samples->GetEnsemble(forest_num)); } [[cpp11::register]] -void remove_sample_forest_container_cpp(cpp11::external_pointer forest_samples, +void remove_sample_forest_container_cpp(cpp11::external_pointer forest_samples, int forest_num) { - forest_samples->DeleteSample(forest_num); + forest_samples->DeleteSample(forest_num); } [[cpp11::register]] cpp11::writable::doubles_matrix<> predict_forest_cpp(cpp11::external_pointer forest_samples, cpp11::external_pointer dataset) { - // Predict from the sampled forests - std::vector output_raw = forest_samples->Predict(*dataset); - - // Convert result to a matrix - int n = dataset->GetCovariates().rows(); - int num_samples = forest_samples->NumSamples(); - cpp11::writable::doubles_matrix<> output(n, num_samples); - for (size_t i = 0; i < n; i++) { - for (int j = 0; j < num_samples; j++) { - output(i, j) = output_raw[n*j + i]; - } + // Predict from the sampled forests + std::vector output_raw = forest_samples->Predict(*dataset); + + // Convert result to a matrix + int n = dataset->GetCovariates().rows(); + int num_samples = forest_samples->NumSamples(); + cpp11::writable::doubles_matrix<> output(n, num_samples); + for (size_t i = 0; i < n; i++) { + for (int j = 0; j < num_samples; j++) { + output(i, j) = output_raw[n * j + i]; } - - return output; + } + + return output; } [[cpp11::register]] cpp11::writable::doubles predict_forest_raw_cpp(cpp11::external_pointer forest_samples, cpp11::external_pointer dataset) { - // Predict from the sampled forests - std::vector output_raw = forest_samples->PredictRaw(*dataset); - - // Unpack / re-arrange results - int n = dataset->GetCovariates().rows(); - int num_samples = forest_samples->NumSamples(); - int output_dimension = forest_samples->OutputDimension(); - cpp11::writable::doubles output(n*output_dimension*num_samples); - for (size_t i = 0; i < n; i++) { - for (int j = 0; j < output_dimension; j++) { - for (int k = 0; k < num_samples; k++) { - // Convert from idiosyncratic C++ storage to "column-major" --- first dimension is data row, second is output column, third is sample number - output.at(k*output_dimension*n + j*n + i) = output_raw[k*output_dimension*n + i*output_dimension + j]; - } - } + // Predict from the sampled forests + std::vector output_raw = forest_samples->PredictRaw(*dataset); + + // Unpack / re-arrange results + int n = dataset->GetCovariates().rows(); + int num_samples = forest_samples->NumSamples(); + int output_dimension = forest_samples->OutputDimension(); + cpp11::writable::doubles output(n * output_dimension * num_samples); + for (size_t i = 0; i < n; i++) { + for (int j = 0; j < output_dimension; j++) { + for (int k = 0; k < num_samples; k++) { + // Convert from idiosyncratic C++ storage to "column-major" --- first dimension is data row, second is output column, third is sample number + output.at(k * output_dimension * n + j * n + i) = output_raw[k * output_dimension * n + i * output_dimension + j]; + } } - - return output; + } + + return output; } [[cpp11::register]] cpp11::writable::doubles_matrix<> predict_forest_raw_single_forest_cpp(cpp11::external_pointer forest_samples, cpp11::external_pointer dataset, int forest_num) { - // Predict from the sampled forests - std::vector output_raw = forest_samples->PredictRaw(*dataset, forest_num); - - // Convert result to a matrix - int n = dataset->GetCovariates().rows(); - int output_dimension = forest_samples->OutputDimension(); - cpp11::writable::doubles_matrix<> output(n, output_dimension); - for (size_t i = 0; i < n; i++) { - for (int j = 0; j < output_dimension; j++) { - output(i, j) = output_raw[i*output_dimension + j]; - } + // Predict from the sampled forests + std::vector output_raw = forest_samples->PredictRaw(*dataset, forest_num); + + // Convert result to a matrix + int n = dataset->GetCovariates().rows(); + int output_dimension = forest_samples->OutputDimension(); + cpp11::writable::doubles_matrix<> output(n, output_dimension); + for (size_t i = 0; i < n; i++) { + for (int j = 0; j < output_dimension; j++) { + output(i, j) = output_raw[i * output_dimension + j]; } - - return output; + } + + return output; } [[cpp11::register]] cpp11::writable::doubles_matrix<> predict_forest_raw_single_tree_cpp(cpp11::external_pointer forest_samples, cpp11::external_pointer dataset, int forest_num, int tree_num) { - // Predict from the sampled forests - std::vector output_raw = forest_samples->PredictRawSingleTree(*dataset, forest_num, tree_num); - - // Convert result to a matrix - int n = dataset->GetCovariates().rows(); - int output_dimension = forest_samples->OutputDimension(); - cpp11::writable::doubles_matrix<> output(n, output_dimension); - for (size_t i = 0; i < n; i++) { - for (int j = 0; j < output_dimension; j++) { - output(i, j) = output_raw[i*output_dimension + j]; - } + // Predict from the sampled forests + std::vector output_raw = forest_samples->PredictRawSingleTree(*dataset, forest_num, tree_num); + + // Convert result to a matrix + int n = dataset->GetCovariates().rows(); + int output_dimension = forest_samples->OutputDimension(); + cpp11::writable::doubles_matrix<> output(n, output_dimension); + for (size_t i = 0; i < n; i++) { + for (int j = 0; j < output_dimension; j++) { + output(i, j) = output_raw[i * output_dimension + j]; } - return output; + } + return output; } [[cpp11::register]] cpp11::writable::doubles predict_active_forest_cpp(cpp11::external_pointer active_forest, cpp11::external_pointer dataset) { - int n = dataset->GetCovariates().rows(); - std::vector output(n); - active_forest->PredictInplace(*dataset, output, 0); - return output; + int n = dataset->GetCovariates().rows(); + std::vector output(n); + active_forest->PredictInplace(*dataset, output, 0); + return output; } [[cpp11::register]] cpp11::writable::doubles predict_raw_active_forest_cpp(cpp11::external_pointer active_forest, cpp11::external_pointer dataset) { - int n = dataset->GetCovariates().rows(); - int output_dimension = active_forest->OutputDimension(); - std::vector output_raw(n*output_dimension); - active_forest->PredictRawInplace(*dataset, output_raw, 0); - - cpp11::writable::doubles output(n*output_dimension); - for (size_t i = 0; i < n; i++) { - for (int j = 0; j < output_dimension; j++) { - // Convert from row-major to column-major - output.at(j*n + i) = output_raw[i*output_dimension + j]; - } + int n = dataset->GetCovariates().rows(); + int output_dimension = active_forest->OutputDimension(); + std::vector output_raw(n * output_dimension); + active_forest->PredictRawInplace(*dataset, output_raw, 0); + + cpp11::writable::doubles output(n * output_dimension); + for (size_t i = 0; i < n; i++) { + for (int j = 0; j < output_dimension; j++) { + // Convert from row-major to column-major + output.at(j * n + i) = output_raw[i * output_dimension + j]; } - - return output; + } + + return output; } [[cpp11::register]] int leaf_dimension_active_forest_cpp(cpp11::external_pointer active_forest) { - return active_forest->OutputDimension(); + return active_forest->OutputDimension(); } [[cpp11::register]] double average_max_depth_active_forest_cpp(cpp11::external_pointer active_forest) { - return active_forest->AverageMaxDepth(); + return active_forest->AverageMaxDepth(); } [[cpp11::register]] int num_trees_active_forest_cpp(cpp11::external_pointer active_forest) { - return active_forest->NumTrees(); + return active_forest->NumTrees(); } [[cpp11::register]] int ensemble_tree_max_depth_active_forest_cpp(cpp11::external_pointer active_forest, int tree_num) { - return active_forest->TreeMaxDepth(tree_num); + return active_forest->TreeMaxDepth(tree_num); } [[cpp11::register]] int is_leaf_constant_active_forest_cpp(cpp11::external_pointer active_forest) { - return active_forest->IsLeafConstant(); + return active_forest->IsLeafConstant(); } [[cpp11::register]] int is_exponentiated_active_forest_cpp(cpp11::external_pointer active_forest) { - return active_forest->IsExponentiated(); + return active_forest->IsExponentiated(); } [[cpp11::register]] bool all_roots_active_forest_cpp(cpp11::external_pointer active_forest) { - return active_forest->AllRoots(); + return active_forest->AllRoots(); } [[cpp11::register]] void set_leaf_value_active_forest_cpp(cpp11::external_pointer active_forest, double leaf_value) { - active_forest->SetLeafValue(leaf_value); + active_forest->SetLeafValue(leaf_value); } [[cpp11::register]] void set_leaf_vector_active_forest_cpp(cpp11::external_pointer active_forest, cpp11::doubles leaf_vector) { - std::vector leaf_vector_cast(leaf_vector.begin(), leaf_vector.end()); - active_forest->SetLeafVector(leaf_vector_cast); + std::vector leaf_vector_cast(leaf_vector.begin(), leaf_vector.end()); + active_forest->SetLeafVector(leaf_vector_cast); } [[cpp11::register]] void add_numeric_split_tree_value_active_forest_cpp(cpp11::external_pointer active_forest, int tree_num, int leaf_num, int feature_num, double split_threshold, double left_leaf_value, double right_leaf_value) { - if (active_forest->OutputDimension() != 1) { - cpp11::stop("leaf_vector must match forest leaf dimension"); - } - StochTree::Tree* tree = active_forest->GetTree(tree_num); - if (!tree->IsLeaf(leaf_num)) { - cpp11::stop("leaf_num is not a leaf"); - } - tree->ExpandNode(leaf_num, feature_num, split_threshold, left_leaf_value, right_leaf_value); + if (active_forest->OutputDimension() != 1) { + cpp11::stop("leaf_vector must match forest leaf dimension"); + } + StochTree::Tree* tree = active_forest->GetTree(tree_num); + if (!tree->IsLeaf(leaf_num)) { + cpp11::stop("leaf_num is not a leaf"); + } + tree->ExpandNode(leaf_num, feature_num, split_threshold, left_leaf_value, right_leaf_value); } [[cpp11::register]] void add_numeric_split_tree_vector_active_forest_cpp(cpp11::external_pointer active_forest, int tree_num, int leaf_num, int feature_num, double split_threshold, cpp11::doubles left_leaf_vector, cpp11::doubles right_leaf_vector) { - if (active_forest->OutputDimension() != left_leaf_vector.size()) { - cpp11::stop("left_leaf_vector must match forest leaf dimension"); - } - if (active_forest->OutputDimension() != right_leaf_vector.size()) { - cpp11::stop("right_leaf_vector must match forest leaf dimension"); - } - std::vector left_leaf_vector_cast(left_leaf_vector.begin(), left_leaf_vector.end()); - std::vector right_leaf_vector_cast(right_leaf_vector.begin(), right_leaf_vector.end()); - StochTree::Tree* tree = active_forest->GetTree(tree_num); - if (!tree->IsLeaf(leaf_num)) { - cpp11::stop("leaf_num is not a leaf"); - } - tree->ExpandNode(leaf_num, feature_num, split_threshold, left_leaf_vector_cast, right_leaf_vector_cast); + if (active_forest->OutputDimension() != left_leaf_vector.size()) { + cpp11::stop("left_leaf_vector must match forest leaf dimension"); + } + if (active_forest->OutputDimension() != right_leaf_vector.size()) { + cpp11::stop("right_leaf_vector must match forest leaf dimension"); + } + std::vector left_leaf_vector_cast(left_leaf_vector.begin(), left_leaf_vector.end()); + std::vector right_leaf_vector_cast(right_leaf_vector.begin(), right_leaf_vector.end()); + StochTree::Tree* tree = active_forest->GetTree(tree_num); + if (!tree->IsLeaf(leaf_num)) { + cpp11::stop("leaf_num is not a leaf"); + } + tree->ExpandNode(leaf_num, feature_num, split_threshold, left_leaf_vector_cast, right_leaf_vector_cast); } [[cpp11::register]] cpp11::writable::integers get_tree_leaves_active_forest_cpp(cpp11::external_pointer active_forest, int tree_num) { - StochTree::Tree* tree = active_forest->GetTree(tree_num); - std::vector leaves_raw = tree->GetLeaves(); - cpp11::writable::integers leaves(leaves_raw.begin(), leaves_raw.end()); - return leaves; + StochTree::Tree* tree = active_forest->GetTree(tree_num); + std::vector leaves_raw = tree->GetLeaves(); + cpp11::writable::integers leaves(leaves_raw.begin(), leaves_raw.end()); + return leaves; } [[cpp11::register]] cpp11::writable::integers get_tree_split_counts_active_forest_cpp(cpp11::external_pointer active_forest, int tree_num, int num_features) { - cpp11::writable::integers output(num_features); - for (int i = 0; i < output.size(); i++) output.at(i) = 0; - StochTree::Tree* tree = active_forest->GetTree(tree_num); - std::vector split_nodes = tree->GetInternalNodes(); - for (int i = 0; i < split_nodes.size(); i++) { - auto node_id = split_nodes.at(i); - auto feature_split = tree->SplitIndex(node_id); - output.at(feature_split)++; - } - return output; + cpp11::writable::integers output(num_features); + for (int i = 0; i < output.size(); i++) output.at(i) = 0; + StochTree::Tree* tree = active_forest->GetTree(tree_num); + std::vector split_nodes = tree->GetInternalNodes(); + for (int i = 0; i < split_nodes.size(); i++) { + auto node_id = split_nodes.at(i); + auto feature_split = tree->SplitIndex(node_id); + output.at(feature_split)++; + } + return output; } [[cpp11::register]] cpp11::writable::integers get_overall_split_counts_active_forest_cpp(cpp11::external_pointer active_forest, int num_features) { - cpp11::writable::integers output(num_features); - for (int i = 0; i < output.size(); i++) output.at(i) = 0; - int num_trees = active_forest->NumTrees(); - for (int i = 0; i < num_trees; i++) { - StochTree::Tree* tree = active_forest->GetTree(i); - std::vector split_nodes = tree->GetInternalNodes(); - for (int j = 0; j < split_nodes.size(); j++) { - auto node_id = split_nodes.at(j); - auto feature_split = tree->SplitIndex(node_id); - output.at(feature_split)++; - } + cpp11::writable::integers output(num_features); + for (int i = 0; i < output.size(); i++) output.at(i) = 0; + int num_trees = active_forest->NumTrees(); + for (int i = 0; i < num_trees; i++) { + StochTree::Tree* tree = active_forest->GetTree(i); + std::vector split_nodes = tree->GetInternalNodes(); + for (int j = 0; j < split_nodes.size(); j++) { + auto node_id = split_nodes.at(j); + auto feature_split = tree->SplitIndex(node_id); + output.at(feature_split)++; } - return output; + } + return output; } [[cpp11::register]] cpp11::writable::integers get_granular_split_count_array_active_forest_cpp(cpp11::external_pointer active_forest, int num_features) { - int num_trees = active_forest->NumTrees(); - cpp11::writable::integers output(num_features*num_trees); - for (int elem = 0; elem < output.size(); elem++) output.at(elem) = 0; - for (int i = 0; i < num_trees; i++) { - StochTree::Tree* tree = active_forest->GetTree(i); - std::vector split_nodes = tree->GetInternalNodes(); - for (int j = 0; j < split_nodes.size(); j++) { - auto node_id = split_nodes.at(j); - auto feature_split = tree->SplitIndex(node_id); - output.at(feature_split*num_trees + i)++; - } - } - return output; -} - -[[cpp11::register]] -void initialize_forest_model_active_forest_cpp(cpp11::external_pointer data, - cpp11::external_pointer residual, - cpp11::external_pointer active_forest, - cpp11::external_pointer tracker, - cpp11::doubles init_values, int leaf_model_int){ - // Convert leaf model type to enum - StochTree::ModelType model_type; - if (leaf_model_int == 0) model_type = StochTree::ModelType::kConstantLeafGaussian; - else if (leaf_model_int == 1) model_type = StochTree::ModelType::kUnivariateRegressionLeafGaussian; - else if (leaf_model_int == 2) model_type = StochTree::ModelType::kMultivariateRegressionLeafGaussian; - else if (leaf_model_int == 3) model_type = StochTree::ModelType::kLogLinearVariance; - else if (leaf_model_int == 4) model_type = StochTree::ModelType::kCloglogOrdinal; - else StochTree::Log::Fatal("Invalid model type"); - - // Unpack initial value - int num_trees = active_forest->NumTrees(); - double init_val; - std::vector init_value_vector; - if ((model_type == StochTree::ModelType::kConstantLeafGaussian) || - (model_type == StochTree::ModelType::kUnivariateRegressionLeafGaussian) || - (model_type == StochTree::ModelType::kLogLinearVariance)) { - init_val = init_values.at(0); - } else if (model_type == StochTree::ModelType::kMultivariateRegressionLeafGaussian) { - int leaf_dim = init_values.size(); - init_value_vector.resize(leaf_dim); - for (int i = 0; i < leaf_dim; i++) { - init_value_vector[i] = init_values[i] / static_cast(num_trees); - } - } - - // Initialize the models accordingly - double leaf_init_val; - if (model_type == StochTree::ModelType::kConstantLeafGaussian) { - leaf_init_val = init_val / static_cast(num_trees); - active_forest->SetLeafValue(leaf_init_val); - UpdateResidualEntireForest(*tracker, *data, *residual, active_forest.get(), false, std::minus()); - tracker->UpdatePredictions(active_forest.get(), *data); - } else if (model_type == StochTree::ModelType::kUnivariateRegressionLeafGaussian) { - leaf_init_val = init_val / static_cast(num_trees); - active_forest->SetLeafValue(leaf_init_val); - UpdateResidualEntireForest(*tracker, *data, *residual, active_forest.get(), true, std::minus()); - tracker->UpdatePredictions(active_forest.get(), *data); - } else if (model_type == StochTree::ModelType::kMultivariateRegressionLeafGaussian) { - active_forest->SetLeafVector(init_value_vector); - UpdateResidualEntireForest(*tracker, *data, *residual, active_forest.get(), true, std::minus()); - tracker->UpdatePredictions(active_forest.get(), *data); - } else if (model_type == StochTree::ModelType::kLogLinearVariance) { - leaf_init_val = std::log(init_val) / static_cast(num_trees); - active_forest->SetLeafValue(leaf_init_val); - tracker->UpdatePredictions(active_forest.get(), *data); - int n = data->NumObservations(); - std::vector initial_preds(n, init_val); - data->AddVarianceWeights(initial_preds.data(), n); - } else if (model_type == StochTree::ModelType::kCloglogOrdinal) { - leaf_init_val = init_val / static_cast(num_trees); - active_forest->SetLeafValue(leaf_init_val); - UpdateResidualEntireForest(*tracker, *data, *residual, active_forest.get(), false, std::minus()); - tracker->UpdatePredictions(active_forest.get(), *data); - } -} - -[[cpp11::register]] -void adjust_residual_active_forest_cpp(cpp11::external_pointer data, - cpp11::external_pointer residual, - cpp11::external_pointer active_forest, - cpp11::external_pointer tracker, + int num_trees = active_forest->NumTrees(); + cpp11::writable::integers output(num_features * num_trees); + for (int elem = 0; elem < output.size(); elem++) output.at(elem) = 0; + for (int i = 0; i < num_trees; i++) { + StochTree::Tree* tree = active_forest->GetTree(i); + std::vector split_nodes = tree->GetInternalNodes(); + for (int j = 0; j < split_nodes.size(); j++) { + auto node_id = split_nodes.at(j); + auto feature_split = tree->SplitIndex(node_id); + output.at(feature_split * num_trees + i)++; + } + } + return output; +} + +[[cpp11::register]] +void initialize_forest_model_active_forest_cpp(cpp11::external_pointer data, + cpp11::external_pointer residual, + cpp11::external_pointer active_forest, + cpp11::external_pointer tracker, + cpp11::doubles init_values, int leaf_model_int) { + // Convert leaf model type to enum + StochTree::ModelType model_type; + if (leaf_model_int == 0) + model_type = StochTree::ModelType::kConstantLeafGaussian; + else if (leaf_model_int == 1) + model_type = StochTree::ModelType::kUnivariateRegressionLeafGaussian; + else if (leaf_model_int == 2) + model_type = StochTree::ModelType::kMultivariateRegressionLeafGaussian; + else if (leaf_model_int == 3) + model_type = StochTree::ModelType::kLogLinearVariance; + else if (leaf_model_int == 4) + model_type = StochTree::ModelType::kCloglogOrdinal; + else + StochTree::Log::Fatal("Invalid model type"); + + // Unpack initial value + int num_trees = active_forest->NumTrees(); + double init_val; + std::vector init_value_vector; + if ((model_type == StochTree::ModelType::kConstantLeafGaussian) || + (model_type == StochTree::ModelType::kUnivariateRegressionLeafGaussian) || + (model_type == StochTree::ModelType::kLogLinearVariance)) { + init_val = init_values.at(0); + } else if (model_type == StochTree::ModelType::kMultivariateRegressionLeafGaussian) { + int leaf_dim = init_values.size(); + init_value_vector.resize(leaf_dim); + for (int i = 0; i < leaf_dim; i++) { + init_value_vector[i] = init_values[i] / static_cast(num_trees); + } + } + + // Initialize the models accordingly + double leaf_init_val; + if (model_type == StochTree::ModelType::kConstantLeafGaussian) { + leaf_init_val = init_val / static_cast(num_trees); + active_forest->SetLeafValue(leaf_init_val); + UpdateResidualEntireForest(*tracker, *data, *residual, active_forest.get(), false, std::minus()); + tracker->UpdatePredictions(active_forest.get(), *data); + } else if (model_type == StochTree::ModelType::kUnivariateRegressionLeafGaussian) { + leaf_init_val = init_val / static_cast(num_trees); + active_forest->SetLeafValue(leaf_init_val); + UpdateResidualEntireForest(*tracker, *data, *residual, active_forest.get(), true, std::minus()); + tracker->UpdatePredictions(active_forest.get(), *data); + } else if (model_type == StochTree::ModelType::kMultivariateRegressionLeafGaussian) { + active_forest->SetLeafVector(init_value_vector); + UpdateResidualEntireForest(*tracker, *data, *residual, active_forest.get(), true, std::minus()); + tracker->UpdatePredictions(active_forest.get(), *data); + } else if (model_type == StochTree::ModelType::kLogLinearVariance) { + leaf_init_val = std::log(init_val) / static_cast(num_trees); + active_forest->SetLeafValue(leaf_init_val); + tracker->UpdatePredictions(active_forest.get(), *data); + int n = data->NumObservations(); + std::vector initial_preds(n, init_val); + data->AddVarianceWeights(initial_preds.data(), n); + } else if (model_type == StochTree::ModelType::kCloglogOrdinal) { + leaf_init_val = init_val / static_cast(num_trees); + active_forest->SetLeafValue(leaf_init_val); + UpdateResidualEntireForest(*tracker, *data, *residual, active_forest.get(), false, std::minus()); + tracker->UpdatePredictions(active_forest.get(), *data); + } +} + +[[cpp11::register]] +void adjust_residual_active_forest_cpp(cpp11::external_pointer data, + cpp11::external_pointer residual, + cpp11::external_pointer active_forest, + cpp11::external_pointer tracker, bool requires_basis, bool add) { - // Determine whether or not we are adding forest predictions to the residuals - std::function op; - if (add) op = std::plus(); - else op = std::minus(); - - // Perform the update (addition / subtraction) operation - StochTree::UpdateResidualEntireForest(*tracker, *data, *residual, active_forest.get(), requires_basis, op); + // Determine whether or not we are adding forest predictions to the residuals + std::function op; + if (add) + op = std::plus(); + else + op = std::minus(); + + // Perform the update (addition / subtraction) operation + StochTree::UpdateResidualEntireForest(*tracker, *data, *residual, active_forest.get(), requires_basis, op); } [[cpp11::register]] -void propagate_basis_update_active_forest_cpp(cpp11::external_pointer data, - cpp11::external_pointer residual, - cpp11::external_pointer active_forest, +void propagate_basis_update_active_forest_cpp(cpp11::external_pointer data, + cpp11::external_pointer residual, + cpp11::external_pointer active_forest, cpp11::external_pointer tracker) { - // Perform the update (addition / subtraction) operation - StochTree::UpdateResidualNewBasis(*tracker, *data, *residual, active_forest.get()); + // Perform the update (addition / subtraction) operation + StochTree::UpdateResidualNewBasis(*tracker, *data, *residual, active_forest.get()); } [[cpp11::register]] -void reset_active_forest_cpp(cpp11::external_pointer active_forest, - cpp11::external_pointer forest_samples, +void reset_active_forest_cpp(cpp11::external_pointer active_forest, + cpp11::external_pointer forest_samples, int forest_num) { - // Extract raw pointer to the forest held at index forest_num - StochTree::TreeEnsemble* forest = forest_samples->GetEnsemble(forest_num); + // Extract raw pointer to the forest held at index forest_num + StochTree::TreeEnsemble* forest = forest_samples->GetEnsemble(forest_num); - // Reset active forest using the forest held at index forest_num - active_forest->ReconstituteFromForest(*forest); + // Reset active forest using the forest held at index forest_num + active_forest->ReconstituteFromForest(*forest); } [[cpp11::register]] -void reset_forest_model_cpp(cpp11::external_pointer forest_tracker, - cpp11::external_pointer forest, - cpp11::external_pointer data, - cpp11::external_pointer residual, +void reset_forest_model_cpp(cpp11::external_pointer forest_tracker, + cpp11::external_pointer forest, + cpp11::external_pointer data, + cpp11::external_pointer residual, bool is_mean_model) { - // Reset forest tracker using the forest held at index forest_num - forest_tracker->ReconstituteFromForest(*forest, *data, *residual, is_mean_model); + // Reset forest tracker using the forest held at index forest_num + forest_tracker->ReconstituteFromForest(*forest, *data, *residual, is_mean_model); } [[cpp11::register]] void root_reset_active_forest_cpp(cpp11::external_pointer active_forest) { - // Reset active forest to root - active_forest->ResetRoot(); + // Reset active forest to root + active_forest->ResetRoot(); } diff --git a/src/kernel.cpp b/src/kernel.cpp index 38fdd35c..1d3309d6 100644 --- a/src/kernel.cpp +++ b/src/kernel.cpp @@ -8,38 +8,37 @@ typedef Eigen::Map forest_container, int forest_num) { - return forest_container->GetEnsemble(forest_num)->GetMaxLeafIndex() - 1; + return forest_container->GetEnsemble(forest_num)->GetMaxLeafIndex() - 1; } [[cpp11::register]] cpp11::writable::integers_matrix<> compute_leaf_indices_cpp( - cpp11::external_pointer forest_container, - cpp11::doubles_matrix<> covariates, cpp11::integers forest_nums -) { - // Wrap an Eigen Map around the raw data of the covariate matrix - StochTree::data_size_t num_obs = covariates.nrow(); - int num_covariates = covariates.ncol(); - double* covariate_data_ptr = REAL(PROTECT(covariates)); - DoubleMatrixType covariates_eigen(covariate_data_ptr, num_obs, num_covariates); - - // Extract other output dimensions - int num_trees = forest_container->NumTrees(); - int num_samples = forest_nums.size(); - - // Declare outputs - cpp11::writable::integers_matrix<> output_matrix(num_obs*num_trees, num_samples); - - // Wrap Eigen Maps around kernel and kernel inverse matrices - int* output_data_ptr = INTEGER(PROTECT(output_matrix)); - IntMatrixType output_eigen(output_data_ptr, num_obs*num_trees, num_samples); - - // Compute leaf indices - std::vector forest_indices(forest_nums.begin(), forest_nums.end()); - forest_container->PredictLeafIndicesInplace(covariates_eigen, output_eigen, forest_indices, num_trees, num_obs); - - // Unprotect pointers to R data - UNPROTECT(2); - - // Return matrix - return output_matrix; + cpp11::external_pointer forest_container, + cpp11::doubles_matrix<> covariates, cpp11::integers forest_nums) { + // Wrap an Eigen Map around the raw data of the covariate matrix + StochTree::data_size_t num_obs = covariates.nrow(); + int num_covariates = covariates.ncol(); + double* covariate_data_ptr = REAL(PROTECT(covariates)); + DoubleMatrixType covariates_eigen(covariate_data_ptr, num_obs, num_covariates); + + // Extract other output dimensions + int num_trees = forest_container->NumTrees(); + int num_samples = forest_nums.size(); + + // Declare outputs + cpp11::writable::integers_matrix<> output_matrix(num_obs * num_trees, num_samples); + + // Wrap Eigen Maps around kernel and kernel inverse matrices + int* output_data_ptr = INTEGER(PROTECT(output_matrix)); + IntMatrixType output_eigen(output_data_ptr, num_obs * num_trees, num_samples); + + // Compute leaf indices + std::vector forest_indices(forest_nums.begin(), forest_nums.end()); + forest_container->PredictLeafIndicesInplace(covariates_eigen, output_eigen, forest_indices, num_trees, num_obs); + + // Unprotect pointers to R data + UNPROTECT(2); + + // Return matrix + return output_matrix; } diff --git a/src/leaf_model.cpp b/src/leaf_model.cpp index c456c29b..5538db9f 100644 --- a/src/leaf_model.cpp +++ b/src/leaf_model.cpp @@ -5,37 +5,31 @@ namespace StochTree { double GaussianConstantLeafModel::SplitLogMarginalLikelihood(GaussianConstantSuffStat& left_stat, GaussianConstantSuffStat& right_stat, double global_variance) { - double left_log_ml = ( - -0.5*std::log(1 + tau_*(left_stat.sum_w/global_variance)) + ((tau_*left_stat.sum_yw*left_stat.sum_yw)/(2.0*global_variance*(tau_*left_stat.sum_w + global_variance))) - ); + double left_log_ml = (-0.5 * std::log(1 + tau_ * (left_stat.sum_w / global_variance)) + ((tau_ * left_stat.sum_yw * left_stat.sum_yw) / (2.0 * global_variance * (tau_ * left_stat.sum_w + global_variance)))); - double right_log_ml = ( - -0.5*std::log(1 + tau_*(right_stat.sum_w/global_variance)) + ((tau_*right_stat.sum_yw*right_stat.sum_yw)/(2.0*global_variance*(tau_*right_stat.sum_w + global_variance))) - ); + double right_log_ml = (-0.5 * std::log(1 + tau_ * (right_stat.sum_w / global_variance)) + ((tau_ * right_stat.sum_yw * right_stat.sum_yw) / (2.0 * global_variance * (tau_ * right_stat.sum_w + global_variance)))); return left_log_ml + right_log_ml; } double GaussianConstantLeafModel::NoSplitLogMarginalLikelihood(GaussianConstantSuffStat& suff_stat, double global_variance) { - double log_ml = ( - -0.5*std::log(1 + tau_*(suff_stat.sum_w/global_variance)) + ((tau_*suff_stat.sum_yw*suff_stat.sum_yw)/(2.0*global_variance*(tau_*suff_stat.sum_w + global_variance))) - ); + double log_ml = (-0.5 * std::log(1 + tau_ * (suff_stat.sum_w / global_variance)) + ((tau_ * suff_stat.sum_yw * suff_stat.sum_yw) / (2.0 * global_variance * (tau_ * suff_stat.sum_w + global_variance)))); return log_ml; } double GaussianConstantLeafModel::PosteriorParameterMean(GaussianConstantSuffStat& suff_stat, double global_variance) { - return (tau_*suff_stat.sum_yw) / (suff_stat.sum_w*tau_ + global_variance); + return (tau_ * suff_stat.sum_yw) / (suff_stat.sum_w * tau_ + global_variance); } double GaussianConstantLeafModel::PosteriorParameterVariance(GaussianConstantSuffStat& suff_stat, double global_variance) { - return (tau_*global_variance) / (suff_stat.sum_w*tau_ + global_variance); + return (tau_ * global_variance) / (suff_stat.sum_w * tau_ + global_variance); } void GaussianConstantLeafModel::SampleLeafParameters(ForestDataset& dataset, ForestTracker& tracker, ColumnVector& residual, Tree* tree, int tree_num, double global_variance, std::mt19937& gen) { // Vector of leaf indices for tree std::vector tree_leaves = tree->GetLeaves(); - + // Initialize sufficient statistics GaussianConstantSuffStat node_suff_stat = GaussianConstantSuffStat(); @@ -49,11 +43,11 @@ void GaussianConstantLeafModel::SampleLeafParameters(ForestDataset& dataset, For leaf_id = tree_leaves[i]; node_suff_stat.ResetSuffStat(); AccumulateSingleNodeSuffStat(node_suff_stat, dataset, tracker, residual, tree_num, leaf_id); - + // Compute posterior mean and variance node_mean = PosteriorParameterMean(node_suff_stat, global_variance); node_variance = PosteriorParameterVariance(node_suff_stat, global_variance); - + // Draw from N(mean, stddev^2) and set the leaf parameter with each draw node_mu = normal_sampler_.Sample(node_mean, node_variance, gen); tree->SetLeaf(leaf_id, node_mu); @@ -69,37 +63,31 @@ void GaussianConstantLeafModel::SetEnsembleRootPredictedValue(ForestDataset& dat } double GaussianUnivariateRegressionLeafModel::SplitLogMarginalLikelihood(GaussianUnivariateRegressionSuffStat& left_stat, GaussianUnivariateRegressionSuffStat& right_stat, double global_variance) { - double left_log_ml = ( - -0.5*std::log(1 + tau_*(left_stat.sum_xxw/global_variance)) + ((tau_*left_stat.sum_yxw*left_stat.sum_yxw)/(2.0*global_variance*(tau_*left_stat.sum_xxw + global_variance))) - ); + double left_log_ml = (-0.5 * std::log(1 + tau_ * (left_stat.sum_xxw / global_variance)) + ((tau_ * left_stat.sum_yxw * left_stat.sum_yxw) / (2.0 * global_variance * (tau_ * left_stat.sum_xxw + global_variance)))); - double right_log_ml = ( - -0.5*std::log(1 + tau_*(right_stat.sum_xxw/global_variance)) + ((tau_*right_stat.sum_yxw*right_stat.sum_yxw)/(2.0*global_variance*(tau_*right_stat.sum_xxw + global_variance))) - ); + double right_log_ml = (-0.5 * std::log(1 + tau_ * (right_stat.sum_xxw / global_variance)) + ((tau_ * right_stat.sum_yxw * right_stat.sum_yxw) / (2.0 * global_variance * (tau_ * right_stat.sum_xxw + global_variance)))); return left_log_ml + right_log_ml; } double GaussianUnivariateRegressionLeafModel::NoSplitLogMarginalLikelihood(GaussianUnivariateRegressionSuffStat& suff_stat, double global_variance) { - double log_ml = ( - -0.5*std::log(1 + tau_*(suff_stat.sum_xxw/global_variance)) + ((tau_*suff_stat.sum_yxw*suff_stat.sum_yxw)/(2.0*global_variance*(tau_*suff_stat.sum_xxw + global_variance))) - ); + double log_ml = (-0.5 * std::log(1 + tau_ * (suff_stat.sum_xxw / global_variance)) + ((tau_ * suff_stat.sum_yxw * suff_stat.sum_yxw) / (2.0 * global_variance * (tau_ * suff_stat.sum_xxw + global_variance)))); return log_ml; } double GaussianUnivariateRegressionLeafModel::PosteriorParameterMean(GaussianUnivariateRegressionSuffStat& suff_stat, double global_variance) { - return (tau_*suff_stat.sum_yxw) / (suff_stat.sum_xxw*tau_ + global_variance); + return (tau_ * suff_stat.sum_yxw) / (suff_stat.sum_xxw * tau_ + global_variance); } double GaussianUnivariateRegressionLeafModel::PosteriorParameterVariance(GaussianUnivariateRegressionSuffStat& suff_stat, double global_variance) { - return (tau_*global_variance) / (suff_stat.sum_xxw*tau_ + global_variance); + return (tau_ * global_variance) / (suff_stat.sum_xxw * tau_ + global_variance); } void GaussianUnivariateRegressionLeafModel::SampleLeafParameters(ForestDataset& dataset, ForestTracker& tracker, ColumnVector& residual, Tree* tree, int tree_num, double global_variance, std::mt19937& gen) { // Vector of leaf indices for tree std::vector tree_leaves = tree->GetLeaves(); - + // Initialize sufficient statistics GaussianUnivariateRegressionSuffStat node_suff_stat = GaussianUnivariateRegressionSuffStat(); @@ -113,11 +101,11 @@ void GaussianUnivariateRegressionLeafModel::SampleLeafParameters(ForestDataset& leaf_id = tree_leaves[i]; node_suff_stat.ResetSuffStat(); AccumulateSingleNodeSuffStat(node_suff_stat, dataset, tracker, residual, tree_num, leaf_id); - + // Compute posterior mean and variance node_mean = PosteriorParameterMean(node_suff_stat, global_variance); node_variance = PosteriorParameterVariance(node_suff_stat, global_variance); - + // Draw from N(mean, stddev^2) and set the leaf parameter with each draw node_mu = normal_sampler_.Sample(node_mean, node_variance, gen); tree->SetLeaf(leaf_id, node_mu); @@ -134,38 +122,32 @@ void GaussianUnivariateRegressionLeafModel::SetEnsembleRootPredictedValue(Forest double GaussianMultivariateRegressionLeafModel::SplitLogMarginalLikelihood(GaussianMultivariateRegressionSuffStat& left_stat, GaussianMultivariateRegressionSuffStat& right_stat, double global_variance) { Eigen::MatrixXd I_p = Eigen::MatrixXd::Identity(left_stat.p, left_stat.p); - double left_log_ml = ( - -0.5*std::log((I_p + (Sigma_0_ * left_stat.XtWX)/global_variance).determinant()) + 0.5*((left_stat.ytWX/global_variance) * (Sigma_0_.inverse() + (left_stat.XtWX/global_variance)).inverse() * (left_stat.ytWX/global_variance).transpose()).value() - ); + double left_log_ml = (-0.5 * std::log((I_p + (Sigma_0_ * left_stat.XtWX) / global_variance).determinant()) + 0.5 * ((left_stat.ytWX / global_variance) * (Sigma_0_.inverse() + (left_stat.XtWX / global_variance)).inverse() * (left_stat.ytWX / global_variance).transpose()).value()); - double right_log_ml = ( - -0.5*std::log((I_p + (Sigma_0_ * right_stat.XtWX)/global_variance).determinant()) + 0.5*((right_stat.ytWX/global_variance) * (Sigma_0_.inverse() + (right_stat.XtWX/global_variance)).inverse() * (right_stat.ytWX/global_variance).transpose()).value() - ); + double right_log_ml = (-0.5 * std::log((I_p + (Sigma_0_ * right_stat.XtWX) / global_variance).determinant()) + 0.5 * ((right_stat.ytWX / global_variance) * (Sigma_0_.inverse() + (right_stat.XtWX / global_variance)).inverse() * (right_stat.ytWX / global_variance).transpose()).value()); return left_log_ml + right_log_ml; } double GaussianMultivariateRegressionLeafModel::NoSplitLogMarginalLikelihood(GaussianMultivariateRegressionSuffStat& suff_stat, double global_variance) { Eigen::MatrixXd I_p = Eigen::MatrixXd::Identity(suff_stat.p, suff_stat.p); - double log_ml = ( - -0.5*std::log((I_p + (Sigma_0_ * suff_stat.XtWX)/global_variance).determinant()) + 0.5*((suff_stat.ytWX/global_variance) * (Sigma_0_.inverse() + (suff_stat.XtWX/global_variance)).inverse() * (suff_stat.ytWX/global_variance).transpose()).value() - ); + double log_ml = (-0.5 * std::log((I_p + (Sigma_0_ * suff_stat.XtWX) / global_variance).determinant()) + 0.5 * ((suff_stat.ytWX / global_variance) * (Sigma_0_.inverse() + (suff_stat.XtWX / global_variance)).inverse() * (suff_stat.ytWX / global_variance).transpose()).value()); return log_ml; } Eigen::VectorXd GaussianMultivariateRegressionLeafModel::PosteriorParameterMean(GaussianMultivariateRegressionSuffStat& suff_stat, double global_variance) { - return (Sigma_0_.inverse() + (suff_stat.XtWX/global_variance)).inverse() * (suff_stat.ytWX/global_variance).transpose(); + return (Sigma_0_.inverse() + (suff_stat.XtWX / global_variance)).inverse() * (suff_stat.ytWX / global_variance).transpose(); } Eigen::MatrixXd GaussianMultivariateRegressionLeafModel::PosteriorParameterVariance(GaussianMultivariateRegressionSuffStat& suff_stat, double global_variance) { - return (Sigma_0_.inverse() + (suff_stat.XtWX/global_variance)).inverse(); + return (Sigma_0_.inverse() + (suff_stat.XtWX / global_variance)).inverse(); } void GaussianMultivariateRegressionLeafModel::SampleLeafParameters(ForestDataset& dataset, ForestTracker& tracker, ColumnVector& residual, Tree* tree, int tree_num, double global_variance, std::mt19937& gen) { // Vector of leaf indices for tree std::vector tree_leaves = tree->GetLeaves(); - + // Initialize sufficient statistics int num_basis = dataset.GetBasis().cols(); GaussianMultivariateRegressionSuffStat node_suff_stat = GaussianMultivariateRegressionSuffStat(num_basis); @@ -180,11 +162,11 @@ void GaussianMultivariateRegressionLeafModel::SampleLeafParameters(ForestDataset leaf_id = tree_leaves[i]; node_suff_stat.ResetSuffStat(); AccumulateSingleNodeSuffStat(node_suff_stat, dataset, tracker, residual, tree_num, leaf_id); - + // Compute posterior mean and variance node_mean = PosteriorParameterMean(node_suff_stat, global_variance); node_variance = PosteriorParameterVariance(node_suff_stat, global_variance); - + // Draw from N(mean, stddev^2) and set the leaf parameter with each draw node_mu = multivariate_normal_sampler_.Sample(node_mean, node_variance, gen); tree->SetLeafVector(leaf_id, node_mu); @@ -194,7 +176,7 @@ void GaussianMultivariateRegressionLeafModel::SampleLeafParameters(ForestDataset void GaussianMultivariateRegressionLeafModel::SetEnsembleRootPredictedValue(ForestDataset& dataset, TreeEnsemble* ensemble, double root_pred_value) { int num_trees = ensemble->NumTrees(); int num_basis = dataset.GetBasis().cols(); - + // Check that root predicted value is close to 0 // TODO: formalize and document this if ((root_pred_value < -0.1) || root_pred_value > 0.1) { @@ -240,7 +222,7 @@ double LogLinearVarianceLeafModel::PosteriorParameterScale(LogLinearVarianceSuff void LogLinearVarianceLeafModel::SampleLeafParameters(ForestDataset& dataset, ForestTracker& tracker, ColumnVector& residual, Tree* tree, int tree_num, double global_variance, std::mt19937& gen) { // Vector of leaf indices for tree std::vector tree_leaves = tree->GetLeaves(); - + // Initialize sufficient statistics LogLinearVarianceSuffStat node_suff_stat = LogLinearVarianceSuffStat(); @@ -254,11 +236,11 @@ void LogLinearVarianceLeafModel::SampleLeafParameters(ForestDataset& dataset, Fo leaf_id = tree_leaves[i]; node_suff_stat.ResetSuffStat(); AccumulateSingleNodeSuffStat(node_suff_stat, dataset, tracker, residual, tree_num, leaf_id); - + // Compute posterior mean and variance node_shape = PosteriorParameterShape(node_suff_stat, global_variance); node_rate = PosteriorParameterScale(node_suff_stat, global_variance); - + // Draw from IG(shape, scale) and set the leaf parameter with each draw node_mu = -std::log(sample_gamma(gen, node_shape, 1.) / node_rate); // node_mu = std::log(gamma_sampler_.Sample(node_shape, node_rate, gen, true)); @@ -306,7 +288,7 @@ double CloglogOrdinalLeafModel::PosteriorParameterRate(CloglogOrdinalSuffStat& s void CloglogOrdinalLeafModel::SampleLeafParameters(ForestDataset& dataset, ForestTracker& tracker, ColumnVector& residual, Tree* tree, int tree_num, double global_variance, std::mt19937& gen) { // Vector of leaf indices for tree std::vector tree_leaves = tree->GetLeaves(); - + // Initialize sufficient statistics CloglogOrdinalSuffStat node_suff_stat = CloglogOrdinalSuffStat(); @@ -320,7 +302,7 @@ void CloglogOrdinalLeafModel::SampleLeafParameters(ForestDataset& dataset, Fores leaf_id = tree_leaves[i]; node_suff_stat.ResetSuffStat(); AccumulateSingleNodeSuffStat(node_suff_stat, dataset, tracker, residual, tree_num, leaf_id); - + // Compute posterior shape and rate node_shape = PosteriorParameterShape(node_suff_stat, global_variance); node_rate = PosteriorParameterRate(node_suff_stat, global_variance); @@ -334,4 +316,4 @@ void CloglogOrdinalLeafModel::SampleLeafParameters(ForestDataset& dataset, Fores } } -} // namespace StochTree +} // namespace StochTree diff --git a/src/ordinal_sampler.cpp b/src/ordinal_sampler.cpp index 81841bde..dcb62ab6 100644 --- a/src/ordinal_sampler.cpp +++ b/src/ordinal_sampler.cpp @@ -17,9 +17,9 @@ double OrdinalSampler::SampleTruncatedExponential(std::mt19937& gen, double rate void OrdinalSampler::UpdateLatentVariables(ForestDataset& dataset, Eigen::VectorXd& outcome, std::mt19937& gen) { // Get auxiliary data vectors - const std::vector& gamma = dataset.GetAuxiliaryDataVector(2); // gamma cutpoints + const std::vector& gamma = dataset.GetAuxiliaryDataVector(2); // gamma cutpoints const std::vector& lambda_hat = dataset.GetAuxiliaryDataVector(1); // forest predictions: lambda_hat_i = sum_t lambda_t(x_i) - std::vector& Z = dataset.GetAuxiliaryDataVector(0); // latent variables: z_i ~ TExp(e^{gamma[y_i] + lambda_hat_i}; 0, 1) + std::vector& Z = dataset.GetAuxiliaryDataVector(0); // latent variables: z_i ~ TExp(e^{gamma[y_i] + lambda_hat_i}; 0, 1) int K = gamma.size() + 1; // Number of ordinal categories int N = dataset.NumObservations(); @@ -41,12 +41,12 @@ void OrdinalSampler::UpdateLatentVariables(ForestDataset& dataset, Eigen::Vector } } -void OrdinalSampler::UpdateGammaParams(ForestDataset& dataset, Eigen::VectorXd& outcome, - double alpha_gamma, double beta_gamma, +void OrdinalSampler::UpdateGammaParams(ForestDataset& dataset, Eigen::VectorXd& outcome, + double alpha_gamma, double beta_gamma, double gamma_0, std::mt19937& gen) { // Get auxiliary data vectors - std::vector& gamma = dataset.GetAuxiliaryDataVector(2); // cutpoints gamma_k's - const std::vector& Z = dataset.GetAuxiliaryDataVector(0); // latent variables z_i's + std::vector& gamma = dataset.GetAuxiliaryDataVector(2); // cutpoints gamma_k's + const std::vector& Z = dataset.GetAuxiliaryDataVector(0); // latent variables z_i's const std::vector& lambda_hat = dataset.GetAuxiliaryDataVector(1); // forest predictions: lambda_hat_i = sum_t lambda_t(x_i) int K = gamma.size() + 1; // Number of ordinal categories @@ -78,23 +78,23 @@ void OrdinalSampler::UpdateGammaParams(ForestDataset& dataset, Eigen::VectorXd& // Set the first gamma parameter to gamma_0 (e.g., 0) for identifiability // if (K > 2) { - gamma[0] = gamma_0; + gamma[0] = gamma_0; // } } void OrdinalSampler::UpdateCumulativeExpSums(ForestDataset& dataset) { // Get auxiliary data vectors const std::vector& gamma = dataset.GetAuxiliaryDataVector(2); // cutpoints gamma_k's - std::vector& seg = dataset.GetAuxiliaryDataVector(3); // seg_k = sum_{j=0}^{k-1} exp(gamma_j) + std::vector& seg = dataset.GetAuxiliaryDataVector(3); // seg_k = sum_{j=0}^{k-1} exp(gamma_j) // Update seg (sum of exponentials of gamma cutpoints) for (int j = 0; j < static_cast(seg.size()); j++) { if (j == 0) { - seg[j] = 0.0; // checked and it is correct + seg[j] = 0.0; // checked and it is correct } else { seg[j] = seg[j - 1] + std::exp(gamma[j - 1]); // checked and it is correct } } } -} // namespace StochTree +} // namespace StochTree diff --git a/src/partition_tracker.cpp b/src/partition_tracker.cpp index 73b37fe8..6cbbf605 100644 --- a/src/partition_tracker.cpp +++ b/src/partition_tracker.cpp @@ -36,7 +36,6 @@ void ForestTracker::ReconstituteFromForest(TreeEnsemble& forest, ForestDataset& // Reconstitute each of the remaining data structures in the tracker based on splits in the ensemble // UnsortedNodeSampleTracker unsorted_node_sample_tracker_->ReconstituteFromForest(forest, dataset); - } void ForestTracker::ResetRoot(Eigen::MatrixXd& covariates, std::vector& feature_types, int32_t tree_num) { @@ -45,19 +44,19 @@ void ForestTracker::ResetRoot(Eigen::MatrixXd& covariates, std::vectorGetNodeId(observation_num, tree_num);} +data_size_t ForestTracker::GetNodeId(int observation_num, int tree_num) { return sample_node_mapper_->GetNodeId(observation_num, tree_num); } -data_size_t ForestTracker::UnsortedNodeBegin(int tree_id, int node_id) {return unsorted_node_sample_tracker_->NodeBegin(tree_id, node_id);} +data_size_t ForestTracker::UnsortedNodeBegin(int tree_id, int node_id) { return unsorted_node_sample_tracker_->NodeBegin(tree_id, node_id); } -data_size_t ForestTracker::UnsortedNodeEnd(int tree_id, int node_id) {return unsorted_node_sample_tracker_->NodeEnd(tree_id, node_id);} +data_size_t ForestTracker::UnsortedNodeEnd(int tree_id, int node_id) { return unsorted_node_sample_tracker_->NodeEnd(tree_id, node_id); } -data_size_t ForestTracker::UnsortedNodeSize(int tree_id, int node_id) {return unsorted_node_sample_tracker_->NodeSize(tree_id, node_id);} +data_size_t ForestTracker::UnsortedNodeSize(int tree_id, int node_id) { return unsorted_node_sample_tracker_->NodeSize(tree_id, node_id); } -data_size_t ForestTracker::SortedNodeBegin(int node_id, int feature_id) {return sorted_node_sample_tracker_->NodeBegin(node_id, feature_id);} +data_size_t ForestTracker::SortedNodeBegin(int node_id, int feature_id) { return sorted_node_sample_tracker_->NodeBegin(node_id, feature_id); } -data_size_t ForestTracker::SortedNodeEnd(int node_id, int feature_id) {return sorted_node_sample_tracker_->NodeEnd(node_id, feature_id);} +data_size_t ForestTracker::SortedNodeEnd(int node_id, int feature_id) { return sorted_node_sample_tracker_->NodeEnd(node_id, feature_id); } -data_size_t ForestTracker::SortedNodeSize(int node_id, int feature_id) {return sorted_node_sample_tracker_->NodeSize(node_id, feature_id);} +data_size_t ForestTracker::SortedNodeSize(int node_id, int feature_id) { return sorted_node_sample_tracker_->NodeSize(node_id, feature_id); } std::vector::iterator ForestTracker::UnsortedNodeBeginIterator(int tree_id, int node_id) { return unsorted_node_sample_tracker_->NodeBeginIterator(tree_id, node_id); @@ -84,7 +83,7 @@ void ForestTracker::AssignAllSamplesToRoot(int32_t tree_num) { void ForestTracker::AssignAllSamplesToConstantPrediction(double value) { for (data_size_t i = 0; i < num_observations_; i++) { - sum_predictions_[i] = value*num_trees_; + sum_predictions_[i] = value * num_trees_; } for (int i = 0; i < num_trees_; i++) { sample_pred_mapper_->AssignAllSamplesToConstantPrediction(i, value); @@ -369,60 +368,66 @@ void FeatureUnsortedPartition::ReconstituteFromTree(Tree& tree, ForestDataset& d data_size_t num_true, num_false; TreeSplit split_rule; int split_index; + // Record deleted (recycled) node slots so the free list matches the tree. for (int i = 0; i < num_nodes_; i++) { - is_deleted = tree.IsDeleted(i); - if (is_deleted) { - deleted_nodes_.push_back(i); + if (tree.IsDeleted(i)) deleted_nodes_.push_back(i); + } + + // Partition the observation indices top-down, processing each node only after its + // parent has set this node's [node_begin_, node_begin_+node_length_) bounds. We must + // traverse in topological (root-to-leaf) order rather than by node id, because the + // tree's node ids may be RECYCLED after prune+regrow, so a child can have a smaller + // id than its parent. Iterating by id (the previous behavior) read stale bounds for + // such nodes and mis-assigned observations, which broke warm-start continuation. + node_begin_[0] = 0; + node_length_[0] = n; + std::vector traversal_queue; + traversal_queue.push_back(0); + for (std::size_t qh = 0; qh < traversal_queue.size(); qh++) { + int i = traversal_queue[qh]; + node_start_idx = node_begin_[i]; + num_node_elements = node_length_[i]; + // Tree node info + parent_nodes_[i] = tree.Parent(i); + node_type = tree.NodeType(i); + left_nodes_[i] = tree.LeftChild(i); + right_nodes_[i] = tree.RightChild(i); + // Only partition / descend if this node is an internal split node + if (node_type == TreeNodeType::kNumericalSplitNode) { + split_rule = TreeSplit(tree.Threshold(i)); + split_index = tree.SplitIndex(i); + } else if (node_type == TreeNodeType::kCategoricalSplitNode) { + std::vector categories = tree.CategoryList(i); + split_rule = TreeSplit(categories); + split_index = tree.SplitIndex(i); } else { - // Node beginning and length in indices_ - if (i == 0) { - node_start_idx = 0; - num_node_elements = n; - } else { - node_start_idx = node_begin_[i]; - num_node_elements = node_length_[i]; - } - // Tree node info - parent_nodes_[i] = tree.Parent(i); - node_type = tree.NodeType(i); - left_nodes_[i] = tree.LeftChild(i); - right_nodes_[i] = tree.RightChild(i); - // Only update indices_, node_begin_ and node_length_ if a split is to be added - if (node_type == TreeNodeType::kNumericalSplitNode) { - // Extract split rule - split_rule = TreeSplit(tree.Threshold(i)); - split_index = tree.SplitIndex(i); - } else if (node_type == TreeNodeType::kCategoricalSplitNode) { - std::vector categories = tree.CategoryList(i); - split_rule = TreeSplit(categories); - split_index = tree.SplitIndex(i); - } else { - continue; - } - // Partition the node indices - auto node_begin = (indices_.begin() + node_begin_[i]); - auto node_end = (indices_.begin() + node_begin_[i] + node_length_[i]); - auto right_node_begin = std::stable_partition(node_begin, node_end, [&](int row) { return split_rule.SplitTrue(covariates(row, split_index)); }); - - // Determine the number of true and false elements - node_begin = (indices_.begin() + node_begin_[i]); - num_true = std::distance(node_begin, right_node_begin); - num_false = num_node_elements - num_true; - - // Add left node tracking information - node_begin_[left_nodes_[i]] = node_start_idx; - node_length_[left_nodes_[i]] = num_true; - parent_nodes_[left_nodes_[i]] = i; - left_nodes_[left_nodes_[i]] = StochTree::Tree::kInvalidNodeId; - left_nodes_[right_nodes_[i]] = StochTree::Tree::kInvalidNodeId; - - // Add right node tracking information - node_begin_[right_nodes_[i]] = node_start_idx + num_true; - node_length_[right_nodes_[i]] = num_false; - parent_nodes_[right_nodes_[i]] = i; - right_nodes_[left_nodes_[i]] = StochTree::Tree::kInvalidNodeId; - right_nodes_[right_nodes_[i]] = StochTree::Tree::kInvalidNodeId; + continue; } + // Partition the node indices + auto node_begin = (indices_.begin() + node_start_idx); + auto node_end = (indices_.begin() + node_start_idx + num_node_elements); + auto right_node_begin = std::stable_partition(node_begin, node_end, [&](int row) { return split_rule.SplitTrue(covariates(row, split_index)); }); + + // Determine the number of true and false elements + num_true = std::distance(node_begin, right_node_begin); + num_false = num_node_elements - num_true; + + int left_id = left_nodes_[i]; + int right_id = right_nodes_[i]; + + // Add left node tracking information + node_begin_[left_id] = node_start_idx; + node_length_[left_id] = num_true; + parent_nodes_[left_id] = i; + + // Add right node tracking information + node_begin_[right_id] = node_start_idx + num_true; + node_length_[right_id] = num_false; + parent_nodes_[right_id] = i; + + // Descend into children (their bounds are now set) + traversal_queue.push_back(left_id); + traversal_queue.push_back(right_id); } } @@ -565,8 +570,7 @@ bool FeatureUnsortedPartition::IsValidNode(int node_id) { if (node_id >= num_nodes_ || node_id < 0) { return false; } - return !(std::find(deleted_nodes_.begin(), deleted_nodes_.end(), node_id) - != deleted_nodes_.end()); + return !(std::find(deleted_nodes_.begin(), deleted_nodes_.end(), node_id) != deleted_nodes_.end()); } bool FeatureUnsortedPartition::LeftNodeIsLeaf(int node_id) { diff --git a/src/prediction.cpp b/src/prediction.cpp new file mode 100644 index 00000000..1b1d5b4a --- /dev/null +++ b/src/prediction.cpp @@ -0,0 +1,746 @@ +/*! Copyright (c) 2026 by stochtree authors */ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "stochtree/log.h" + +namespace StochTree { + +void location_scale_adjust_predictions(std::vector& predictions, double location, double scale) { + for (double& pred : predictions) { + pred = pred * scale + location; + } +} + +void probability_transform_probit_predictions(std::vector& predictions) { + for (double& pred : predictions) { + pred = norm_cdf(pred); + } +} + +void probability_transform_binary_cloglog_predictions(std::vector& predictions) { + for (double& pred : predictions) { + pred = std::exp(-std::exp(pred)); + } +} + +void probability_transform_ordinal_cloglog_predictions(std::vector& predictions, std::vector& output, double* cutpoints, int num_obs, int num_classes, int num_samples) { + // Sequential ordinal cloglog model: P(Y=k) = prod_{j& predictions) { + for (double& pred : predictions) { + pred = pred >= 0.5 ? 1.0 : 0.0; + } +} + +/*! + * \brief Assumes that class probabilities are stored in a column-major format in `predictions`, with dimensions num_obs x num_classes x num_samples, and transforms these into class predictions by taking the class with the highest predicted probability for each observation and sample. + * The output is stored in `output`, which should be a pre-allocated vector of size num_obs x num_samples, also in column-major format (i.e., all predictions for the first sample are stored contiguously, followed by all predictions for the second sample, etc.). + * + * \param predictions Vector of class probabilities in column-major format with dimensions num_obs x num_classes x num_samples + * \param output Pre-allocated vector to store class predictions in column-major format with dimensions num_obs x num_samples + * \param num_obs Number of observations + * \param num_classes Number of classes + * \param num_samples Number of samples + */ +void class_transform_multiclass_outcome_predictions(std::vector& predictions, std::vector& output, int num_obs, int num_classes, int num_samples) { + for (int j = 0; j < num_samples; j++) { + for (int i = 0; i < num_obs; i++) { + int predicted_class = 0; + double max_prob = predictions[j * num_classes * num_obs + 0 * num_obs + i]; + for (int k = 1; k < num_classes; k++) { + double prob_k = predictions[j * num_classes * num_obs + k * num_obs + i]; + if (prob_k > max_prob) { + max_prob = prob_k; + predicted_class = k; + } + } + output[j * num_obs + i] = static_cast(predicted_class); + } + } +} + +/*! + * \brief Internal helper function to average over the columns of a column-major 3d array. Works similarly to `np.mean(..., axis=1)` in numpy. + * + * \param input Vector representation of a 2d array stored in column-major order + * \param output Empty vector allocated to store the 1d averaged array + * \param num_rows Number of rows in `input` + * \param num_cols Number of columns in `input` + */ + +void average_col_major_2d(std::vector& input, std::vector& output, int num_rows, int num_cols) { + for (int i = 0; i < num_rows; i++) { + double sum = 0; + for (int j = 0; j < num_cols; j++) { + sum += input[j * num_rows + i]; + } + output[i] = sum / num_cols; + } +} + +/*! + * \brief Internal helper function to average over a specific axis of a column-major 3d array. + * input_idx_scale1, input_idx_scale2, input_idx_scale3 refer to the scaling factors for the three dimensions of the input array, which are used to compute the correct index in the flattened input vector. + * output_idx_scale1 and output_idx_scale2 refer to the scaling factors for the two dimensions of the output array, which are used to compute the correct index in the flattened output vector. + * + * \param input Vector representation of a 3d array stored in column-major order + * \param output Empty vector allocated to store the 2d averaged array in column-major order + * \param dim1 Size of the first dimension of the loop (NOTE: ~not~ necessarily the first dimension of the array) + * \param dim2 Size of the second dimension of the loop (NOTE: ~not~ necessarily the second dimension of the array) + * \param dim3 Size of the third dimension of the loop (NOTE: ~not~ necessarily the third dimension of the array) + * \param input_idx_scale1 Scaling factor for the index corresponding to the first dimension of the loop, which ensures that the correct array index is computed when accessing data from `input` + * \param input_idx_scale2 Scaling factor for the index corresponding to the second dimension of the loop, which ensures that the correct array index is computed when accessing data from `input` + * \param input_idx_scale3 Scaling factor for the index corresponding to the third dimension of the loop, which ensures that the correct array index is computed when accessing data from `input` + * \param output_idx_scale1 Scaling factor for the index corresponding to the first dimension of the loop, which ensures that the correct array index is computed when storing data in `output` + * \param output_idx_scale2 Scaling factor for the index corresponding to the second dimension of the loop, which ensures that the correct array index is computed when storing data in `output` + */ +void _average_col_major_3d_helper(std::vector& input, std::vector& output, int dim1, int dim2, int dim3, + int input_idx_scale1, int input_idx_scale2, int input_idx_scale3, + int output_idx_scale1, int output_idx_scale2) { + for (int i = 0; i < dim1; i++) { + for (int j = 0; j < dim2; j++) { + double sum = 0; + for (int k = 0; k < dim3; k++) { + const int input_idx = k * input_idx_scale3 + j * input_idx_scale2 + i * input_idx_scale1; + sum += input[input_idx]; + } + const int output_idx = j * output_idx_scale2 + i * output_idx_scale1; + output[output_idx] = sum / dim3; + } + } +} + +/*! + * \brief Internal helper function to average over a specific axis of a column-major 3d array. Works similarly to `np.mean(..., axis=dim_average)` in numpy. The `dim_average` parameter specifies which dimension to average over. + * Assumes `input` has size dim1 x dim2 x dim3 and is stored in column-major order and that `output` has size corresponding to the two dimensions that are not being averaged over (i.e. if dim_average==0, output should have size dim2 x dim3; if dim_average==1, output should have size dim1 x dim3; if dim_average==2, output should have size dim1 x dim2). + * input_idx_scale1, input_idx_scale2, input_idx_scale3 refer to the scaling factors for the three dimensions of the input array, which are used to compute the correct index in the flattened input vector. + * + * \param input Vector representation of a 3d array stored in column-major order + * \param output Empty vector allocated to store the 2d averaged array in column-major order + * \param dim1 Size of the first dimension of the loop (NOTE: ~not~ necessarily the first dimension of the array) + * \param dim2 Size of the second dimension of the loop (NOTE: ~not~ necessarily the second dimension of the array) + * \param dim3 Size of the third dimension of the loop (NOTE: ~not~ necessarily the third dimension of the array) + * \param dim_average Dimension of the input array over which to average (must be 0, 1, or 2) + */ +void average_col_major_3d(std::vector& input, std::vector& output, int dim1, int dim2, int dim3, int dim_average) { + if (dim_average == 2) { + _average_col_major_3d_helper(input, output, dim1, dim2, dim3, 1, dim1, dim1 * dim2, 1, dim1); + } else if (dim_average == 1) { + _average_col_major_3d_helper(input, output, dim1, dim3, dim2, 1, dim1 * dim2, dim1, 1, dim1); + } else if (dim_average == 0) { + _average_col_major_3d_helper(input, output, dim2, dim3, dim1, dim1, dim1 * dim2, 1, 1, dim2); + } else { + Log::Fatal("dim_average must be in {0, 1, 2}."); + } +} + +/*! + * The return value, BARTPredictionResult, is a struct that contains many optional data fields + * stored as std::vectors that are left empty if a model term is not requested by the prediction call. + * + * In some cases, model terms need to be computed even if not directly requested. + * For example, the conditional outcome mean (y_hat) requires mean forest and any random effects predictions. + * In the case that a term is needed as an intermediate computation but not requested as an output, we + * compute it internally and not return it. + */ +BARTPredictionResult predict_bart_model(BARTData& data, BARTSamples& samples, BARTPredictionMetadata& metadata) { + // Initialize a prediction result object + BARTPredictionResult output{}; + + // Key input / output dimensions + const int num_samples = metadata.num_samples; + const int num_obs = metadata.num_obs; + + // Key model components + const bool has_mean_forest = samples.mean_forests != nullptr; + const bool has_variance_forest = samples.variance_forests != nullptr; + const bool has_rfx = metadata.has_rfx; + const bool rfx_custom = metadata.rfx_model_spec == BARTRFXModelSpec::Custom; + + // Input data / config checks + if (has_rfx) { + if (rfx_custom && data.rfx_basis_test == nullptr) { + Log::Fatal("Model includes random effects with custom basis, but no random effect basis was provided in the test data for prediction"); + } + } + if (metadata.pred_scale == PredScale::kClass && metadata.pred_type == PredType::kMean) { + Log::Fatal("Taking the posterior mean of integer-valued class predictions is not an informative quantity, so this combination of pred_scale and pred_type is not supported directly by stochtree's prediction capabilities. If you do wish to obtain a posterior mean of class label predictions, we recommend predicting the class label posterior and then taking the average across MCMC samples in the resulting array"); + } + + // Model output details: + // - num_samples_output refers to the posterior sample dimension, which is num_samples for posterior predictions and 1 for posterior mean transformations + // - each of the need_* fields are true if a term needs to be computed en route to the user's requested outputs + int num_samples_output = 1; + if (metadata.pred_type == PredType::kPosterior) { + num_samples_output = num_samples; + } + bool need_mean = has_mean_forest && (metadata.pred_terms.y_hat || metadata.pred_terms.mean_forest); + bool need_rfx = has_rfx && (metadata.pred_terms.y_hat || metadata.pred_terms.random_effects); + bool need_variance_forest = has_variance_forest && metadata.pred_terms.variance_forest; + + // Resize any output vectors to be returned to users + const bool probability_scale = metadata.pred_scale == PredScale::kProbability; + const bool class_scale = metadata.pred_scale == PredScale::kClass; + const bool ordinal_cloglog_prob_scale = probability_scale && metadata.link_function == LinkFunction::Cloglog && metadata.outcome_type == OutcomeType::Ordinal; + if (metadata.pred_terms.y_hat) output.y_hat.resize(num_obs * (ordinal_cloglog_prob_scale ? metadata.cloglog_num_classes : 1) * num_samples_output); + if (metadata.pred_terms.mean_forest) output.mean_forest_predictions.resize(num_obs * (ordinal_cloglog_prob_scale ? metadata.cloglog_num_classes : 1) * num_samples_output); + if (metadata.pred_terms.variance_forest) output.variance_forest_predictions.resize(num_obs * num_samples_output); + if (metadata.pred_terms.random_effects) output.rfx_predictions.resize(num_obs * num_samples_output); + + // Initialize temporary containers needed to compute the requested predictions + std::vector mean_forest; + std::vector rfx; + std::vector variance_forest; + std::vector y_hat; + if (need_mean) { + mean_forest.resize(num_obs * num_samples); + } + if (need_rfx) { + rfx.resize(num_obs * num_samples); + } + if (need_variance_forest) { + variance_forest.resize(num_obs * num_samples); + } + if (metadata.pred_terms.y_hat) { + y_hat.resize(num_obs * num_samples); + } + + // Construct ForestDataset -- use the "test" fields + ForestDataset forest_dataset{}; + forest_dataset.AddCovariates(data.X_test, data.n_test, data.p, /*row_major=*/false); + if (data.basis_test != nullptr) { + forest_dataset.AddBasis(data.basis_test, data.n_test, data.basis_dim, /*row_major=*/false); + } + + if (need_mean) { + // Predict from mean forest + mean_forest = samples.mean_forests->Predict(forest_dataset); + } + + // Compute overall random effects predictions + if (need_rfx) { + RandomEffectsDataset rfx_dataset; + rfx_dataset.AddGroupLabels(data.rfx_group_ids_test, num_obs); + if (data.rfx_basis_test != nullptr) { + rfx_dataset.AddBasis(data.rfx_basis_test, num_obs, data.rfx_basis_dim, /*row_major=*/false); + } else if (metadata.rfx_model_spec == BARTRFXModelSpec::InterceptOnly) { + std::vector rfx_basis(data.n_test, 1.0); + rfx_dataset.AddBasis(rfx_basis.data(), num_obs, 1, /*row_major=*/false); + } else { + Log::Fatal("BART model random effects term was not sampled with intercept_only or intercept_plus_treatment specification, but not random effect basis was provided for prediction"); + } + samples.rfx_container->Predict(rfx_dataset, *samples.rfx_label_mapper, rfx); + } + + if (need_variance_forest) { + variance_forest = samples.variance_forests->Predict(forest_dataset); + } + if (metadata.pred_terms.y_hat) { + // y_hat is default initialized to 0, so we can just add the mean forest and random effects predictions as needed + for (int i = 0; i < num_obs; i++) { + if (need_mean) { + for (int j = 0; j < num_samples; j++) { + y_hat[j * num_obs + i] += mean_forest[j * num_obs + i]; + } + } + if (need_rfx) { + for (int j = 0; j < num_samples; j++) { + y_hat[j * num_obs + i] += rfx[j * num_obs + i]; + } + } + } + } + + // Scale the outputs + if (metadata.pred_terms.mean_forest) { + location_scale_adjust_predictions(mean_forest, samples.y_bar, samples.y_std); + } + if (metadata.pred_terms.random_effects) { + location_scale_adjust_predictions(rfx, 0.0, samples.y_std); + } + if (metadata.pred_terms.y_hat) { + location_scale_adjust_predictions(y_hat, samples.y_bar, samples.y_std); + } + if (need_variance_forest) { + location_scale_adjust_predictions(variance_forest, 0.0, samples.y_std * samples.y_std); + } + + // Transform if necessary (e.g. for probit models) + if (metadata.link_function == LinkFunction::Probit) { + if (metadata.pred_terms.mean_forest && probability_scale) { + probability_transform_probit_predictions(mean_forest); + } + if (metadata.pred_terms.random_effects && probability_scale) { + probability_transform_probit_predictions(rfx); + } + if (metadata.pred_terms.y_hat && (probability_scale || class_scale)) { + probability_transform_probit_predictions(y_hat); + if (class_scale) { + class_transform_binary_outcome_predictions(y_hat); + } + } + } else if (metadata.link_function == LinkFunction::Cloglog) { + if (metadata.outcome_type == OutcomeType::Binary) { + if (metadata.pred_terms.mean_forest && probability_scale) { + probability_transform_binary_cloglog_predictions(mean_forest); + } + // NOTE: RFX not compatible with cloglog link, so we skip RFX transformation + if (metadata.pred_terms.y_hat && (probability_scale || class_scale)) { + probability_transform_binary_cloglog_predictions(y_hat); + if (class_scale) { + class_transform_binary_outcome_predictions(y_hat); + } + } + } else if (metadata.outcome_type == OutcomeType::Ordinal) { + if (metadata.pred_terms.mean_forest && probability_scale) { + std::vector mean_forest_prob(num_obs * num_samples * metadata.cloglog_num_classes); + probability_transform_ordinal_cloglog_predictions(mean_forest, mean_forest_prob, samples.cloglog_cutpoint_samples.data(), num_obs, metadata.cloglog_num_classes, num_samples); + mean_forest = std::move(mean_forest_prob); + } + if (metadata.pred_terms.y_hat && (probability_scale || class_scale)) { + std::vector y_hat_prob(num_obs * num_samples * metadata.cloglog_num_classes); + probability_transform_ordinal_cloglog_predictions(y_hat, y_hat_prob, samples.cloglog_cutpoint_samples.data(), num_obs, metadata.cloglog_num_classes, num_samples); + if (metadata.pred_scale == PredScale::kClass) { + class_transform_multiclass_outcome_predictions(y_hat_prob, y_hat, num_obs, metadata.cloglog_num_classes, num_samples); + } else { + y_hat = std::move(y_hat_prob); + } + } + } + } + + // Unpack into returned outputs, aggregating if necessary + if (metadata.pred_terms.mean_forest) { + if (metadata.pred_type == PredType::kMean) { + if (metadata.pred_scale == PredScale::kProbability && metadata.outcome_type == OutcomeType::Ordinal && metadata.link_function == LinkFunction::Cloglog) { + average_col_major_3d(mean_forest, output.mean_forest_predictions, /*dim1=*/num_obs, /*dim2=*/metadata.cloglog_num_classes, /*dim3=*/num_samples, /*dim_average=*/2); + } else { + average_col_major_2d(mean_forest, output.mean_forest_predictions, /*num_rows=*/num_obs, /*num_cols=*/num_samples); + } + } else { + output.mean_forest_predictions = std::move(mean_forest); + } + } + if (need_variance_forest) { + if (metadata.pred_type == PredType::kMean) { + // NOTE: variance forest not compatible with ordinal cloglog model so we don't need to worry about 3d averaging here + average_col_major_2d(variance_forest, output.variance_forest_predictions, /*num_rows=*/num_obs, /*num_cols=*/num_samples); + } else { + output.variance_forest_predictions = std::move(variance_forest); + } + } + if (metadata.pred_terms.random_effects) { + if (metadata.pred_type == PredType::kMean) { + // NOTE: random effects not compatible with ordinal cloglog model so we don't need to worry about 3d averaging here + average_col_major_2d(rfx, output.rfx_predictions, /*num_rows=*/num_obs, /*num_cols=*/num_samples); + } else { + output.rfx_predictions = std::move(rfx); + } + } + if (metadata.pred_terms.y_hat) { + if (metadata.pred_type == PredType::kMean) { + if (metadata.pred_scale == PredScale::kProbability && metadata.outcome_type == OutcomeType::Ordinal && metadata.link_function == LinkFunction::Cloglog) { + average_col_major_3d(y_hat, output.y_hat, /*dim1=*/num_obs, /*dim2=*/metadata.cloglog_num_classes, /*dim3=*/num_samples, /*dim_average=*/2); + } else { + average_col_major_2d(y_hat, output.y_hat, /*num_rows=*/num_obs, /*num_cols=*/num_samples); + } + } else { + output.y_hat = std::move(y_hat); + } + } + + return output; +} + +/*! + * The return value, BCFPRedictionResult, is a struct that contains many optional data fields + * stored as std::vectors that are left empty if a model term is not requested by the prediction call. + * + * In some cases, model terms need to be computed even if not directly requested. + * For example, the conditional outcome mean (y_hat) requires mu_x, tau_x and any random effects predictions. + * In the case that a term is needed as an intermediate computation but not requested as an output, we + * compute it internally and not return it. + */ +BCFPredictionResult predict_bcf_model(BCFData& data, BCFSamples& samples, BCFPredictionMetadata& metadata) { + // Initialize a prediction result object + BCFPredictionResult output{}; + + // Key input / output dimensions + const int num_samples = metadata.num_samples; + const int num_obs = metadata.num_obs; + const int num_treatment = metadata.treatment_dim; + + // Key model components + const bool has_mu_forest = samples.mu_forests != nullptr; + const bool has_tau_forest = samples.tau_forests != nullptr; + const bool has_variance_forest = samples.variance_forests != nullptr; + const bool has_rfx = samples.rfx_container != nullptr; + const bool rfx_custom = metadata.rfx_model_spec == BCFRFXModelSpec::Custom; + const bool rfx_intercept = metadata.rfx_model_spec == BCFRFXModelSpec::InterceptOnly || metadata.rfx_model_spec == BCFRFXModelSpec::InterceptPlusTreatment; + const bool rfx_treatment = metadata.rfx_model_spec == BCFRFXModelSpec::InterceptPlusTreatment; + + // Model output details: + // - num_samples_output refers to the posterior sample dimension, which is num_samples for posterior predictions and 1 for posterior mean transformations + // - each of the need_* fields are true if a term needs to be computed en route to the user's requested outputs + int num_samples_output = 1; + if (metadata.pred_type == PredType::kPosterior) { + num_samples_output = num_samples; + } + bool need_tau_interm = metadata.pred_terms.y_hat || metadata.pred_terms.tau_x || metadata.pred_terms.cate; + bool need_mu = metadata.pred_terms.y_hat || metadata.pred_terms.mu_x || metadata.pred_terms.prognostic_function || (metadata.adaptive_coding && need_tau_interm); + bool need_tau = need_tau_interm || (metadata.adaptive_coding && need_mu); + bool need_rfx = has_rfx && (metadata.pred_terms.y_hat || metadata.pred_terms.random_effects); + bool need_rfx_intercept = has_rfx && rfx_intercept && metadata.pred_terms.prognostic_function; + bool need_rfx_treatment = has_rfx && rfx_treatment && metadata.pred_terms.cate; + bool need_variance_forest = has_variance_forest && metadata.pred_terms.conditional_variance; + + // Resize any output vectors to be returned to users + if (metadata.pred_terms.y_hat) output.y_hat.resize(num_obs * num_samples_output); + if (metadata.pred_terms.mu_x) output.mu_x.resize(num_obs * num_samples_output); + if (metadata.pred_terms.tau_x) output.tau_x.resize(num_obs * num_treatment * num_samples_output); + if (metadata.pred_terms.prognostic_function) output.prognostic_function.resize(num_obs * num_samples_output); + if (metadata.pred_terms.cate) output.cate.resize(num_obs * num_treatment * num_samples_output); + if (metadata.pred_terms.conditional_variance) output.conditional_variance.resize(num_obs * num_samples_output); + if (metadata.pred_terms.random_effects) output.random_effects.resize(num_obs * num_samples_output); + + // Initialize temporary containers needed to compute the requested predictions + std::vector mu_x; + std::vector prognostic_function; + std::vector tau_x; + std::vector cate; + std::vector rfx_mu; + std::vector rfx_tau; + std::vector rfx; + std::vector variance_forest; + std::vector y_hat; + if (need_mu) { + mu_x.resize(num_obs * num_samples); + } + if (metadata.pred_terms.prognostic_function) { + prognostic_function.resize(num_obs * num_samples); + } + if (need_tau) { + tau_x.resize(num_obs * num_treatment * num_samples); + } + if (metadata.pred_terms.cate) { + cate.resize(num_obs * num_treatment * num_samples); + } + if (need_rfx_intercept) { + rfx_mu.resize(num_obs * num_samples); + } + if (need_rfx_treatment) { + rfx_tau.resize(num_obs * num_treatment * num_samples); + } + if (need_rfx) { + rfx.resize(num_obs * num_samples); + } + if (need_variance_forest) { + variance_forest.resize(num_obs * num_samples); + } + if (metadata.pred_terms.y_hat) { + y_hat.resize(num_obs * num_samples); + } + + // Construct ForestDataset -- use the "test" fields + ForestDataset forest_dataset{}; + forest_dataset.AddCovariates(data.X_test, data.n_test, data.p, /*row_major=*/false); + // NOTE: not adding treatment as basis to forest_dataset since we always predict the raw treatment effect forest values and multiply by + // either the raw or recoded treatment (if adaptive coding) + if (data.obs_weights_test != nullptr) { + forest_dataset.AddVarianceWeights(data.obs_weights_test, data.n_test); + } + + if (need_mu) { + // Predict from mu forest + mu_x = samples.mu_forests->Predict(forest_dataset); + } + + if (need_tau) { + // Predict from tau forest. We use PredictRaw for the tau forest because we + // don't want to pre-multiply by the treatment / basis at this stage -- we want to be + // able to return the treatment effect itself, not the treatment effect times Z (or recoded Z) + tau_x = samples.tau_forests->PredictRaw(forest_dataset, /*row_major=*/false); + // Add tau_0 to the treatment effect function predictions if it was sampled. + // tau_0_samples layout: col-major (treatment dim k, sample j) -> j * treatment_dim + k. + // For treatment_dim==1 this collapses to samples.tau_0_samples[j]. + // NOTE: tau_0_samples is stored in STANDARDIZED scale (matching the sampler and the R main-branch + // convention), same as tau_x from PredictRaw, so we add it directly here; the y_std scale step + // applied later (location_scale_adjust_predictions) then scales the full CATE tau_0 + tau(x). + if (metadata.sample_tau_0) { + for (int j = 0; j < num_samples; j++) { + for (int k = 0; k < num_treatment; k++) { + for (int i = 0; i < num_obs; i++) { + const int idx = j * num_obs * num_treatment + k * num_obs + i; + tau_x[idx] += samples.tau_0_samples[j * num_treatment + k]; + } + } + } + } + // Handle adaptive coding correctly: + // When treatment is b_0 (1-Z) + b_1 Z, the conditional mean model: + // mu(x) + [tau_0 + tau(x)] * (b_0 * (1-Z) + b_1 * Z) + // turns into + // [mu(x) + b_0 * (tau_0 + tau(x))] + (tau_0 + tau(x)) * (b_1 - b_0) * Z + // So the treatment effect function that gets multiplied by Z is actually (b_1 - b_0) * (tau_0 + tau(x)) + // and the prognostic function has an added contribution of b_0 * (tau_0 + tau(x)) + // NOTE: adaptive coding is only supported for a univariate binary treatment, so we construct our indices as if tau_x is 2d because whenever adaptive coding is true, treatment_dim must be 1 and the array is effectively 2d. + if (metadata.adaptive_coding) { + for (int i = 0; i < num_samples; i++) { + double b_0 = samples.b0_samples[i]; + double b_1 = samples.b1_samples[i]; + for (int j = 0; j < num_obs; j++) { + const int idx = i * num_obs + j; + // Add b_0 * (tau_0 + tau(x)) to the prognostic function predictions + mu_x[idx] += b_0 * tau_x[idx]; + // Scale tau_x by (b_1 - b_0) + tau_x[idx] *= (b_1 - b_0); + } + } + } + } + + // Add random effects contribution to prognostic function if needed + if (need_rfx_intercept) { + // Extract just the random intercept effects from the RFX model + int group_ind; + const int k = 0; // We only want the first column from the RFX parameters + std::vector& beta = samples.rfx_container->GetBeta(); + int num_components = samples.rfx_container->NumComponents(); + int num_groups = samples.rfx_container->NumGroups(); + for (int i = 0; i < num_obs; i++) { + group_ind = samples.rfx_label_mapper->CategoryNumber(data.rfx_group_ids_test[i]); + for (int j = 0; j < num_samples; j++) { + const int idx = j * num_obs + i; + rfx_mu[idx] = beta.at(j * num_groups * num_components + group_ind * num_components + k); + } + } + } + + // Add random effects contribution to CATE function if needed + if (need_rfx_treatment) { + // Extract just the random treatment effects from the RFX model + int group_ind; + std::vector& beta = samples.rfx_container->GetBeta(); + int num_components = samples.rfx_container->NumComponents(); + int num_groups = samples.rfx_container->NumGroups(); + for (int i = 0; i < num_obs; i++) { + group_ind = samples.rfx_label_mapper->CategoryNumber(data.rfx_group_ids_test[i]); + for (int j = 0; j < num_samples; j++) { + // In the "intercept_plus_treatment" RFX specification, the random intercept is in column 0 and the random treatment effect(s) start from column 1, + // so we loop from k=1 to num_components to extract the treatment effect contribution(s) + for (int k = 1; k < num_components; k++) { + // NOTE, the "treatment" index is k - 1 since k starts at 1 in this loop + const int idx = j * num_obs * num_treatment + (k - 1) * num_obs + i; + rfx_tau[idx] = beta.at(j * num_groups * num_components + group_ind * num_components + k); + } + } + } + } + + // Compute overall random effects predictions + if (need_rfx) { + RandomEffectsDataset rfx_dataset; + rfx_dataset.AddGroupLabels(data.rfx_group_ids_test, num_obs); + if (data.rfx_basis_test != nullptr) { + rfx_dataset.AddBasis(data.rfx_basis_test, num_obs, data.rfx_basis_dim, /*row_major=*/false); + } else if (metadata.rfx_model_spec == BCFRFXModelSpec::InterceptOnly) { + std::vector rfx_basis(data.n_test, 1.0); + rfx_dataset.AddBasis(rfx_basis.data(), num_obs, 1, /*row_major=*/false); + } else if (metadata.rfx_model_spec == BCFRFXModelSpec::InterceptPlusTreatment) { + // Column-major rfx basis + std::vector rfx_basis(data.n_test * (1 + num_treatment)); + for (int i = 0; i < num_obs; i++) { + rfx_basis[i] = 1.0; + } + for (int j = 0; j < num_treatment; j++) { + for (int i = 0; i < num_obs; i++) { + rfx_basis[(j + 1) * num_obs + i] = data.treatment_test[j * num_obs + i]; + } + } + rfx_dataset.AddBasis(rfx_basis.data(), num_obs, 1 + num_treatment, /*row_major=*/false); + } else { + Log::Fatal("BCF model random effects term was not sampled with intercept_only or intercept_plus_treatment specification, but not random effect basis was provided for prediction"); + } + samples.rfx_container->Predict(rfx_dataset, *samples.rfx_label_mapper, rfx); + } + + // Unpack into returned outputs + if (metadata.pred_terms.prognostic_function) { + for (int i = 0; i < mu_x.size(); i++) { + prognostic_function[i] = mu_x[i]; + if (need_rfx_intercept) { + prognostic_function[i] += rfx_mu[i]; + } + } + } + if (metadata.pred_terms.cate) { + for (int i = 0; i < tau_x.size(); i++) { + cate[i] = tau_x[i]; + if (need_rfx_treatment) { + cate[i] += rfx_tau[i]; + } + } + } + if (need_variance_forest) { + variance_forest = samples.variance_forests->Predict(forest_dataset); + } + if (metadata.pred_terms.y_hat) { + for (int i = 0; i < num_obs; i++) { + for (int j = 0; j < num_samples; j++) { + y_hat[j * num_obs + i] = mu_x[j * num_obs + i]; + for (int k = 0; k < num_treatment; k++) { + y_hat[j * num_obs + i] += tau_x[j * num_obs * num_treatment + k * num_obs + i] * data.treatment_test[k * num_obs + i]; + } + } + if (need_rfx) { + for (int j = 0; j < num_samples; j++) { + y_hat[j * num_obs + i] += rfx[j * num_obs + i]; + } + } + } + } + + // Scale the outputs + if (metadata.pred_terms.mu_x) { + location_scale_adjust_predictions(mu_x, samples.y_bar, samples.y_std); + } + if (metadata.pred_terms.prognostic_function) { + location_scale_adjust_predictions(prognostic_function, samples.y_bar, samples.y_std); + } + if (metadata.pred_terms.tau_x) { + location_scale_adjust_predictions(tau_x, 0.0, samples.y_std); + } + if (metadata.pred_terms.cate) { + location_scale_adjust_predictions(cate, 0.0, samples.y_std); + } + if (metadata.pred_terms.random_effects) { + location_scale_adjust_predictions(rfx, 0.0, samples.y_std); + } + if (metadata.pred_terms.y_hat) { + location_scale_adjust_predictions(y_hat, samples.y_bar, samples.y_std); + } + if (need_variance_forest) { + location_scale_adjust_predictions(variance_forest, 0.0, samples.y_std * samples.y_std); + } + + // Transform if necessary (e.g. for probit models) + // NOTE: if we support cloglog or ordinal probit BCF in the future (likely), + // we must add more link function guards to this block of code + const bool probability_scale = metadata.pred_scale == PredScale::kProbability; + const bool class_scale = metadata.pred_scale == PredScale::kClass; + if (metadata.pred_terms.mu_x && probability_scale) { + probability_transform_probit_predictions(mu_x); + } + if (metadata.pred_terms.prognostic_function && probability_scale) { + probability_transform_probit_predictions(prognostic_function); + } + if (metadata.pred_terms.tau_x && probability_scale) { + probability_transform_probit_predictions(tau_x); + } + if (metadata.pred_terms.cate && probability_scale) { + probability_transform_probit_predictions(cate); + } + if (metadata.pred_terms.y_hat && (probability_scale || class_scale)) { + probability_transform_probit_predictions(y_hat); + if (class_scale) { + class_transform_binary_outcome_predictions(y_hat); + } + } + + // Unpack into returned outputs, aggregating if necessary + if (metadata.pred_terms.mu_x) { + if (metadata.pred_type == PredType::kMean) { + average_col_major_2d(mu_x, output.mu_x, num_obs, num_samples); + } else { + output.mu_x = std::move(mu_x); + } + } + if (metadata.pred_terms.tau_x) { + if (metadata.pred_type == PredType::kMean) { + if (num_treatment == 1) { + // If only one treatment, tau_x is num_obs by num_samples, so average across samples in columns + average_col_major_2d(tau_x, output.tau_x, /*num_rows=*/num_obs, /*num_cols=*/num_samples); + } else { + // If multiple treatments, tau_x is num_obs by num_treatment by num_samples in column-major order, so average across samples while keeping treatment dimension separate + average_col_major_3d(tau_x, output.tau_x, /*dim1=*/num_obs, /*dim2=*/num_treatment, /*dim3=*/num_samples, /*dim_average=*/2); + } + } else { + output.tau_x = std::move(tau_x); + } + } + if (metadata.pred_terms.prognostic_function) { + if (metadata.pred_type == PredType::kMean) { + average_col_major_2d(prognostic_function, output.prognostic_function, num_obs, num_samples); + } else { + output.prognostic_function = std::move(prognostic_function); + } + } + if (metadata.pred_terms.cate) { + if (metadata.pred_type == PredType::kMean) { + if (num_treatment == 1) { + // If only one treatment, cate is num_obs by num_samples, so average across samples in columns + average_col_major_2d(cate, output.cate, /*num_rows=*/num_obs, /*num_cols=*/num_samples); + } else { + // If multiple treatments, cate is num_obs by num_treatment by num_samples in column-major order, so average across samples while keeping treatment dimension separate + average_col_major_3d(cate, output.cate, /*dim1=*/num_obs, /*dim2=*/num_treatment, /*dim3=*/num_samples, /*dim_average=*/2); + } + } else { + output.cate = std::move(cate); + } + } + if (need_variance_forest) { + if (metadata.pred_type == PredType::kMean) { + average_col_major_2d(variance_forest, output.conditional_variance, num_obs, num_samples); + } else { + output.conditional_variance = std::move(variance_forest); + } + } + if (metadata.pred_terms.random_effects) { + if (metadata.pred_type == PredType::kMean) { + average_col_major_2d(rfx, output.random_effects, num_obs, num_samples); + } else { + output.random_effects = std::move(rfx); + } + } + if (metadata.pred_terms.y_hat) { + if (metadata.pred_type == PredType::kMean) { + average_col_major_2d(y_hat, output.y_hat, num_obs, num_samples); + } else { + output.y_hat = std::move(y_hat); + } + } + + return output; +} + +} // namespace StochTree \ No newline at end of file diff --git a/src/py_stochtree.cpp b/src/py_stochtree.cpp index a477b4ff..4ad5f43d 100644 --- a/src/py_stochtree.cpp +++ b/src/py_stochtree.cpp @@ -1,12 +1,18 @@ #include #include +#include #include #include +#include +#include +#include +#include #include #include #include #include #include +#include #include #include #include @@ -34,7 +40,7 @@ class ForestDatasetCpp { void AddCovariates(py::array_t covariate_matrix, data_size_t num_row, int num_col, bool row_major) { // Extract pointer to contiguous block of memory double* data_ptr = static_cast(covariate_matrix.mutable_data()); - + // Load covariates dataset_->AddCovariates(data_ptr, num_row, num_col, row_major); } @@ -50,7 +56,7 @@ class ForestDatasetCpp { void UpdateBasis(py::array_t basis_matrix, data_size_t num_row, int num_col, bool row_major) { // Extract pointer to contiguous block of memory double* data_ptr = static_cast(basis_matrix.mutable_data()); - + // Load covariates dataset_->UpdateBasis(data_ptr, num_row, num_col, row_major); } @@ -66,7 +72,7 @@ class ForestDatasetCpp { void UpdateVarianceWeights(py::array_t weight_vector, data_size_t num_row, bool exponentiate) { // Extract pointer to contiguous block of memory double* data_ptr = static_cast(weight_vector.mutable_data()); - + // Load covariates dataset_->UpdateVarWeights(data_ptr, num_row, exponentiate); } @@ -79,7 +85,7 @@ class ForestDatasetCpp { auto accessor = result.mutable_unchecked<2>(); for (size_t i = 0; i < n; i++) { for (int j = 0; j < num_covariates; j++) { - accessor(i,j) = dataset_->CovariateValue(i,j); + accessor(i, j) = dataset_->CovariateValue(i, j); } } @@ -94,7 +100,7 @@ class ForestDatasetCpp { auto accessor = result.mutable_unchecked<2>(); for (size_t i = 0; i < n; i++) { for (int j = 0; j < num_basis; j++) { - accessor(i,j) = dataset_->BasisValue(i,j); + accessor(i, j) = dataset_->BasisValue(i, j); } } @@ -169,7 +175,7 @@ class ResidualCpp { ResidualCpp(py::array_t residual_array, data_size_t num_row) { // Extract pointer to contiguous block of memory double* data_ptr = static_cast(residual_array.mutable_data()); - + // Initialize pointer to C++ ColumnVector class residual_ = std::make_unique(data_ptr, num_row); } @@ -182,13 +188,13 @@ class ResidualCpp { py::array_t GetResidualArray() { // Obtain a reference to the underlying Eigen::VectorXd Eigen::VectorXd& resid_vector = residual_->GetData(); - + // Initialize n x 1 numpy array to store the residual data_size_t n = residual_->NumRows(); auto result = py::array_t(py::detail::any_container({n, 1})); auto accessor = result.mutable_unchecked<2>(); for (size_t i = 0; i < n; i++) { - accessor(i,0) = resid_vector(i); + accessor(i, 0) = resid_vector(i); } return result; @@ -253,8 +259,17 @@ class ForestContainerCpp { is_leaf_constant_ = is_leaf_constant; is_exponentiated_ = is_exponentiated; } + ForestContainerCpp(std::unique_ptr forest_samples, int num_trees, int output_dimension = 1, bool is_leaf_constant = true, bool is_exponentiated = false) { + forest_samples_ = std::move(forest_samples); + num_trees_ = num_trees; + output_dimension_ = output_dimension; + is_leaf_constant_ = is_leaf_constant; + is_exponentiated_ = is_exponentiated; + } ~ForestContainerCpp() {} + StochTree::ForestContainer* GetPtr() { return forest_samples_.get(); } + void CombineForests(py::array_t forest_inds) { int num_forests = forest_inds.size(); for (int j = 1; j < num_forests; j++) { @@ -311,7 +326,7 @@ class ForestContainerCpp { for (size_t i = 0; i < n; i++) { for (int j = 0; j < num_samples; j++) { // NOTE: converting from "column-major" to "row-major" here - accessor(i,j) = output_raw[j*n + i]; + accessor(i, j) = output_raw[j * n + i]; // ptr[i*num_samples + j] = output_raw[j*n + i]; } } @@ -335,7 +350,7 @@ class ForestContainerCpp { for (size_t i = 0; i < n; i++) { for (int j = 0; j < output_dim; j++) { for (int k = 0; k < num_samples; k++) { - accessor(i,k,j) = output_raw[k*(output_dim*n) + i*output_dim + j]; + accessor(i, k, j) = output_raw[k * (output_dim * n) + i * output_dim + j]; // ptr[i*(output_dim*num_samples) + j*output_dim + k] = output_raw[k*(output_dim*n) + i*output_dim + j]; } } @@ -359,7 +374,7 @@ class ForestContainerCpp { // double *ptr = static_cast(buf.ptr); for (size_t i = 0; i < n; i++) { for (int j = 0; j < output_dim; j++) { - accessor(i,j) = output_raw[i*output_dim + j]; + accessor(i, j) = output_raw[i * output_dim + j]; // ptr[i*output_dim + j] = output_raw[i*output_dim + j]; } } @@ -382,7 +397,7 @@ class ForestContainerCpp { // double *ptr = static_cast(buf.ptr); for (size_t i = 0; i < n; i++) { for (int j = 0; j < output_dim; j++) { - accessor(i,j) = output_raw[i*output_dim + j]; + accessor(i, j) = output_raw[i * output_dim + j]; // ptr[i*output_dim + j] = output_raw[i*output_dim + j]; } } @@ -397,7 +412,7 @@ class ForestContainerCpp { void SetRootVector(int forest_num, py::array_t& leaf_vector, int leaf_size) { std::vector leaf_vector_converted(leaf_size); for (int i = 0; i < leaf_size; i++) { - leaf_vector_converted[i] = leaf_vector.at(i); + leaf_vector_converted[i] = leaf_vector.at(i); } forest_samples_->InitializeRoot(leaf_vector_converted); } @@ -453,8 +468,8 @@ class ForestContainerCpp { StochTree::TreeEnsemble* ensemble = forest_samples_->GetEnsemble(num_samples); int num_trees = ensemble->NumTrees(); for (int i = 0; i < num_trees; i++) { - StochTree::Tree* tree = ensemble->GetTree(i); - tree->SetLeaf(0, leaf_value); + StochTree::Tree* tree = ensemble->GetTree(i); + tree->SetLeaf(0, leaf_value); } } @@ -469,8 +484,8 @@ class ForestContainerCpp { StochTree::TreeEnsemble* ensemble = forest_samples_->GetEnsemble(num_samples); int num_trees = ensemble->NumTrees(); for (int i = 0; i < num_trees; i++) { - StochTree::Tree* tree = ensemble->GetTree(i); - tree->SetLeafVector(0, leaf_vector_cast); + StochTree::Tree* tree = ensemble->GetTree(i); + tree->SetLeafVector(0, leaf_vector_cast); } } @@ -495,8 +510,8 @@ class ForestContainerCpp { tree->ExpandNode(leaf_num, feature_num, split_threshold, left_leaf_vector_cast, right_leaf_vector_cast); } - void AddNumericSplitValue(int forest_num, int tree_num, int leaf_num, int feature_num, - double split_threshold, double left_leaf_value, double right_leaf_value) { + void AddNumericSplitValue(int forest_num, int tree_num, int leaf_num, int feature_num, + double split_threshold, double left_leaf_value, double right_leaf_value) { if (forest_samples_->OutputDimension() != 1) { StochTree::Log::Fatal("left_leaf_value must match forest leaf dimension"); } @@ -534,9 +549,9 @@ class ForestContainerCpp { StochTree::Tree* tree = ensemble->GetTree(tree_num); std::vector split_nodes = tree->GetInternalNodes(); for (int i = 0; i < split_nodes.size(); i++) { - auto node_id = split_nodes.at(i); - auto split_feature = tree->SplitIndex(node_id); - accessor(split_feature)++; + auto node_id = split_nodes.at(i); + auto split_feature = tree->SplitIndex(node_id); + accessor(split_feature)++; } return result; } @@ -587,12 +602,12 @@ class ForestContainerCpp { py::array_t GetGranularSplitCounts(int num_features) { int num_samples = forest_samples_->NumSamples(); int num_trees = forest_samples_->NumTrees(); - auto result = py::array_t(py::detail::any_container({num_samples,num_trees,num_features})); + auto result = py::array_t(py::detail::any_container({num_samples, num_trees, num_features})); auto accessor = result.mutable_unchecked<3>(); for (int i = 0; i < num_samples; i++) { for (int j = 0; j < num_trees; j++) { for (int k = 0; k < num_features; k++) { - accessor(i,j,k) = 0; + accessor(i, j, k) = 0; } } } @@ -604,7 +619,7 @@ class ForestContainerCpp { for (int k = 0; k < split_nodes.size(); k++) { auto node_id = split_nodes.at(k); auto split_feature = tree->SplitIndex(node_id); - accessor(i,j,split_feature)++; + accessor(i, j, split_feature)++; } } } @@ -760,7 +775,7 @@ class ForestCpp { } ~ForestCpp() {} - StochTree::TreeEnsemble* GetForestPtr() {return forest_.get();} + StochTree::TreeEnsemble* GetForestPtr() { return forest_.get(); } void MergeForest(ForestCpp& outbound_forest) { forest_->MergeForest(*outbound_forest.GetForestPtr()); @@ -813,7 +828,7 @@ class ForestCpp { auto result = py::array_t(py::detail::any_container({n})); auto accessor = result.mutable_unchecked<1>(); for (size_t i = 0; i < n; i++) { - accessor(i) = output_raw[i]; + accessor(i) = output_raw[i]; } return result; @@ -831,7 +846,7 @@ class ForestCpp { auto accessor = result.mutable_unchecked<2>(); for (size_t i = 0; i < n; i++) { for (int j = 0; j < output_dim; j++) { - accessor(i,j) = output_raw[i*output_dim + j]; + accessor(i, j) = output_raw[i * output_dim + j]; } } @@ -845,7 +860,7 @@ class ForestCpp { void SetRootVector(py::array_t& leaf_vector, int leaf_size) { std::vector leaf_vector_converted(leaf_size); for (int i = 0; i < leaf_size; i++) { - leaf_vector_converted[i] = leaf_vector.at(i); + leaf_vector_converted[i] = leaf_vector.at(i); } forest_->SetLeafVector(leaf_vector_converted); } @@ -856,7 +871,7 @@ class ForestCpp { return forest_.get(); } - void AddNumericSplitValue(int tree_num, int leaf_num, int feature_num, double split_threshold, + void AddNumericSplitValue(int tree_num, int leaf_num, int feature_num, double split_threshold, double left_leaf_value, double right_leaf_value) { if (forest_->OutputDimension() != 1) { StochTree::Log::Fatal("left_leaf_value must match forest leaf dimension"); @@ -913,9 +928,9 @@ class ForestCpp { StochTree::Tree* tree = forest_->GetTree(tree_num); std::vector split_nodes = tree->GetInternalNodes(); for (int i = 0; i < split_nodes.size(); i++) { - auto node_id = split_nodes.at(i); - auto split_feature = tree->SplitIndex(node_id); - accessor(split_feature)++; + auto node_id = split_nodes.at(i); + auto split_feature = tree->SplitIndex(node_id); + accessor(split_feature)++; } return result; } @@ -941,11 +956,11 @@ class ForestCpp { py::array_t GetGranularSplitCounts(int num_features) { int num_trees = forest_->NumTrees(); - auto result = py::array_t(py::detail::any_container({num_trees,num_features})); + auto result = py::array_t(py::detail::any_container({num_trees, num_features})); auto accessor = result.mutable_unchecked<2>(); for (int i = 0; i < num_trees; i++) { for (int j = 0; j < num_features; j++) { - accessor(i,j) = 0; + accessor(i, j) = 0; } } for (int i = 0; i < num_trees; i++) { @@ -954,7 +969,7 @@ class ForestCpp { for (int j = 0; j < split_nodes.size(); j++) { auto node_id = split_nodes.at(i); auto split_feature = tree->SplitIndex(node_id); - accessor(i,split_feature)++; + accessor(i, split_feature)++; } } return result; @@ -1086,9 +1101,9 @@ class ForestSamplerCpp { // Convert vector of integers to std::vector of enum FeatureType std::vector feature_types_(feature_types.size()); for (int i = 0; i < feature_types.size(); i++) { - feature_types_[i] = static_cast(feature_types.at(i)); + feature_types_[i] = static_cast(feature_types.at(i)); } - + // Initialize pointer to C++ ForestTracker and TreePrior classes StochTree::ForestDataset* dataset_ptr = dataset.GetDataset(); tracker_ = std::make_unique(dataset_ptr->GetCovariates(), feature_types_, num_trees, num_obs); @@ -1096,14 +1111,14 @@ class ForestSamplerCpp { } ~ForestSamplerCpp() {} - StochTree::ForestTracker* GetTracker() {return tracker_.get();} + StochTree::ForestTracker* GetTracker() { return tracker_.get(); } void ReconstituteTrackerFromForest(ForestCpp& forest, ForestDatasetCpp& dataset, ResidualCpp& residual, bool is_mean_model) { // Extract raw pointer to the forest and dataset StochTree::TreeEnsemble* forest_ptr = forest.GetEnsemble(); StochTree::ForestDataset* data_ptr = dataset.GetDataset(); StochTree::ColumnVector* residual_ptr = residual.GetData(); - + // Reset forest tracker using the forest held at index forest_num tracker_->ReconstituteFromForest(*forest_ptr, *data_ptr, *residual_ptr, is_mean_model); } @@ -1115,56 +1130,61 @@ class ForestSamplerCpp { // Refactoring completely out of the Python interface. // Intention to refactor out of the C++ interface in the future. bool pre_initialized = true; - + // Unpack feature types std::vector feature_types_(feature_types.size()); for (int i = 0; i < feature_types.size(); i++) { feature_types_[i] = static_cast(feature_types.at(i)); } - + // Unpack sweep indices std::vector sweep_update_indices_; if (sweep_update_indices.size() > 0) { - sweep_update_indices_.resize(sweep_update_indices.size()); - for (int i = 0; i < sweep_update_indices.size(); i++) { - sweep_update_indices_[i] = sweep_update_indices.at(i); - } + sweep_update_indices_.resize(sweep_update_indices.size()); + for (int i = 0; i < sweep_update_indices.size(); i++) { + sweep_update_indices_[i] = sweep_update_indices.at(i); + } } // Convert leaf model type to enum StochTree::ModelType model_type; - if (leaf_model_int == 0) model_type = StochTree::ModelType::kConstantLeafGaussian; - else if (leaf_model_int == 1) model_type = StochTree::ModelType::kUnivariateRegressionLeafGaussian; - else if (leaf_model_int == 2) model_type = StochTree::ModelType::kMultivariateRegressionLeafGaussian; - else if (leaf_model_int == 3) model_type = StochTree::ModelType::kLogLinearVariance; - else if (leaf_model_int == 4) model_type = StochTree::ModelType::kCloglogOrdinal; + if (leaf_model_int == 0) + model_type = StochTree::ModelType::kConstantLeafGaussian; + else if (leaf_model_int == 1) + model_type = StochTree::ModelType::kUnivariateRegressionLeafGaussian; + else if (leaf_model_int == 2) + model_type = StochTree::ModelType::kMultivariateRegressionLeafGaussian; + else if (leaf_model_int == 3) + model_type = StochTree::ModelType::kLogLinearVariance; + else if (leaf_model_int == 4) + model_type = StochTree::ModelType::kCloglogOrdinal; // Unpack leaf model parameters double leaf_scale; Eigen::MatrixXd leaf_scale_matrix; if ((model_type == StochTree::ModelType::kConstantLeafGaussian) || (model_type == StochTree::ModelType::kUnivariateRegressionLeafGaussian)) { - leaf_scale = leaf_model_scale_input.at(0,0); + leaf_scale = leaf_model_scale_input.at(0, 0); } else if (model_type == StochTree::ModelType::kMultivariateRegressionLeafGaussian) { - int num_row = leaf_model_scale_input.shape(0); - int num_col = leaf_model_scale_input.shape(1); - leaf_scale_matrix.resize(num_row, num_col); - for (int i = 0; i < num_row; i++) { - for (int j = 0; j < num_col; j++) { - leaf_scale_matrix(i,j) = leaf_model_scale_input.at(i,j); - } + int num_row = leaf_model_scale_input.shape(0); + int num_col = leaf_model_scale_input.shape(1); + leaf_scale_matrix.resize(num_row, num_col); + for (int i = 0; i < num_row; i++) { + for (int j = 0; j < num_col; j++) { + leaf_scale_matrix(i, j) = leaf_model_scale_input.at(i, j); } + } } // Convert variable weights to std::vector std::vector var_weights_vector(variable_weights.size()); for (int i = 0; i < variable_weights.size(); i++) { - var_weights_vector[i] = variable_weights.at(i); + var_weights_vector[i] = variable_weights.at(i); } // Prepare the samplers StochTree::LeafModelVariant leaf_model = StochTree::leafModelFactory(model_type, leaf_scale, leaf_scale_matrix, a_forest, b_forest); - + // Run one iteration of the sampler StochTree::ForestContainer* forest_sample_ptr = forest_samples.GetContainer(); StochTree::TreeEnsemble* active_forest_ptr = forest.GetEnsemble(); @@ -1203,12 +1223,18 @@ class ForestSamplerCpp { int leaf_model_int, py::array_t initial_values) { // Convert leaf model type to enum StochTree::ModelType model_type; - if (leaf_model_int == 0) model_type = StochTree::ModelType::kConstantLeafGaussian; - else if (leaf_model_int == 1) model_type = StochTree::ModelType::kUnivariateRegressionLeafGaussian; - else if (leaf_model_int == 2) model_type = StochTree::ModelType::kMultivariateRegressionLeafGaussian; - else if (leaf_model_int == 3) model_type = StochTree::ModelType::kLogLinearVariance; - else if (leaf_model_int == 4) model_type = StochTree::ModelType::kCloglogOrdinal; - else StochTree::Log::Fatal("Invalid model type"); + if (leaf_model_int == 0) + model_type = StochTree::ModelType::kConstantLeafGaussian; + else if (leaf_model_int == 1) + model_type = StochTree::ModelType::kUnivariateRegressionLeafGaussian; + else if (leaf_model_int == 2) + model_type = StochTree::ModelType::kMultivariateRegressionLeafGaussian; + else if (leaf_model_int == 3) + model_type = StochTree::ModelType::kLogLinearVariance; + else if (leaf_model_int == 4) + model_type = StochTree::ModelType::kCloglogOrdinal; + else + StochTree::Log::Fatal("Invalid model type"); // Unpack initial value StochTree::TreeEnsemble* forest_ptr = forest.GetEnsemble(); @@ -1221,43 +1247,43 @@ class ForestSamplerCpp { (model_type == StochTree::ModelType::kUnivariateRegressionLeafGaussian) || (model_type == StochTree::ModelType::kLogLinearVariance) || (model_type == StochTree::ModelType::kCloglogOrdinal)) { - init_val = initial_values.at(0); + init_val = initial_values.at(0); } else if (model_type == StochTree::ModelType::kMultivariateRegressionLeafGaussian) { - int leaf_dim = initial_values.size(); - init_value_vector.resize(leaf_dim); - for (int i = 0; i < leaf_dim; i++) { - init_value_vector[i] = initial_values.at(i) / static_cast(num_trees); - } + int leaf_dim = initial_values.size(); + init_value_vector.resize(leaf_dim); + for (int i = 0; i < leaf_dim; i++) { + init_value_vector[i] = initial_values.at(i) / static_cast(num_trees); + } } - + // Initialize the models accordingly double leaf_init_val; if (model_type == StochTree::ModelType::kConstantLeafGaussian) { - leaf_init_val = init_val / static_cast(num_trees); - forest_ptr->SetLeafValue(leaf_init_val); - StochTree::UpdateResidualEntireForest(*tracker_, *forest_data_ptr, *residual_data_ptr, forest_ptr, false, std::minus()); - tracker_->UpdatePredictions(forest_ptr, *forest_data_ptr); + leaf_init_val = init_val / static_cast(num_trees); + forest_ptr->SetLeafValue(leaf_init_val); + StochTree::UpdateResidualEntireForest(*tracker_, *forest_data_ptr, *residual_data_ptr, forest_ptr, false, std::minus()); + tracker_->UpdatePredictions(forest_ptr, *forest_data_ptr); } else if (model_type == StochTree::ModelType::kUnivariateRegressionLeafGaussian) { - leaf_init_val = init_val / static_cast(num_trees); - forest_ptr->SetLeafValue(leaf_init_val); - StochTree::UpdateResidualEntireForest(*tracker_, *forest_data_ptr, *residual_data_ptr, forest_ptr, true, std::minus()); - tracker_->UpdatePredictions(forest_ptr, *forest_data_ptr); + leaf_init_val = init_val / static_cast(num_trees); + forest_ptr->SetLeafValue(leaf_init_val); + StochTree::UpdateResidualEntireForest(*tracker_, *forest_data_ptr, *residual_data_ptr, forest_ptr, true, std::minus()); + tracker_->UpdatePredictions(forest_ptr, *forest_data_ptr); } else if (model_type == StochTree::ModelType::kMultivariateRegressionLeafGaussian) { - forest_ptr->SetLeafVector(init_value_vector); - StochTree::UpdateResidualEntireForest(*tracker_, *forest_data_ptr, *residual_data_ptr, forest_ptr, true, std::minus()); - tracker_->UpdatePredictions(forest_ptr, *forest_data_ptr); + forest_ptr->SetLeafVector(init_value_vector); + StochTree::UpdateResidualEntireForest(*tracker_, *forest_data_ptr, *residual_data_ptr, forest_ptr, true, std::minus()); + tracker_->UpdatePredictions(forest_ptr, *forest_data_ptr); } else if (model_type == StochTree::ModelType::kLogLinearVariance) { - leaf_init_val = std::log(init_val) / static_cast(num_trees); - forest_ptr->SetLeafValue(leaf_init_val); - tracker_->UpdatePredictions(forest_ptr, *forest_data_ptr); - int n = forest_data_ptr->NumObservations(); - std::vector initial_preds(n, init_val); - forest_data_ptr->AddVarianceWeights(initial_preds.data(), n); + leaf_init_val = std::log(init_val) / static_cast(num_trees); + forest_ptr->SetLeafValue(leaf_init_val); + tracker_->UpdatePredictions(forest_ptr, *forest_data_ptr); + int n = forest_data_ptr->NumObservations(); + std::vector initial_preds(n, init_val); + forest_data_ptr->AddVarianceWeights(initial_preds.data(), n); } else if (model_type == StochTree::ModelType::kCloglogOrdinal) { - leaf_init_val = init_val / static_cast(num_trees); - forest_ptr->SetLeafValue(leaf_init_val); - StochTree::UpdateResidualEntireForest(*tracker_, *forest_data_ptr, *residual_data_ptr, forest_ptr, false, std::minus()); - tracker_->UpdatePredictions(forest_ptr, *forest_data_ptr); + leaf_init_val = init_val / static_cast(num_trees); + forest_ptr->SetLeafValue(leaf_init_val); + StochTree::UpdateResidualEntireForest(*tracker_, *forest_data_ptr, *residual_data_ptr, forest_ptr, false, std::minus()); + tracker_->UpdatePredictions(forest_ptr, *forest_data_ptr); } } @@ -1331,7 +1357,7 @@ class GlobalVarianceModelCpp { StochTree::ColumnVector* residual_ptr = residual.GetData(); std::mt19937* rng_ptr = rng.GetRng(); return var_model_.SampleVarianceParameter(residual_ptr->GetData(), a, b, *rng_ptr); - } + } private: StochTree::GlobalHomoskedasticVarianceModel var_model_; @@ -1391,7 +1417,7 @@ class OrdinalSamplerCpp { class RandomEffectsDatasetCpp { public: - RandomEffectsDatasetCpp() { + RandomEffectsDatasetCpp() { rfx_dataset_ = std::make_unique(); } ~RandomEffectsDatasetCpp() {} @@ -1443,7 +1469,7 @@ class RandomEffectsDatasetCpp { auto accessor = result.mutable_unchecked<2>(); for (py::ssize_t i = 0; i < num_row; i++) { for (int j = 0; j < num_col; j++) { - accessor(i,j) = rfx_dataset_->BasisValue(i,j); + accessor(i, j) = rfx_dataset_->BasisValue(i, j); } } return result; @@ -1466,9 +1492,9 @@ class RandomEffectsDatasetCpp { } return result; } - bool HasGroupLabels() {return rfx_dataset_->HasGroupLabels();} - bool HasBasis() {return rfx_dataset_->HasBasis();} - bool HasVarianceWeights() {return rfx_dataset_->HasVarWeights();} + bool HasGroupLabels() { return rfx_dataset_->HasGroupLabels(); } + bool HasBasis() { return rfx_dataset_->HasBasis(); } + bool HasVarianceWeights() { return rfx_dataset_->HasVarWeights(); } private: std::unique_ptr rfx_dataset_; @@ -1483,7 +1509,13 @@ class RandomEffectsContainerCpp { RandomEffectsContainerCpp() { rfx_container_ = std::make_unique(); } + explicit RandomEffectsContainerCpp(std::unique_ptr ptr) + : rfx_container_(std::move(ptr)) {} + ~RandomEffectsContainerCpp() {} + + StochTree::RandomEffectsContainer* GetPtr() { return rfx_container_.get(); } + void SetComponentsAndGroups(int num_components, int num_groups) { rfx_container_->SetNumComponents(num_components); rfx_container_->SetNumGroups(num_groups); @@ -1508,7 +1540,7 @@ class RandomEffectsContainerCpp { for (int i = 0; i < num_components; i++) { for (int j = 0; j < num_groups; j++) { for (int k = 0; k < num_samples; k++) { - accessor(i,j,k) = beta_raw[k*num_groups*num_components + j*num_components + i]; + accessor(i, j, k) = beta_raw[k * num_groups * num_components + j * num_components + i]; } } } @@ -1524,7 +1556,7 @@ class RandomEffectsContainerCpp { for (int i = 0; i < num_components; i++) { for (int j = 0; j < num_groups; j++) { for (int k = 0; k < num_samples; k++) { - accessor(i,j,k) = xi_raw[k*num_groups*num_components + j*num_components + i]; + accessor(i, j, k) = xi_raw[k * num_groups * num_components + j * num_components + i]; } } } @@ -1538,7 +1570,7 @@ class RandomEffectsContainerCpp { auto accessor = result.mutable_unchecked<2>(); for (int i = 0; i < num_components; i++) { for (int j = 0; j < num_samples; j++) { - accessor(i,j) = alpha_raw[j*num_components + i]; + accessor(i, j) = alpha_raw[j * num_components + i]; } } return result; @@ -1551,7 +1583,7 @@ class RandomEffectsContainerCpp { auto accessor = result.mutable_unchecked<2>(); for (int i = 0; i < num_components; i++) { for (int j = 0; j < num_samples; j++) { - accessor(i,j) = sigma_raw[j*num_components + i]; + accessor(i, j) = sigma_raw[j * num_components + i]; } } return result; @@ -1577,7 +1609,7 @@ class RandomEffectsContainerCpp { StochTree::RandomEffectsContainer* GetRandomEffectsContainer() { return rfx_container_.get(); } - + private: std::unique_ptr rfx_container_; }; @@ -1618,7 +1650,13 @@ class RandomEffectsLabelMapperCpp { RandomEffectsLabelMapperCpp() { rfx_label_mapper_ = std::make_unique(); } + explicit RandomEffectsLabelMapperCpp(std::unique_ptr ptr) + : rfx_label_mapper_(std::move(ptr)) {} + ~RandomEffectsLabelMapperCpp() {} + + StochTree::LabelMapper* GetPtr() { return rfx_label_mapper_.get(); } + void LoadFromTracker(RandomEffectsTrackerCpp& rfx_tracker) { StochTree::RandomEffectsTracker* internal_tracker = rfx_tracker.GetTracker(); rfx_label_mapper_->LoadFromLabelMap(internal_tracker->GetLabelMap()); @@ -1651,6 +1689,13 @@ class RandomEffectsLabelMapperCpp { } return result; } + py::array_t GetUniqueGroupIds() { + std::vector& keys = rfx_label_mapper_->Keys(); + auto result = py::array_t(py::detail::any_container({(py::ssize_t)keys.size()})); + auto accessor = result.mutable_unchecked<1>(); + for (int i = 0; i < (int)keys.size(); i++) accessor(i) = keys[i]; + return result; + } private: std::unique_ptr rfx_label_mapper_; @@ -1665,8 +1710,8 @@ class RandomEffectsModelCpp { StochTree::MultivariateRegressionRandomEffectsModel* GetModel() { return rfx_model_.get(); } - void SampleRandomEffects(RandomEffectsDatasetCpp& rfx_dataset, ResidualCpp& residual, - RandomEffectsTrackerCpp& rfx_tracker, RandomEffectsContainerCpp& rfx_container, + void SampleRandomEffects(RandomEffectsDatasetCpp& rfx_dataset, ResidualCpp& residual, + RandomEffectsTrackerCpp& rfx_tracker, RandomEffectsContainerCpp& rfx_container, bool keep_sample, double global_variance, RngCpp& rng); py::array_t Predict(RandomEffectsDatasetCpp& rfx_dataset, RandomEffectsTrackerCpp& rfx_tracker) { std::vector output = rfx_model_->Predict(*rfx_dataset.GetDataset(), *rfx_tracker.GetTracker()); @@ -1691,7 +1736,7 @@ class RandomEffectsModelCpp { Eigen::MatrixXd group_params_eigen(nrow, ncol); for (py::ssize_t i = 0; i < nrow; i++) { for (int j = 0; j < ncol; j++) { - group_params_eigen(i,j) = group_params.at(i,j); + group_params_eigen(i, j) = group_params.at(i, j); } } rfx_model_->SetGroupParameters(group_params_eigen); @@ -1702,7 +1747,7 @@ class RandomEffectsModelCpp { Eigen::MatrixXd working_param_cov_eigen(nrow, ncol); for (int i = 0; i < nrow; i++) { for (int j = 0; j < ncol; j++) { - working_param_cov_eigen(i,j) = working_param_cov.at(i,j); + working_param_cov_eigen(i, j) = working_param_cov.at(i, j); } } rfx_model_->SetWorkingParameterCovariance(working_param_cov_eigen); @@ -1713,7 +1758,7 @@ class RandomEffectsModelCpp { Eigen::MatrixXd group_param_cov_eigen(nrow, ncol); for (int i = 0; i < nrow; i++) { for (int j = 0; j < ncol; j++) { - group_param_cov_eigen(i,j) = group_param_cov.at(i,j); + group_param_cov_eigen(i, j) = group_param_cov.at(i, j); } } rfx_model_->SetGroupParameterCovariance(group_param_cov_eigen); @@ -1764,9 +1809,11 @@ class JsonCpp { return json_->dump(); } - std::string AddForest(ForestContainerCpp& forest_samples) { + std::string AddForest(ForestContainerCpp& forest_samples, std::string forest_label = "") { int forest_num = json_->at("num_forests"); - std::string forest_label = "forest_" + std::to_string(forest_num); + if (forest_label.empty()) { + forest_label = "forest_" + std::to_string(forest_num); + } nlohmann::json forest_json = forest_samples.ToJson(); json_->at("forests").emplace(forest_label, forest_json); json_->at("num_forests") = forest_num + 1; @@ -1795,7 +1842,7 @@ class JsonCpp { nlohmann::json groupids_json = nlohmann::json::array(); for (int i = 0; i < rfx_group_ids.size(); i++) { groupids_json.emplace_back(rfx_group_ids.at(i)); - } + } json_->at("random_effects").emplace(rfx_label, groupids_json); return rfx_label; } @@ -2027,6 +2074,32 @@ class JsonCpp { } } + void EraseField(std::string field_name) { + if (json_->contains(field_name)) { + json_->erase(field_name); + } + } + + void EraseFieldSubfolder(std::string subfolder_name, std::string field_name) { + if (json_->contains(subfolder_name) && json_->at(subfolder_name).contains(field_name)) { + json_->at(subfolder_name).erase(field_name); + } + } + + void RenameField(std::string old_name, std::string new_name) { + if (json_->contains(old_name)) { + (*json_)[new_name] = json_->at(old_name); + json_->erase(old_name); + } + } + + void RenameFieldSubfolder(std::string subfolder_name, std::string old_name, std::string new_name) { + if (json_->contains(subfolder_name) && json_->at(subfolder_name).contains(old_name)) { + json_->at(subfolder_name)[new_name] = json_->at(subfolder_name).at(old_name); + json_->at(subfolder_name).erase(old_name); + } + } + double ExtractDouble(std::string field_name) { return json_->at(field_name); } @@ -2131,10 +2204,1702 @@ class JsonCpp { return json_->at("random_effects"); } + // Direct access to the underlying nlohmann object so the samples wrappers can read/write the + // samples-owned subtree in place (BARTSamplesCpp::LoadFromJson / AddToJson). + nlohmann::json& GetJson() { return *json_; } + private: std::unique_ptr json_; }; +// Shared helpers for the single-owner samples wrappers (BART + BCF): marshal a C++ vector to a numpy +// array, and materialize a standalone deep copy of a forest container (via to_json/from_json) wrapped +// in ForestContainerCpp for the deprecated direct forest accessor (returns nullptr -> None if absent). +inline py::array_t samples_vec_to_numpy(const std::vector& v) { + py::array_t arr(static_cast(v.size())); + std::copy(v.begin(), v.end(), arr.mutable_data()); + return arr; +} + +inline std::vector numpy_to_samples_vec(py::object obj) { + if (obj.is_none()) return {}; + auto arr = obj.cast>(); + return std::vector(arr.data(), arr.data() + arr.size()); +} + +// Deep-copy a ForestContainer sample-by-sample (matching dims), without a JSON round-trip. +inline std::unique_ptr copy_forest_container(StochTree::ForestContainer* src) { + auto copy = std::make_unique( + src->NumTrees(), src->OutputDimension(), src->IsLeafConstant(), src->IsExponentiated()); + for (int i = 0; i < src->NumSamples(); i++) copy->AddSample(*src->GetEnsemble(i)); + return copy; +} + +inline std::unique_ptr materialize_forest_container( + const std::unique_ptr& src) { + if (src == nullptr) return nullptr; + auto copy = copy_forest_container(src.get()); + // Read dims before moving the unique_ptr (argument evaluation order is unspecified). + int num_trees = copy->NumTrees(); + int output_dim = copy->OutputDimension(); + bool leaf_constant = copy->IsLeafConstant(); + bool exponentiated = copy->IsExponentiated(); + return std::make_unique(std::move(copy), num_trees, output_dim, leaf_constant, exponentiated); +} + +// Thin owning wrapper around a single StochTree::BARTSamples, held as ONE external pointer in the +// Python model (single-owner design). It owns the unique_ptr and forwards to the core methods, +// doing only Python type marshalling: parameter traces -> numpy, and a materialize-on-demand deep +// copy of a forest container wrapped in ForestContainerCpp for the (deprecated) direct forest +// accessor. Serialization/merge logic lives in core BARTSamples; this class adds no model logic. +class BARTSamplesCpp { + public: + BARTSamplesCpp() { samples_ = std::make_unique(); } + BARTSamplesCpp(std::unique_ptr samples) { samples_ = std::move(samples); } + ~BARTSamplesCpp() {} + + StochTree::BARTSamples* GetPtr() { return samples_.get(); } + + // Populate the samples-owned subtree from a parsed JSON object held by a JsonCpp (envelope already + // resolved/migrated by the per-language caller). Inverse: AddToJson merges the subtree back in. + void LoadFromJson(JsonCpp& json) { samples_->FromJson(json.GetJson()); } + void AddToJson(JsonCpp& json) { json.GetJson().update(samples_->ToJson()); } + + // Convenience string forms (used for isolated testing; the real save/load path uses JsonCpp). + static std::unique_ptr FromJsonString(std::string json_string) { + auto wrapper = std::make_unique(); + nlohmann::json obj = nlohmann::json::parse(json_string); + wrapper->samples_->FromJson(obj); + return wrapper; + } + std::string ToJsonString() { return samples_->ToJson().dump(); } + + // Build a BARTSamples by deep-copying existing forest containers and parameter arrays (the + // "construct from components" path, #406). Lets the model assemble the single owned object from + // the current sampler outputs without restructuring the binding; variance_forest/global_var/ + // leaf_scale are optional (pass None when absent). + static std::unique_ptr FromComponents( + py::object mean_forest, py::object variance_forest, + py::object global_var_samples, py::object leaf_scale_samples, + double y_bar, double y_std, int num_samples) { + auto wrapper = std::make_unique(); + StochTree::BARTSamples* s = wrapper->samples_.get(); + // BART supports mean-only, variance-only, or both -- both forests are optional. + if (!mean_forest.is_none()) { + s->mean_forests = copy_forest_container(mean_forest.cast()->GetPtr()); + } + if (!variance_forest.is_none()) { + s->variance_forests = copy_forest_container(variance_forest.cast()->GetPtr()); + } + s->global_error_variance_samples = numpy_to_samples_vec(global_var_samples); + s->leaf_scale_samples = numpy_to_samples_vec(leaf_scale_samples); + s->y_bar = y_bar; + s->y_std = y_std; + s->num_samples = num_samples; + return wrapper; + } + + // Append another chain's draws (multi-chain combine); forwards to core BARTSamples::Merge. + void Merge(BARTSamplesCpp& other) { samples_->Merge(*other.samples_); } + + int NumSamples() { return samples_->num_samples; } + double YBar() { return samples_->y_bar; } + double YStd() { return samples_->y_std; } + int NumTrain() { return samples_->num_train; } + int NumTest() { return samples_->num_test; } + + py::array_t GlobalVarSamples() { return samples_vec_to_numpy(samples_->global_error_variance_samples); } + py::array_t LeafScaleSamples() { return samples_vec_to_numpy(samples_->leaf_scale_samples); } + + bool HasMeanForest() { return samples_->mean_forests != nullptr; } + bool HasVarianceForest() { return samples_->variance_forests != nullptr; } + + // Materialize a standalone deep copy of a forest container, wrapped in ForestContainerCpp, for the + // (deprecated) direct forest accessor. Returns None when the forest is absent. + std::unique_ptr MaterializeMeanForest() { return materialize_forest_container(samples_->mean_forests); } + std::unique_ptr MaterializeVarianceForest() { return materialize_forest_container(samples_->variance_forests); } + + private: + std::unique_ptr samples_; +}; + +// Thin owning wrapper around a single StochTree::BCFSamples (BCF mirror of BARTSamplesCpp). Owns the +// unique_ptr and forwards to core; adds only Python marshalling. BCF carries three forests +// (prognostic/treatment/variance), extra parameter traces (leaf_scale_mu/tau, tau_0, b0/b1), and a +// treatment_dim scalar. +class BCFSamplesCpp { + public: + BCFSamplesCpp() { samples_ = std::make_unique(); } + BCFSamplesCpp(std::unique_ptr samples) { samples_ = std::move(samples); } + ~BCFSamplesCpp() {} + + StochTree::BCFSamples* GetPtr() { return samples_.get(); } + + void LoadFromJson(JsonCpp& json) { samples_->FromJson(json.GetJson()); } + void AddToJson(JsonCpp& json) { json.GetJson().update(samples_->ToJson()); } + static std::unique_ptr FromJsonString(std::string json_string) { + auto wrapper = std::make_unique(); + nlohmann::json obj = nlohmann::json::parse(json_string); + wrapper->samples_->FromJson(obj); + return wrapper; + } + std::string ToJsonString() { return samples_->ToJson().dump(); } + + // Build a BCFSamples by deep-copying existing forest containers and parameter arrays (#406 BCF + // mirror). mu/tau forests are required; variance_forest and the parameter arrays are optional. + static std::unique_ptr FromComponents( + ForestContainerCpp& mu_forest, ForestContainerCpp& tau_forest, py::object variance_forest, + py::object global_var_samples, py::object leaf_scale_mu_samples, py::object leaf_scale_tau_samples, + py::object tau_0_samples, py::object b0_samples, py::object b1_samples, + double y_bar, double y_std, int num_samples, int treatment_dim) { + auto wrapper = std::make_unique(); + StochTree::BCFSamples* s = wrapper->samples_.get(); + s->mu_forests = copy_forest_container(mu_forest.GetPtr()); + s->tau_forests = copy_forest_container(tau_forest.GetPtr()); + if (!variance_forest.is_none()) { + s->variance_forests = copy_forest_container(variance_forest.cast()->GetPtr()); + } + s->global_error_variance_samples = numpy_to_samples_vec(global_var_samples); + s->leaf_scale_mu_samples = numpy_to_samples_vec(leaf_scale_mu_samples); + s->leaf_scale_tau_samples = numpy_to_samples_vec(leaf_scale_tau_samples); + s->tau_0_samples = numpy_to_samples_vec(tau_0_samples); + s->b0_samples = numpy_to_samples_vec(b0_samples); + s->b1_samples = numpy_to_samples_vec(b1_samples); + s->y_bar = y_bar; + s->y_std = y_std; + s->num_samples = num_samples; + s->treatment_dim = treatment_dim; + return wrapper; + } + + void Merge(BCFSamplesCpp& other) { samples_->Merge(*other.samples_); } + + int NumSamples() { return samples_->num_samples; } + double YBar() { return samples_->y_bar; } + double YStd() { return samples_->y_std; } + int NumTrain() { return samples_->num_train; } + int NumTest() { return samples_->num_test; } + int TreatmentDim() { return samples_->treatment_dim; } + + py::array_t GlobalVarSamples() { return samples_vec_to_numpy(samples_->global_error_variance_samples); } + py::array_t LeafScaleMuSamples() { return samples_vec_to_numpy(samples_->leaf_scale_mu_samples); } + py::array_t LeafScaleTauSamples() { return samples_vec_to_numpy(samples_->leaf_scale_tau_samples); } + py::array_t Tau0Samples() { return samples_vec_to_numpy(samples_->tau_0_samples); } + py::array_t B0Samples() { return samples_vec_to_numpy(samples_->b0_samples); } + py::array_t B1Samples() { return samples_vec_to_numpy(samples_->b1_samples); } + + bool HasMuForest() { return samples_->mu_forests != nullptr; } + bool HasTauForest() { return samples_->tau_forests != nullptr; } + bool HasVarianceForest() { return samples_->variance_forests != nullptr; } + + std::unique_ptr MaterializeMuForest() { return materialize_forest_container(samples_->mu_forests); } + std::unique_ptr MaterializeTauForest() { return materialize_forest_container(samples_->tau_forests); } + std::unique_ptr MaterializeVarianceForest() { return materialize_forest_container(samples_->variance_forests); } + + private: + std::unique_ptr samples_; +}; + +template +T get_config_scalar_default(py::dict& config_dict, const char* config_key, T default_value) { + return config_dict.contains(config_key) ? config_dict[config_key].cast() : default_value; +} + +inline StochTree::BARTConfig convert_dict_to_bart_config(py::dict config_dict) { + StochTree::BARTConfig output; + + // Global model parameters + output.standardize_outcome = get_config_scalar_default(config_dict, "standardize_outcome", true); + output.num_threads = get_config_scalar_default(config_dict, "num_threads", 1); + output.verbose = get_config_scalar_default(config_dict, "verbose", false); + output.cutpoint_grid_size = get_config_scalar_default(config_dict, "cutpoint_grid_size", 100); + output.link_function = static_cast(get_config_scalar_default(config_dict, "link_function", 0)); + output.outcome_type = static_cast(get_config_scalar_default(config_dict, "outcome_type", 0)); + output.random_seed = get_config_scalar_default(config_dict, "random_seed", 1); + output.keep_gfr = get_config_scalar_default(config_dict, "keep_gfr", 1); + output.keep_burnin = get_config_scalar_default(config_dict, "keep_burnin", 1); + + // Global error variance parameters + output.a_sigma2_global = get_config_scalar_default(config_dict, "a_sigma2_global", 0.0); + output.b_sigma2_global = get_config_scalar_default(config_dict, "b_sigma2_global", 0.0); + output.sigma2_global_init = get_config_scalar_default(config_dict, "sigma2_global_init", 1.0); + output.sample_sigma2_global = get_config_scalar_default(config_dict, "sample_sigma2_global", true); + + // Mean forest parameters + output.num_trees_mean = get_config_scalar_default(config_dict, "num_trees_mean", 200); + output.alpha_mean = get_config_scalar_default(config_dict, "alpha_mean", 0.95); + output.beta_mean = get_config_scalar_default(config_dict, "beta_mean", 2.0); + output.min_samples_leaf_mean = get_config_scalar_default(config_dict, "min_samples_leaf_mean", 5); + output.max_depth_mean = get_config_scalar_default(config_dict, "max_depth_mean", -1); + output.leaf_constant_mean = get_config_scalar_default(config_dict, "leaf_constant_mean", true); + output.leaf_dim_mean = get_config_scalar_default(config_dict, "leaf_dim_mean", 1); + output.exponentiated_leaf_mean = get_config_scalar_default(config_dict, "exponentiated_leaf_mean", true); + output.num_features_subsample_mean = get_config_scalar_default(config_dict, "num_features_subsample_mean", 0); + output.a_sigma2_mean = get_config_scalar_default(config_dict, "a_sigma2_mean", 3.0); + output.b_sigma2_mean = get_config_scalar_default(config_dict, "b_sigma2_mean", -1.0); + output.sigma2_mean_init = get_config_scalar_default(config_dict, "sigma2_mean_init", -1.0); + output.sample_sigma2_leaf_mean = get_config_scalar_default(config_dict, "sample_sigma2_leaf_mean", false); + output.mean_leaf_model_type = static_cast(get_config_scalar_default(config_dict, "mean_leaf_model_type", 0)); + output.num_classes_cloglog = get_config_scalar_default(config_dict, "num_classes_cloglog", 0); + output.cloglog_leaf_prior_shape = get_config_scalar_default(config_dict, "cloglog_leaf_prior_shape", 2.0); + output.cloglog_leaf_prior_scale = get_config_scalar_default(config_dict, "cloglog_leaf_prior_scale", 2.0); + output.cloglog_cutpoint_0 = get_config_scalar_default(config_dict, "cloglog_cutpoint_0", 0.0); + + // Variance forest parameters + output.num_trees_variance = get_config_scalar_default(config_dict, "num_trees_variance", 0); + output.leaf_prior_calibration_param = get_config_scalar_default(config_dict, "leaf_prior_calibration_param", 1.5); + output.shape_variance_forest = get_config_scalar_default(config_dict, "shape_variance_forest", -1.0); + output.scale_variance_forest = get_config_scalar_default(config_dict, "scale_variance_forest", -1.0); + output.variance_forest_leaf_init = get_config_scalar_default(config_dict, "variance_forest_leaf_init", -1.0); + output.alpha_variance = get_config_scalar_default(config_dict, "alpha_variance", 0.5); + output.beta_variance = get_config_scalar_default(config_dict, "beta_variance", 2.0); + output.min_samples_leaf_variance = get_config_scalar_default(config_dict, "min_samples_leaf_variance", 5); + output.max_depth_variance = get_config_scalar_default(config_dict, "max_depth_variance", -1); + output.leaf_constant_variance = get_config_scalar_default(config_dict, "leaf_constant_variance", true); + output.leaf_dim_variance = get_config_scalar_default(config_dict, "leaf_dim_variance", 1); + output.exponentiated_leaf_variance = get_config_scalar_default(config_dict, "exponentiated_leaf_variance", true); + output.num_features_subsample_variance = get_config_scalar_default(config_dict, "num_features_subsample_variance", 0); + + // Random effects parameters + output.has_random_effects = get_config_scalar_default(config_dict, "has_random_effects", false); + output.rfx_model_spec = static_cast(get_config_scalar_default(config_dict, "rfx_model_spec", 0)); + output.rfx_variance_prior_shape = get_config_scalar_default(config_dict, "rfx_variance_prior_shape", 1.0); + output.rfx_variance_prior_scale = get_config_scalar_default(config_dict, "rfx_variance_prior_scale", 1.0); + + // Handle vector conversions separately + if (config_dict.contains("feature_types")) { + std::vector feature_types_vector = config_dict["feature_types"].cast>(); + for (auto item : feature_types_vector) { + output.feature_types.push_back(static_cast(item)); + } + } + if (config_dict.contains("sweep_update_indices_mean")) { + output.sweep_update_indices_mean = config_dict["sweep_update_indices_mean"].cast>(); + } + if (config_dict.contains("sweep_update_indices_variance")) { + output.sweep_update_indices_variance = config_dict["sweep_update_indices_variance"].cast>(); + } + if (config_dict.contains("var_weights_mean")) { + output.var_weights_mean = config_dict["var_weights_mean"].cast>(); + } + if (config_dict.contains("sigma2_leaf_mean_matrix")) { + output.sigma2_leaf_mean_matrix = config_dict["sigma2_leaf_mean_matrix"].cast>(); + } + if (config_dict.contains("var_weights_variance")) { + output.var_weights_variance = config_dict["var_weights_variance"].cast>(); + } + if (config_dict.contains("rfx_working_parameter_mean_prior")) { + py::array_t arr = + config_dict["rfx_working_parameter_mean_prior"].cast>(); + output.rfx_working_parameter_mean_prior = std::vector( + arr.data(), arr.data() + arr.size()); + } + if (config_dict.contains("rfx_group_parameter_mean_prior")) { + py::array_t arr = + config_dict["rfx_group_parameter_mean_prior"].cast>(); + output.rfx_group_parameter_mean_prior = std::vector( + arr.data(), arr.data() + arr.size()); + } + if (config_dict.contains("rfx_working_parameter_cov_prior")) { + py::array_t arr = + config_dict["rfx_working_parameter_cov_prior"].cast>(); + output.rfx_working_parameter_cov_prior = std::vector( + arr.data(), arr.data() + arr.size()); + } + if (config_dict.contains("rfx_group_parameter_cov_prior")) { + py::array_t arr = + config_dict["rfx_group_parameter_cov_prior"].cast>(); + output.rfx_group_parameter_cov_prior = std::vector( + arr.data(), arr.data() + arr.size()); + } + return output; +} + +inline StochTree::BARTData convert_numpy_to_bart_data( + py::object X_train, + py::object y_train, + py::object X_test, + int n_train, + int n_test, + int p, + py::object basis_train, + py::object basis_test, + int basis_dim, + py::object obs_weights_train, + py::object obs_weights_test, + py::object rfx_group_ids_train, + py::object rfx_group_ids_test, + py::object rfx_basis_train, + py::object rfx_basis_test, + int rfx_num_groups, + int rfx_basis_dim) { + StochTree::BARTData output; + if (!X_train.is_none()) { + py::array_t X_train_array = X_train.cast>(); + output.X_train = static_cast(X_train_array.mutable_data()); + } + if (!y_train.is_none()) { + py::array_t y_train_array = y_train.cast>(); + output.y_train = static_cast(y_train_array.mutable_data()); + } + if (!X_test.is_none()) { + py::array_t X_test_array = X_test.cast>(); + output.X_test = static_cast(X_test_array.mutable_data()); + } + if (!basis_train.is_none()) { + py::array_t basis_train_array = basis_train.cast>(); + output.basis_train = static_cast(basis_train_array.mutable_data()); + } + if (!basis_test.is_none()) { + py::array_t basis_test_array = basis_test.cast>(); + output.basis_test = static_cast(basis_test_array.mutable_data()); + } + if (!obs_weights_train.is_none()) { + py::array_t obs_weights_train_array = obs_weights_train.cast>(); + output.obs_weights_train = static_cast(obs_weights_train_array.mutable_data()); + } + if (!obs_weights_test.is_none()) { + py::array_t obs_weights_test_array = obs_weights_test.cast>(); + output.obs_weights_test = static_cast(obs_weights_test_array.mutable_data()); + } + if (!rfx_group_ids_train.is_none()) { + py::array_t rfx_group_ids_train_array = rfx_group_ids_train.cast>(); + output.rfx_group_ids_train = static_cast(rfx_group_ids_train_array.mutable_data()); + } + if (!rfx_group_ids_test.is_none()) { + py::array_t rfx_group_ids_test_array = rfx_group_ids_test.cast>(); + output.rfx_group_ids_test = static_cast(rfx_group_ids_test_array.mutable_data()); + } + if (!rfx_basis_train.is_none()) { + py::array_t rfx_basis_train_array = rfx_basis_train.cast>(); + output.rfx_basis_train = static_cast(rfx_basis_train_array.mutable_data()); + } + if (!rfx_basis_test.is_none()) { + py::array_t rfx_basis_test_array = rfx_basis_test.cast>(); + output.rfx_basis_test = static_cast(rfx_basis_test_array.mutable_data()); + } + output.n_train = n_train; + output.n_test = n_test; + output.p = p; + output.basis_dim = basis_dim; + output.rfx_num_groups = rfx_num_groups; + output.rfx_basis_dim = rfx_basis_dim; + return output; +} + +inline py::dict convert_bart_results_to_dict( + StochTree::BARTSamples& results_raw, StochTree::BARTConfig& config) { + py::dict output; + + // Transfer ownership of mean forest pointers + if (results_raw.mean_forests != nullptr) { + output["forest_container_mean"] = py::cast(std::make_unique(std::move(results_raw.mean_forests), config.num_trees_mean, config.leaf_dim_mean, config.leaf_constant_mean, config.exponentiated_leaf_mean)); + } else { + output["forest_container_mean"] = py::none(); + } + + // Transfer ownership of variance forest pointers + if (results_raw.variance_forests != nullptr) { + output["forest_container_variance"] = py::cast(std::make_unique(std::move(results_raw.variance_forests), config.num_trees_variance, config.leaf_dim_variance, config.leaf_constant_variance, config.exponentiated_leaf_variance)); + } else { + output["forest_container_variance"] = py::none(); + } + + // Move parameter vector samples + + // Train set mean forest predictions + if (results_raw.mean_forest_predictions_train.empty()) { + output["mean_forest_predictions_train"] = py::none(); + } else { + auto input_vec = results_raw.mean_forest_predictions_train; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["mean_forest_predictions_train"] = array; + } + + // Test set mean forest predictions + if (results_raw.mean_forest_predictions_test.empty()) { + output["mean_forest_predictions_test"] = py::none(); + } else { + auto input_vec = results_raw.mean_forest_predictions_test; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["mean_forest_predictions_test"] = array; + } + + // Train set variance forest predictions + if (results_raw.variance_forest_predictions_train.empty()) { + output["variance_forest_predictions_train"] = py::none(); + } else { + auto input_vec = results_raw.variance_forest_predictions_train; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["variance_forest_predictions_train"] = array; + } + + // Test set variance forest predictions + if (results_raw.variance_forest_predictions_test.empty()) { + output["variance_forest_predictions_test"] = py::none(); + } else { + auto input_vec = results_raw.variance_forest_predictions_test; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["variance_forest_predictions_test"] = array; + } + + // Global error variance samples + if (results_raw.global_error_variance_samples.empty()) { + output["global_var_samples"] = py::none(); + } else { + auto input_vec = results_raw.global_error_variance_samples; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["global_var_samples"] = array; + } + + // Leaf scale samples + if (results_raw.leaf_scale_samples.empty()) { + output["leaf_scale_samples"] = py::none(); + } else { + auto input_vec = results_raw.leaf_scale_samples; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["leaf_scale_samples"] = array; + } + + // Cloglog cutpoint samples + if (results_raw.cloglog_cutpoint_samples.empty()) { + output["cloglog_cutpoint_samples"] = py::none(); + } else { + auto input_vec = results_raw.cloglog_cutpoint_samples; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["cloglog_cutpoint_samples"] = array; + } + + // Unpack RFX predictions + if (!results_raw.rfx_predictions_train.empty()) { + auto& v = results_raw.rfx_predictions_train; + py::array_t array(v.size()); + std::copy(v.begin(), v.end(), array.mutable_data()); + output["rfx_predictions_train"] = array; + } else { + output["rfx_predictions_train"] = py::none(); + } + if (!results_raw.rfx_predictions_test.empty()) { + auto& v = results_raw.rfx_predictions_test; + py::array_t array(v.size()); + std::copy(v.begin(), v.end(), array.mutable_data()); + output["rfx_predictions_test"] = array; + } else { + output["rfx_predictions_test"] = py::none(); + } + + // Transfer ownership of random effects container pointers + if (results_raw.rfx_container != nullptr) { + output["rfx_container"] = py::cast(std::make_unique(std::move(results_raw.rfx_container))); + } else { + output["rfx_container"] = py::none(); + } + + // Transfer ownership of random effects label mapper pointer + if (results_raw.rfx_label_mapper != nullptr) { + output["rfx_label_mapper"] = py::cast(std::make_unique(std::move(results_raw.rfx_label_mapper))); + } else { + output["rfx_label_mapper"] = py::none(); + } + + // Unpack scalars + output["y_bar"] = results_raw.y_bar; + output["y_std"] = results_raw.y_std; + output["num_samples"] = results_raw.num_samples; + output["num_train"] = results_raw.num_train; + output["num_test"] = results_raw.num_test; + + return output; +} + +void add_config_to_bart_result_dict(py::dict& result, StochTree::BARTConfig& config) { + // Unpack more metadata about the model that was sampled + result["sigma2_init"] = config.sigma2_global_init; + result["sigma2_mean_init"] = config.sigma2_mean_init; + result["b_sigma2_mean"] = config.b_sigma2_mean; + result["shape_variance_forest"] = config.shape_variance_forest; + result["scale_variance_forest"] = config.scale_variance_forest; +} + +py::dict bart_sample_cpp( + py::object X_train, + py::object y_train, + py::object X_test, + int n_train, + int n_test, + int p, + py::object basis_train, + py::object basis_test, + int basis_dim, + py::object obs_weights_train, + py::object obs_weights_test, + py::object rfx_group_ids_train, + py::object rfx_group_ids_test, + py::object rfx_basis_train, + py::object rfx_basis_test, + int rfx_num_groups, + int rfx_basis_dim, + int num_gfr, + int num_burnin, + int keep_every, + int num_mcmc, + int num_chains, + py::dict config_input) { + // Convert config dict to BARTConfig struct + StochTree::BARTConfig bart_config = convert_dict_to_bart_config(config_input); + + // Unpack pointers to input data to BARTData object + StochTree::BARTData bart_data = convert_numpy_to_bart_data(X_train, y_train, X_test, n_train, n_test, p, basis_train, basis_test, basis_dim, obs_weights_train, obs_weights_test, rfx_group_ids_train, rfx_group_ids_test, rfx_basis_train, rfx_basis_test, rfx_num_groups, rfx_basis_dim); + + // Create outcome object + StochTree::BARTSamples bart_results_raw = StochTree::BARTSamples(); + + // Initialize a BART sampler + StochTree::BARTSampler bart_sampler(bart_results_raw, bart_config, bart_data); + + // Run the sampler + bart_sampler.run_gfr(bart_results_raw, num_gfr, bart_config.keep_gfr, num_chains); + if (num_chains > 1) { + bart_sampler.run_mcmc_chains(bart_results_raw, num_chains, num_burnin, keep_every, num_mcmc); + } else { + bart_sampler.run_mcmc(bart_results_raw, num_burnin, keep_every, num_mcmc); + } + bart_sampler.postprocess_samples(bart_results_raw); + + // Convert results to Python dictionary + py::dict bart_results = convert_bart_results_to_dict(bart_results_raw, bart_config); + add_config_to_bart_result_dict(bart_results, bart_config); + // Persist the final RNG state so continue_sampling() can resume the random stream + // for bit-identical continuation. Only meaningful for single-chain runs; multi-chain + // runs use one RNG per chain, so there is no single stream to resume. + if (num_chains <= 1) { + bart_results["rng_state"] = bart_sampler.GetRngState(); + bart_results["leaf_normal_cache"] = bart_sampler.GetLeafNormalCache(); + } else { + bart_results["rng_state"] = py::none(); + bart_results["leaf_normal_cache"] = py::none(); + } + return bart_results; +} + +// Continue (warm-start) sampling from an already-fit BART model. +// +// Reconstruct-on-demand stopgap (RFC 0005 / #408): rather than transferring +// ownership of the model's forest container, this DEEP-COPIES it into a fresh +// BARTSamples so the caller's container remains the source of truth. The scalar +// sample histories are pre-populated in STANDARDIZED space (forward sampling +// keeps them standardized until postprocess_samples), then the sampler is +// constructed in continuation mode (warm-starts the active forest from the last +// retained sample and appends new draws to the copied container). +// +// Returns the same dict shape as bart_sample_cpp, with the forest container and +// parameter arrays extended to (history + new) samples. Predictions in the +// returned dict are recomputed by the Python wrapper post-hoc. +py::dict bart_continue_sample_cpp( + py::object X_train, + py::object y_train, + py::object X_test, + int n_train, + int n_test, + int p, + py::object basis_train, + py::object basis_test, + int basis_dim, + py::object obs_weights_train, + py::object obs_weights_test, + py::object rfx_group_ids_train, + py::object rfx_group_ids_test, + py::object rfx_basis_train, + py::object rfx_basis_test, + int rfx_num_groups, + int rfx_basis_dim, + int num_burnin, + int keep_every, + int num_mcmc, + ForestContainerCpp& mean_forest_container, + py::object global_var_samples, + py::object leaf_scale_samples, + double y_bar, + double y_std, + std::string rng_state_in, + bool override_seed, + std::string leaf_normal_cache_in, + py::dict config_input) { + // Convert config dict to BARTConfig struct + StochTree::BARTConfig bart_config = convert_dict_to_bart_config(config_input); + + // Unpack pointers to (re-supplied) input data to BARTData object + StochTree::BARTData bart_data = convert_numpy_to_bart_data(X_train, y_train, X_test, n_train, n_test, p, basis_train, basis_test, basis_dim, obs_weights_train, obs_weights_test, rfx_group_ids_train, rfx_group_ids_test, rfx_basis_train, rfx_basis_test, rfx_num_groups, rfx_basis_dim); + + // Create samples object and deep-copy the model's mean forest container into it. + // The caller's container is only read (via GetPtr), never moved, so it stays valid. + StochTree::BARTSamples bart_results_raw = StochTree::BARTSamples(); + StochTree::ForestContainer* src_mean = mean_forest_container.GetPtr(); + int num_history = src_mean->NumSamples(); + bart_results_raw.mean_forests = std::make_unique( + bart_config.num_trees_mean, bart_config.leaf_dim_mean, bart_config.leaf_constant_mean, bart_config.exponentiated_leaf_mean); + for (int i = 0; i < num_history; i++) { + bart_results_raw.mean_forests->AddSample(*src_mean->GetEnsemble(i)); + } + + // History counts + standardization scalars + bart_results_raw.y_bar = y_bar; + bart_results_raw.y_std = y_std; + bart_results_raw.num_samples = num_history; + + // Pre-populate scalar sample histories in standardized space so that + // postprocess_samples re-scales the whole (history + new) array consistently. + if (!global_var_samples.is_none()) { + auto arr = global_var_samples.cast>(); + auto r = arr.unchecked<1>(); + double y_std2 = y_std * y_std; + bart_results_raw.global_error_variance_samples.reserve(static_cast(r.shape(0)) + num_mcmc); + for (py::ssize_t i = 0; i < r.shape(0); i++) { + // Stored values are post-processed (x y_std^2); divide back out to standardized space. + bart_results_raw.global_error_variance_samples.push_back(r(i) / y_std2); + } + } + if (!leaf_scale_samples.is_none()) { + auto arr = leaf_scale_samples.cast>(); + auto r = arr.unchecked<1>(); + bart_results_raw.leaf_scale_samples.reserve(static_cast(r.shape(0)) + num_mcmc); + for (py::ssize_t i = 0; i < r.shape(0); i++) { + // leaf_scale samples are already standardized (not post-processed). + bart_results_raw.leaf_scale_samples.push_back(r(i)); + } + } + + // Initialize a BART sampler in continuation mode (warm-start from last sample) + StochTree::BARTSampler bart_sampler(bart_results_raw, bart_config, bart_data, /*continuation=*/true); + + // Resume the RNG stream from where the prior run left off (bit-identical continuation), + // unless the user supplied a new seed (override_seed) -- in which case the fresh seed set + // by InitializeState stands. The warm-start init consumes no RNG draws, so the restored + // state is positioned exactly at the next draw. + if (!override_seed && !rng_state_in.empty()) { + bart_sampler.SetRngState(rng_state_in); + } + // Restore the leaf normal sampler's cached spare value (unless re-seeding from scratch). + if (!override_seed) { + bart_sampler.SetLeafNormalCache(leaf_normal_cache_in); + } + + // Append new MCMC samples (continuation does not run GFR) + bart_sampler.run_mcmc(bart_results_raw, num_burnin, keep_every, num_mcmc); + bart_sampler.postprocess_samples(bart_results_raw); + + // Convert results to Python dictionary + py::dict bart_results = convert_bart_results_to_dict(bart_results_raw, bart_config); + add_config_to_bart_result_dict(bart_results, bart_config); + // Carry the new final RNG + leaf-cache state forward so chained continuations stay bit-identical. + bart_results["rng_state"] = bart_sampler.GetRngState(); + bart_results["leaf_normal_cache"] = bart_sampler.GetLeafNormalCache(); + return bart_results; +} + +inline StochTree::BCFConfig convert_dict_to_bcf_config(py::dict config_dict) { + StochTree::BCFConfig output; + + // Global model parameters + output.standardize_outcome = get_config_scalar_default(config_dict, "standardize_outcome", true); + output.num_threads = get_config_scalar_default(config_dict, "num_threads", 1); + output.verbose = get_config_scalar_default(config_dict, "verbose", false); + output.cutpoint_grid_size = get_config_scalar_default(config_dict, "cutpoint_grid_size", 100); + output.link_function = static_cast(get_config_scalar_default(config_dict, "link_function", 0)); + output.outcome_type = static_cast(get_config_scalar_default(config_dict, "outcome_type", 0)); + output.random_seed = get_config_scalar_default(config_dict, "random_seed", 1); + output.keep_gfr = get_config_scalar_default(config_dict, "keep_gfr", 0); + output.keep_burnin = get_config_scalar_default(config_dict, "keep_burnin", 0); + output.adaptive_coding = get_config_scalar_default(config_dict, "adaptive_coding", 0); + output.b_0_init = get_config_scalar_default(config_dict, "b_0_init", 0.0); + output.b_1_init = get_config_scalar_default(config_dict, "b_1_init", 1.0); + + // Global error variance parameters + output.a_sigma2_global = get_config_scalar_default(config_dict, "a_sigma2_global", 0.0); + output.b_sigma2_global = get_config_scalar_default(config_dict, "b_sigma2_global", 0.0); + output.sigma2_global_init = get_config_scalar_default(config_dict, "sigma2_global_init", 1.0); + output.sample_sigma2_global = get_config_scalar_default(config_dict, "sample_sigma2_global", true); + + // Mu forest parameters + output.num_trees_mu = get_config_scalar_default(config_dict, "num_trees_mu", 200); + output.alpha_mu = get_config_scalar_default(config_dict, "alpha_mu", 0.95); + output.beta_mu = get_config_scalar_default(config_dict, "beta_mu", 2.0); + output.min_samples_leaf_mu = get_config_scalar_default(config_dict, "min_samples_leaf_mu", 5); + output.max_depth_mu = get_config_scalar_default(config_dict, "max_depth_mu", -1); + output.leaf_constant_mu = get_config_scalar_default(config_dict, "leaf_constant_mu", true); + output.leaf_dim_mu = get_config_scalar_default(config_dict, "leaf_dim_mu", 1); + output.exponentiated_leaf_mu = get_config_scalar_default(config_dict, "exponentiated_leaf_mu", false); + output.num_features_subsample_mu = get_config_scalar_default(config_dict, "num_features_subsample_mu", 0); + output.a_sigma2_mu = get_config_scalar_default(config_dict, "a_sigma2_mu", 3.0); + output.b_sigma2_mu = get_config_scalar_default(config_dict, "b_sigma2_mu", -1.0); + output.sigma2_mu_init = get_config_scalar_default(config_dict, "sigma2_mu_init", -1.0); + output.sample_sigma2_leaf_mu = get_config_scalar_default(config_dict, "sample_sigma2_leaf_mu", false); + + // Tau forest parameters + output.num_trees_tau = get_config_scalar_default(config_dict, "num_trees_tau", 50); + output.alpha_tau = get_config_scalar_default(config_dict, "alpha_tau", 0.95); + output.beta_tau = get_config_scalar_default(config_dict, "beta_tau", 2.0); + output.min_samples_leaf_tau = get_config_scalar_default(config_dict, "min_samples_leaf_tau", 5); + output.max_depth_tau = get_config_scalar_default(config_dict, "max_depth_tau", -1); + output.leaf_constant_tau = get_config_scalar_default(config_dict, "leaf_constant_tau", false); + output.leaf_dim_tau = get_config_scalar_default(config_dict, "leaf_dim_tau", 1); + output.exponentiated_leaf_tau = get_config_scalar_default(config_dict, "exponentiated_leaf_tau", false); + output.num_features_subsample_tau = get_config_scalar_default(config_dict, "num_features_subsample_tau", 0); + output.a_sigma2_tau = get_config_scalar_default(config_dict, "a_sigma2_tau", 3.0); + output.b_sigma2_tau = get_config_scalar_default(config_dict, "b_sigma2_tau", -1.0); + output.sigma2_tau_init = get_config_scalar_default(config_dict, "sigma2_tau_init", -1.0); + output.sample_sigma2_leaf_tau = get_config_scalar_default(config_dict, "sample_sigma2_leaf_tau", false); + output.tau_leaf_model_type = static_cast(get_config_scalar_default(config_dict, "tau_leaf_model_type", 0)); + output.sample_tau_0 = get_config_scalar_default(config_dict, "sample_tau_0", true); + output.tau_0_prior_var_scalar = get_config_scalar_default(config_dict, "tau_0_prior_var_scalar", -1.0); + + // Variance forest parameters + output.num_trees_variance = get_config_scalar_default(config_dict, "num_trees_variance", 0); + output.leaf_prior_calibration_param = get_config_scalar_default(config_dict, "leaf_prior_calibration_param", 1.5); + output.shape_variance_forest = get_config_scalar_default(config_dict, "shape_variance_forest", -1.0); + output.scale_variance_forest = get_config_scalar_default(config_dict, "scale_variance_forest", -1.0); + output.variance_forest_leaf_init = get_config_scalar_default(config_dict, "variance_forest_leaf_init", -1.0); + output.alpha_variance = get_config_scalar_default(config_dict, "alpha_variance", 0.5); + output.beta_variance = get_config_scalar_default(config_dict, "beta_variance", 2.0); + output.min_samples_leaf_variance = get_config_scalar_default(config_dict, "min_samples_leaf_variance", 5); + output.max_depth_variance = get_config_scalar_default(config_dict, "max_depth_variance", -1); + output.leaf_constant_variance = get_config_scalar_default(config_dict, "leaf_constant_variance", true); + output.leaf_dim_variance = get_config_scalar_default(config_dict, "leaf_dim_variance", 1); + output.exponentiated_leaf_variance = get_config_scalar_default(config_dict, "exponentiated_leaf_variance", true); + output.num_features_subsample_variance = get_config_scalar_default(config_dict, "num_features_subsample_variance", 0); + + // Random effects parameters + output.has_random_effects = get_config_scalar_default(config_dict, "has_random_effects", false); + output.rfx_model_spec = static_cast(get_config_scalar_default(config_dict, "rfx_model_spec", 0)); + output.rfx_variance_prior_shape = get_config_scalar_default(config_dict, "rfx_variance_prior_shape", 1.0); + output.rfx_variance_prior_scale = get_config_scalar_default(config_dict, "rfx_variance_prior_scale", 1.0); + + // Handle vector conversions separately + if (config_dict.contains("feature_types")) { + std::vector feature_types_vector = config_dict["feature_types"].cast>(); + for (auto item : feature_types_vector) { + output.feature_types.push_back(static_cast(item)); + } + } + if (config_dict.contains("sweep_update_indices_mu")) { + output.sweep_update_indices_mu = config_dict["sweep_update_indices_mu"].cast>(); + } + if (config_dict.contains("sweep_update_indices_tau")) { + output.sweep_update_indices_tau = config_dict["sweep_update_indices_tau"].cast>(); + } + if (config_dict.contains("sweep_update_indices_variance")) { + output.sweep_update_indices_variance = config_dict["sweep_update_indices_variance"].cast>(); + } + if (config_dict.contains("var_weights_mu")) { + output.var_weights_mu = config_dict["var_weights_mu"].cast>(); + } + if (config_dict.contains("var_weights_tau")) { + output.var_weights_tau = config_dict["var_weights_tau"].cast>(); + } + if (config_dict.contains("var_weights_variance")) { + output.var_weights_variance = config_dict["var_weights_variance"].cast>(); + } + if (config_dict.contains("sigma2_leaf_tau_matrix")) { + output.sigma2_leaf_tau_matrix = config_dict["sigma2_leaf_tau_matrix"].cast>(); + } + if (config_dict.contains("tau_0_prior_var_multivariate")) { + output.tau_0_prior_var_multivariate = config_dict["tau_0_prior_var_multivariate"].cast>(); + } + if (config_dict.contains("rfx_working_parameter_mean_prior")) { + py::array_t arr = + config_dict["rfx_working_parameter_mean_prior"].cast>(); + output.rfx_working_parameter_mean_prior = std::vector( + arr.data(), arr.data() + arr.size()); + } + if (config_dict.contains("rfx_group_parameter_mean_prior")) { + py::array_t arr = + config_dict["rfx_group_parameter_mean_prior"].cast>(); + output.rfx_group_parameter_mean_prior = std::vector( + arr.data(), arr.data() + arr.size()); + } + if (config_dict.contains("rfx_working_parameter_cov_prior")) { + py::array_t arr = + config_dict["rfx_working_parameter_cov_prior"].cast>(); + output.rfx_working_parameter_cov_prior = std::vector( + arr.data(), arr.data() + arr.size()); + } + if (config_dict.contains("rfx_group_parameter_cov_prior")) { + py::array_t arr = + config_dict["rfx_group_parameter_cov_prior"].cast>(); + output.rfx_group_parameter_cov_prior = std::vector( + arr.data(), arr.data() + arr.size()); + } + return output; +} + +inline StochTree::BCFData convert_numpy_to_bcf_data( + py::object X_train, + py::object Z_train, + py::object y_train, + py::object X_test, + py::object Z_test, + int n_train, + int n_test, + int p, + int treatment_dim, + py::object obs_weights_train, + py::object obs_weights_test, + py::object rfx_group_ids_train, + py::object rfx_group_ids_test, + py::object rfx_basis_train, + py::object rfx_basis_test, + int rfx_num_groups, + int rfx_basis_dim) { + StochTree::BCFData output; + if (!X_train.is_none()) { + py::array_t X_train_array = X_train.cast>(); + output.X_train = static_cast(X_train_array.mutable_data()); + } + if (!Z_train.is_none()) { + py::array_t Z_train_array = Z_train.cast>(); + output.treatment_train = static_cast(Z_train_array.mutable_data()); + } + if (!y_train.is_none()) { + py::array_t y_train_array = y_train.cast>(); + output.y_train = static_cast(y_train_array.mutable_data()); + } + if (!X_test.is_none()) { + py::array_t X_test_array = X_test.cast>(); + output.X_test = static_cast(X_test_array.mutable_data()); + } + if (!Z_test.is_none()) { + py::array_t Z_test_array = Z_test.cast>(); + output.treatment_test = static_cast(Z_test_array.mutable_data()); + } + if (!obs_weights_train.is_none()) { + py::array_t obs_weights_train_array = obs_weights_train.cast>(); + output.obs_weights_train = static_cast(obs_weights_train_array.mutable_data()); + } + if (!obs_weights_test.is_none()) { + py::array_t obs_weights_test_array = obs_weights_test.cast>(); + output.obs_weights_test = static_cast(obs_weights_test_array.mutable_data()); + } + if (!rfx_group_ids_train.is_none()) { + py::array_t rfx_group_ids_train_array = rfx_group_ids_train.cast>(); + output.rfx_group_ids_train = static_cast(rfx_group_ids_train_array.mutable_data()); + } + if (!rfx_group_ids_test.is_none()) { + py::array_t rfx_group_ids_test_array = rfx_group_ids_test.cast>(); + output.rfx_group_ids_test = static_cast(rfx_group_ids_test_array.mutable_data()); + } + if (!rfx_basis_train.is_none()) { + py::array_t rfx_basis_train_array = rfx_basis_train.cast>(); + output.rfx_basis_train = static_cast(rfx_basis_train_array.mutable_data()); + } + if (!rfx_basis_test.is_none()) { + py::array_t rfx_basis_test_array = rfx_basis_test.cast>(); + output.rfx_basis_test = static_cast(rfx_basis_test_array.mutable_data()); + } + output.n_train = n_train; + output.n_test = n_test; + output.p = p; + output.treatment_dim = treatment_dim; + output.rfx_num_groups = rfx_num_groups; + output.rfx_basis_dim = rfx_basis_dim; + return output; +} + +inline py::dict convert_bcf_results_to_dict( + StochTree::BCFSamples& results_raw, StochTree::BCFConfig& config) { + py::dict output; + + // Transfer ownership of mean forest pointers + if (results_raw.mu_forests != nullptr) { + output["forest_container_mu"] = py::cast(std::make_unique(std::move(results_raw.mu_forests), config.num_trees_mu, config.leaf_dim_mu, config.leaf_constant_mu, config.exponentiated_leaf_mu)); + } else { + output["forest_container_mu"] = py::none(); + } + if (results_raw.tau_forests != nullptr) { + output["forest_container_tau"] = py::cast(std::make_unique(std::move(results_raw.tau_forests), config.num_trees_tau, config.leaf_dim_tau, config.leaf_constant_tau, config.exponentiated_leaf_tau)); + } else { + output["forest_container_tau"] = py::none(); + } + + // Transfer ownership of variance forest pointers + if (results_raw.variance_forests != nullptr) { + output["forest_container_variance"] = py::cast(std::make_unique(std::move(results_raw.variance_forests), config.num_trees_variance, config.leaf_dim_variance, config.leaf_constant_variance, config.exponentiated_leaf_variance)); + } else { + output["forest_container_variance"] = py::none(); + } + + // Move parameter vector samples + + // Train set prognostic forest predictions + if (results_raw.mu_forest_predictions_train.empty()) { + output["mu_forest_predictions_train"] = py::none(); + } else { + auto input_vec = results_raw.mu_forest_predictions_train; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["mu_forest_predictions_train"] = array; + } + + // Test set prognostic forest predictions + if (results_raw.mu_forest_predictions_test.empty()) { + output["mu_forest_predictions_test"] = py::none(); + } else { + auto input_vec = results_raw.mu_forest_predictions_test; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["mu_forest_predictions_test"] = array; + } + + // Train set treatment effect forest predictions + if (results_raw.tau_forest_predictions_train.empty()) { + output["tau_forest_predictions_train"] = py::none(); + } else { + auto input_vec = results_raw.tau_forest_predictions_train; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["tau_forest_predictions_train"] = array; + } + + // Test set treatment effect forest predictions + if (results_raw.tau_forest_predictions_test.empty()) { + output["tau_forest_predictions_test"] = py::none(); + } else { + auto input_vec = results_raw.tau_forest_predictions_test; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["tau_forest_predictions_test"] = array; + } + + // Train set outcome predictions + if (results_raw.y_hat_train.empty()) { + output["y_hat_train"] = py::none(); + } else { + auto input_vec = results_raw.y_hat_train; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["y_hat_train"] = array; + } + + // Test set outcome predictions + if (results_raw.y_hat_test.empty()) { + output["y_hat_test"] = py::none(); + } else { + auto input_vec = results_raw.y_hat_test; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["y_hat_test"] = array; + } + + // Train set variance forest predictions + if (results_raw.variance_forest_predictions_train.empty()) { + output["variance_forest_predictions_train"] = py::none(); + } else { + auto input_vec = results_raw.variance_forest_predictions_train; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["variance_forest_predictions_train"] = array; + } + + // Test set variance forest predictions + if (results_raw.variance_forest_predictions_test.empty()) { + output["variance_forest_predictions_test"] = py::none(); + } else { + auto input_vec = results_raw.variance_forest_predictions_test; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["variance_forest_predictions_test"] = array; + } + + // Global error variance samples + if (results_raw.global_error_variance_samples.empty()) { + output["global_var_samples"] = py::none(); + } else { + auto input_vec = results_raw.global_error_variance_samples; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["global_var_samples"] = array; + } + + // Leaf scale samples -- prognostic forest + if (results_raw.leaf_scale_mu_samples.empty()) { + output["leaf_scale_mu_samples"] = py::none(); + } else { + auto input_vec = results_raw.leaf_scale_mu_samples; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["leaf_scale_mu_samples"] = array; + } + + // Leaf scale samples -- treatment effect forest + if (results_raw.leaf_scale_tau_samples.empty()) { + output["leaf_scale_tau_samples"] = py::none(); + } else { + auto input_vec = results_raw.leaf_scale_tau_samples; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["leaf_scale_tau_samples"] = array; + } + + // tau_0 samples + if (results_raw.tau_0_samples.empty()) { + output["tau_0_samples"] = py::none(); + } else { + auto input_vec = results_raw.tau_0_samples; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["tau_0_samples"] = array; + } + + // adaptive coding samples + // b0 + if (results_raw.b0_samples.empty()) { + output["b0_samples"] = py::none(); + } else { + auto input_vec = results_raw.b0_samples; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["b0_samples"] = array; + } + // b1 + if (results_raw.b1_samples.empty()) { + output["b1_samples"] = py::none(); + } else { + auto input_vec = results_raw.b1_samples; + py::array_t array(input_vec.size()); + std::copy(input_vec.begin(), input_vec.end(), array.mutable_data()); + output["b1_samples"] = array; + } + + // Unpack RFX predictions + if (!results_raw.rfx_predictions_train.empty()) { + auto& v = results_raw.rfx_predictions_train; + py::array_t array(v.size()); + std::copy(v.begin(), v.end(), array.mutable_data()); + output["rfx_predictions_train"] = array; + } else { + output["rfx_predictions_train"] = py::none(); + } + if (!results_raw.rfx_predictions_test.empty()) { + auto& v = results_raw.rfx_predictions_test; + py::array_t array(v.size()); + std::copy(v.begin(), v.end(), array.mutable_data()); + output["rfx_predictions_test"] = array; + } else { + output["rfx_predictions_test"] = py::none(); + } + + // Transfer ownership of random effects container pointers + if (results_raw.rfx_container != nullptr) { + output["rfx_container"] = py::cast(std::make_unique(std::move(results_raw.rfx_container))); + } else { + output["rfx_container"] = py::none(); + } + + // Transfer ownership of random effects label mapper pointer + if (results_raw.rfx_label_mapper != nullptr) { + output["rfx_label_mapper"] = py::cast(std::make_unique(std::move(results_raw.rfx_label_mapper))); + } else { + output["rfx_label_mapper"] = py::none(); + } + + // Unpack scalars + output["y_bar"] = results_raw.y_bar; + output["y_std"] = results_raw.y_std; + output["num_samples"] = results_raw.num_samples; + output["num_train"] = results_raw.num_train; + output["num_test"] = results_raw.num_test; + + return output; +} + +inline py::dict convert_bart_preds_to_dict(StochTree::BARTPredictionResult& results_raw) { + py::dict output; + + // Move prediction samples + + if (results_raw.y_hat.empty()) { + output["y_hat"] = py::none(); + } else { + py::array_t array(results_raw.y_hat.size()); + std::copy(results_raw.y_hat.begin(), results_raw.y_hat.end(), array.mutable_data()); + output["y_hat"] = array; + } + + if (results_raw.mean_forest_predictions.empty()) { + output["mean_forest_predictions"] = py::none(); + } else { + py::array_t array(results_raw.mean_forest_predictions.size()); + std::copy(results_raw.mean_forest_predictions.begin(), results_raw.mean_forest_predictions.end(), array.mutable_data()); + output["mean_forest_predictions"] = array; + } + + if (results_raw.variance_forest_predictions.empty()) { + output["variance_forest_predictions"] = py::none(); + } else { + py::array_t array(results_raw.variance_forest_predictions.size()); + std::copy(results_raw.variance_forest_predictions.begin(), results_raw.variance_forest_predictions.end(), array.mutable_data()); + output["variance_forest_predictions"] = array; + } + + if (results_raw.rfx_predictions.empty()) { + output["rfx_predictions"] = py::none(); + } else { + py::array_t array(results_raw.rfx_predictions.size()); + std::copy(results_raw.rfx_predictions.begin(), results_raw.rfx_predictions.end(), array.mutable_data()); + output["rfx_predictions"] = array; + } + + return output; +} + +py::dict bart_predict_cpp( + py::dict bart_model_dict, + py::object X, + py::object leaf_basis, + int n, + int p, + int num_basis, + py::object obs_weights, + py::object rfx_group_ids, + py::object rfx_basis, + int rfx_num_groups, + int rfx_basis_dim, + bool posterior, + int scale, + bool predict_y_hat, + bool predict_mean_forest, + bool predict_variance_forest, + bool predict_random_effects) { + // Pre-convert all numpy inputs to F-contiguous at function scope so the raw pointers + // stored in BARTData outlive the convert_numpy_to_bart_data call. + using FArray = py::array_t; + using IArray = py::array_t; + FArray X_farr, leaf_basis_farr, obs_weights_farr, rfx_basis_farr; + IArray rfx_group_ids_iarr; + if (!X.is_none()) X_farr = X.cast(); + if (!leaf_basis.is_none()) leaf_basis_farr = leaf_basis.cast(); + if (!obs_weights.is_none()) obs_weights_farr = obs_weights.cast(); + if (!rfx_group_ids.is_none()) rfx_group_ids_iarr = rfx_group_ids.cast(); + if (!rfx_basis.is_none()) rfx_basis_farr = rfx_basis.cast(); + + // Build BARTData with test-only fields + StochTree::BARTData bart_data = convert_numpy_to_bart_data( + /*X_train=*/py::none(), + /*y_train=*/py::none(), + /*X_test=*/X.is_none() ? py::object(py::none()) : py::object(X_farr), + /*n_train=*/0, + /*n_test=*/n, + /*p=*/p, + /*basis_train=*/py::none(), + /*basis_test=*/leaf_basis.is_none() ? py::object(py::none()) : py::object(leaf_basis_farr), + /*basis_dim=*/num_basis, + /*obs_weights_train=*/py::none(), + /*obs_weights_test=*/obs_weights.is_none() ? py::object(py::none()) : py::object(obs_weights_farr), + /*rfx_group_ids_train=*/py::none(), + /*rfx_group_ids_test=*/rfx_group_ids.is_none() ? py::object(py::none()) : py::object(rfx_group_ids_iarr), + /*rfx_basis_train=*/py::none(), + /*rfx_basis_test=*/rfx_basis.is_none() ? py::object(py::none()) : py::object(rfx_basis_farr), + /*rfx_num_groups=*/rfx_num_groups, + /*rfx_basis_dim=*/rfx_basis_dim); + + // Build BARTPredictionInput from the model dict + StochTree::BARTPredictionInput pred_input; + + py::array_t global_var_arr, leaf_scale_arr, cloglog_cutpoints_arr; + if (bart_model_dict.contains("sigma2_global_samples") && !bart_model_dict["sigma2_global_samples"].is_none()) { + global_var_arr = bart_model_dict["sigma2_global_samples"].cast>(); + pred_input.global_error_variance_samples = static_cast(global_var_arr.mutable_data()); + } + if (bart_model_dict.contains("sigma2_leaf_samples") && !bart_model_dict["sigma2_leaf_samples"].is_none()) { + leaf_scale_arr = bart_model_dict["sigma2_leaf_samples"].cast>(); + pred_input.leaf_scale_samples = static_cast(leaf_scale_arr.mutable_data()); + } + if (bart_model_dict.contains("mean_forests") && !bart_model_dict["mean_forests"].is_none()) { + pred_input.mean_forests = bart_model_dict["mean_forests"].cast()->GetPtr(); + } + if (bart_model_dict.contains("variance_forests") && !bart_model_dict["variance_forests"].is_none()) { + pred_input.variance_forests = bart_model_dict["variance_forests"].cast()->GetPtr(); + } + if (bart_model_dict.contains("rfx_container") && !bart_model_dict["rfx_container"].is_none()) { + pred_input.rfx_container = bart_model_dict["rfx_container"].cast()->GetPtr(); + } + if (bart_model_dict.contains("rfx_label_mapper") && !bart_model_dict["rfx_label_mapper"].is_none()) { + pred_input.rfx_label_mapper = bart_model_dict["rfx_label_mapper"].cast()->GetPtr(); + } + if (bart_model_dict.contains("cloglog_cutpoint_samples") && !bart_model_dict["cloglog_cutpoint_samples"].is_none()) { + cloglog_cutpoints_arr = bart_model_dict["cloglog_cutpoint_samples"].cast>(); + pred_input.cloglog_cutpoint_samples = static_cast(cloglog_cutpoints_arr.mutable_data()); + } + pred_input.cloglog_num_classes = bart_model_dict.contains("cloglog_num_classes") ? bart_model_dict["cloglog_num_classes"].cast() : 0; + pred_input.num_samples = bart_model_dict["num_samples"].cast(); + pred_input.num_obs = n; + pred_input.num_basis = num_basis; + pred_input.y_bar = bart_model_dict["y_bar"].cast(); + pred_input.y_std = bart_model_dict["y_std"].cast(); + pred_input.has_variance_forest = bart_model_dict["include_variance_forest"].cast(); + pred_input.has_rfx = bart_model_dict["has_rfx"].cast(); + { + std::string rfx_spec_str = ""; + if (bart_model_dict.contains("rfx_model_spec") && !bart_model_dict["rfx_model_spec"].is_none()) { + rfx_spec_str = bart_model_dict["rfx_model_spec"].cast(); + } + pred_input.rfx_model_spec = (rfx_spec_str == "intercept_only") + ? StochTree::BARTRFXModelSpec::InterceptOnly + : StochTree::BARTRFXModelSpec::Custom; + } + { + std::string link_str = bart_model_dict.contains("link_function") ? bart_model_dict["link_function"].cast() : "identity"; + if (link_str == "probit") + pred_input.link_function = StochTree::LinkFunction::Probit; + else if (link_str == "cloglog") + pred_input.link_function = StochTree::LinkFunction::Cloglog; + else + pred_input.link_function = StochTree::LinkFunction::Identity; + } + { + std::string outcome_str = bart_model_dict.contains("outcome_type") ? bart_model_dict["outcome_type"].cast() : "continuous"; + if (outcome_str == "binary") + pred_input.outcome_type = StochTree::OutcomeType::Binary; + else if (outcome_str == "ordinal") + pred_input.outcome_type = StochTree::OutcomeType::Ordinal; + else + pred_input.outcome_type = StochTree::OutcomeType::Continuous; + } + pred_input.pred_type = posterior ? StochTree::PredType::kPosterior : StochTree::PredType::kMean; + if (scale == 0) + pred_input.pred_scale = StochTree::PredScale::kLinear; + else if (scale == 1) + pred_input.pred_scale = StochTree::PredScale::kProbability; + else + pred_input.pred_scale = StochTree::PredScale::kClass; + pred_input.pred_terms.y_hat = predict_y_hat; + pred_input.pred_terms.mean_forest = predict_mean_forest; + pred_input.pred_terms.variance_forest = predict_variance_forest; + pred_input.pred_terms.random_effects = predict_random_effects; + + // Run prediction + StochTree::BARTPredictionResult pred_results = predict_bart_model(bart_data, pred_input); + + return convert_bart_preds_to_dict(pred_results); +} + +inline py::dict convert_bcf_preds_to_dict(StochTree::BCFPredictionResult& results_raw) { + py::dict output; + + // Move prediction samples + + if (results_raw.y_hat.empty()) { + output["y_hat"] = py::none(); + } else { + py::array_t array(results_raw.y_hat.size()); + std::copy(results_raw.y_hat.begin(), results_raw.y_hat.end(), array.mutable_data()); + output["y_hat"] = array; + } + + if (results_raw.mu_x.empty()) { + output["mu_x"] = py::none(); + } else { + py::array_t array(results_raw.mu_x.size()); + std::copy(results_raw.mu_x.begin(), results_raw.mu_x.end(), array.mutable_data()); + output["mu_x"] = array; + } + + if (results_raw.tau_x.empty()) { + output["tau_x"] = py::none(); + } else { + py::array_t array(results_raw.tau_x.size()); + std::copy(results_raw.tau_x.begin(), results_raw.tau_x.end(), array.mutable_data()); + output["tau_x"] = array; + } + + if (results_raw.prognostic_function.empty()) { + output["prognostic_function"] = py::none(); + } else { + py::array_t array(results_raw.prognostic_function.size()); + std::copy(results_raw.prognostic_function.begin(), results_raw.prognostic_function.end(), array.mutable_data()); + output["prognostic_function"] = array; + } + + if (results_raw.cate.empty()) { + output["cate"] = py::none(); + } else { + py::array_t array(results_raw.cate.size()); + std::copy(results_raw.cate.begin(), results_raw.cate.end(), array.mutable_data()); + output["cate"] = array; + } + + if (results_raw.conditional_variance.empty()) { + output["conditional_variance"] = py::none(); + } else { + py::array_t array(results_raw.conditional_variance.size()); + std::copy(results_raw.conditional_variance.begin(), results_raw.conditional_variance.end(), array.mutable_data()); + output["conditional_variance"] = array; + } + + if (results_raw.random_effects.empty()) { + output["random_effects"] = py::none(); + } else { + py::array_t array(results_raw.random_effects.size()); + std::copy(results_raw.random_effects.begin(), results_raw.random_effects.end(), array.mutable_data()); + output["random_effects"] = array; + } + + return output; +} + +void add_config_to_bcf_result_dict(py::dict& result, StochTree::BCFConfig& config) { + // Unpack more metadata about the model that was sampled + result["sigma2_init"] = config.sigma2_global_init; + result["sigma2_mu_init"] = config.sigma2_mu_init; + result["sigma2_tau_init"] = config.sigma2_tau_init; + result["b_sigma2_mu"] = config.b_sigma2_mu; + result["b_sigma2_tau"] = config.b_sigma2_tau; + result["shape_variance_forest"] = config.shape_variance_forest; + result["scale_variance_forest"] = config.scale_variance_forest; +} + +py::dict bcf_sample_cpp( + py::object X_train, + py::object Z_train, + py::object y_train, + py::object X_test, + py::object Z_test, + int n_train, + int n_test, + int p, + int treatment_dim, + py::object obs_weights_train, + py::object obs_weights_test, + py::object rfx_group_ids_train, + py::object rfx_group_ids_test, + py::object rfx_basis_train, + py::object rfx_basis_test, + int rfx_num_groups, + int rfx_basis_dim, + int num_gfr, + int num_burnin, + int keep_every, + int num_mcmc, + int num_chains, + bool adaptive_coding, + py::dict config_input) { + // Convert config dict to BCFConfig struct + StochTree::BCFConfig bcf_config = convert_dict_to_bcf_config(config_input); + + // Unpack pointers to input data to BCFData object + StochTree::BCFData bcf_data = convert_numpy_to_bcf_data(X_train, Z_train, y_train, X_test, Z_test, n_train, n_test, p, treatment_dim, obs_weights_train, obs_weights_test, rfx_group_ids_train, rfx_group_ids_test, rfx_basis_train, rfx_basis_test, rfx_num_groups, rfx_basis_dim); + + // Create outcome object + StochTree::BCFSamples bcf_results_raw = StochTree::BCFSamples(); + + // Initialize a BCF sampler + StochTree::BCFSampler bcf_sampler(bcf_results_raw, bcf_config, bcf_data); + + // Run the sampler + bcf_sampler.run_gfr(bcf_results_raw, num_gfr, bcf_config.keep_gfr, num_chains); + if (num_chains > 1) { + bcf_sampler.run_mcmc_chains(bcf_results_raw, num_chains, num_burnin, keep_every, num_mcmc); + } else { + bcf_sampler.run_mcmc(bcf_results_raw, num_burnin, keep_every, num_mcmc); + } + bcf_sampler.postprocess_samples(bcf_results_raw); + + // Convert results to Python dictionary + py::dict bcf_results = convert_bcf_results_to_dict(bcf_results_raw, bcf_config); + add_config_to_bcf_result_dict(bcf_results, bcf_config); + // Persist the final RNG + leaf-cache state so continue_sampling() can resume the random stream + // for bit-identical continuation. Only meaningful for single-chain runs; multi-chain runs use + // one RNG per chain, so there is no single stream to resume. + if (num_chains <= 1) { + bcf_results["rng_state"] = bcf_sampler.GetRngState(); + bcf_results["leaf_normal_cache"] = bcf_sampler.GetLeafNormalCache(); + } else { + bcf_results["rng_state"] = py::none(); + bcf_results["leaf_normal_cache"] = py::none(); + } + return bcf_results; +} + +// Continue (warm-start) sampling from an already-fit BCF model. +// +// Reconstruct-on-demand stopgap (RFC 0005 / #408), mirroring bart_continue_sample_cpp: +// DEEP-COPIES the model's mu and tau forest containers into a fresh BCFSamples so the caller's +// containers remain the source of truth, pre-populates the scalar sample histories in +// STANDARDIZED space, constructs the sampler in continuation mode (warm-starts the active +// forests from the last retained samples and appends new draws), then returns the same dict +// shape as bcf_sample_cpp with the containers and parameter arrays extended. +// +// Supports identity-link, univariate-treatment Gaussian BCF without a variance forest, random +// effects, or adaptive coding (the sampler hard-errors on those). Predictions in the returned +// dict are recomputed by the Python wrapper post-hoc, so postprocess_samples is skipped here +// (it would index the per-iteration prediction arrays, which only hold the new samples). +py::dict bcf_continue_sample_cpp( + py::object X_train, + py::object Z_train, + py::object y_train, + int n_train, + int p, + int treatment_dim, + py::object obs_weights_train, + int num_burnin, + int keep_every, + int num_mcmc, + ForestContainerCpp& mu_forest_container, + ForestContainerCpp& tau_forest_container, + py::object global_var_samples, + py::object leaf_scale_mu_samples, + py::object leaf_scale_tau_samples, + py::object tau_0_samples, + double y_bar, + double y_std, + std::string rng_state_in, + bool override_seed, + std::string leaf_normal_cache_in, + py::dict config_input) { + // Convert config dict to BCFConfig struct + StochTree::BCFConfig bcf_config = convert_dict_to_bcf_config(config_input); + + // Unpack pointers to (re-supplied) input data to BCFData object (no test data; the Python + // wrapper recomputes predictions via predict()). + StochTree::BCFData bcf_data = convert_numpy_to_bcf_data( + X_train, Z_train, y_train, /*X_test=*/py::none(), /*Z_test=*/py::none(), + n_train, /*n_test=*/0, p, treatment_dim, obs_weights_train, /*obs_weights_test=*/py::none(), + /*rfx_group_ids_train=*/py::none(), /*rfx_group_ids_test=*/py::none(), + /*rfx_basis_train=*/py::none(), /*rfx_basis_test=*/py::none(), + /*rfx_num_groups=*/0, /*rfx_basis_dim=*/0); + + // Create samples object and deep-copy the model's mu/tau forest containers into it. + // The caller's containers are only read (via GetPtr), never moved, so they stay valid. + StochTree::BCFSamples bcf_results_raw = StochTree::BCFSamples(); + StochTree::ForestContainer* src_mu = mu_forest_container.GetPtr(); + StochTree::ForestContainer* src_tau = tau_forest_container.GetPtr(); + int num_history = src_mu->NumSamples(); + bcf_results_raw.mu_forests = std::make_unique( + bcf_config.num_trees_mu, bcf_config.leaf_dim_mu, bcf_config.leaf_constant_mu, bcf_config.exponentiated_leaf_mu); + for (int i = 0; i < num_history; i++) { + bcf_results_raw.mu_forests->AddSample(*src_mu->GetEnsemble(i)); + } + bcf_results_raw.tau_forests = std::make_unique( + bcf_config.num_trees_tau, bcf_config.leaf_dim_tau, bcf_config.leaf_constant_tau, bcf_config.exponentiated_leaf_tau); + for (int i = 0; i < num_history; i++) { + bcf_results_raw.tau_forests->AddSample(*src_tau->GetEnsemble(i)); + } + + // History counts + standardization scalars + bcf_results_raw.y_bar = y_bar; + bcf_results_raw.y_std = y_std; + bcf_results_raw.num_samples = num_history; + bcf_results_raw.treatment_dim = treatment_dim; + + // Pre-populate scalar sample histories in STANDARDIZED space. The model stores global error + // variance post-processed (x y_std^2) and tau_0 scaled by y_std; leaf scales are stored + // standardized. Un-postprocess so the warm-start reads the correct standardized values and + // forward sampling stays consistent. + double y_std2 = y_std * y_std; + if (!global_var_samples.is_none()) { + auto arr = global_var_samples.cast>(); + auto r = arr.unchecked<1>(); + bcf_results_raw.global_error_variance_samples.reserve(static_cast(r.shape(0)) + num_mcmc); + for (py::ssize_t i = 0; i < r.shape(0); i++) { + bcf_results_raw.global_error_variance_samples.push_back(r(i) / y_std2); + } + } + if (!leaf_scale_mu_samples.is_none()) { + auto arr = leaf_scale_mu_samples.cast>(); + auto r = arr.unchecked<1>(); + bcf_results_raw.leaf_scale_mu_samples.reserve(static_cast(r.shape(0)) + num_mcmc); + for (py::ssize_t i = 0; i < r.shape(0); i++) { + bcf_results_raw.leaf_scale_mu_samples.push_back(r(i)); + } + } + if (!leaf_scale_tau_samples.is_none()) { + auto arr = leaf_scale_tau_samples.cast>(); + auto r = arr.unchecked<1>(); + bcf_results_raw.leaf_scale_tau_samples.reserve(static_cast(r.shape(0)) + num_mcmc); + for (py::ssize_t i = 0; i < r.shape(0); i++) { + bcf_results_raw.leaf_scale_tau_samples.push_back(r(i)); + } + } + if (!tau_0_samples.is_none()) { + auto arr = tau_0_samples.cast>(); + auto r = arr.unchecked<1>(); + bcf_results_raw.tau_0_samples.reserve(static_cast(r.shape(0)) + num_mcmc); + for (py::ssize_t i = 0; i < r.shape(0); i++) { + // Model stores tau_0 scaled by y_std; divide back out to standardized space. + bcf_results_raw.tau_0_samples.push_back(r(i) / y_std); + } + } + + // Initialize a BCF sampler in continuation mode (warm-start from last samples) + StochTree::BCFSampler bcf_sampler(bcf_results_raw, bcf_config, bcf_data, /*continuation=*/true); + + // Resume the RNG stream and leaf-sampler caches unless the user re-seeded (override_seed). + if (!override_seed && !rng_state_in.empty()) { + bcf_sampler.SetRngState(rng_state_in); + } + if (!override_seed) { + bcf_sampler.SetLeafNormalCache(leaf_normal_cache_in); + } + + // Append new MCMC samples (continuation does not run GFR) + bcf_sampler.run_mcmc(bcf_results_raw, num_burnin, keep_every, num_mcmc); + + // Skip postprocess_samples (it would OOB-index the per-iteration prediction arrays, which + // hold only the new samples). The Python wrapper recomputes predictions from the extended + // forest containers. Manually rescale global error variance to the original outcome scale to + // match what postprocess_samples / convert_bcf_results_to_dict normally produce. + for (double& v : bcf_results_raw.global_error_variance_samples) v *= y_std2; + + // Convert results to Python dictionary + py::dict bcf_results = convert_bcf_results_to_dict(bcf_results_raw, bcf_config); + add_config_to_bcf_result_dict(bcf_results, bcf_config); + // Carry the new final RNG + leaf-cache state forward so chained continuations stay bit-identical. + bcf_results["rng_state"] = bcf_sampler.GetRngState(); + bcf_results["leaf_normal_cache"] = bcf_sampler.GetLeafNormalCache(); + return bcf_results; +} + +py::dict bcf_predict_cpp( + py::dict bcf_model_dict, + py::object X, + py::object Z, + int n, + int p, + int treatment_dim, + py::object obs_weights, + py::object rfx_group_ids, + py::object rfx_basis, + int rfx_num_groups, + int rfx_basis_dim, + bool posterior, + int scale, + bool predict_y_hat, + bool predict_mu_x, + bool predict_tau_x, + bool predict_prognostic_function, + bool predict_cate, + bool predict_conditional_variance, + bool predict_random_effects) { + // Pre-convert test data to F-contiguous at function scope so the buffers outlive bcf_data + // and the predict_bcf_model call. convert_numpy_to_bcf_data casts inside if-blocks, so its + // temporaries are freed before predict_bcf_model runs -- these function-scope arrays prevent that. + using FArray = py::array_t; + using IArray = py::array_t; + FArray X_farr, Z_farr, obs_weights_farr, rfx_basis_farr; + IArray rfx_group_ids_iarr; + if (!X.is_none()) X_farr = X.cast(); + if (!Z.is_none()) Z_farr = Z.cast(); + if (!obs_weights.is_none()) obs_weights_farr = obs_weights.cast(); + if (!rfx_group_ids.is_none()) rfx_group_ids_iarr = rfx_group_ids.cast(); + if (!rfx_basis.is_none()) rfx_basis_farr = rfx_basis.cast(); + + // Unpack pointers to input data to BCFData object -- use only the "test" data fields as this is what the predict function expects + StochTree::BCFData bcf_data = convert_numpy_to_bcf_data( + /*X_train=*/py::none(), /*Z_train=*/py::none(), /*y_train=*/py::none(), + /*X_test=*/X.is_none() ? py::object(py::none()) : py::object(X_farr), + /*Z_test=*/Z.is_none() ? py::object(py::none()) : py::object(Z_farr), + /*n_train=*/0, /*n_test=*/n, /*p=*/p, /*treatment_dim=*/treatment_dim, + /*obs_weights_train=*/py::none(), + /*obs_weights_test=*/obs_weights.is_none() ? py::object(py::none()) : py::object(obs_weights_farr), + /*rfx_group_ids_train=*/py::none(), + /*rfx_group_ids_test=*/rfx_group_ids.is_none() ? py::object(py::none()) : py::object(rfx_group_ids_iarr), + /*rfx_basis_train=*/py::none(), + /*rfx_basis_test=*/rfx_basis.is_none() ? py::object(py::none()) : py::object(rfx_basis_farr), + /*rfx_num_groups=*/rfx_num_groups, /*rfx_basis_dim=*/rfx_basis_dim); + + // Load the BCF model and config from the model list + StochTree::BCFPredictionInput pred_input; + + py::array_t global_error_variance_array, leaf_mu_variance_array, leaf_tau_variance_array, b0_samples_array, b1_samples_array, tau_0_samples_array; + if (bcf_model_dict.contains("sigma2_global_samples") && !bcf_model_dict["sigma2_global_samples"].is_none()) { + global_error_variance_array = bcf_model_dict["sigma2_global_samples"].cast>(); + pred_input.global_error_variance_samples = static_cast(global_error_variance_array.mutable_data()); + } + if (bcf_model_dict.contains("sigma2_leaf_mu_samples") && !bcf_model_dict["sigma2_leaf_mu_samples"].is_none()) { + leaf_mu_variance_array = bcf_model_dict["sigma2_leaf_mu_samples"].cast>(); + pred_input.leaf_scale_mu_samples = static_cast(leaf_mu_variance_array.mutable_data()); + } + if (bcf_model_dict.contains("sigma2_leaf_tau_samples") && !bcf_model_dict["sigma2_leaf_tau_samples"].is_none()) { + leaf_tau_variance_array = bcf_model_dict["sigma2_leaf_tau_samples"].cast>(); + pred_input.leaf_scale_tau_samples = static_cast(leaf_tau_variance_array.mutable_data()); + } + if (bcf_model_dict.contains("b0_samples") && !bcf_model_dict["b0_samples"].is_none()) { + b0_samples_array = bcf_model_dict["b0_samples"].cast>(); + pred_input.b0_samples = static_cast(b0_samples_array.mutable_data()); + } + if (bcf_model_dict.contains("b1_samples") && !bcf_model_dict["b1_samples"].is_none()) { + b1_samples_array = bcf_model_dict["b1_samples"].cast>(); + pred_input.b1_samples = static_cast(b1_samples_array.mutable_data()); + } + if (bcf_model_dict.contains("tau_0_samples") && !bcf_model_dict["tau_0_samples"].is_none()) { + tau_0_samples_array = bcf_model_dict["tau_0_samples"].cast>(); + pred_input.tau_0_samples = static_cast(tau_0_samples_array.mutable_data()); + } + if (bcf_model_dict.contains("mu_forests") && !bcf_model_dict["mu_forests"].is_none()) { + pred_input.mu_forests = bcf_model_dict["mu_forests"].cast()->GetPtr(); + } + if (bcf_model_dict.contains("tau_forests") && !bcf_model_dict["tau_forests"].is_none()) { + pred_input.tau_forests = bcf_model_dict["tau_forests"].cast()->GetPtr(); + } + if (bcf_model_dict.contains("variance_forests") && !bcf_model_dict["variance_forests"].is_none()) { + pred_input.variance_forests = bcf_model_dict["variance_forests"].cast()->GetPtr(); + } + if (bcf_model_dict.contains("rfx_container") && !bcf_model_dict["rfx_container"].is_none()) { + pred_input.rfx_container = bcf_model_dict["rfx_container"].cast()->GetPtr(); + } + if (bcf_model_dict.contains("rfx_label_mapper") && !bcf_model_dict["rfx_label_mapper"].is_none()) { + pred_input.rfx_label_mapper = bcf_model_dict["rfx_label_mapper"].cast()->GetPtr(); + } + pred_input.num_samples = bcf_model_dict["num_samples"].cast(); + pred_input.num_obs = n; + pred_input.treatment_dim = treatment_dim; + pred_input.y_bar = bcf_model_dict["y_bar"].cast(); + pred_input.y_std = bcf_model_dict["y_std"].cast(); + pred_input.has_variance_forest = bcf_model_dict["include_variance_forest"].cast(); + pred_input.has_rfx = bcf_model_dict["has_rfx"].cast(); + std::string rfx_model_spec_str = ""; + if (bcf_model_dict.contains("rfx_model_spec") && !bcf_model_dict["rfx_model_spec"].is_none()) { + rfx_model_spec_str = bcf_model_dict["rfx_model_spec"].cast(); + } + if (rfx_model_spec_str == "intercept_only") { + pred_input.rfx_model_spec = StochTree::BCFRFXModelSpec::InterceptOnly; + } else if (rfx_model_spec_str == "intercept_plus_treatment") { + pred_input.rfx_model_spec = StochTree::BCFRFXModelSpec::InterceptPlusTreatment; + } else { + pred_input.rfx_model_spec = StochTree::BCFRFXModelSpec::Custom; + } + pred_input.adaptive_coding = bcf_model_dict["adaptive_coding"].cast(); + pred_input.sample_tau_0 = bcf_model_dict["sample_tau_0"].cast(); + pred_input.pred_type = posterior ? StochTree::PredType::kPosterior : StochTree::PredType::kMean; + if (scale == 0) { + pred_input.pred_scale = StochTree::PredScale::kLinear; + } else if (scale == 1) { + pred_input.pred_scale = StochTree::PredScale::kProbability; + } else { + pred_input.pred_scale = StochTree::PredScale::kClass; + } + pred_input.pred_terms.y_hat = predict_y_hat; + pred_input.pred_terms.mu_x = predict_mu_x; + pred_input.pred_terms.tau_x = predict_tau_x; + pred_input.pred_terms.prognostic_function = predict_prognostic_function; + pred_input.pred_terms.cate = predict_cate; + pred_input.pred_terms.conditional_variance = predict_conditional_variance; + pred_input.pred_terms.random_effects = predict_random_effects; + + // Run the prediction function + StochTree::BCFPredictionResult pred_results = predict_bcf_model(bcf_data, pred_input); + + // Unpack outputs + py::dict output = convert_bcf_preds_to_dict(pred_results); + return output; +} + py::array_t cppComputeForestContainerLeafIndices(ForestContainerCpp& forest_container, py::array_t& covariates, py::array_t& forest_nums) { // Wrap an Eigen Map around the raw data of the covariate matrix StochTree::data_size_t num_obs = covariates.shape(0); @@ -2153,9 +3918,9 @@ py::array_t cppComputeForestContainerLeafIndices(ForestContainerCpp& forest } // Compute leaf indices - auto result = py::array_t(py::detail::any_container({num_obs*num_trees, num_samples})); + auto result = py::array_t(py::detail::any_container({num_obs * num_trees, num_samples})); int* output_data_ptr = static_cast(result.mutable_data()); - Eigen::Map> output_eigen(output_data_ptr, num_obs*num_trees, num_samples); + Eigen::Map> output_eigen(output_data_ptr, num_obs * num_trees, num_samples); forest_container.GetContainer()->PredictLeafIndicesInplace(covariates_eigen, output_eigen, forest_indices, num_trees, num_obs); // Return matrix @@ -2180,9 +3945,11 @@ void ForestContainerCpp::AppendFromJson(JsonCpp& json, std::string forest_label) void ForestContainerCpp::AdjustResidual(ForestDatasetCpp& dataset, ResidualCpp& residual, ForestSamplerCpp& sampler, bool requires_basis, int forest_num, bool add) { // Determine whether or not we are adding forest_num to the residuals std::function op; - if (add) op = std::plus(); - else op = std::minus(); - + if (add) + op = std::plus(); + else + op = std::minus(); + // Perform the update (addition / subtraction) operation StochTree::UpdateResidualEntireForest(*(sampler.GetTracker()), *(dataset.GetDataset()), *(residual.GetData()), forest_samples_->GetEnsemble(forest_num), requires_basis, op); } @@ -2190,9 +3957,11 @@ void ForestContainerCpp::AdjustResidual(ForestDatasetCpp& dataset, ResidualCpp& void ForestCpp::AdjustResidual(ForestDatasetCpp& dataset, ResidualCpp& residual, ForestSamplerCpp& sampler, bool requires_basis, bool add) { // Determine whether or not we are adding forest predictions to the residuals std::function op; - if (add) op = std::plus(); - else op = std::minus(); - + if (add) + op = std::plus(); + else + op = std::minus(); + // Perform the update (addition / subtraction) operation StochTree::UpdateResidualEntireForest(*(sampler.GetTracker()), *(dataset.GetDataset()), *(residual.GetData()), forest_.get(), requires_basis, op); } @@ -2215,13 +3984,13 @@ void RandomEffectsContainerCpp::AddSample(RandomEffectsModelCpp& rfx_model) { py::array_t RandomEffectsContainerCpp::Predict(RandomEffectsDatasetCpp& rfx_dataset, RandomEffectsLabelMapperCpp& label_mapper) { py::ssize_t num_observations = rfx_dataset.NumObservations(); int num_samples = rfx_container_->NumSamples(); - std::vector output(num_observations*num_samples); + std::vector output(num_observations * num_samples); rfx_container_->Predict(*rfx_dataset.GetDataset(), *label_mapper.GetLabelMapper(), output); auto result = py::array_t(py::detail::any_container({num_observations, num_samples})); auto accessor = result.mutable_unchecked<2>(); for (size_t i = 0; i < num_observations; i++) { for (int j = 0; j < num_samples; j++) { - accessor(i, j) = output.at(j*num_observations + i); + accessor(i, j) = output.at(j * num_observations + i); } } return result; @@ -2233,10 +4002,10 @@ void RandomEffectsLabelMapperCpp::LoadFromJson(JsonCpp& json, std::string rfx_la rfx_label_mapper_->from_json(rfx_json); } -void RandomEffectsModelCpp::SampleRandomEffects(RandomEffectsDatasetCpp& rfx_dataset, ResidualCpp& residual, - RandomEffectsTrackerCpp& rfx_tracker, RandomEffectsContainerCpp& rfx_container, +void RandomEffectsModelCpp::SampleRandomEffects(RandomEffectsDatasetCpp& rfx_dataset, ResidualCpp& residual, + RandomEffectsTrackerCpp& rfx_tracker, RandomEffectsContainerCpp& rfx_container, bool keep_sample, double global_variance, RngCpp& rng) { - rfx_model_->SampleRandomEffects(*rfx_dataset.GetDataset(), *residual.GetData(), + rfx_model_->SampleRandomEffects(*rfx_dataset.GetDataset(), *residual.GetData(), *rfx_tracker.GetTracker(), global_variance, *rng.GetRng()); if (keep_sample) rfx_container.AddSample(*this); } @@ -2252,193 +4021,390 @@ void RandomEffectsTrackerCpp::RootReset(RandomEffectsModelCpp& rfx_model, Random PYBIND11_MODULE(stochtree_cpp, m) { m.def("cppComputeForestContainerLeafIndices", &cppComputeForestContainerLeafIndices, "Compute leaf indices of the forests in a forest container"); m.def("cppComputeForestMaxLeafIndex", &cppComputeForestMaxLeafIndex, "Compute max leaf index of a forest in a forest container"); + m.def("bart_sample_cpp", &bart_sample_cpp, "Run BART sampler in C++ implementation", + py::arg("X_train"), + py::arg("y_train"), + py::arg("X_test") = py::none(), + py::arg("n_train"), + py::arg("n_test"), + py::arg("p"), + py::arg("basis_train") = py::none(), + py::arg("basis_test") = py::none(), + py::arg("basis_dim"), + py::arg("obs_weights_train") = py::none(), + py::arg("obs_weights_test") = py::none(), + py::arg("rfx_group_ids_train") = py::none(), + py::arg("rfx_group_ids_test") = py::none(), + py::arg("rfx_basis_train") = py::none(), + py::arg("rfx_basis_test") = py::none(), + py::arg("rfx_num_groups"), + py::arg("rfx_basis_dim"), + py::arg("num_gfr"), + py::arg("num_burnin"), + py::arg("keep_every"), + py::arg("num_mcmc"), + py::arg("num_chains"), + py::arg("config_input")); + + m.def("bart_continue_sample_cpp", &bart_continue_sample_cpp, "Continue (warm-start) BART sampling from an existing model in C++", + py::arg("X_train"), + py::arg("y_train"), + py::arg("X_test") = py::none(), + py::arg("n_train"), + py::arg("n_test"), + py::arg("p"), + py::arg("basis_train") = py::none(), + py::arg("basis_test") = py::none(), + py::arg("basis_dim"), + py::arg("obs_weights_train") = py::none(), + py::arg("obs_weights_test") = py::none(), + py::arg("rfx_group_ids_train") = py::none(), + py::arg("rfx_group_ids_test") = py::none(), + py::arg("rfx_basis_train") = py::none(), + py::arg("rfx_basis_test") = py::none(), + py::arg("rfx_num_groups"), + py::arg("rfx_basis_dim"), + py::arg("num_burnin"), + py::arg("keep_every"), + py::arg("num_mcmc"), + py::arg("mean_forest_container"), + py::arg("global_var_samples") = py::none(), + py::arg("leaf_scale_samples") = py::none(), + py::arg("y_bar"), + py::arg("y_std"), + py::arg("rng_state_in") = std::string(), + py::arg("override_seed") = false, + py::arg("leaf_normal_cache_in") = std::string(), + py::arg("config_input")); + + m.def("bcf_sample_cpp", &bcf_sample_cpp, "Run BCF sampler in C++ implementation", + py::arg("X_train"), + py::arg("Z_train"), + py::arg("y_train"), + py::arg("X_test") = py::none(), + py::arg("Z_test") = py::none(), + py::arg("n_train"), + py::arg("n_test"), + py::arg("p"), + py::arg("treatment_dim"), + py::arg("obs_weights_train") = py::none(), + py::arg("obs_weights_test") = py::none(), + py::arg("rfx_group_ids_train") = py::none(), + py::arg("rfx_group_ids_test") = py::none(), + py::arg("rfx_basis_train") = py::none(), + py::arg("rfx_basis_test") = py::none(), + py::arg("rfx_num_groups"), + py::arg("rfx_basis_dim"), + py::arg("num_gfr"), + py::arg("num_burnin"), + py::arg("keep_every"), + py::arg("num_mcmc"), + py::arg("num_chains"), + py::arg("adaptive_coding"), + py::arg("config_input")); + + m.def("bcf_continue_sample_cpp", &bcf_continue_sample_cpp, "Continue (warm-start) BCF sampling from an existing model in C++", + py::arg("X_train"), + py::arg("Z_train"), + py::arg("y_train"), + py::arg("n_train"), + py::arg("p"), + py::arg("treatment_dim"), + py::arg("obs_weights_train") = py::none(), + py::arg("num_burnin"), + py::arg("keep_every"), + py::arg("num_mcmc"), + py::arg("mu_forest_container"), + py::arg("tau_forest_container"), + py::arg("global_var_samples") = py::none(), + py::arg("leaf_scale_mu_samples") = py::none(), + py::arg("leaf_scale_tau_samples") = py::none(), + py::arg("tau_0_samples") = py::none(), + py::arg("y_bar"), + py::arg("y_std"), + py::arg("rng_state_in") = std::string(), + py::arg("override_seed") = false, + py::arg("leaf_normal_cache_in") = std::string(), + py::arg("config_input")); + + m.def("bart_predict_cpp", &bart_predict_cpp, "Run BART predictions in C++", + py::arg("bart_model_dict"), + py::arg("X"), + py::arg("leaf_basis") = py::none(), + py::arg("n"), + py::arg("p"), + py::arg("num_basis") = 0, + py::arg("obs_weights") = py::none(), + py::arg("rfx_group_ids") = py::none(), + py::arg("rfx_basis") = py::none(), + py::arg("rfx_num_groups") = 0, + py::arg("rfx_basis_dim") = 0, + py::arg("posterior"), + py::arg("scale"), + py::arg("predict_y_hat"), + py::arg("predict_mean_forest"), + py::arg("predict_variance_forest"), + py::arg("predict_random_effects")); + + m.def("bcf_predict_cpp", &bcf_predict_cpp, "Run BCF predictions in C++", + py::arg("bcf_model_dict"), + py::arg("X"), + py::arg("Z"), + py::arg("n"), + py::arg("p"), + py::arg("treatment_dim"), + py::arg("obs_weights") = py::none(), + py::arg("rfx_group_ids") = py::none(), + py::arg("rfx_basis") = py::none(), + py::arg("rfx_num_groups"), + py::arg("rfx_basis_dim"), + py::arg("posterior"), + py::arg("scale"), + py::arg("predict_y_hat"), + py::arg("predict_mu_x"), + py::arg("predict_tau_x"), + py::arg("predict_prognostic_function"), + py::arg("predict_cate"), + py::arg("predict_conditional_variance"), + py::arg("predict_random_effects")); py::class_(m, "JsonCpp") - .def(py::init<>()) - .def("LoadFile", &JsonCpp::LoadFile) - .def("SaveFile", &JsonCpp::SaveFile) - .def("LoadFromString", &JsonCpp::LoadFromString) - .def("DumpJson", &JsonCpp::DumpJson) - .def("AddDouble", &JsonCpp::AddDouble) - .def("AddDoubleSubfolder", &JsonCpp::AddDoubleSubfolder) - .def("AddInteger", &JsonCpp::AddInteger) - .def("AddIntegerSubfolder", &JsonCpp::AddIntegerSubfolder) - .def("AddBool", &JsonCpp::AddBool) - .def("AddBoolSubfolder", &JsonCpp::AddBoolSubfolder) - .def("AddString", &JsonCpp::AddString) - .def("AddStringSubfolder", &JsonCpp::AddStringSubfolder) - .def("AddDoubleVector", &JsonCpp::AddDoubleVector) - .def("AddDoubleVectorSubfolder", &JsonCpp::AddDoubleVectorSubfolder) - .def("AddIntegerVector", &JsonCpp::AddIntegerVector) - .def("AddIntegerVectorSubfolder", &JsonCpp::AddIntegerVectorSubfolder) - .def("AddStringVector", &JsonCpp::AddStringVector) - .def("AddStringVectorSubfolder", &JsonCpp::AddStringVectorSubfolder) - .def("AddForest", &JsonCpp::AddForest) - .def("AddRandomEffectsContainer", &JsonCpp::AddRandomEffectsContainer) - .def("AddRandomEffectsLabelMapper", &JsonCpp::AddRandomEffectsLabelMapper) - .def("AddRandomEffectsGroupIDs", &JsonCpp::AddRandomEffectsGroupIDs) - .def("ContainsField", &JsonCpp::ContainsField) - .def("ContainsFieldSubfolder", &JsonCpp::ContainsFieldSubfolder) - .def("ExtractDouble", &JsonCpp::ExtractDouble) - .def("ExtractDoubleSubfolder", &JsonCpp::ExtractDoubleSubfolder) - .def("ExtractInteger", &JsonCpp::ExtractInteger) - .def("ExtractIntegerSubfolder", &JsonCpp::ExtractIntegerSubfolder) - .def("ExtractBool", &JsonCpp::ExtractBool) - .def("ExtractBoolSubfolder", &JsonCpp::ExtractBoolSubfolder) - .def("ExtractString", &JsonCpp::ExtractString) - .def("ExtractStringSubfolder", &JsonCpp::ExtractStringSubfolder) - .def("ExtractDoubleVector", &JsonCpp::ExtractDoubleVector) - .def("ExtractDoubleVectorSubfolder", &JsonCpp::ExtractDoubleVectorSubfolder) - .def("ExtractIntegerVector", &JsonCpp::ExtractIntegerVector) - .def("ExtractIntegerVectorSubfolder", &JsonCpp::ExtractIntegerVectorSubfolder) - .def("ExtractStringVector", &JsonCpp::ExtractStringVector) - .def("ExtractStringVectorSubfolder", &JsonCpp::ExtractStringVectorSubfolder) - .def("IncrementRandomEffectsCount", &JsonCpp::IncrementRandomEffectsCount) - .def("SubsetJsonForest", &JsonCpp::SubsetJsonForest) - .def("SubsetJsonRFX", &JsonCpp::SubsetJsonRFX); + .def(py::init<>()) + .def("LoadFile", &JsonCpp::LoadFile) + .def("SaveFile", &JsonCpp::SaveFile) + .def("LoadFromString", &JsonCpp::LoadFromString) + .def("DumpJson", &JsonCpp::DumpJson) + .def("AddDouble", &JsonCpp::AddDouble) + .def("AddDoubleSubfolder", &JsonCpp::AddDoubleSubfolder) + .def("AddInteger", &JsonCpp::AddInteger) + .def("AddIntegerSubfolder", &JsonCpp::AddIntegerSubfolder) + .def("AddBool", &JsonCpp::AddBool) + .def("AddBoolSubfolder", &JsonCpp::AddBoolSubfolder) + .def("AddString", &JsonCpp::AddString) + .def("AddStringSubfolder", &JsonCpp::AddStringSubfolder) + .def("AddDoubleVector", &JsonCpp::AddDoubleVector) + .def("AddDoubleVectorSubfolder", &JsonCpp::AddDoubleVectorSubfolder) + .def("AddIntegerVector", &JsonCpp::AddIntegerVector) + .def("AddIntegerVectorSubfolder", &JsonCpp::AddIntegerVectorSubfolder) + .def("AddStringVector", &JsonCpp::AddStringVector) + .def("AddStringVectorSubfolder", &JsonCpp::AddStringVectorSubfolder) + .def("AddForest", &JsonCpp::AddForest, py::arg("forest_samples"), py::arg("forest_label") = std::string("")) + .def("AddRandomEffectsContainer", &JsonCpp::AddRandomEffectsContainer) + .def("AddRandomEffectsLabelMapper", &JsonCpp::AddRandomEffectsLabelMapper) + .def("AddRandomEffectsGroupIDs", &JsonCpp::AddRandomEffectsGroupIDs) + .def("ContainsField", &JsonCpp::ContainsField) + .def("ContainsFieldSubfolder", &JsonCpp::ContainsFieldSubfolder) + .def("EraseField", &JsonCpp::EraseField) + .def("EraseFieldSubfolder", &JsonCpp::EraseFieldSubfolder) + .def("RenameField", &JsonCpp::RenameField) + .def("RenameFieldSubfolder", &JsonCpp::RenameFieldSubfolder) + .def("ExtractDouble", &JsonCpp::ExtractDouble) + .def("ExtractDoubleSubfolder", &JsonCpp::ExtractDoubleSubfolder) + .def("ExtractInteger", &JsonCpp::ExtractInteger) + .def("ExtractIntegerSubfolder", &JsonCpp::ExtractIntegerSubfolder) + .def("ExtractBool", &JsonCpp::ExtractBool) + .def("ExtractBoolSubfolder", &JsonCpp::ExtractBoolSubfolder) + .def("ExtractString", &JsonCpp::ExtractString) + .def("ExtractStringSubfolder", &JsonCpp::ExtractStringSubfolder) + .def("ExtractDoubleVector", &JsonCpp::ExtractDoubleVector) + .def("ExtractDoubleVectorSubfolder", &JsonCpp::ExtractDoubleVectorSubfolder) + .def("ExtractIntegerVector", &JsonCpp::ExtractIntegerVector) + .def("ExtractIntegerVectorSubfolder", &JsonCpp::ExtractIntegerVectorSubfolder) + .def("ExtractStringVector", &JsonCpp::ExtractStringVector) + .def("ExtractStringVectorSubfolder", &JsonCpp::ExtractStringVectorSubfolder) + .def("IncrementRandomEffectsCount", &JsonCpp::IncrementRandomEffectsCount) + .def("SubsetJsonForest", &JsonCpp::SubsetJsonForest) + .def("SubsetJsonRFX", &JsonCpp::SubsetJsonRFX); py::class_(m, "ForestDatasetCpp") - .def(py::init<>()) - .def("AddCovariates", &ForestDatasetCpp::AddCovariates) - .def("AddBasis", &ForestDatasetCpp::AddBasis) - .def("UpdateBasis", &ForestDatasetCpp::UpdateBasis) - .def("AddVarianceWeights", &ForestDatasetCpp::AddVarianceWeights) - .def("UpdateVarianceWeights", &ForestDatasetCpp::UpdateVarianceWeights) - .def("NumRows", &ForestDatasetCpp::NumRows) - .def("NumCovariates", &ForestDatasetCpp::NumCovariates) - .def("NumBasis", &ForestDatasetCpp::NumBasis) - .def("GetCovariates", &ForestDatasetCpp::GetCovariates) - .def("GetBasis", &ForestDatasetCpp::GetBasis) - .def("GetVarianceWeights", &ForestDatasetCpp::GetVarianceWeights) - .def("HasBasis", &ForestDatasetCpp::HasBasis) - .def("HasVarianceWeights", &ForestDatasetCpp::HasVarianceWeights) - .def("AddAuxiliaryDimension", &ForestDatasetCpp::AddAuxiliaryDimension) - .def("SetAuxiliaryDataValue", &ForestDatasetCpp::SetAuxiliaryDataValue) - .def("GetAuxiliaryDataValue", &ForestDatasetCpp::GetAuxiliaryDataValue) - .def("GetAuxiliaryDataVector", &ForestDatasetCpp::GetAuxiliaryDataVector); + .def(py::init<>()) + .def("AddCovariates", &ForestDatasetCpp::AddCovariates) + .def("AddBasis", &ForestDatasetCpp::AddBasis) + .def("UpdateBasis", &ForestDatasetCpp::UpdateBasis) + .def("AddVarianceWeights", &ForestDatasetCpp::AddVarianceWeights) + .def("UpdateVarianceWeights", &ForestDatasetCpp::UpdateVarianceWeights) + .def("NumRows", &ForestDatasetCpp::NumRows) + .def("NumCovariates", &ForestDatasetCpp::NumCovariates) + .def("NumBasis", &ForestDatasetCpp::NumBasis) + .def("GetCovariates", &ForestDatasetCpp::GetCovariates) + .def("GetBasis", &ForestDatasetCpp::GetBasis) + .def("GetVarianceWeights", &ForestDatasetCpp::GetVarianceWeights) + .def("HasBasis", &ForestDatasetCpp::HasBasis) + .def("HasVarianceWeights", &ForestDatasetCpp::HasVarianceWeights) + .def("AddAuxiliaryDimension", &ForestDatasetCpp::AddAuxiliaryDimension) + .def("SetAuxiliaryDataValue", &ForestDatasetCpp::SetAuxiliaryDataValue) + .def("GetAuxiliaryDataValue", &ForestDatasetCpp::GetAuxiliaryDataValue) + .def("GetAuxiliaryDataVector", &ForestDatasetCpp::GetAuxiliaryDataVector); py::class_(m, "ResidualCpp") - .def(py::init,data_size_t>()) - .def("GetResidualArray", &ResidualCpp::GetResidualArray) - .def("ReplaceData", &ResidualCpp::ReplaceData) - .def("AddToData", &ResidualCpp::AddToData) - .def("SubtractFromData", &ResidualCpp::SubtractFromData); + .def(py::init, data_size_t>()) + .def("GetResidualArray", &ResidualCpp::GetResidualArray) + .def("ReplaceData", &ResidualCpp::ReplaceData) + .def("AddToData", &ResidualCpp::AddToData) + .def("SubtractFromData", &ResidualCpp::SubtractFromData); py::class_(m, "RngCpp") - .def(py::init()); - + .def(py::init()); + + py::class_(m, "BARTSamplesCpp") + .def(py::init<>()) + .def_static("from_json_string", &BARTSamplesCpp::FromJsonString) + .def_static("from_components", &BARTSamplesCpp::FromComponents) + .def("to_json_string", &BARTSamplesCpp::ToJsonString) + .def("load_from_json", &BARTSamplesCpp::LoadFromJson) + .def("add_to_json", &BARTSamplesCpp::AddToJson) + .def("merge", &BARTSamplesCpp::Merge) + .def("num_samples", &BARTSamplesCpp::NumSamples) + .def("y_bar", &BARTSamplesCpp::YBar) + .def("y_std", &BARTSamplesCpp::YStd) + .def("num_train", &BARTSamplesCpp::NumTrain) + .def("num_test", &BARTSamplesCpp::NumTest) + .def("global_var_samples", &BARTSamplesCpp::GlobalVarSamples) + .def("leaf_scale_samples", &BARTSamplesCpp::LeafScaleSamples) + .def("has_mean_forest", &BARTSamplesCpp::HasMeanForest) + .def("has_variance_forest", &BARTSamplesCpp::HasVarianceForest) + .def("materialize_mean_forest", &BARTSamplesCpp::MaterializeMeanForest) + .def("materialize_variance_forest", &BARTSamplesCpp::MaterializeVarianceForest); + + py::class_(m, "BCFSamplesCpp") + .def(py::init<>()) + .def_static("from_json_string", &BCFSamplesCpp::FromJsonString) + .def_static("from_components", &BCFSamplesCpp::FromComponents) + .def("to_json_string", &BCFSamplesCpp::ToJsonString) + .def("load_from_json", &BCFSamplesCpp::LoadFromJson) + .def("add_to_json", &BCFSamplesCpp::AddToJson) + .def("merge", &BCFSamplesCpp::Merge) + .def("num_samples", &BCFSamplesCpp::NumSamples) + .def("y_bar", &BCFSamplesCpp::YBar) + .def("y_std", &BCFSamplesCpp::YStd) + .def("num_train", &BCFSamplesCpp::NumTrain) + .def("num_test", &BCFSamplesCpp::NumTest) + .def("treatment_dim", &BCFSamplesCpp::TreatmentDim) + .def("global_var_samples", &BCFSamplesCpp::GlobalVarSamples) + .def("leaf_scale_mu_samples", &BCFSamplesCpp::LeafScaleMuSamples) + .def("leaf_scale_tau_samples", &BCFSamplesCpp::LeafScaleTauSamples) + .def("tau_0_samples", &BCFSamplesCpp::Tau0Samples) + .def("b0_samples", &BCFSamplesCpp::B0Samples) + .def("b1_samples", &BCFSamplesCpp::B1Samples) + .def("has_mu_forest", &BCFSamplesCpp::HasMuForest) + .def("has_tau_forest", &BCFSamplesCpp::HasTauForest) + .def("has_variance_forest", &BCFSamplesCpp::HasVarianceForest) + .def("materialize_mu_forest", &BCFSamplesCpp::MaterializeMuForest) + .def("materialize_tau_forest", &BCFSamplesCpp::MaterializeTauForest) + .def("materialize_variance_forest", &BCFSamplesCpp::MaterializeVarianceForest); + py::class_(m, "ForestContainerCpp") - .def(py::init()) - .def("CombineForests", &ForestContainerCpp::CombineForests) - .def("AddToForest", &ForestContainerCpp::AddToForest) - .def("MultiplyForest", &ForestContainerCpp::MultiplyForest) - .def("OutputDimension", &ForestContainerCpp::OutputDimension) - .def("NumTrees", &ForestContainerCpp::NumTrees) - .def("NumSamples", &ForestContainerCpp::NumSamples) - .def("DeleteSample", &ForestContainerCpp::DeleteSample) - .def("Predict", &ForestContainerCpp::Predict) - .def("PredictRaw", &ForestContainerCpp::PredictRaw) - .def("PredictRawSingleForest", &ForestContainerCpp::PredictRawSingleForest) - .def("SetRootValue", &ForestContainerCpp::SetRootValue) - .def("SetRootVector", &ForestContainerCpp::SetRootVector) - .def("AdjustResidual", &ForestContainerCpp::AdjustResidual) - .def("SaveToJsonFile", &ForestContainerCpp::SaveToJsonFile) - .def("LoadFromJsonFile", &ForestContainerCpp::LoadFromJsonFile) - .def("LoadFromJson", &ForestContainerCpp::LoadFromJson) - .def("AppendFromJson", &ForestContainerCpp::AppendFromJson) - .def("DumpJsonString", &ForestContainerCpp::DumpJsonString) - .def("LoadFromJsonString", &ForestContainerCpp::LoadFromJsonString) - .def("AddSampleValue", &ForestContainerCpp::AddSampleValue) - .def("AddSampleVector", &ForestContainerCpp::AddSampleVector) - .def("AddNumericSplitValue", &ForestContainerCpp::AddNumericSplitValue) - .def("AddNumericSplitVector", &ForestContainerCpp::AddNumericSplitVector) - .def("GetTreeLeaves", &ForestContainerCpp::GetTreeLeaves) - .def("GetTreeSplitCounts", &ForestContainerCpp::GetTreeSplitCounts) - .def("GetForestSplitCounts", &ForestContainerCpp::GetForestSplitCounts) - .def("GetOverallSplitCounts", &ForestContainerCpp::GetOverallSplitCounts) - .def("GetGranularSplitCounts", &ForestContainerCpp::GetGranularSplitCounts) - .def("NumLeavesForest", &ForestContainerCpp::NumLeavesForest) - .def("SumLeafSquared", &ForestContainerCpp::SumLeafSquared) - .def("IsLeafNode", &ForestContainerCpp::IsLeafNode) - .def("IsNumericSplitNode", &ForestContainerCpp::IsNumericSplitNode) - .def("IsCategoricalSplitNode", &ForestContainerCpp::IsCategoricalSplitNode) - .def("ParentNode", &ForestContainerCpp::ParentNode) - .def("LeftChildNode", &ForestContainerCpp::LeftChildNode) - .def("RightChildNode", &ForestContainerCpp::RightChildNode) - .def("SplitIndex", &ForestContainerCpp::SplitIndex) - .def("NodeDepth", &ForestContainerCpp::NodeDepth) - .def("SplitThreshold", &ForestContainerCpp::SplitThreshold) - .def("SplitCategories", &ForestContainerCpp::SplitCategories) - .def("NodeLeafValues", &ForestContainerCpp::NodeLeafValues) - .def("NumNodes", &ForestContainerCpp::NumNodes) - .def("NumLeaves", &ForestContainerCpp::NumLeaves) - .def("NumLeafParents", &ForestContainerCpp::NumLeafParents) - .def("NumSplitNodes", &ForestContainerCpp::NumSplitNodes) - .def("Nodes", &ForestContainerCpp::Nodes) - .def("Leaves", &ForestContainerCpp::Leaves); + .def(py::init()) + .def("CombineForests", &ForestContainerCpp::CombineForests) + .def("AddToForest", &ForestContainerCpp::AddToForest) + .def("MultiplyForest", &ForestContainerCpp::MultiplyForest) + .def("OutputDimension", &ForestContainerCpp::OutputDimension) + .def("NumTrees", &ForestContainerCpp::NumTrees) + .def("NumSamples", &ForestContainerCpp::NumSamples) + .def("DeleteSample", &ForestContainerCpp::DeleteSample) + .def("Predict", &ForestContainerCpp::Predict) + .def("PredictRaw", &ForestContainerCpp::PredictRaw) + .def("PredictRawSingleForest", &ForestContainerCpp::PredictRawSingleForest) + .def("SetRootValue", &ForestContainerCpp::SetRootValue) + .def("SetRootVector", &ForestContainerCpp::SetRootVector) + .def("AdjustResidual", &ForestContainerCpp::AdjustResidual) + .def("SaveToJsonFile", &ForestContainerCpp::SaveToJsonFile) + .def("LoadFromJsonFile", &ForestContainerCpp::LoadFromJsonFile) + .def("LoadFromJson", &ForestContainerCpp::LoadFromJson) + .def("AppendFromJson", &ForestContainerCpp::AppendFromJson) + .def("DumpJsonString", &ForestContainerCpp::DumpJsonString) + .def("LoadFromJsonString", &ForestContainerCpp::LoadFromJsonString) + .def("AddSampleValue", &ForestContainerCpp::AddSampleValue) + .def("AddSampleVector", &ForestContainerCpp::AddSampleVector) + .def("AddNumericSplitValue", &ForestContainerCpp::AddNumericSplitValue) + .def("AddNumericSplitVector", &ForestContainerCpp::AddNumericSplitVector) + .def("GetTreeLeaves", &ForestContainerCpp::GetTreeLeaves) + .def("GetTreeSplitCounts", &ForestContainerCpp::GetTreeSplitCounts) + .def("GetForestSplitCounts", &ForestContainerCpp::GetForestSplitCounts) + .def("GetOverallSplitCounts", &ForestContainerCpp::GetOverallSplitCounts) + .def("GetGranularSplitCounts", &ForestContainerCpp::GetGranularSplitCounts) + .def("NumLeavesForest", &ForestContainerCpp::NumLeavesForest) + .def("SumLeafSquared", &ForestContainerCpp::SumLeafSquared) + .def("IsLeafNode", &ForestContainerCpp::IsLeafNode) + .def("IsNumericSplitNode", &ForestContainerCpp::IsNumericSplitNode) + .def("IsCategoricalSplitNode", &ForestContainerCpp::IsCategoricalSplitNode) + .def("ParentNode", &ForestContainerCpp::ParentNode) + .def("LeftChildNode", &ForestContainerCpp::LeftChildNode) + .def("RightChildNode", &ForestContainerCpp::RightChildNode) + .def("SplitIndex", &ForestContainerCpp::SplitIndex) + .def("NodeDepth", &ForestContainerCpp::NodeDepth) + .def("SplitThreshold", &ForestContainerCpp::SplitThreshold) + .def("SplitCategories", &ForestContainerCpp::SplitCategories) + .def("NodeLeafValues", &ForestContainerCpp::NodeLeafValues) + .def("NumNodes", &ForestContainerCpp::NumNodes) + .def("NumLeaves", &ForestContainerCpp::NumLeaves) + .def("NumLeafParents", &ForestContainerCpp::NumLeafParents) + .def("NumSplitNodes", &ForestContainerCpp::NumSplitNodes) + .def("Nodes", &ForestContainerCpp::Nodes) + .def("Leaves", &ForestContainerCpp::Leaves); py::class_(m, "ForestCpp") - .def(py::init()) - .def("GetForestPtr", &ForestCpp::GetForestPtr) - .def("MergeForest", &ForestCpp::MergeForest) - .def("AddConstant", &ForestCpp::AddConstant) - .def("MultiplyConstant", &ForestCpp::MultiplyConstant) - .def("OutputDimension", &ForestCpp::OutputDimension) - .def("NumTrees", &ForestCpp::NumTrees) - .def("NumLeavesForest", &ForestCpp::NumLeavesForest) - .def("SumLeafSquared", &ForestCpp::SumLeafSquared) - .def("ResetRoot", &ForestCpp::ResetRoot) - .def("Reset", &ForestCpp::Reset) - .def("Predict", &ForestCpp::Predict) - .def("PredictRaw", &ForestCpp::PredictRaw) - .def("SetRootValue", &ForestCpp::SetRootValue) - .def("SetRootVector", &ForestCpp::SetRootVector) - .def("AdjustResidual", &ForestCpp::AdjustResidual) - .def("AddNumericSplitValue", &ForestCpp::AddNumericSplitValue) - .def("AddNumericSplitVector", &ForestCpp::AddNumericSplitVector) - .def("GetEnsemble", &ForestCpp::GetEnsemble) - .def("GetTreeLeaves", &ForestCpp::GetTreeLeaves) - .def("GetTreeSplitCounts", &ForestCpp::GetTreeSplitCounts) - .def("GetOverallSplitCounts", &ForestCpp::GetOverallSplitCounts) - .def("GetGranularSplitCounts", &ForestCpp::GetGranularSplitCounts) - .def("NumLeavesForest", &ForestCpp::NumLeavesForest) - .def("SumLeafSquared", &ForestCpp::SumLeafSquared) - .def("IsLeafNode", &ForestCpp::IsLeafNode) - .def("IsNumericSplitNode", &ForestCpp::IsNumericSplitNode) - .def("IsCategoricalSplitNode", &ForestCpp::IsCategoricalSplitNode) - .def("ParentNode", &ForestCpp::ParentNode) - .def("LeftChildNode", &ForestCpp::LeftChildNode) - .def("RightChildNode", &ForestCpp::RightChildNode) - .def("SplitIndex", &ForestCpp::SplitIndex) - .def("NodeDepth", &ForestCpp::NodeDepth) - .def("SplitThreshold", &ForestCpp::SplitThreshold) - .def("SplitCategories", &ForestCpp::SplitCategories) - .def("NodeLeafValues", &ForestCpp::NodeLeafValues) - .def("NumNodes", &ForestCpp::NumNodes) - .def("NumLeaves", &ForestCpp::NumLeaves) - .def("NumLeafParents", &ForestCpp::NumLeafParents) - .def("NumSplitNodes", &ForestCpp::NumSplitNodes) - .def("Nodes", &ForestCpp::Nodes) - .def("Leaves", &ForestCpp::Leaves); - + .def(py::init()) + .def("GetForestPtr", &ForestCpp::GetForestPtr) + .def("MergeForest", &ForestCpp::MergeForest) + .def("AddConstant", &ForestCpp::AddConstant) + .def("MultiplyConstant", &ForestCpp::MultiplyConstant) + .def("OutputDimension", &ForestCpp::OutputDimension) + .def("NumTrees", &ForestCpp::NumTrees) + .def("NumLeavesForest", &ForestCpp::NumLeavesForest) + .def("SumLeafSquared", &ForestCpp::SumLeafSquared) + .def("ResetRoot", &ForestCpp::ResetRoot) + .def("Reset", &ForestCpp::Reset) + .def("Predict", &ForestCpp::Predict) + .def("PredictRaw", &ForestCpp::PredictRaw) + .def("SetRootValue", &ForestCpp::SetRootValue) + .def("SetRootVector", &ForestCpp::SetRootVector) + .def("AdjustResidual", &ForestCpp::AdjustResidual) + .def("AddNumericSplitValue", &ForestCpp::AddNumericSplitValue) + .def("AddNumericSplitVector", &ForestCpp::AddNumericSplitVector) + .def("GetEnsemble", &ForestCpp::GetEnsemble) + .def("GetTreeLeaves", &ForestCpp::GetTreeLeaves) + .def("GetTreeSplitCounts", &ForestCpp::GetTreeSplitCounts) + .def("GetOverallSplitCounts", &ForestCpp::GetOverallSplitCounts) + .def("GetGranularSplitCounts", &ForestCpp::GetGranularSplitCounts) + .def("NumLeavesForest", &ForestCpp::NumLeavesForest) + .def("SumLeafSquared", &ForestCpp::SumLeafSquared) + .def("IsLeafNode", &ForestCpp::IsLeafNode) + .def("IsNumericSplitNode", &ForestCpp::IsNumericSplitNode) + .def("IsCategoricalSplitNode", &ForestCpp::IsCategoricalSplitNode) + .def("ParentNode", &ForestCpp::ParentNode) + .def("LeftChildNode", &ForestCpp::LeftChildNode) + .def("RightChildNode", &ForestCpp::RightChildNode) + .def("SplitIndex", &ForestCpp::SplitIndex) + .def("NodeDepth", &ForestCpp::NodeDepth) + .def("SplitThreshold", &ForestCpp::SplitThreshold) + .def("SplitCategories", &ForestCpp::SplitCategories) + .def("NodeLeafValues", &ForestCpp::NodeLeafValues) + .def("NumNodes", &ForestCpp::NumNodes) + .def("NumLeaves", &ForestCpp::NumLeaves) + .def("NumLeafParents", &ForestCpp::NumLeafParents) + .def("NumSplitNodes", &ForestCpp::NumSplitNodes) + .def("Nodes", &ForestCpp::Nodes) + .def("Leaves", &ForestCpp::Leaves); + py::class_(m, "ForestSamplerCpp") - .def(py::init, int, data_size_t, double, double, int, int>()) - .def("ReconstituteTrackerFromForest", &ForestSamplerCpp::ReconstituteTrackerFromForest) - .def("SampleOneIteration", &ForestSamplerCpp::SampleOneIteration) - .def("InitializeForestModel", &ForestSamplerCpp::InitializeForestModel) - .def("GetCachedForestPredictions", &ForestSamplerCpp::GetCachedForestPredictions) - .def("PropagateBasisUpdate", &ForestSamplerCpp::PropagateBasisUpdate) - .def("PropagateResidualUpdate", &ForestSamplerCpp::PropagateResidualUpdate) - .def("UpdateAlpha", &ForestSamplerCpp::UpdateAlpha) - .def("UpdateBeta", &ForestSamplerCpp::UpdateBeta) - .def("UpdateMinSamplesLeaf", &ForestSamplerCpp::UpdateMinSamplesLeaf) - .def("UpdateMaxDepth", &ForestSamplerCpp::UpdateMaxDepth) - .def("GetAlpha", &ForestSamplerCpp::GetAlpha) - .def("GetBeta", &ForestSamplerCpp::GetBeta) - .def("GetMinSamplesLeaf", &ForestSamplerCpp::GetMinSamplesLeaf) - .def("GetMaxDepth", &ForestSamplerCpp::GetMaxDepth); - - py::class_(m, "RandomEffectsDatasetCpp") + .def(py::init, int, data_size_t, double, double, int, int>()) + .def("ReconstituteTrackerFromForest", &ForestSamplerCpp::ReconstituteTrackerFromForest) + .def("SampleOneIteration", &ForestSamplerCpp::SampleOneIteration) + .def("InitializeForestModel", &ForestSamplerCpp::InitializeForestModel) + .def("GetCachedForestPredictions", &ForestSamplerCpp::GetCachedForestPredictions) + .def("PropagateBasisUpdate", &ForestSamplerCpp::PropagateBasisUpdate) + .def("PropagateResidualUpdate", &ForestSamplerCpp::PropagateResidualUpdate) + .def("UpdateAlpha", &ForestSamplerCpp::UpdateAlpha) + .def("UpdateBeta", &ForestSamplerCpp::UpdateBeta) + .def("UpdateMinSamplesLeaf", &ForestSamplerCpp::UpdateMinSamplesLeaf) + .def("UpdateMaxDepth", &ForestSamplerCpp::UpdateMaxDepth) + .def("GetAlpha", &ForestSamplerCpp::GetAlpha) + .def("GetBeta", &ForestSamplerCpp::GetBeta) + .def("GetMinSamplesLeaf", &ForestSamplerCpp::GetMinSamplesLeaf) + .def("GetMaxDepth", &ForestSamplerCpp::GetMaxDepth); + + py::class_(m, "RandomEffectsDatasetCpp") .def(py::init<>()) .def("GetDataset", &RandomEffectsDatasetCpp::GetDataset) .def("NumObservations", &RandomEffectsDatasetCpp::NumObservations) @@ -2457,71 +4423,73 @@ PYBIND11_MODULE(stochtree_cpp, m) { .def("HasVarianceWeights", &RandomEffectsDatasetCpp::HasVarianceWeights); py::class_(m, "RandomEffectsContainerCpp") - .def(py::init<>()) - .def("SetComponentsAndGroups", &RandomEffectsContainerCpp::SetComponentsAndGroups) - .def("AddSample", &RandomEffectsContainerCpp::AddSample) - .def("NumSamples", &RandomEffectsContainerCpp::NumSamples) - .def("NumComponents", &RandomEffectsContainerCpp::NumComponents) - .def("NumGroups", &RandomEffectsContainerCpp::NumGroups) - .def("GetBeta", &RandomEffectsContainerCpp::GetBeta) - .def("GetXi", &RandomEffectsContainerCpp::GetXi) - .def("GetAlpha", &RandomEffectsContainerCpp::GetAlpha) - .def("GetSigma", &RandomEffectsContainerCpp::GetSigma) - .def("DeleteSample", &RandomEffectsContainerCpp::DeleteSample) - .def("Predict", &RandomEffectsContainerCpp::Predict) - .def("SaveToJsonFile", &RandomEffectsContainerCpp::SaveToJsonFile) - .def("LoadFromJsonFile", &RandomEffectsContainerCpp::LoadFromJsonFile) - .def("DumpJsonString", &RandomEffectsContainerCpp::DumpJsonString) - .def("LoadFromJsonString", &RandomEffectsContainerCpp::LoadFromJsonString) - .def("LoadFromJson", &RandomEffectsContainerCpp::LoadFromJson) - .def("AppendFromJson", &RandomEffectsContainerCpp::AppendFromJson) - .def("GetRandomEffectsContainer", &RandomEffectsContainerCpp::GetRandomEffectsContainer); + .def(py::init<>()) + .def("SetComponentsAndGroups", &RandomEffectsContainerCpp::SetComponentsAndGroups) + .def("AddSample", &RandomEffectsContainerCpp::AddSample) + .def("NumSamples", &RandomEffectsContainerCpp::NumSamples) + .def("NumComponents", &RandomEffectsContainerCpp::NumComponents) + .def("NumGroups", &RandomEffectsContainerCpp::NumGroups) + .def("GetBeta", &RandomEffectsContainerCpp::GetBeta) + .def("GetXi", &RandomEffectsContainerCpp::GetXi) + .def("GetAlpha", &RandomEffectsContainerCpp::GetAlpha) + .def("GetSigma", &RandomEffectsContainerCpp::GetSigma) + .def("DeleteSample", &RandomEffectsContainerCpp::DeleteSample) + .def("Predict", &RandomEffectsContainerCpp::Predict) + .def("SaveToJsonFile", &RandomEffectsContainerCpp::SaveToJsonFile) + .def("LoadFromJsonFile", &RandomEffectsContainerCpp::LoadFromJsonFile) + .def("DumpJsonString", &RandomEffectsContainerCpp::DumpJsonString) + .def("LoadFromJsonString", &RandomEffectsContainerCpp::LoadFromJsonString) + .def("LoadFromJson", &RandomEffectsContainerCpp::LoadFromJson) + .def("AppendFromJson", &RandomEffectsContainerCpp::AppendFromJson) + .def("GetRandomEffectsContainer", &RandomEffectsContainerCpp::GetRandomEffectsContainer); py::class_(m, "RandomEffectsTrackerCpp") - .def(py::init>()) - .def("GetUniqueGroupIds", &RandomEffectsTrackerCpp::GetUniqueGroupIds) - .def("GetTracker", &RandomEffectsTrackerCpp::GetTracker) - .def("Reset", &RandomEffectsTrackerCpp::Reset) - .def("RootReset", &RandomEffectsTrackerCpp::RootReset); + .def(py::init>()) + .def("GetUniqueGroupIds", &RandomEffectsTrackerCpp::GetUniqueGroupIds) + .def("GetTracker", &RandomEffectsTrackerCpp::GetTracker) + .def("Reset", &RandomEffectsTrackerCpp::Reset) + .def("RootReset", &RandomEffectsTrackerCpp::RootReset); py::class_(m, "RandomEffectsLabelMapperCpp") - .def(py::init<>()) - .def("LoadFromTracker", &RandomEffectsLabelMapperCpp::LoadFromTracker) - .def("SaveToJsonFile", &RandomEffectsLabelMapperCpp::SaveToJsonFile) - .def("LoadFromJsonFile", &RandomEffectsLabelMapperCpp::LoadFromJsonFile) - .def("DumpJsonString", &RandomEffectsLabelMapperCpp::DumpJsonString) - .def("LoadFromJsonString", &RandomEffectsLabelMapperCpp::LoadFromJsonString) - .def("LoadFromJson", &RandomEffectsLabelMapperCpp::LoadFromJson) - .def("GetLabelMapper", &RandomEffectsLabelMapperCpp::GetLabelMapper) - .def("MapGroupIdToArrayIndex", &RandomEffectsLabelMapperCpp::MapGroupIdToArrayIndex) - .def("MapMultipleGroupIdsToArrayIndices", &RandomEffectsLabelMapperCpp::MapMultipleGroupIdsToArrayIndices); + .def(py::init<>()) + .def("LoadFromTracker", &RandomEffectsLabelMapperCpp::LoadFromTracker) + .def("SaveToJsonFile", &RandomEffectsLabelMapperCpp::SaveToJsonFile) + .def("LoadFromJsonFile", &RandomEffectsLabelMapperCpp::LoadFromJsonFile) + .def("DumpJsonString", &RandomEffectsLabelMapperCpp::DumpJsonString) + .def("LoadFromJsonString", &RandomEffectsLabelMapperCpp::LoadFromJsonString) + .def("LoadFromJson", &RandomEffectsLabelMapperCpp::LoadFromJson) + .def("GetLabelMapper", &RandomEffectsLabelMapperCpp::GetLabelMapper) + .def("MapGroupIdToArrayIndex", &RandomEffectsLabelMapperCpp::MapGroupIdToArrayIndex) + .def("MapMultipleGroupIdsToArrayIndices", &RandomEffectsLabelMapperCpp::MapMultipleGroupIdsToArrayIndices) + .def("GetUniqueGroupIds", &RandomEffectsLabelMapperCpp::GetUniqueGroupIds); py::class_(m, "RandomEffectsModelCpp") - .def(py::init()) - .def("GetModel", &RandomEffectsModelCpp::GetModel) - .def("SampleRandomEffects", &RandomEffectsModelCpp::SampleRandomEffects) - .def("Predict", &RandomEffectsModelCpp::Predict) - .def("SetWorkingParameter", &RandomEffectsModelCpp::SetWorkingParameter) - .def("SetGroupParameters", &RandomEffectsModelCpp::SetGroupParameters) - .def("SetWorkingParameterCovariance", &RandomEffectsModelCpp::SetWorkingParameterCovariance) - .def("SetGroupParameterCovariance", &RandomEffectsModelCpp::SetGroupParameterCovariance) - .def("SetVariancePriorShape", &RandomEffectsModelCpp::SetVariancePriorShape) - .def("SetVariancePriorScale", &RandomEffectsModelCpp::SetVariancePriorScale) - .def("Reset", &RandomEffectsModelCpp::Reset); + .def(py::init()) + .def("GetModel", &RandomEffectsModelCpp::GetModel) + .def("SampleRandomEffects", &RandomEffectsModelCpp::SampleRandomEffects) + .def("Predict", &RandomEffectsModelCpp::Predict) + .def("SetWorkingParameter", &RandomEffectsModelCpp::SetWorkingParameter) + .def("SetGroupParameters", &RandomEffectsModelCpp::SetGroupParameters) + .def("SetWorkingParameterCovariance", &RandomEffectsModelCpp::SetWorkingParameterCovariance) + .def("SetGroupParameterCovariance", &RandomEffectsModelCpp::SetGroupParameterCovariance) + .def("SetVariancePriorShape", &RandomEffectsModelCpp::SetVariancePriorShape) + .def("SetVariancePriorScale", &RandomEffectsModelCpp::SetVariancePriorScale) + .def("Reset", &RandomEffectsModelCpp::Reset); py::class_(m, "GlobalVarianceModelCpp") - .def(py::init<>()) - .def("SampleOneIteration", &GlobalVarianceModelCpp::SampleOneIteration); + .def(py::init<>()) + .def("SampleOneIteration", &GlobalVarianceModelCpp::SampleOneIteration); py::class_(m, "LeafVarianceModelCpp") - .def(py::init<>()) - .def("SampleOneIteration", &LeafVarianceModelCpp::SampleOneIteration); + .def(py::init<>()) + .def("SampleOneIteration", &LeafVarianceModelCpp::SampleOneIteration); py::class_(m, "OrdinalSamplerCpp") - .def(py::init<>()) - .def("UpdateLatentVariables", &OrdinalSamplerCpp::UpdateLatentVariables) - .def("UpdateGammaParams", &OrdinalSamplerCpp::UpdateGammaParams) - .def("UpdateCumulativeExpSums", &OrdinalSamplerCpp::UpdateCumulativeExpSums); + .def(py::init<>()) + .def("UpdateLatentVariables", &OrdinalSamplerCpp::UpdateLatentVariables) + .def("UpdateGammaParams", &OrdinalSamplerCpp::UpdateGammaParams) + .def("UpdateCumulativeExpSums", &OrdinalSamplerCpp::UpdateCumulativeExpSums); + ; #ifdef VERSION_INFO m.attr("__version__") = MACRO_STRINGIFY(VERSION_INFO); diff --git a/src/random_effects.cpp b/src/random_effects.cpp index 40c828ba..32a62c68 100644 --- a/src/random_effects.cpp +++ b/src/random_effects.cpp @@ -3,7 +3,7 @@ namespace StochTree { -RandomEffectsTracker::RandomEffectsTracker(std::vector& group_indices) { +RandomEffectsTracker::RandomEffectsTracker(std::vector& group_indices) { sample_category_mapper_ = std::make_unique(group_indices); category_sample_tracker_ = std::make_unique(group_indices); num_categories_ = category_sample_tracker_->NumCategories(); @@ -11,6 +11,14 @@ RandomEffectsTracker::RandomEffectsTracker(std::vector& group_indices) rfx_predictions_.resize(num_observations_, 0.); } +RandomEffectsTracker::RandomEffectsTracker(int* group_indices, int num_observations) { + sample_category_mapper_ = std::make_unique(group_indices, num_observations); + category_sample_tracker_ = std::make_unique(group_indices, num_observations); + num_categories_ = category_sample_tracker_->NumCategories(); + num_observations_ = num_observations; + rfx_predictions_.resize(num_observations_, 0.); +} + nlohmann::json LabelMapper::to_json() { json output_obj; // Initialize a map with names of the node vectors and empty json arrays @@ -32,17 +40,17 @@ void LabelMapper::from_json(const nlohmann::json& rfx_label_mapper_json) { int num_values = rfx_label_mapper_json.at("values").size(); CHECK_EQ(num_keys, num_values); for (int i = 0; i < num_keys; i++) { - int32_t key = rfx_label_mapper_json.at("keys").at(i); - int32_t value = rfx_label_mapper_json.at("values").at(i); + int key = rfx_label_mapper_json.at("keys").at(i); + int value = rfx_label_mapper_json.at("values").at(i); keys_.push_back(key); label_map_.insert({key, value}); } } -void RandomEffectsTracker::ResetFromSample(MultivariateRegressionRandomEffectsModel& rfx_model, +void RandomEffectsTracker::ResetFromSample(MultivariateRegressionRandomEffectsModel& rfx_model, RandomEffectsDataset& rfx_dataset, ColumnVector& residual) { Eigen::MatrixXd X = rfx_dataset.GetBasis(); - std::vector group_labels = rfx_dataset.GetGroupLabels(); + std::vector group_labels = rfx_dataset.GetGroupLabels(); CHECK_EQ(X.rows(), group_labels.size()); int n = X.rows(); double prev_pred; @@ -50,7 +58,7 @@ void RandomEffectsTracker::ResetFromSample(MultivariateRegressionRandomEffectsMo double new_resid; Eigen::MatrixXd alpha_diag = rfx_model.GetWorkingParameter().asDiagonal().toDenseMatrix(); Eigen::MatrixXd xi = rfx_model.GetGroupParameters(); - std::int32_t group_ind; + int group_ind; for (int i = 0; i < n; i++) { group_ind = CategoryNumber(group_labels[i]); prev_pred = GetPrediction(i); @@ -61,7 +69,7 @@ void RandomEffectsTracker::ResetFromSample(MultivariateRegressionRandomEffectsMo } } -void RandomEffectsTracker::RootReset(MultivariateRegressionRandomEffectsModel& rfx_model, +void RandomEffectsTracker::RootReset(MultivariateRegressionRandomEffectsModel& rfx_model, RandomEffectsDataset& rfx_dataset, ColumnVector& residual) { int n = rfx_dataset.NumObservations(); CHECK_EQ(n, num_observations_); @@ -78,26 +86,26 @@ void RandomEffectsTracker::RootReset(MultivariateRegressionRandomEffectsModel& r } void MultivariateRegressionRandomEffectsModel::ResetFromSample(RandomEffectsContainer& rfx_container, int sample_num) { - // Extract parameter vectors - std::vector& alpha = rfx_container.GetAlpha(); - std::vector& xi = rfx_container.GetXi(); - std::vector& sigma = rfx_container.GetSigma(); - - // Unpack parameters - for (int i = 0; i < num_components_; i++) { - working_parameter_(i) = alpha.at(sample_num*num_components_ + i); - group_parameter_covariance_(i, i) = sigma.at(sample_num*num_components_ + i); - for (int j = 0; j < num_groups_; j++) { - group_parameters_(i,j) = xi.at(sample_num*num_groups_*num_components_ + j*num_components_ + i); - } + // Extract parameter vectors + std::vector& alpha = rfx_container.GetAlpha(); + std::vector& xi = rfx_container.GetXi(); + std::vector& sigma = rfx_container.GetSigma(); + + // Unpack parameters + for (int i = 0; i < num_components_; i++) { + working_parameter_(i) = alpha.at(sample_num * num_components_ + i); + group_parameter_covariance_(i, i) = sigma.at(sample_num * num_components_ + i); + for (int j = 0; j < num_groups_; j++) { + group_parameters_(i, j) = xi.at(sample_num * num_groups_ * num_components_ + j * num_components_ + i); } } +} -void MultivariateRegressionRandomEffectsModel::SampleRandomEffects(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, +void MultivariateRegressionRandomEffectsModel::SampleRandomEffects(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, std::mt19937& gen) { // Update partial residual to add back in the random effects AddCurrentPredictionToResidual(dataset, rfx_tracker, residual); - + // Sample random effects SampleGroupParameters(dataset, residual, rfx_tracker, global_variance, gen); SampleWorkingParameter(dataset, residual, rfx_tracker, global_variance, gen); @@ -107,16 +115,16 @@ void MultivariateRegressionRandomEffectsModel::SampleRandomEffects(RandomEffects SubtractNewPredictionFromResidual(dataset, rfx_tracker, residual); } -void MultivariateRegressionRandomEffectsModel::SampleWorkingParameter(RandomEffectsDataset& dataset, ColumnVector& residual, +void MultivariateRegressionRandomEffectsModel::SampleWorkingParameter(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, std::mt19937& gen) { Eigen::VectorXd posterior_mean = WorkingParameterMean(dataset, residual, rfx_tracker, global_variance); Eigen::MatrixXd posterior_covariance = WorkingParameterVariance(dataset, residual, rfx_tracker, global_variance); working_parameter_ = normal_sampler_.SampleEigen(posterior_mean, posterior_covariance, gen); } -void MultivariateRegressionRandomEffectsModel::SampleGroupParameters(RandomEffectsDataset& dataset, ColumnVector& residual, +void MultivariateRegressionRandomEffectsModel::SampleGroupParameters(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, std::mt19937& gen) { - int32_t num_groups = num_groups_; + int num_groups = num_groups_; Eigen::VectorXd posterior_mean; Eigen::MatrixXd posterior_covariance; Eigen::VectorXd output; @@ -124,12 +132,12 @@ void MultivariateRegressionRandomEffectsModel::SampleGroupParameters(RandomEffec posterior_mean = GroupParameterMean(dataset, residual, rfx_tracker, global_variance, i); posterior_covariance = GroupParameterVariance(dataset, residual, rfx_tracker, global_variance, i); group_parameters_(Eigen::all, i) = normal_sampler_.SampleEigen(posterior_mean, posterior_covariance, gen); - } + } } -void MultivariateRegressionRandomEffectsModel::SampleVarianceComponents(RandomEffectsDataset& dataset, ColumnVector& residual, +void MultivariateRegressionRandomEffectsModel::SampleVarianceComponents(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, std::mt19937& gen) { - int32_t num_components = num_components_; + int num_components = num_components_; double posterior_shape; double posterior_scale; double output; @@ -140,10 +148,10 @@ void MultivariateRegressionRandomEffectsModel::SampleVarianceComponents(RandomEf } } -Eigen::VectorXd MultivariateRegressionRandomEffectsModel::WorkingParameterMean(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, - double global_variance){ - int32_t num_components = num_components_; - int32_t num_groups = num_groups_; +Eigen::VectorXd MultivariateRegressionRandomEffectsModel::WorkingParameterMean(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, + double global_variance) { + int num_components = num_components_; + int num_groups = num_groups_; std::vector observation_indices; Eigen::MatrixXd X_group; Eigen::VectorXd y_group; @@ -164,9 +172,9 @@ Eigen::VectorXd MultivariateRegressionRandomEffectsModel::WorkingParameterMean(R return posterior_denominator.inverse() * posterior_numerator; } -Eigen::MatrixXd MultivariateRegressionRandomEffectsModel::WorkingParameterVariance(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance){ - int32_t num_components = num_components_; - int32_t num_groups = num_groups_; +Eigen::MatrixXd MultivariateRegressionRandomEffectsModel::WorkingParameterVariance(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance) { + int num_components = num_components_; + int num_groups = num_groups_; std::vector observation_indices; Eigen::MatrixXd X_group; Eigen::VectorXd y_group; @@ -186,9 +194,9 @@ Eigen::MatrixXd MultivariateRegressionRandomEffectsModel::WorkingParameterVarian return posterior_denominator.inverse(); } -Eigen::VectorXd MultivariateRegressionRandomEffectsModel::GroupParameterMean(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, int32_t group_id) { - int32_t num_components = num_components_; - int32_t num_groups = num_groups_; +Eigen::VectorXd MultivariateRegressionRandomEffectsModel::GroupParameterMean(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, int group_id) { + int num_components = num_components_; + int num_groups = num_groups_; Eigen::MatrixXd X = dataset.GetBasis(); Eigen::VectorXd y = residual.GetData(); Eigen::VectorXd alpha = working_parameter_; @@ -202,80 +210,80 @@ Eigen::VectorXd MultivariateRegressionRandomEffectsModel::GroupParameterMean(Ran return posterior_denominator.inverse() * posterior_numerator; } -Eigen::MatrixXd MultivariateRegressionRandomEffectsModel::GroupParameterVariance(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, int32_t group_id){ - int32_t num_components = num_components_; - int32_t num_groups = num_groups_; +Eigen::MatrixXd MultivariateRegressionRandomEffectsModel::GroupParameterVariance(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, int group_id) { + int num_components = num_components_; + int num_groups = num_groups_; Eigen::MatrixXd X = dataset.GetBasis(); Eigen::VectorXd y = residual.GetData(); Eigen::VectorXd alpha = working_parameter_; Eigen::MatrixXd posterior_denominator = group_parameter_covariance_.inverse(); -// Eigen::VectorXd posterior_numerator = Eigen::VectorXd::Zero(num_components); + // Eigen::VectorXd posterior_numerator = Eigen::VectorXd::Zero(num_components); std::vector observation_indices = rfx_tracker.NodeIndicesInternalIndex(group_id); Eigen::MatrixXd X_group = X(observation_indices, Eigen::all); -// Eigen::VectorXd y_group = y(observation_indices, Eigen::all); + // Eigen::VectorXd y_group = y(observation_indices, Eigen::all); posterior_denominator += ((alpha).asDiagonal() * X_group.transpose() * X_group * (alpha).asDiagonal()) / (global_variance); -// posterior_numerator += (alpha).asDiagonal() * X_group.transpose() * y_group; + // posterior_numerator += (alpha).asDiagonal() * X_group.transpose() * y_group; return posterior_denominator.inverse(); } -double MultivariateRegressionRandomEffectsModel::VarianceComponentShape(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, int32_t component_id) { +double MultivariateRegressionRandomEffectsModel::VarianceComponentShape(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, int component_id) { return static_cast(variance_prior_shape_ + num_groups_); } -double MultivariateRegressionRandomEffectsModel::VarianceComponentScale(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, int32_t component_id) { - int32_t num_groups = num_groups_; +double MultivariateRegressionRandomEffectsModel::VarianceComponentScale(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, int component_id) { + int num_groups = num_groups_; Eigen::MatrixXd xi = group_parameters_; double output = variance_prior_scale_; for (int i = 0; i < num_groups; i++) { - output += xi(component_id, i)*xi(component_id, i); + output += xi(component_id, i) * xi(component_id, i); } return output; } -void RandomEffectsContainer::AddSample(MultivariateRegressionRandomEffectsModel& model){ +void RandomEffectsContainer::AddSample(MultivariateRegressionRandomEffectsModel& model) { // Increment number of samples int sample_ind = num_samples_; num_samples_++; // Add alpha - alpha_.resize(num_samples_*num_components_); + alpha_.resize(num_samples_ * num_components_); for (int i = 0; i < num_components_; i++) { - alpha_.at(sample_ind*num_components_ + i) = model.GetWorkingParameter()(i); + alpha_.at(sample_ind * num_components_ + i) = model.GetWorkingParameter()(i); } // Add xi and beta - xi_.resize(num_samples_*num_components_*num_groups_); - beta_.resize(num_samples_*num_components_*num_groups_); + xi_.resize(num_samples_ * num_components_ * num_groups_); + beta_.resize(num_samples_ * num_components_ * num_groups_); for (int i = 0; i < num_components_; i++) { for (int j = 0; j < num_groups_; j++) { - xi_.at(sample_ind*num_groups_*num_components_ + j*num_components_ + i) = model.GetGroupParameters()(i,j); - beta_.at(sample_ind*num_groups_*num_components_ + j*num_components_ + i) = xi_.at(sample_ind*num_groups_*num_components_ + j*num_components_ + i) * alpha_.at(sample_ind*num_components_ + i); + xi_.at(sample_ind * num_groups_ * num_components_ + j * num_components_ + i) = model.GetGroupParameters()(i, j); + beta_.at(sample_ind * num_groups_ * num_components_ + j * num_components_ + i) = xi_.at(sample_ind * num_groups_ * num_components_ + j * num_components_ + i) * alpha_.at(sample_ind * num_components_ + i); } } // Add sigma - sigma_xi_.resize(num_samples_*num_components_); + sigma_xi_.resize(num_samples_ * num_components_); for (int i = 0; i < num_components_; i++) { - sigma_xi_.at(sample_ind*num_components_ + i) = model.GetGroupParameterCovariance()(i,i); + sigma_xi_.at(sample_ind * num_components_ + i) = model.GetGroupParameterCovariance()(i, i); } } void RandomEffectsContainer::Predict(RandomEffectsDataset& dataset, LabelMapper& label_mapper, std::vector& output) { Eigen::MatrixXd X = dataset.GetBasis(); - std::vector group_labels = dataset.GetGroupLabels(); + std::vector group_labels = dataset.GetGroupLabels(); CHECK_EQ(X.rows(), group_labels.size()); int n = X.rows(); - CHECK_EQ(n*num_samples_, output.size()); - std::int32_t group_ind; + CHECK_EQ(n * num_samples_, output.size()); + int group_ind; double pred; for (int i = 0; i < n; i++) { group_ind = label_mapper.CategoryNumber(group_labels[i]); for (int j = 0; j < num_samples_; j++) { pred = 0; for (int k = 0; k < num_components_; k++) { - pred += X(i,k) * beta_.at(j*num_groups_*num_components_ + group_ind*num_components_ + k); + pred += X(i, k) * beta_.at(j * num_groups_ * num_components_ + group_ind * num_components_ + k); } - output.at(j*n + i) = pred; + output.at(j * n + i) = pred; } } } @@ -288,8 +296,8 @@ nlohmann::json RandomEffectsContainer::to_json() { result_obj.emplace("num_groups", num_groups_); // Store some meta-level information about the containers - int beta_size = num_groups_*num_components_*num_samples_; - int alpha_size = num_components_*num_samples_; + int beta_size = num_groups_ * num_components_ * num_samples_; + int alpha_size = num_components_ * num_samples_; result_obj.emplace("beta_size", beta_size); result_obj.emplace("alpha_size", alpha_size); @@ -317,39 +325,39 @@ nlohmann::json RandomEffectsContainer::to_json() { result_obj.emplace(pair); } -return result_obj; + return result_obj; } -void RandomEffectsContainer::DeleteSample(int sample_num){ +void RandomEffectsContainer::DeleteSample(int sample_num) { // Decrement number of samples num_samples_--; // Remove sample_num from alpha // ---------------------------- - // This code works because the data are stored in a "column-major" format, - // with components comprising rows and and samples comprising columns, so that - // element `sample_num*num_components_ + i` will contain the "i"-th component of the - // sample indexed by sample_num. Erasing the `sample_num*num_components_ + 0` - // element of the vector will move the element that was previously in position + // This code works because the data are stored in a "column-major" format, + // with components comprising rows and and samples comprising columns, so that + // element `sample_num*num_components_ + i` will contain the "i"-th component of the + // sample indexed by sample_num. Erasing the `sample_num*num_components_ + 0` + // element of the vector will move the element that was previously in position // `sample_num*num_components_ + 1` into the position `sample_num*num_components_ + 0` // and thus we can repeat `alpha_.erase(alpha_.begin() + sample_num*num_components_);` // exactly `num_components_` times to erase each component pertaining to this sample. for (int i = 0; i < num_components_; i++) { - alpha_.erase(alpha_.begin() + sample_num*num_components_); + alpha_.erase(alpha_.begin() + sample_num * num_components_); } // Remove sample_num from xi and beta // ---------------------------------- - // This code works as above, with the added nuance of the three-dimensional (Fortran-aligned) array, - // in which sample number is the third dimension, group number is the second dimension, and component - // number is the third dimension. The nested loop assembles all `num_groups_*num_components_` offsets, - // expressed as `j*num_components_ + i`. In order to remove each of the elements stored in these offsets - // from `sample_num*num_groups_*num_components_`, we simply need to erase the + // This code works as above, with the added nuance of the three-dimensional (Fortran-aligned) array, + // in which sample number is the third dimension, group number is the second dimension, and component + // number is the third dimension. The nested loop assembles all `num_groups_*num_components_` offsets, + // expressed as `j*num_components_ + i`. In order to remove each of the elements stored in these offsets + // from `sample_num*num_groups_*num_components_`, we simply need to erase the // `sample_num*num_groups_*num_components_` element, exactly `num_groups_*num_components_` times. for (int i = 0; i < num_components_; i++) { for (int j = 0; j < num_groups_; j++) { - xi_.erase(xi_.begin() + sample_num*num_groups_*num_components_); - beta_.erase(beta_.begin() + sample_num*num_groups_*num_components_); + xi_.erase(xi_.begin() + sample_num * num_groups_ * num_components_); + beta_.erase(beta_.begin() + sample_num * num_groups_ * num_components_); } } @@ -357,7 +365,7 @@ void RandomEffectsContainer::DeleteSample(int sample_num){ // ---------------------------- // This code works as with alpha for (int i = 0; i < num_components_; i++) { - sigma_xi_.erase(sigma_xi_.begin() + sample_num*num_components_); + sigma_xi_.erase(sigma_xi_.begin() + sample_num * num_components_); } } @@ -375,13 +383,13 @@ void RandomEffectsContainer::from_json(const nlohmann::json& rfx_container_json) this->num_samples_ = rfx_container_json.at("num_samples"); this->num_components_ = rfx_container_json.at("num_components"); this->num_groups_ = rfx_container_json.at("num_groups"); - + // Unpack beta and xi for (int i = 0; i < beta_size; i++) { beta_.push_back(rfx_container_json.at("beta").at(i)); xi_.push_back(rfx_container_json.at("xi").at(i)); } - + // Unpack alpha and sigma_xi for (int i = 0; i < alpha_size; i++) { alpha_.push_back(rfx_container_json.at("alpha").at(i)); @@ -392,19 +400,19 @@ void RandomEffectsContainer::from_json(const nlohmann::json& rfx_container_json) void RandomEffectsContainer::append_from_json(const nlohmann::json& rfx_container_json) { CHECK_EQ(this->num_components_, rfx_container_json.at("num_components")); CHECK_EQ(this->num_groups_, rfx_container_json.at("num_groups")); - + // Update internal sample count and extract size of parameter vectors int new_num_samples = rfx_container_json.at("num_samples"); this->num_samples_ += new_num_samples; int beta_size = rfx_container_json.at("beta_size"); int alpha_size = rfx_container_json.at("alpha_size"); - + // Unpack beta and xi for (int i = 0; i < beta_size; i++) { beta_.push_back(rfx_container_json.at("beta").at(i)); xi_.push_back(rfx_container_json.at("xi").at(i)); } - + // Unpack alpha and sigma_xi for (int i = 0; i < alpha_size; i++) { alpha_.push_back(rfx_container_json.at("alpha").at(i)); diff --git a/src/sampler.cpp b/src/sampler.cpp index f356d968..84df1cd6 100644 --- a/src/sampler.cpp +++ b/src/sampler.cpp @@ -24,76 +24,81 @@ void sample_gfr_one_iteration_cpp(cpp11::external_pointer feature_types_(feature_types.size()); - for (int i = 0; i < feature_types.size(); i++) { - feature_types_[i] = static_cast(feature_types[i]); - } - - // Unpack sweep indices - std::vector sweep_indices_(sweep_indices.size()); - // if (sweep_indices.size() > 0) { - // sweep_indices_.resize(sweep_indices.size()); - for (int i = 0; i < sweep_indices.size(); i++) { - sweep_indices_[i] = sweep_indices[i]; - } - // } - - // Convert leaf model type to enum - StochTree::ModelType model_type; - if (leaf_model_int == 0) model_type = StochTree::ModelType::kConstantLeafGaussian; - else if (leaf_model_int == 1) model_type = StochTree::ModelType::kUnivariateRegressionLeafGaussian; - else if (leaf_model_int == 2) model_type = StochTree::ModelType::kMultivariateRegressionLeafGaussian; - else if (leaf_model_int == 3) model_type = StochTree::ModelType::kLogLinearVariance; - else if (leaf_model_int == 4) model_type = StochTree::ModelType::kCloglogOrdinal; - else StochTree::Log::Fatal("Invalid model type"); - - // Unpack leaf model parameters - double leaf_scale; - Eigen::MatrixXd leaf_scale_matrix; - if ((model_type == StochTree::ModelType::kConstantLeafGaussian) || - (model_type == StochTree::ModelType::kUnivariateRegressionLeafGaussian)) { - leaf_scale = leaf_model_scale_input(0,0); - } else if (model_type == StochTree::ModelType::kMultivariateRegressionLeafGaussian) { - int num_row = leaf_model_scale_input.nrow(); - int num_col = leaf_model_scale_input.ncol(); - leaf_scale_matrix.resize(num_row, num_col); - for (int i = 0; i < num_row; i++) { - for (int j = 0; j < num_col; j++) { - leaf_scale_matrix(i,j) = leaf_model_scale_input(i,j); - } - } - } - - // Convert variable weights to std::vector - std::vector var_weights_vector(variable_weights.size()); - for (int i = 0; i < variable_weights.size(); i++) { - var_weights_vector[i] = variable_weights[i]; - } - - // Prepare the samplers - StochTree::LeafModelVariant leaf_model = StochTree::leafModelFactory(model_type, leaf_scale, leaf_scale_matrix, a_forest, b_forest); - int num_basis = data->NumBasis(); - - // Run one iteration of the sampler - if (model_type == StochTree::ModelType::kConstantLeafGaussian) { - StochTree::GFRSampleOneIter(*active_forest, *tracker, *forest_samples, std::get(leaf_model), *data, *residual, *split_prior, *rng, var_weights_vector, sweep_indices_, global_variance, feature_types_, cutpoint_grid_size, keep_forest, pre_initialized, true, num_features_subsample, num_threads); - } else if (model_type == StochTree::ModelType::kUnivariateRegressionLeafGaussian) { - StochTree::GFRSampleOneIter(*active_forest, *tracker, *forest_samples, std::get(leaf_model), *data, *residual, *split_prior, *rng, var_weights_vector, sweep_indices_, global_variance, feature_types_, cutpoint_grid_size, keep_forest, pre_initialized, true, num_features_subsample, num_threads); - } else if (model_type == StochTree::ModelType::kMultivariateRegressionLeafGaussian) { - StochTree::GFRSampleOneIter(*active_forest, *tracker, *forest_samples, std::get(leaf_model), *data, *residual, *split_prior, *rng, var_weights_vector, sweep_indices_, global_variance, feature_types_, cutpoint_grid_size, keep_forest, pre_initialized, true, num_features_subsample, num_threads, num_basis); - } else if (model_type == StochTree::ModelType::kLogLinearVariance) { - StochTree::GFRSampleOneIter(*active_forest, *tracker, *forest_samples, std::get(leaf_model), *data, *residual, *split_prior, *rng, var_weights_vector, sweep_indices_, global_variance, feature_types_, cutpoint_grid_size, keep_forest, pre_initialized, false, num_features_subsample, num_threads); - } else if (model_type == StochTree::ModelType::kCloglogOrdinal) { - StochTree::GFRSampleOneIter(*active_forest, *tracker, *forest_samples, std::get(leaf_model), *data, *residual, *split_prior, *rng, var_weights_vector, sweep_indices_, global_variance, feature_types_, cutpoint_grid_size, keep_forest, pre_initialized, false, num_features_subsample, num_threads); + bool keep_forest, int num_features_subsample, + int num_threads) { + // Refactoring completely out of the R interface. + // Intention to refactor out of the C++ interface in the future. + bool pre_initialized = true; + + // Unpack feature types + std::vector feature_types_(feature_types.size()); + for (int i = 0; i < feature_types.size(); i++) { + feature_types_[i] = static_cast(feature_types[i]); + } + + // Unpack sweep indices + std::vector sweep_indices_(sweep_indices.size()); + // if (sweep_indices.size() > 0) { + // sweep_indices_.resize(sweep_indices.size()); + for (int i = 0; i < sweep_indices.size(); i++) { + sweep_indices_[i] = sweep_indices[i]; + } + // } + + // Convert leaf model type to enum + StochTree::ModelType model_type; + if (leaf_model_int == 0) + model_type = StochTree::ModelType::kConstantLeafGaussian; + else if (leaf_model_int == 1) + model_type = StochTree::ModelType::kUnivariateRegressionLeafGaussian; + else if (leaf_model_int == 2) + model_type = StochTree::ModelType::kMultivariateRegressionLeafGaussian; + else if (leaf_model_int == 3) + model_type = StochTree::ModelType::kLogLinearVariance; + else if (leaf_model_int == 4) + model_type = StochTree::ModelType::kCloglogOrdinal; + else + StochTree::Log::Fatal("Invalid model type"); + + // Unpack leaf model parameters + double leaf_scale; + Eigen::MatrixXd leaf_scale_matrix; + if ((model_type == StochTree::ModelType::kConstantLeafGaussian) || + (model_type == StochTree::ModelType::kUnivariateRegressionLeafGaussian)) { + leaf_scale = leaf_model_scale_input(0, 0); + } else if (model_type == StochTree::ModelType::kMultivariateRegressionLeafGaussian) { + int num_row = leaf_model_scale_input.nrow(); + int num_col = leaf_model_scale_input.ncol(); + leaf_scale_matrix.resize(num_row, num_col); + for (int i = 0; i < num_row; i++) { + for (int j = 0; j < num_col; j++) { + leaf_scale_matrix(i, j) = leaf_model_scale_input(i, j); + } } + } + + // Convert variable weights to std::vector + std::vector var_weights_vector(variable_weights.size()); + for (int i = 0; i < variable_weights.size(); i++) { + var_weights_vector[i] = variable_weights[i]; + } + + // Prepare the samplers + StochTree::LeafModelVariant leaf_model = StochTree::leafModelFactory(model_type, leaf_scale, leaf_scale_matrix, a_forest, b_forest); + int num_basis = data->NumBasis(); + + // Run one iteration of the sampler + if (model_type == StochTree::ModelType::kConstantLeafGaussian) { + StochTree::GFRSampleOneIter(*active_forest, *tracker, *forest_samples, std::get(leaf_model), *data, *residual, *split_prior, *rng, var_weights_vector, sweep_indices_, global_variance, feature_types_, cutpoint_grid_size, keep_forest, pre_initialized, true, num_features_subsample, num_threads); + } else if (model_type == StochTree::ModelType::kUnivariateRegressionLeafGaussian) { + StochTree::GFRSampleOneIter(*active_forest, *tracker, *forest_samples, std::get(leaf_model), *data, *residual, *split_prior, *rng, var_weights_vector, sweep_indices_, global_variance, feature_types_, cutpoint_grid_size, keep_forest, pre_initialized, true, num_features_subsample, num_threads); + } else if (model_type == StochTree::ModelType::kMultivariateRegressionLeafGaussian) { + StochTree::GFRSampleOneIter(*active_forest, *tracker, *forest_samples, std::get(leaf_model), *data, *residual, *split_prior, *rng, var_weights_vector, sweep_indices_, global_variance, feature_types_, cutpoint_grid_size, keep_forest, pre_initialized, true, num_features_subsample, num_threads, num_basis); + } else if (model_type == StochTree::ModelType::kLogLinearVariance) { + StochTree::GFRSampleOneIter(*active_forest, *tracker, *forest_samples, std::get(leaf_model), *data, *residual, *split_prior, *rng, var_weights_vector, sweep_indices_, global_variance, feature_types_, cutpoint_grid_size, keep_forest, pre_initialized, false, num_features_subsample, num_threads); + } else if (model_type == StochTree::ModelType::kCloglogOrdinal) { + StochTree::GFRSampleOneIter(*active_forest, *tracker, *forest_samples, std::get(leaf_model), *data, *residual, *split_prior, *rng, var_weights_vector, sweep_indices_, global_variance, feature_types_, cutpoint_grid_size, keep_forest, pre_initialized, false, num_features_subsample, num_threads); + } } [[cpp11::register]] @@ -110,234 +115,235 @@ void sample_mcmc_one_iteration_cpp(cpp11::external_pointer feature_types_(feature_types.size()); - for (int i = 0; i < feature_types.size(); i++) { - feature_types_[i] = static_cast(feature_types[i]); - } - - // Unpack sweep indices - std::vector sweep_indices_; - if (sweep_indices.size() > 0) { - sweep_indices_.resize(sweep_indices.size()); - for (int i = 0; i < sweep_indices.size(); i++) { - sweep_indices_[i] = sweep_indices[i]; - } - } - - // Convert leaf model type to enum - StochTree::ModelType model_type; - if (leaf_model_int == 0) model_type = StochTree::ModelType::kConstantLeafGaussian; - else if (leaf_model_int == 1) model_type = StochTree::ModelType::kUnivariateRegressionLeafGaussian; - else if (leaf_model_int == 2) model_type = StochTree::ModelType::kMultivariateRegressionLeafGaussian; - else if (leaf_model_int == 3) model_type = StochTree::ModelType::kLogLinearVariance; - else if (leaf_model_int == 4) model_type = StochTree::ModelType::kCloglogOrdinal; - else StochTree::Log::Fatal("Invalid model type"); - - // Unpack leaf model parameters - double leaf_scale; - Eigen::MatrixXd leaf_scale_matrix; - if ((model_type == StochTree::ModelType::kConstantLeafGaussian) || - (model_type == StochTree::ModelType::kUnivariateRegressionLeafGaussian)) { - leaf_scale = leaf_model_scale_input(0,0); - } else if (model_type == StochTree::ModelType::kMultivariateRegressionLeafGaussian) { - int num_row = leaf_model_scale_input.nrow(); - int num_col = leaf_model_scale_input.ncol(); - leaf_scale_matrix.resize(num_row, num_col); - for (int i = 0; i < num_row; i++) { - for (int j = 0; j < num_col; j++) { - leaf_scale_matrix(i,j) = leaf_model_scale_input(i,j); - } - } - } - - // Convert variable weights to std::vector - std::vector var_weights_vector(variable_weights.size()); - for (int i = 0; i < variable_weights.size(); i++) { - var_weights_vector[i] = variable_weights[i]; + bool keep_forest, int num_threads) { + // Refactoring completely out of the R interface. + // Intention to refactor out of the C++ interface in the future. + bool pre_initialized = true; + + // Unpack feature types + std::vector feature_types_(feature_types.size()); + for (int i = 0; i < feature_types.size(); i++) { + feature_types_[i] = static_cast(feature_types[i]); + } + + // Unpack sweep indices + std::vector sweep_indices_; + if (sweep_indices.size() > 0) { + sweep_indices_.resize(sweep_indices.size()); + for (int i = 0; i < sweep_indices.size(); i++) { + sweep_indices_[i] = sweep_indices[i]; } - - // Prepare the samplers - StochTree::LeafModelVariant leaf_model = StochTree::leafModelFactory(model_type, leaf_scale, leaf_scale_matrix, a_forest, b_forest); - int num_basis = data->NumBasis(); - - // Run one iteration of the sampler - if (model_type == StochTree::ModelType::kConstantLeafGaussian) { - StochTree::MCMCSampleOneIter(*active_forest, *tracker, *forest_samples, std::get(leaf_model), *data, *residual, *split_prior, *rng, var_weights_vector, sweep_indices_, global_variance, keep_forest, pre_initialized, true, num_threads); - } else if (model_type == StochTree::ModelType::kUnivariateRegressionLeafGaussian) { - StochTree::MCMCSampleOneIter(*active_forest, *tracker, *forest_samples, std::get(leaf_model), *data, *residual, *split_prior, *rng, var_weights_vector, sweep_indices_, global_variance, keep_forest, pre_initialized, true, num_threads); - } else if (model_type == StochTree::ModelType::kMultivariateRegressionLeafGaussian) { - StochTree::MCMCSampleOneIter(*active_forest, *tracker, *forest_samples, std::get(leaf_model), *data, *residual, *split_prior, *rng, var_weights_vector, sweep_indices_, global_variance, keep_forest, pre_initialized, true, num_threads, num_basis); - } else if (model_type == StochTree::ModelType::kLogLinearVariance) { - StochTree::MCMCSampleOneIter(*active_forest, *tracker, *forest_samples, std::get(leaf_model), *data, *residual, *split_prior, *rng, var_weights_vector, sweep_indices_, global_variance, keep_forest, pre_initialized, false, num_threads); - } else if (model_type == StochTree::ModelType::kCloglogOrdinal) { - StochTree::MCMCSampleOneIter(*active_forest, *tracker, *forest_samples, std::get(leaf_model), *data, *residual, *split_prior, *rng, var_weights_vector, sweep_indices_, global_variance, keep_forest, pre_initialized, false, num_threads); + } + + // Convert leaf model type to enum + StochTree::ModelType model_type; + if (leaf_model_int == 0) + model_type = StochTree::ModelType::kConstantLeafGaussian; + else if (leaf_model_int == 1) + model_type = StochTree::ModelType::kUnivariateRegressionLeafGaussian; + else if (leaf_model_int == 2) + model_type = StochTree::ModelType::kMultivariateRegressionLeafGaussian; + else if (leaf_model_int == 3) + model_type = StochTree::ModelType::kLogLinearVariance; + else if (leaf_model_int == 4) + model_type = StochTree::ModelType::kCloglogOrdinal; + else + StochTree::Log::Fatal("Invalid model type"); + + // Unpack leaf model parameters + double leaf_scale; + Eigen::MatrixXd leaf_scale_matrix; + if ((model_type == StochTree::ModelType::kConstantLeafGaussian) || + (model_type == StochTree::ModelType::kUnivariateRegressionLeafGaussian)) { + leaf_scale = leaf_model_scale_input(0, 0); + } else if (model_type == StochTree::ModelType::kMultivariateRegressionLeafGaussian) { + int num_row = leaf_model_scale_input.nrow(); + int num_col = leaf_model_scale_input.ncol(); + leaf_scale_matrix.resize(num_row, num_col); + for (int i = 0; i < num_row; i++) { + for (int j = 0; j < num_col; j++) { + leaf_scale_matrix(i, j) = leaf_model_scale_input(i, j); + } } + } + + // Convert variable weights to std::vector + std::vector var_weights_vector(variable_weights.size()); + for (int i = 0; i < variable_weights.size(); i++) { + var_weights_vector[i] = variable_weights[i]; + } + + // Prepare the samplers + StochTree::LeafModelVariant leaf_model = StochTree::leafModelFactory(model_type, leaf_scale, leaf_scale_matrix, a_forest, b_forest); + int num_basis = data->NumBasis(); + + // Run one iteration of the sampler + if (model_type == StochTree::ModelType::kConstantLeafGaussian) { + StochTree::MCMCSampleOneIter(*active_forest, *tracker, *forest_samples, std::get(leaf_model), *data, *residual, *split_prior, *rng, var_weights_vector, sweep_indices_, global_variance, keep_forest, pre_initialized, true, num_threads); + } else if (model_type == StochTree::ModelType::kUnivariateRegressionLeafGaussian) { + StochTree::MCMCSampleOneIter(*active_forest, *tracker, *forest_samples, std::get(leaf_model), *data, *residual, *split_prior, *rng, var_weights_vector, sweep_indices_, global_variance, keep_forest, pre_initialized, true, num_threads); + } else if (model_type == StochTree::ModelType::kMultivariateRegressionLeafGaussian) { + StochTree::MCMCSampleOneIter(*active_forest, *tracker, *forest_samples, std::get(leaf_model), *data, *residual, *split_prior, *rng, var_weights_vector, sweep_indices_, global_variance, keep_forest, pre_initialized, true, num_threads, num_basis); + } else if (model_type == StochTree::ModelType::kLogLinearVariance) { + StochTree::MCMCSampleOneIter(*active_forest, *tracker, *forest_samples, std::get(leaf_model), *data, *residual, *split_prior, *rng, var_weights_vector, sweep_indices_, global_variance, keep_forest, pre_initialized, false, num_threads); + } else if (model_type == StochTree::ModelType::kCloglogOrdinal) { + StochTree::MCMCSampleOneIter(*active_forest, *tracker, *forest_samples, std::get(leaf_model), *data, *residual, *split_prior, *rng, var_weights_vector, sweep_indices_, global_variance, keep_forest, pre_initialized, false, num_threads); + } } [[cpp11::register]] double sample_sigma2_one_iteration_cpp(cpp11::external_pointer residual, cpp11::external_pointer dataset, cpp11::external_pointer rng, - double a, double b -) { - // Run one iteration of the sampler - StochTree::GlobalHomoskedasticVarianceModel var_model = StochTree::GlobalHomoskedasticVarianceModel(); - if (dataset->HasVarWeights()) { - return var_model.SampleVarianceParameter(residual->GetData(), dataset->GetVarWeights(), a, b, *rng); - } else { - return var_model.SampleVarianceParameter(residual->GetData(), a, b, *rng); - } + double a, double b) { + // Run one iteration of the sampler + StochTree::GlobalHomoskedasticVarianceModel var_model = StochTree::GlobalHomoskedasticVarianceModel(); + if (dataset->HasVarWeights()) { + return var_model.SampleVarianceParameter(residual->GetData(), dataset->GetVarWeights(), a, b, *rng); + } else { + return var_model.SampleVarianceParameter(residual->GetData(), a, b, *rng); + } } [[cpp11::register]] double sample_tau_one_iteration_cpp(cpp11::external_pointer active_forest, cpp11::external_pointer rng, - double a, double b -) { - // Run one iteration of the sampler - StochTree::LeafNodeHomoskedasticVarianceModel var_model = StochTree::LeafNodeHomoskedasticVarianceModel(); - return var_model.SampleVarianceParameter(active_forest.get(), a, b, *rng); + double a, double b) { + // Run one iteration of the sampler + StochTree::LeafNodeHomoskedasticVarianceModel var_model = StochTree::LeafNodeHomoskedasticVarianceModel(); + return var_model.SampleVarianceParameter(active_forest.get(), a, b, *rng); } [[cpp11::register]] cpp11::external_pointer rng_cpp(int random_seed = -1) { - std::unique_ptr rng_; - if (random_seed == -1) { - std::random_device rd; - rng_ = std::make_unique(rd()); - } else { - rng_ = std::make_unique(random_seed); - } + std::unique_ptr rng_; + if (random_seed == -1) { + std::random_device rd; + rng_ = std::make_unique(rd()); + } else { + rng_ = std::make_unique(random_seed); + } - // Release management of the pointer to R session - return cpp11::external_pointer(rng_.release()); + // Release management of the pointer to R session + return cpp11::external_pointer(rng_.release()); } [[cpp11::register]] cpp11::external_pointer tree_prior_cpp(double alpha, double beta, int min_samples_leaf, int max_depth = -1) { - // Create smart pointer to newly allocated object - std::unique_ptr prior_ptr_ = std::make_unique(alpha, beta, min_samples_leaf, max_depth); + // Create smart pointer to newly allocated object + std::unique_ptr prior_ptr_ = std::make_unique(alpha, beta, min_samples_leaf, max_depth); - // Release management of the pointer to R session - return cpp11::external_pointer(prior_ptr_.release()); + // Release management of the pointer to R session + return cpp11::external_pointer(prior_ptr_.release()); } [[cpp11::register]] void update_alpha_tree_prior_cpp(cpp11::external_pointer tree_prior_ptr, double alpha) { - // Update alpha - tree_prior_ptr->SetAlpha(alpha); + // Update alpha + tree_prior_ptr->SetAlpha(alpha); } [[cpp11::register]] void update_beta_tree_prior_cpp(cpp11::external_pointer tree_prior_ptr, double beta) { - // Update beta - tree_prior_ptr->SetBeta(beta); + // Update beta + tree_prior_ptr->SetBeta(beta); } [[cpp11::register]] void update_min_samples_leaf_tree_prior_cpp(cpp11::external_pointer tree_prior_ptr, int min_samples_leaf) { - // Update min_samples_leaf - tree_prior_ptr->SetMinSamplesLeaf(min_samples_leaf); + // Update min_samples_leaf + tree_prior_ptr->SetMinSamplesLeaf(min_samples_leaf); } [[cpp11::register]] void update_max_depth_tree_prior_cpp(cpp11::external_pointer tree_prior_ptr, int max_depth) { - // Update max_depth - tree_prior_ptr->SetMaxDepth(max_depth); + // Update max_depth + tree_prior_ptr->SetMaxDepth(max_depth); } [[cpp11::register]] double get_alpha_tree_prior_cpp(cpp11::external_pointer tree_prior_ptr) { - return tree_prior_ptr->GetAlpha(); + return tree_prior_ptr->GetAlpha(); } [[cpp11::register]] double get_beta_tree_prior_cpp(cpp11::external_pointer tree_prior_ptr) { - // Update beta - return tree_prior_ptr->GetBeta(); + // Update beta + return tree_prior_ptr->GetBeta(); } [[cpp11::register]] int get_min_samples_leaf_tree_prior_cpp(cpp11::external_pointer tree_prior_ptr) { - return tree_prior_ptr->GetMinSamplesLeaf(); + return tree_prior_ptr->GetMinSamplesLeaf(); } [[cpp11::register]] int get_max_depth_tree_prior_cpp(cpp11::external_pointer tree_prior_ptr) { - return tree_prior_ptr->GetMaxDepth(); + return tree_prior_ptr->GetMaxDepth(); } [[cpp11::register]] cpp11::external_pointer forest_tracker_cpp(cpp11::external_pointer data, cpp11::integers feature_types, int num_trees, StochTree::data_size_t n) { - // Convert vector of integers to std::vector of enum FeatureType - std::vector feature_types_(feature_types.size()); - for (int i = 0; i < feature_types.size(); i++) { - feature_types_[i] = static_cast(feature_types[i]); - } + // Convert vector of integers to std::vector of enum FeatureType + std::vector feature_types_(feature_types.size()); + for (int i = 0; i < feature_types.size(); i++) { + feature_types_[i] = static_cast(feature_types[i]); + } - // Create smart pointer to newly allocated object - std::unique_ptr tracker_ptr_ = std::make_unique(data->GetCovariates(), feature_types_, num_trees, n); + // Create smart pointer to newly allocated object + std::unique_ptr tracker_ptr_ = std::make_unique(data->GetCovariates(), feature_types_, num_trees, n); - // Release management of the pointer to R session - return cpp11::external_pointer(tracker_ptr_.release()); + // Release management of the pointer to R session + return cpp11::external_pointer(tracker_ptr_.release()); } [[cpp11::register]] cpp11::writable::doubles get_cached_forest_predictions_cpp(cpp11::external_pointer tracker_ptr) { - int n_train = tracker_ptr->GetNumObservations(); - cpp11::writable::doubles output(n_train); - for (int i = 0; i < n_train; i++) { - output[i] = tracker_ptr->GetSamplePrediction(i); - } - return output; + int n_train = tracker_ptr->GetNumObservations(); + cpp11::writable::doubles output(n_train); + for (int i = 0; i < n_train; i++) { + output[i] = tracker_ptr->GetSamplePrediction(i); + } + return output; } [[cpp11::register]] cpp11::writable::integers sample_without_replacement_integer_cpp( cpp11::integers population_vector, cpp11::doubles sampling_probs, - int sample_size -) { - // Unpack pointer to population vector - int population_size = population_vector.size(); - int* population_vector_ptr = INTEGER(PROTECT(population_vector)); + int sample_size) { + // Unpack pointer to population vector + int population_size = population_vector.size(); + int* population_vector_ptr = INTEGER(PROTECT(population_vector)); - // Unpack pointer to sampling probabilities - double* sampling_probs_ptr = REAL(PROTECT(sampling_probs)); + // Unpack pointer to sampling probabilities + double* sampling_probs_ptr = REAL(PROTECT(sampling_probs)); - // Create output vector - cpp11::writable::integers output(sample_size); + // Create output vector + cpp11::writable::integers output(sample_size); - // Unpack pointer to output vector - int* output_ptr = INTEGER(PROTECT(output)); + // Unpack pointer to output vector + int* output_ptr = INTEGER(PROTECT(output)); - // Create C++ RNG - std::random_device rd; - std::mt19937 gen(rd()); + // Create C++ RNG + std::random_device rd; + std::mt19937 gen(rd()); - // Run the sampler - StochTree::sample_without_replacement( - output_ptr, sampling_probs_ptr, population_vector_ptr, population_size, sample_size, gen - ); + // Run the sampler + StochTree::sample_without_replacement( + output_ptr, sampling_probs_ptr, population_vector_ptr, population_size, sample_size, gen); - // Unprotect raw pointers - UNPROTECT(3); + // Unprotect raw pointers + UNPROTECT(3); - // Return result - return(output); + // Return result + return (output); } [[cpp11::register]] cpp11::external_pointer ordinal_sampler_cpp() { - std::unique_ptr sampler_ptr = std::make_unique(); - return cpp11::external_pointer(sampler_ptr.release()); + std::unique_ptr sampler_ptr = std::make_unique(); + return cpp11::external_pointer(sampler_ptr.release()); } [[cpp11::register]] @@ -345,9 +351,8 @@ void ordinal_sampler_update_latent_variables_cpp( cpp11::external_pointer sampler_ptr, cpp11::external_pointer data_ptr, cpp11::external_pointer outcome_ptr, - cpp11::external_pointer rng_ptr -) { - sampler_ptr->UpdateLatentVariables(*data_ptr, outcome_ptr->GetData(), *rng_ptr); + cpp11::external_pointer rng_ptr) { + sampler_ptr->UpdateLatentVariables(*data_ptr, outcome_ptr->GetData(), *rng_ptr); } [[cpp11::register]] @@ -358,15 +363,13 @@ void ordinal_sampler_update_gamma_params_cpp( double alpha_gamma, double beta_gamma, double gamma_0, - cpp11::external_pointer rng_ptr -) { - sampler_ptr->UpdateGammaParams(*data_ptr, outcome_ptr->GetData(), alpha_gamma, beta_gamma, gamma_0, *rng_ptr); + cpp11::external_pointer rng_ptr) { + sampler_ptr->UpdateGammaParams(*data_ptr, outcome_ptr->GetData(), alpha_gamma, beta_gamma, gamma_0, *rng_ptr); } [[cpp11::register]] void ordinal_sampler_update_cumsum_exp_cpp( cpp11::external_pointer sampler_ptr, - cpp11::external_pointer data_ptr -) { - sampler_ptr->UpdateCumulativeExpSums(*data_ptr); + cpp11::external_pointer data_ptr) { + sampler_ptr->UpdateCumulativeExpSums(*data_ptr); } diff --git a/src/serialization.cpp b/src/serialization.cpp index fb248f62..a1d5aa1d 100644 --- a/src/serialization.cpp +++ b/src/serialization.cpp @@ -11,405 +11,451 @@ [[cpp11::register]] cpp11::external_pointer init_json_cpp() { - std::unique_ptr json_ptr = std::make_unique(); - json forests = nlohmann::json::object(); - json rfx = nlohmann::json::object(); - json parameters = nlohmann::json::object(); - json_ptr->emplace("forests", forests); - json_ptr->emplace("random_effects", rfx); - json_ptr->emplace("num_forests", 0); - json_ptr->emplace("num_random_effects", 0); - return cpp11::external_pointer(json_ptr.release()); + std::unique_ptr json_ptr = std::make_unique(); + json forests = nlohmann::json::object(); + json rfx = nlohmann::json::object(); + json parameters = nlohmann::json::object(); + json_ptr->emplace("forests", forests); + json_ptr->emplace("random_effects", rfx); + json_ptr->emplace("num_forests", 0); + json_ptr->emplace("num_random_effects", 0); + return cpp11::external_pointer(json_ptr.release()); } [[cpp11::register]] void json_add_double_subfolder_cpp(cpp11::external_pointer json_ptr, std::string subfolder_name, std::string field_name, double field_value) { - if (json_ptr->contains(subfolder_name)) { - if (json_ptr->at(subfolder_name).contains(field_name)) { - json_ptr->at(subfolder_name).at(field_name) = field_value; - } else { - json_ptr->at(subfolder_name).emplace(std::pair(field_name, field_value)); - } + if (json_ptr->contains(subfolder_name)) { + if (json_ptr->at(subfolder_name).contains(field_name)) { + json_ptr->at(subfolder_name).at(field_name) = field_value; } else { - json_ptr->emplace(std::pair(subfolder_name, nlohmann::json::object())); - json_ptr->at(subfolder_name).emplace(std::pair(field_name, field_value)); + json_ptr->at(subfolder_name).emplace(std::pair(field_name, field_value)); } + } else { + json_ptr->emplace(std::pair(subfolder_name, nlohmann::json::object())); + json_ptr->at(subfolder_name).emplace(std::pair(field_name, field_value)); + } } [[cpp11::register]] void json_add_double_cpp(cpp11::external_pointer json_ptr, std::string field_name, double field_value) { - if (json_ptr->contains(field_name)) { - json_ptr->at(field_name) = field_value; - } else { - json_ptr->emplace(std::pair(field_name, field_value)); - } + if (json_ptr->contains(field_name)) { + json_ptr->at(field_name) = field_value; + } else { + json_ptr->emplace(std::pair(field_name, field_value)); + } } [[cpp11::register]] void json_add_integer_subfolder_cpp(cpp11::external_pointer json_ptr, std::string subfolder_name, std::string field_name, int field_value) { - if (json_ptr->contains(subfolder_name)) { - if (json_ptr->at(subfolder_name).contains(field_name)) { - json_ptr->at(subfolder_name).at(field_name) = field_value; - } else { - json_ptr->at(subfolder_name).emplace(std::pair(field_name, field_value)); - } + if (json_ptr->contains(subfolder_name)) { + if (json_ptr->at(subfolder_name).contains(field_name)) { + json_ptr->at(subfolder_name).at(field_name) = field_value; } else { - json_ptr->emplace(std::pair(subfolder_name, nlohmann::json::object())); - json_ptr->at(subfolder_name).emplace(std::pair(field_name, field_value)); + json_ptr->at(subfolder_name).emplace(std::pair(field_name, field_value)); } + } else { + json_ptr->emplace(std::pair(subfolder_name, nlohmann::json::object())); + json_ptr->at(subfolder_name).emplace(std::pair(field_name, field_value)); + } } [[cpp11::register]] void json_add_integer_cpp(cpp11::external_pointer json_ptr, std::string field_name, int field_value) { - if (json_ptr->contains(field_name)) { - json_ptr->at(field_name) = field_value; - } else { - json_ptr->emplace(std::pair(field_name, field_value)); - } + if (json_ptr->contains(field_name)) { + json_ptr->at(field_name) = field_value; + } else { + json_ptr->emplace(std::pair(field_name, field_value)); + } } [[cpp11::register]] void json_add_bool_subfolder_cpp(cpp11::external_pointer json_ptr, std::string subfolder_name, std::string field_name, bool field_value) { - if (json_ptr->contains(subfolder_name)) { - if (json_ptr->at(subfolder_name).contains(field_name)) { - json_ptr->at(subfolder_name).at(field_name) = field_value; - } else { - json_ptr->at(subfolder_name).emplace(std::pair(field_name, field_value)); - } + if (json_ptr->contains(subfolder_name)) { + if (json_ptr->at(subfolder_name).contains(field_name)) { + json_ptr->at(subfolder_name).at(field_name) = field_value; } else { - json_ptr->emplace(std::pair(subfolder_name, nlohmann::json::object())); - json_ptr->at(subfolder_name).emplace(std::pair(field_name, field_value)); + json_ptr->at(subfolder_name).emplace(std::pair(field_name, field_value)); } + } else { + json_ptr->emplace(std::pair(subfolder_name, nlohmann::json::object())); + json_ptr->at(subfolder_name).emplace(std::pair(field_name, field_value)); + } } [[cpp11::register]] void json_add_bool_cpp(cpp11::external_pointer json_ptr, std::string field_name, bool field_value) { - if (json_ptr->contains(field_name)) { - json_ptr->at(field_name) = field_value; - } else { - json_ptr->emplace(std::pair(field_name, field_value)); - } + if (json_ptr->contains(field_name)) { + json_ptr->at(field_name) = field_value; + } else { + json_ptr->emplace(std::pair(field_name, field_value)); + } } [[cpp11::register]] void json_add_vector_subfolder_cpp(cpp11::external_pointer json_ptr, std::string subfolder_name, std::string field_name, cpp11::doubles field_vector) { - int vec_length = field_vector.size(); - if (json_ptr->contains(subfolder_name)) { - if (json_ptr->at(subfolder_name).contains(field_name)) { - json_ptr->at(subfolder_name).at(field_name).clear(); - for (int i = 0; i < vec_length; i++) { - json_ptr->at(subfolder_name).at(field_name).emplace_back(field_vector.at(i)); - } - } else { - json_ptr->at(subfolder_name).emplace(std::pair(field_name, nlohmann::json::array())); - for (int i = 0; i < vec_length; i++) { - json_ptr->at(subfolder_name).at(field_name).emplace_back(field_vector.at(i)); - } - } + int vec_length = field_vector.size(); + if (json_ptr->contains(subfolder_name)) { + if (json_ptr->at(subfolder_name).contains(field_name)) { + json_ptr->at(subfolder_name).at(field_name).clear(); + for (int i = 0; i < vec_length; i++) { + json_ptr->at(subfolder_name).at(field_name).emplace_back(field_vector.at(i)); + } } else { - json_ptr->emplace(std::pair(subfolder_name, nlohmann::json::object())); - json_ptr->at(subfolder_name).emplace(std::pair(field_name, nlohmann::json::array())); - for (int i = 0; i < vec_length; i++) { - json_ptr->at(subfolder_name).at(field_name).emplace_back(field_vector.at(i)); - } + json_ptr->at(subfolder_name).emplace(std::pair(field_name, nlohmann::json::array())); + for (int i = 0; i < vec_length; i++) { + json_ptr->at(subfolder_name).at(field_name).emplace_back(field_vector.at(i)); + } + } + } else { + json_ptr->emplace(std::pair(subfolder_name, nlohmann::json::object())); + json_ptr->at(subfolder_name).emplace(std::pair(field_name, nlohmann::json::array())); + for (int i = 0; i < vec_length; i++) { + json_ptr->at(subfolder_name).at(field_name).emplace_back(field_vector.at(i)); } + } } [[cpp11::register]] void json_add_vector_cpp(cpp11::external_pointer json_ptr, std::string field_name, cpp11::doubles field_vector) { - int vec_length = field_vector.size(); - if (json_ptr->contains(field_name)) { - json_ptr->at(field_name).clear(); - for (int i = 0; i < vec_length; i++) { - json_ptr->at(field_name).emplace_back(field_vector.at(i)); - } - } else { - json_ptr->emplace(std::pair(field_name, nlohmann::json::array())); - for (int i = 0; i < vec_length; i++) { - json_ptr->at(field_name).emplace_back(field_vector.at(i)); - } + int vec_length = field_vector.size(); + if (json_ptr->contains(field_name)) { + json_ptr->at(field_name).clear(); + for (int i = 0; i < vec_length; i++) { + json_ptr->at(field_name).emplace_back(field_vector.at(i)); + } + } else { + json_ptr->emplace(std::pair(field_name, nlohmann::json::array())); + for (int i = 0; i < vec_length; i++) { + json_ptr->at(field_name).emplace_back(field_vector.at(i)); } + } } [[cpp11::register]] void json_add_integer_vector_subfolder_cpp(cpp11::external_pointer json_ptr, std::string subfolder_name, std::string field_name, cpp11::integers field_vector) { - int vec_length = field_vector.size(); - if (json_ptr->contains(subfolder_name)) { - if (json_ptr->at(subfolder_name).contains(field_name)) { - json_ptr->at(subfolder_name).at(field_name).clear(); - for (int i = 0; i < vec_length; i++) { - json_ptr->at(subfolder_name).at(field_name).emplace_back(field_vector.at(i)); - } - } else { - json_ptr->at(subfolder_name).emplace(std::pair(field_name, nlohmann::json::array())); - for (int i = 0; i < vec_length; i++) { - json_ptr->at(subfolder_name).at(field_name).emplace_back(field_vector.at(i)); - } - } + int vec_length = field_vector.size(); + if (json_ptr->contains(subfolder_name)) { + if (json_ptr->at(subfolder_name).contains(field_name)) { + json_ptr->at(subfolder_name).at(field_name).clear(); + for (int i = 0; i < vec_length; i++) { + json_ptr->at(subfolder_name).at(field_name).emplace_back(field_vector.at(i)); + } } else { - json_ptr->emplace(std::pair(subfolder_name, nlohmann::json::object())); - json_ptr->at(subfolder_name).emplace(std::pair(field_name, nlohmann::json::array())); - for (int i = 0; i < vec_length; i++) { - json_ptr->at(subfolder_name).at(field_name).emplace_back(field_vector.at(i)); - } + json_ptr->at(subfolder_name).emplace(std::pair(field_name, nlohmann::json::array())); + for (int i = 0; i < vec_length; i++) { + json_ptr->at(subfolder_name).at(field_name).emplace_back(field_vector.at(i)); + } } + } else { + json_ptr->emplace(std::pair(subfolder_name, nlohmann::json::object())); + json_ptr->at(subfolder_name).emplace(std::pair(field_name, nlohmann::json::array())); + for (int i = 0; i < vec_length; i++) { + json_ptr->at(subfolder_name).at(field_name).emplace_back(field_vector.at(i)); + } + } } [[cpp11::register]] void json_add_integer_vector_cpp(cpp11::external_pointer json_ptr, std::string field_name, cpp11::integers field_vector) { - int vec_length = field_vector.size(); - if (json_ptr->contains(field_name)) { - json_ptr->at(field_name).clear(); - for (int i = 0; i < vec_length; i++) { - json_ptr->at(field_name).emplace_back(field_vector.at(i)); - } - } else { - json_ptr->emplace(std::pair(field_name, nlohmann::json::array())); - for (int i = 0; i < vec_length; i++) { - json_ptr->at(field_name).emplace_back(field_vector.at(i)); - } + int vec_length = field_vector.size(); + if (json_ptr->contains(field_name)) { + json_ptr->at(field_name).clear(); + for (int i = 0; i < vec_length; i++) { + json_ptr->at(field_name).emplace_back(field_vector.at(i)); + } + } else { + json_ptr->emplace(std::pair(field_name, nlohmann::json::array())); + for (int i = 0; i < vec_length; i++) { + json_ptr->at(field_name).emplace_back(field_vector.at(i)); } + } } [[cpp11::register]] void json_add_string_vector_subfolder_cpp(cpp11::external_pointer json_ptr, std::string subfolder_name, std::string field_name, cpp11::strings field_vector) { - int vec_length = field_vector.size(); - if (json_ptr->contains(subfolder_name)) { - if (json_ptr->at(subfolder_name).contains(field_name)) { - json_ptr->at(subfolder_name).at(field_name).clear(); - for (int i = 0; i < vec_length; i++) { - json_ptr->at(subfolder_name).at(field_name).emplace_back(field_vector.at(i)); - } - } else { - json_ptr->at(subfolder_name).emplace(std::pair(field_name, nlohmann::json::array())); - for (int i = 0; i < vec_length; i++) { - json_ptr->at(subfolder_name).at(field_name).emplace_back(field_vector.at(i)); - } - } + int vec_length = field_vector.size(); + if (json_ptr->contains(subfolder_name)) { + if (json_ptr->at(subfolder_name).contains(field_name)) { + json_ptr->at(subfolder_name).at(field_name).clear(); + for (int i = 0; i < vec_length; i++) { + json_ptr->at(subfolder_name).at(field_name).emplace_back(field_vector.at(i)); + } } else { - json_ptr->emplace(std::pair(subfolder_name, nlohmann::json::object())); - json_ptr->at(subfolder_name).emplace(std::pair(field_name, nlohmann::json::array())); - for (int i = 0; i < vec_length; i++) { - json_ptr->at(subfolder_name).at(field_name).emplace_back(field_vector.at(i)); - } + json_ptr->at(subfolder_name).emplace(std::pair(field_name, nlohmann::json::array())); + for (int i = 0; i < vec_length; i++) { + json_ptr->at(subfolder_name).at(field_name).emplace_back(field_vector.at(i)); + } } + } else { + json_ptr->emplace(std::pair(subfolder_name, nlohmann::json::object())); + json_ptr->at(subfolder_name).emplace(std::pair(field_name, nlohmann::json::array())); + for (int i = 0; i < vec_length; i++) { + json_ptr->at(subfolder_name).at(field_name).emplace_back(field_vector.at(i)); + } + } } [[cpp11::register]] void json_add_string_vector_cpp(cpp11::external_pointer json_ptr, std::string field_name, cpp11::strings field_vector) { - int vec_length = field_vector.size(); - if (json_ptr->contains(field_name)) { - json_ptr->at(field_name).clear(); - for (int i = 0; i < vec_length; i++) { - json_ptr->at(field_name).emplace_back(field_vector.at(i)); - } - } else { - json_ptr->emplace(std::pair(field_name, nlohmann::json::array())); - for (int i = 0; i < vec_length; i++) { - json_ptr->at(field_name).emplace_back(field_vector.at(i)); - } + int vec_length = field_vector.size(); + if (json_ptr->contains(field_name)) { + json_ptr->at(field_name).clear(); + for (int i = 0; i < vec_length; i++) { + json_ptr->at(field_name).emplace_back(field_vector.at(i)); } + } else { + json_ptr->emplace(std::pair(field_name, nlohmann::json::array())); + for (int i = 0; i < vec_length; i++) { + json_ptr->at(field_name).emplace_back(field_vector.at(i)); + } + } } [[cpp11::register]] void json_add_string_subfolder_cpp(cpp11::external_pointer json_ptr, std::string subfolder_name, std::string field_name, std::string field_value) { - if (json_ptr->contains(subfolder_name)) { - if (json_ptr->at(subfolder_name).contains(field_name)) { - json_ptr->at(subfolder_name).at(field_name) = field_value; - } else { - json_ptr->at(subfolder_name).emplace(std::pair(field_name, field_value)); - } + if (json_ptr->contains(subfolder_name)) { + if (json_ptr->at(subfolder_name).contains(field_name)) { + json_ptr->at(subfolder_name).at(field_name) = field_value; } else { - json_ptr->emplace(std::pair(subfolder_name, nlohmann::json::object())); - json_ptr->at(subfolder_name).emplace(std::pair(field_name, field_value)); + json_ptr->at(subfolder_name).emplace(std::pair(field_name, field_value)); } + } else { + json_ptr->emplace(std::pair(subfolder_name, nlohmann::json::object())); + json_ptr->at(subfolder_name).emplace(std::pair(field_name, field_value)); + } } [[cpp11::register]] void json_add_string_cpp(cpp11::external_pointer json_ptr, std::string field_name, std::string field_value) { - if (json_ptr->contains(field_name)) { - json_ptr->at(field_name) = field_value; - } else { - json_ptr->emplace(std::pair(field_name, field_value)); - } + if (json_ptr->contains(field_name)) { + json_ptr->at(field_name) = field_value; + } else { + json_ptr->emplace(std::pair(field_name, field_value)); + } } [[cpp11::register]] bool json_contains_field_subfolder_cpp(cpp11::external_pointer json_ptr, std::string subfolder_name, std::string field_name) { - if (json_ptr->contains(subfolder_name)) { - if (json_ptr->at(subfolder_name).contains(field_name)) { - return true; - } else { - return false; - } + if (json_ptr->contains(subfolder_name)) { + if (json_ptr->at(subfolder_name).contains(field_name)) { + return true; } else { - return false; + return false; } + } else { + return false; + } } [[cpp11::register]] bool json_contains_field_cpp(cpp11::external_pointer json_ptr, std::string field_name) { - if (json_ptr->contains(field_name)) { - return true; - } else { - return false; - } + if (json_ptr->contains(field_name)) { + return true; + } else { + return false; + } +} + +[[cpp11::register]] +void json_erase_field_cpp(cpp11::external_pointer json_ptr, std::string field_name) { + if (json_ptr->contains(field_name)) { + json_ptr->erase(field_name); + } +} + +[[cpp11::register]] +void json_erase_field_subfolder_cpp(cpp11::external_pointer json_ptr, std::string subfolder_name, std::string field_name) { + if (json_ptr->contains(subfolder_name) && json_ptr->at(subfolder_name).contains(field_name)) { + json_ptr->at(subfolder_name).erase(field_name); + } +} + +[[cpp11::register]] +void json_rename_field_cpp(cpp11::external_pointer json_ptr, std::string old_name, std::string new_name) { + if (json_ptr->contains(old_name)) { + (*json_ptr)[new_name] = json_ptr->at(old_name); + json_ptr->erase(old_name); + } +} + +[[cpp11::register]] +void json_rename_field_subfolder_cpp(cpp11::external_pointer json_ptr, std::string subfolder_name, std::string old_name, std::string new_name) { + if (json_ptr->contains(subfolder_name) && json_ptr->at(subfolder_name).contains(old_name)) { + json_ptr->at(subfolder_name)[new_name] = json_ptr->at(subfolder_name).at(old_name); + json_ptr->at(subfolder_name).erase(old_name); + } } [[cpp11::register]] double json_extract_double_subfolder_cpp(cpp11::external_pointer json_ptr, std::string subfolder_name, std::string field_name) { - return json_ptr->at(subfolder_name).at(field_name); + return json_ptr->at(subfolder_name).at(field_name); } [[cpp11::register]] double json_extract_double_cpp(cpp11::external_pointer json_ptr, std::string field_name) { - return json_ptr->at(field_name); + return json_ptr->at(field_name); } [[cpp11::register]] int json_extract_integer_subfolder_cpp(cpp11::external_pointer json_ptr, std::string subfolder_name, std::string field_name) { - return json_ptr->at(subfolder_name).at(field_name); + return json_ptr->at(subfolder_name).at(field_name); } [[cpp11::register]] int json_extract_integer_cpp(cpp11::external_pointer json_ptr, std::string field_name) { - return json_ptr->at(field_name); + return json_ptr->at(field_name); } [[cpp11::register]] bool json_extract_bool_subfolder_cpp(cpp11::external_pointer json_ptr, std::string subfolder_name, std::string field_name) { - return json_ptr->at(subfolder_name).at(field_name); + return json_ptr->at(subfolder_name).at(field_name); } [[cpp11::register]] bool json_extract_bool_cpp(cpp11::external_pointer json_ptr, std::string field_name) { - return json_ptr->at(field_name); + return json_ptr->at(field_name); } [[cpp11::register]] std::string json_extract_string_subfolder_cpp(cpp11::external_pointer json_ptr, std::string subfolder_name, std::string field_name) { - return json_ptr->at(subfolder_name).at(field_name); + return json_ptr->at(subfolder_name).at(field_name); } [[cpp11::register]] std::string json_extract_string_cpp(cpp11::external_pointer json_ptr, std::string field_name) { - return json_ptr->at(field_name); + return json_ptr->at(field_name); } [[cpp11::register]] cpp11::writable::doubles json_extract_vector_subfolder_cpp(cpp11::external_pointer json_ptr, std::string subfolder_name, std::string field_name) { - cpp11::writable::doubles output; - int vec_length = json_ptr->at(subfolder_name).at(field_name).size(); - for (int i = 0; i < vec_length; i++) output.push_back((json_ptr->at(subfolder_name).at(field_name).at(i))); - return output; + cpp11::writable::doubles output; + int vec_length = json_ptr->at(subfolder_name).at(field_name).size(); + for (int i = 0; i < vec_length; i++) output.push_back((json_ptr->at(subfolder_name).at(field_name).at(i))); + return output; } [[cpp11::register]] cpp11::writable::doubles json_extract_vector_cpp(cpp11::external_pointer json_ptr, std::string field_name) { - cpp11::writable::doubles output; - int vec_length = json_ptr->at(field_name).size(); - for (int i = 0; i < vec_length; i++) output.push_back((json_ptr->at(field_name).at(i))); - return output; + cpp11::writable::doubles output; + int vec_length = json_ptr->at(field_name).size(); + for (int i = 0; i < vec_length; i++) output.push_back((json_ptr->at(field_name).at(i))); + return output; } [[cpp11::register]] cpp11::writable::integers json_extract_integer_vector_subfolder_cpp(cpp11::external_pointer json_ptr, std::string subfolder_name, std::string field_name) { - cpp11::writable::integers output; - int vec_length = json_ptr->at(subfolder_name).at(field_name).size(); - for (int i = 0; i < vec_length; i++) output.push_back((json_ptr->at(subfolder_name).at(field_name).at(i))); - return output; + cpp11::writable::integers output; + int vec_length = json_ptr->at(subfolder_name).at(field_name).size(); + for (int i = 0; i < vec_length; i++) output.push_back((json_ptr->at(subfolder_name).at(field_name).at(i))); + return output; } [[cpp11::register]] cpp11::writable::integers json_extract_integer_vector_cpp(cpp11::external_pointer json_ptr, std::string field_name) { - cpp11::writable::integers output; - int vec_length = json_ptr->at(field_name).size(); - for (int i = 0; i < vec_length; i++) output.push_back((json_ptr->at(field_name).at(i))); - return output; + cpp11::writable::integers output; + int vec_length = json_ptr->at(field_name).size(); + for (int i = 0; i < vec_length; i++) output.push_back((json_ptr->at(field_name).at(i))); + return output; } [[cpp11::register]] cpp11::writable::strings json_extract_string_vector_subfolder_cpp(cpp11::external_pointer json_ptr, std::string subfolder_name, std::string field_name) { - int vec_length = json_ptr->at(subfolder_name).at(field_name).size(); - std::vector output(vec_length); - for (int i = 0; i < vec_length; i++) output.at(i) = json_ptr->at(subfolder_name).at(field_name).at(i); - return output; + int vec_length = json_ptr->at(subfolder_name).at(field_name).size(); + std::vector output(vec_length); + for (int i = 0; i < vec_length; i++) output.at(i) = json_ptr->at(subfolder_name).at(field_name).at(i); + return output; } [[cpp11::register]] cpp11::writable::strings json_extract_string_vector_cpp(cpp11::external_pointer json_ptr, std::string field_name) { - int vec_length = json_ptr->at(field_name).size(); - std::vector output(vec_length); - for (int i = 0; i < vec_length; i++) output.at(i) = json_ptr->at(field_name).at(i); - return output; + int vec_length = json_ptr->at(field_name).size(); + std::vector output(vec_length); + for (int i = 0; i < vec_length; i++) output.at(i) = json_ptr->at(field_name).at(i); + return output; } [[cpp11::register]] -std::string json_add_forest_cpp(cpp11::external_pointer json_ptr, cpp11::external_pointer forest_samples) { - int forest_num = json_ptr->at("num_forests"); - std::string forest_label = "forest_" + std::to_string(forest_num); - nlohmann::json forest_json = forest_samples->to_json(); - json_ptr->at("forests").emplace(forest_label, forest_json); - json_ptr->at("num_forests") = forest_num + 1; - return forest_label; +std::string json_add_forest_cpp(cpp11::external_pointer json_ptr, cpp11::external_pointer forest_samples, std::string forest_label) { + int forest_num = json_ptr->at("num_forests"); + if (forest_label.empty()) { + forest_label = "forest_" + std::to_string(forest_num); + } + nlohmann::json forest_json = forest_samples->to_json(); + json_ptr->at("forests").emplace(forest_label, forest_json); + json_ptr->at("num_forests") = forest_num + 1; + return forest_label; } [[cpp11::register]] void json_increment_rfx_count_cpp(cpp11::external_pointer json_ptr) { - int rfx_num = json_ptr->at("num_random_effects"); - json_ptr->at("num_random_effects") = rfx_num + 1; + int rfx_num = json_ptr->at("num_random_effects"); + json_ptr->at("num_random_effects") = rfx_num + 1; } [[cpp11::register]] std::string json_add_rfx_container_cpp(cpp11::external_pointer json_ptr, cpp11::external_pointer rfx_samples) { - int rfx_num = json_ptr->at("num_random_effects"); - std::string rfx_label = "random_effect_container_" + std::to_string(rfx_num); - nlohmann::json rfx_json = rfx_samples->to_json(); - json_ptr->at("random_effects").emplace(rfx_label, rfx_json); - return rfx_label; + int rfx_num = json_ptr->at("num_random_effects"); + std::string rfx_label = "random_effect_container_" + std::to_string(rfx_num); + nlohmann::json rfx_json = rfx_samples->to_json(); + json_ptr->at("random_effects").emplace(rfx_label, rfx_json); + return rfx_label; } [[cpp11::register]] std::string json_add_rfx_label_mapper_cpp(cpp11::external_pointer json_ptr, cpp11::external_pointer label_mapper) { - int rfx_num = json_ptr->at("num_random_effects"); - std::string rfx_label = "random_effect_label_mapper_" + std::to_string(rfx_num); - nlohmann::json rfx_json = label_mapper->to_json(); - json_ptr->at("random_effects").emplace(rfx_label, rfx_json); - return rfx_label; + int rfx_num = json_ptr->at("num_random_effects"); + std::string rfx_label = "random_effect_label_mapper_" + std::to_string(rfx_num); + nlohmann::json rfx_json = label_mapper->to_json(); + json_ptr->at("random_effects").emplace(rfx_label, rfx_json); + return rfx_label; +} + +// Unique (sorted) integer group ids tracked by an rfx label mapper. Used to +// reconstruct R's `rfx_unique_group_ids` levels when loading a model written +// without them (e.g. a cross-platform Python model). Mirrors Python's +// RandomEffectsLabelMapperCpp::GetUniqueGroupIds. +[[cpp11::register]] +cpp11::writable::integers rfx_label_mapper_unique_group_ids_cpp(cpp11::external_pointer label_mapper) { + std::vector& keys = label_mapper->Keys(); + cpp11::writable::integers result(static_cast(keys.size())); + for (size_t i = 0; i < keys.size(); i++) { + result[i] = keys[i]; + } + return result; } [[cpp11::register]] std::string json_add_rfx_groupids_cpp(cpp11::external_pointer json_ptr, cpp11::integers groupids) { - int rfx_num = json_ptr->at("num_random_effects"); - std::string rfx_label = "random_effect_groupids_" + std::to_string(rfx_num); - nlohmann::json groupids_json = nlohmann::json::array(); - for (int i = 0; i < groupids.size(); i++) { - groupids_json.emplace_back(groupids.at(i)); - } - json_ptr->at("random_effects").emplace(rfx_label, groupids_json); - return rfx_label; + int rfx_num = json_ptr->at("num_random_effects"); + std::string rfx_label = "random_effect_groupids_" + std::to_string(rfx_num); + nlohmann::json groupids_json = nlohmann::json::array(); + for (int i = 0; i < groupids.size(); i++) { + groupids_json.emplace_back(groupids.at(i)); + } + json_ptr->at("random_effects").emplace(rfx_label, groupids_json); + return rfx_label; } [[cpp11::register]] std::string get_json_string_cpp(cpp11::external_pointer json_ptr) { - return json_ptr->dump(); + return json_ptr->dump(); } [[cpp11::register]] void json_save_file_cpp(cpp11::external_pointer json_ptr, std::string filename) { - std::ofstream output_file(filename); - output_file << *json_ptr << std::endl; + std::ofstream output_file(filename); + output_file << *json_ptr << std::endl; } [[cpp11::register]] void json_load_file_cpp(cpp11::external_pointer json_ptr, std::string filename) { - std::ifstream f(filename); - // nlohmann::json file_json = nlohmann::json::parse(f); - *json_ptr = nlohmann::json::parse(f); - // json_ptr.reset(&file_json); + std::ifstream f(filename); + // nlohmann::json file_json = nlohmann::json::parse(f); + *json_ptr = nlohmann::json::parse(f); + // json_ptr.reset(&file_json); } [[cpp11::register]] void json_load_string_cpp(cpp11::external_pointer json_ptr, std::string json_string) { - *json_ptr = nlohmann::json::parse(json_string); + *json_ptr = nlohmann::json::parse(json_string); } diff --git a/src/stochtree_types.h b/src/stochtree_types.h index 0e17038f..5a9f0b06 100644 --- a/src/stochtree_types.h +++ b/src/stochtree_types.h @@ -1,3 +1,5 @@ +#include +#include #include #include #include @@ -6,5 +8,47 @@ #include #include #include +#include #include #include +#include + +static void check_numeric(cpp11::sexp input, const char* input_name) { + if (TYPEOF(input) != REALSXP && !Rf_isInteger(input)) { + cpp11::stop("Parameter %s must be a numeric array (integer or floating point)", input_name); + } +} + +static double* extract_numeric_pointer(cpp11::sexp input, const char* input_name, int& protect_count) { + if (input == R_NilValue) return nullptr; + check_numeric(input, input_name); + cpp11::sexp input_converted = PROTECT(Rf_coerceVector(input, REALSXP)); + protect_count++; + return REAL(input_converted); +} + +static void check_integer(cpp11::sexp input, const char* input_name) { + if (!Rf_isInteger(input)) { + cpp11::stop("Parameter %s must be an integer array", input_name); + } +} + +static int* extract_integer_pointer(cpp11::sexp input, const char* input_name, int& protect_count) { + if (input == R_NilValue) return nullptr; + check_integer(input, input_name); + return INTEGER(input); +} + +template +T get_config_scalar_default(cpp11::list& config_list, const char* config_key, T default_value) { + cpp11::sexp val = config_list[config_key]; + if (Rf_isNull(val)) return default_value; + return cpp11::as_cpp(val); +} + +template <> +inline int get_config_scalar_default(cpp11::list& config_list, const char* config_key, int default_value) { + cpp11::sexp val = config_list[config_key]; + if (Rf_isNull(val)) return default_value; + return Rf_asInteger(val); +} diff --git a/src/tree.cpp b/src/tree.cpp index 32c51475..43debc57 100644 --- a/src/tree.cpp +++ b/src/tree.cpp @@ -1,5 +1,5 @@ /*! - * Inspired by the design of the tree in the xgboost and treelite package, both released under the Apache license + * Inspired by the design of the tree in the xgboost and treelite package, both released under the Apache license * with the following copyright: * Copyright 2015-2023 by XGBoost Contributors * Copyright 2017-2021 by [treelite] Contributors @@ -24,14 +24,14 @@ std::int32_t Tree::NumLeafParents() const { } std::int32_t Tree::NumSplitNodes() const { - std::int32_t splits { 0 }; + std::int32_t splits{0}; auto const& self = *this; this->WalkTree([&splits, &self](std::int32_t nidx) { - if (!self.IsLeaf(nidx)){ - splits++; - } - return true; - }); + if (!self.IsLeaf(nidx)) { + splits++; + } + return true; + }); return splits; } @@ -110,10 +110,10 @@ std::int32_t Tree::AllocNode() { --num_deleted_nodes; return nid; } - + std::int32_t nd = num_nodes++; CHECK_LT(num_nodes, std::numeric_limits::max()); - + node_type_.push_back(TreeNodeType::kLeafNode); cleft_.push_back(kInvalidNodeId); cright_.push_back(kInvalidNodeId); @@ -121,7 +121,7 @@ std::int32_t Tree::AllocNode() { leaf_value_.push_back(static_cast(0)); threshold_.push_back(static_cast(0)); node_deleted_.push_back(false); - // THIS is a placeholder, currently set after AllocNode is called ... + // THIS is a placeholder, currently set after AllocNode is called ... // ... to be refactored ... parent_.push_back(static_cast(0)); @@ -169,7 +169,7 @@ void Tree::ExpandNode(std::int32_t nid, int split_index, double split_value, dou internal_nodes_.push_back(nid); // Remove nid's parent node (if applicable) from leaf parents - if (!IsRoot(nid)){ + if (!IsRoot(nid)) { std::int32_t parent_idx = Parent(nid); leaf_parents_.erase(std::remove(leaf_parents_.begin(), leaf_parents_.end(), parent_idx), leaf_parents_.end()); } @@ -195,7 +195,7 @@ void Tree::ExpandNode(std::int32_t nid, int split_index, std::vector& void Tree::PredictLeafIndexInplace(Eigen::MatrixXd& covariates, std::vector& output, int32_t offset, int32_t max_leaf) { int n = covariates.rows(); CHECK_GE(output.size(), offset + n); - std::map renumber_map; + std::map renumber_map; for (int i = 0; i < leaves_.size(); i++) { renumber_map.insert({leaves_[i], i}); } @@ -437,7 +437,7 @@ void Tree::PredictLeafIndexInplace(Eigen::MatrixXd& covariates, std::vector>& covariates, std::vector& output, int32_t offset, int32_t max_leaf) { int n = covariates.rows(); CHECK_GE(output.size(), offset + n); - std::map renumber_map; + std::map renumber_map; for (int i = 0; i < leaves_.size(); i++) { renumber_map.insert({leaves_[i], i}); } @@ -448,12 +448,12 @@ void Tree::PredictLeafIndexInplace(Eigen::Map>& covariates, - Eigen::Map>& output, +void Tree::PredictLeafIndexInplace(Eigen::Map>& covariates, + Eigen::Map>& output, int column_ind, int32_t offset, int32_t max_leaf) { int n = covariates.rows(); CHECK_GE(output.size(), offset + n); - std::map renumber_map; + std::map renumber_map; for (int i = 0; i < leaves_.size(); i++) { renumber_map.insert({leaves_[i], i}); } @@ -481,26 +481,26 @@ void TreeNodeVectorsToJson(json& obj, Tree* tree) { tree_array_map.emplace(std::pair("category_list_end", json::array())); // Extract only the non-deleted nodes into tree_array_map -// bool node_deleted; + // bool node_deleted; for (int i = 0; i < tree->NumNodes(); i++) { -// node_deleted = (std::find(tree->deleted_nodes_.begin(), tree->deleted_nodes_.end(), i) -// != tree->deleted_nodes_.end()); -// if (!node_deleted) { - tree_array_map["node_type"].emplace_back(static_cast(tree->node_type_[i])); - tree_array_map["parent"].emplace_back(tree->parent_[i]); - tree_array_map["left"].emplace_back(tree->cleft_[i]); - tree_array_map["right"].emplace_back(tree->cright_[i]); - tree_array_map["split_index"].emplace_back(tree->split_index_[i]); - tree_array_map["leaf_value"].emplace_back(tree->leaf_value_[i]); - tree_array_map["threshold"].emplace_back(tree->threshold_[i]); - tree_array_map["node_deleted"].emplace_back(tree->node_deleted_[i]); - tree_array_map["leaf_vector_begin"].emplace_back(static_cast(tree->leaf_vector_begin_[i])); - tree_array_map["leaf_vector_end"].emplace_back(static_cast(tree->leaf_vector_end_[i])); - tree_array_map["category_list_begin"].emplace_back(static_cast(tree->category_list_begin_[i])); - tree_array_map["category_list_end"].emplace_back(static_cast(tree->category_list_end_[i])); -// } - } - + // node_deleted = (std::find(tree->deleted_nodes_.begin(), tree->deleted_nodes_.end(), i) + // != tree->deleted_nodes_.end()); + // if (!node_deleted) { + tree_array_map["node_type"].emplace_back(static_cast(tree->node_type_[i])); + tree_array_map["parent"].emplace_back(tree->parent_[i]); + tree_array_map["left"].emplace_back(tree->cleft_[i]); + tree_array_map["right"].emplace_back(tree->cright_[i]); + tree_array_map["split_index"].emplace_back(tree->split_index_[i]); + tree_array_map["leaf_value"].emplace_back(tree->leaf_value_[i]); + tree_array_map["threshold"].emplace_back(tree->threshold_[i]); + tree_array_map["node_deleted"].emplace_back(tree->node_deleted_[i]); + tree_array_map["leaf_vector_begin"].emplace_back(static_cast(tree->leaf_vector_begin_[i])); + tree_array_map["leaf_vector_end"].emplace_back(static_cast(tree->leaf_vector_end_[i])); + tree_array_map["category_list_begin"].emplace_back(static_cast(tree->category_list_begin_[i])); + tree_array_map["category_list_end"].emplace_back(static_cast(tree->category_list_end_[i])); + // } + } + // Unpack the map into the reference JSON object for (auto& pair : tree_array_map) { obj.emplace(pair); @@ -532,7 +532,7 @@ void NodeListsToJson(json& obj, Tree* tree) { json vec_leaf_parents = json::array(); json vec_leaves = json::array(); json vec_deleted_nodes = json::array(); - + if (tree->internal_nodes_.size() > 0) { for (int i = 0; i < tree->internal_nodes_.size(); i++) { vec_internal_nodes.emplace_back(tree->internal_nodes_[i]); @@ -556,7 +556,7 @@ void NodeListsToJson(json& obj, Tree* tree) { vec_deleted_nodes.emplace_back(tree->deleted_nodes_[i]); } } - + obj.emplace("internal_nodes", vec_internal_nodes); obj.emplace("leaf_parents", vec_leaf_parents); obj.emplace("leaves", vec_leaves); @@ -577,7 +577,7 @@ json Tree::to_json() { MultivariateLeafVectorToJson(result_obj, this); SplitCategoryVectorToJson(result_obj, this); NodeListsToJson(result_obj, this); - + // Initialize Json from Json::object map and return result return result_obj; } @@ -603,8 +603,10 @@ void JsonToTreeNodeVectors(const json& tree_json, Tree* tree) { tree->cleft_.push_back(tree_json.at("left").at(i)); tree->cright_.push_back(tree_json.at("right").at(i)); tree->split_index_.push_back(tree_json.at("split_index").at(i)); - if (is_univariate) tree->leaf_value_.push_back(tree_json.at("leaf_value").at(i)); - else tree->leaf_value_.push_back(0.); + if (is_univariate) + tree->leaf_value_.push_back(tree_json.at("leaf_value").at(i)); + else + tree->leaf_value_.push_back(0.); tree->threshold_.push_back(tree_json.at("threshold").at(i)); tree->node_deleted_.push_back(tree_json.at("node_deleted").at(i)); // Handle type conversions for node_type, leaf_vector_begin/end, and category_list_begin/end @@ -665,7 +667,7 @@ void Tree::from_json(const json& tree_json) { tree_json.at("has_categorical_split").get_to(this->has_categorical_split_); tree_json.at("output_dimension").get_to(this->output_dimension_); tree_json.at("is_log_scale").get_to(this->is_log_scale_); - + // Unpack the array based fields JsonToTreeNodeVectors(tree_json, this); JsonToMultivariateLeafVector(tree_json, this); @@ -673,4 +675,4 @@ void Tree::from_json(const json& tree_json) { JsonToNodeLists(tree_json, this); } -} // namespace StochTree +} // namespace StochTree diff --git a/stochtree/bart.py b/stochtree/bart.py index 2c352a11..9d5e3fcb 100644 --- a/stochtree/bart.py +++ b/stochtree/bart.py @@ -1,31 +1,22 @@ import json import warnings -from math import log, floor +from math import floor from numbers import Integral from typing import Any, Dict, Optional, Union import numpy as np import pandas as pd -from scipy.stats import norm -from .config import ForestModelConfig, GlobalModelConfig -from .data import Dataset, Residual -from .forest import Forest, ForestContainer +from .forest import ForestContainer from .preprocessing import CovariatePreprocessor, _preprocess_params -from .random_effects import ( - RandomEffectsContainer, - RandomEffectsDataset, - RandomEffectsModel, - RandomEffectsTracker, +from .random_effects import RandomEffectsContainer +from .serialization import ( + JSONSerializer, + SCHEMA_VERSION, + resolve_schema_version, + infer_platform_v0, + enforce_cross_platform_gate, ) -from .sampler import ( - RNG, - ForestSampler, - GlobalVarianceModel, - LeafVarianceModel, - OrdinalSampler, -) -from .serialization import JSONSerializer from .utils import ( OutcomeModel, NotSampledError, @@ -39,6 +30,30 @@ _posterior_predictive_heuristic_multiplier, _summarize_interval, ) +from stochtree_cpp import bart_sample_cpp, bart_continue_sample_cpp, bart_predict_cpp, BARTSamplesCpp + + +def _migrate_bart_v0_to_v1(serializer: JSONSerializer, loaded_version: int) -> None: + """In-place v0 -> v1 migration for a BART model envelope. + + Stamps the writer ``platform`` (inferred from structural fingerprints) and + renames positional forest keys (``forests/forest_0``, ...) to named keys + (``mean_forest`` / ``variance_forest``), driven by the ``include_*_forest`` + flags (unchanged across v0/v1). + """ + serializer.add_string("platform", infer_platform_v0(serializer, "python")) + include_mean = serializer.get_boolean_or_default("include_mean_forest", False) + include_variance = serializer.get_boolean_or_default( + "include_variance_forest", False + ) + if include_mean: + serializer.rename_field("forest_0", "mean_forest", subfolder_name="forests") + if include_variance: + serializer.rename_field( + "forest_1", "variance_forest", subfolder_name="forests" + ) + elif include_variance: + serializer.rename_field("forest_0", "variance_forest", subfolder_name="forests") class BARTModel: @@ -79,6 +94,117 @@ class BARTModel: def __init__(self) -> None: # Internal flag for whether the sample() method has been run self.sampled = False + # Single source of truth for the sampled forests + parameter traces (BARTSamplesCpp). + # The public forest_container_* / global_var_samples / leaf_scale_samples / num_samples + # attributes are properties backed by this object; _fc_*_cache hold the materialized + # (deep-copied) ForestContainer views for the deprecated direct forest accessor. + self._samples = None + self._fc_mean_cache = None + self._fc_variance_cache = None + + def _set_samples(self, samples) -> None: + """Install a new BARTSamplesCpp as the single source of truth and invalidate the + internal materialized forest-container caches so the next access re-derives them.""" + self._samples = samples + self._fc_mean_cache = None + self._fc_variance_cache = None + + @property + def samples(self): + """The single-owner ``BARTSamplesCpp`` holding the sampled forests and parameter traces.""" + return self._samples + + def extract_forest(self, forest: str = "mean"): + """Return a standalone deep copy of a sampled forest as a ``ForestContainer``. + + Parameters + ---------- + forest : str + Which forest to extract: ``"mean"`` or ``"variance"``. + + Returns + ------- + ForestContainer or None + A deep copy independent of the model, or ``None`` if that forest was not sampled. + """ + if self._samples is None: + raise RuntimeError("Model has not been sampled; no forests to extract.") + if forest == "mean": + if not self._samples.has_mean_forest(): + return None + cpp = self._samples.materialize_mean_forest() + output_dim = self.num_basis if self.has_basis else 1 + leaf_constant = not self.has_basis + fc = ForestContainer(cpp.NumTrees(), output_dim, leaf_constant, False) + fc.forest_container_cpp = cpp + return fc + elif forest == "variance": + if not self._samples.has_variance_forest(): + return None + cpp = self._samples.materialize_variance_forest() + fc = ForestContainer(cpp.NumTrees(), 1, True, True) + fc.forest_container_cpp = cpp + return fc + raise ValueError(f"Unknown forest '{forest}'; expected 'mean' or 'variance'.") + + @property + def _forest_container_mean(self): + # Internal cached deep copy for prediction / serialization. The Python samples + # wrapper exposes no borrowed-pointer accessor, so we materialize once and reuse. + if self._samples is None or not self._samples.has_mean_forest(): + return None + if self._fc_mean_cache is None: + cpp = self._samples.materialize_mean_forest() + output_dim = self.num_basis if self.has_basis else 1 + leaf_constant = not self.has_basis + fc = ForestContainer(cpp.NumTrees(), output_dim, leaf_constant, False) + fc.forest_container_cpp = cpp + self._fc_mean_cache = fc + return self._fc_mean_cache + + @property + def _forest_container_variance(self): + if self._samples is None or not self._samples.has_variance_forest(): + return None + if self._fc_variance_cache is None: + cpp = self._samples.materialize_variance_forest() + fc = ForestContainer(cpp.NumTrees(), 1, True, True) + fc.forest_container_cpp = cpp + self._fc_variance_cache = fc + return self._fc_variance_cache + + @property + def forest_container_mean(self): + raise AttributeError( + "`BARTModel.forest_container_mean` has been removed. The sampled forests are owned by " + "`model.samples`; extract a standalone copy with `model.extract_forest('mean')`." + ) + + @property + def forest_container_variance(self): + raise AttributeError( + "`BARTModel.forest_container_variance` has been removed. The sampled forests are owned " + "by `model.samples`; extract a standalone copy with `model.extract_forest('variance')`." + ) + + @property + def global_var_samples(self): + # None when unsampled or not sampled in this model (preserves getattr(..., None) semantics). + if self._samples is None: + return None + arr = self._samples.global_var_samples() + return arr if arr.size else None + + @property + def leaf_scale_samples(self): + if self._samples is None: + return None + arr = self._samples.leaf_scale_samples() + return arr if arr.size else None + + @property + def num_samples(self): + return self._samples.num_samples() if self._samples is not None else None def sample( self, @@ -91,6 +217,7 @@ def sample( leaf_basis_test: Optional[np.ndarray] = None, rfx_group_ids_test: Optional[np.ndarray] = None, rfx_basis_test: Optional[np.ndarray] = None, + observation_weights_train: Optional[np.ndarray] = None, observation_weights: Optional[np.ndarray] = None, num_gfr: int = 5, num_burnin: int = 0, @@ -127,7 +254,7 @@ def sample( test set evaluation for group labels that were not in the training set. rfx_basis_test : np.array, optional Optional test set basis for "random-slope" regression in additive random effects model. - observation_weights : np.array, optional + observation_weights_train : np.array, optional Optional vector of observation weights of length ``n_train``. Weights are applied as ``y_i | - ~ N(mu(X_i), sigma^2 / w_i)``, so larger weights increase an observation's influence on the fit. All weights must be non-negative. Defaults to ``None`` (all @@ -135,6 +262,8 @@ def sample( probit outcome models; not compatible with cloglog link functions. Note: these are referred to internally in the C++ layer as "variance weights" (``var_weights``), since they scale the residual variance. + observation_weights : np.array, optional + Deprecated alias for ``observation_weights_train``; will be removed in a future release. num_gfr : int, optional Number of "warm-start" iterations run using the grow-from-root algorithm (He and Hahn, 2021). Defaults to `5`. num_burnin : int, optional @@ -176,6 +305,7 @@ def sample( than there are samples in `previous_model_json`, a warning will be raised and only the last sample will be used. + Returns ------- self : BARTModel @@ -200,6 +330,7 @@ def sample( - **outcome_model** (*stochtree.OutcomeModel*): An object of class ``OutcomeModel`` specifying the outcome model. Default: ``OutcomeModel(outcome="continuous", link="identity")``. Pre-empts the deprecated ``probit_outcome_model`` parameter if specified. - **probit_outcome_model** (*bool*): Deprecated in favor of ``outcome_model``. Whether or not the outcome should be modeled as explicitly binary via a probit link. If ``True``, ``y`` must only contain the values ``0`` and ``1``. Default: ``False``. - **num_threads** (*int*): Number of threads to use in the GFR and MCMC algorithms, as well as prediction. Defaults to ``1`` (single-threaded). Set to ``-1`` to use the maximum number of available threads, or a positive integer for a specific count. OpenMP must be available for values other than ``1``. + - **verbose** (*bool*): Whether to print sampler progress (GFR / MCMC iteration updates) to the console during sampling. Defaults to ``False``. **mean_forest_params keys** @@ -215,6 +346,8 @@ def sample( - **keep_vars** (*list* or *np.array*): Variable names or column indices to include in the mean forest. Defaults to ``None``. - **drop_vars** (*list* or *np.array*): Variable names or column indices to exclude from the mean forest. Defaults to ``None``. Ignored if ``keep_vars`` is also set. - **num_features_subsample** (*int*): How many features to subsample when growing each tree for the GFR algorithm. Defaults to the number of features in the training dataset. + - **cloglog_leaf_prior_shape** (*float*): Shape parameter for the prior on leaf parameters in a cloglog ordinal leaf model. Defaults to ``2.0``. + - **cloglog_leaf_prior_scale** (*float*): Scale parameter for the prior on leaf parameters in a cloglog ordinal leaf model. Defaults to ``2.0``. **variance_forest_params keys** @@ -258,6 +391,7 @@ def sample( "outcome_model": OutcomeModel(outcome="continuous", link="identity"), "probit_outcome_model": False, "num_threads": 1, + "verbose": False, } general_params_updated = _preprocess_params( general_params_default, general_params @@ -277,6 +411,8 @@ def sample( "keep_vars": None, "drop_vars": None, "num_features_subsample": None, + "cloglog_leaf_prior_shape": 2.0, + "cloglog_leaf_prior_scale": 2.0, } mean_forest_params_updated = _preprocess_params( mean_forest_params_default, mean_forest_params @@ -332,6 +468,7 @@ def sample( self.probit_outcome_model = general_params_updated["probit_outcome_model"] self.outcome_model = general_params_updated["outcome_model"] num_threads = general_params_updated["num_threads"] + verbose = general_params_updated["verbose"] # 2. Mean forest parameters num_trees_mean = mean_forest_params_updated["num_trees"] @@ -348,6 +485,12 @@ def sample( num_features_subsample_mean = mean_forest_params_updated[ "num_features_subsample" ] + cloglog_leaf_prior_shape = mean_forest_params_updated[ + "cloglog_leaf_prior_shape" + ] + cloglog_leaf_prior_scale = mean_forest_params_updated[ + "cloglog_leaf_prior_scale" + ] # 3. Variance forest parameters num_trees_variance = variance_forest_params_updated["num_trees"] @@ -453,15 +596,28 @@ def sample( self.include_mean_forest = True if num_trees_mean > 0 else False self.include_variance_forest = True if num_trees_variance > 0 else False - # observation_weights compatibility checks + # `observation_weights` was renamed to `observation_weights_train`; honor the + # deprecated argument for one release cycle. if observation_weights is not None: + warnings.warn( + "`observation_weights` is deprecated and will be removed in a future " + "release; use `observation_weights_train` instead.", + DeprecationWarning, + stacklevel=2, + ) + if observation_weights_train is None: + observation_weights_train = observation_weights + + # observation_weights_train compatibility checks + if observation_weights_train is not None: if link_is_cloglog: raise ValueError( - "observation_weights are not compatible with cloglog link functions." + "observation_weights_train are not compatible with cloglog link functions." ) if self.include_variance_forest: - warnings.warn( - "Results may be unreliable when observation_weights are deployed alongside a variance forest model." + raise ValueError( + "observation_weights_train are not compatible with a variance forest model." + "Use either observation_weights_train or a variance forest, not both." ) # Check data inputs @@ -502,24 +658,24 @@ def sample( if rfx_basis_test is not None: if not isinstance(rfx_basis_test, np.ndarray): raise ValueError("rfx_basis_test must be a numpy array") - if observation_weights is not None: - if not isinstance(observation_weights, np.ndarray): - raise ValueError("observation_weights must be a numpy array") - observation_weights_ = np.squeeze(observation_weights) - if observation_weights_.ndim != 1: + if observation_weights_train is not None: + if not isinstance(observation_weights_train, np.ndarray): + raise ValueError("observation_weights_train must be a numpy array") + observation_weights_train_ = np.squeeze(observation_weights_train) + if observation_weights_train_.ndim != 1: raise ValueError( - "observation_weights must be a 1-dimensional numpy array" + "observation_weights_train must be a 1-dimensional numpy array" ) - if np.any(observation_weights_ < 0): - raise ValueError("observation_weights cannot have any negative values") + if np.any(observation_weights_train_ < 0): + raise ValueError("observation_weights_train cannot have any negative values") - # Validate that observation_weights are not all-zero when num_gfr > 0 - if observation_weights is not None and num_gfr > 0: - if np.all(observation_weights == 0): + # Validate that observation_weights_train are not all-zero when num_gfr > 0 + if observation_weights_train is not None and num_gfr > 0: + if np.all(observation_weights_train == 0): raise ValueError( - "observation_weights are all zero (prior sampling mode) but num_gfr > 0. " + "observation_weights_train are all zero (prior sampling mode) but num_gfr > 0. " "GFR warm-start is data-dependent and ill-defined with zero weights. " - "Set num_gfr=0 when using all-zero observation_weights." + "Set num_gfr=0 when using all-zero observation_weights_train." ) # Convert everything to standard shape (2-dimensional) @@ -685,16 +841,15 @@ def sample( # Covariate preprocessing self._covariate_preprocessor = CovariatePreprocessor() self._covariate_preprocessor.fit(X_train) - X_train_processed = self._covariate_preprocessor.transform(X_train) + X_train_processed = self._covariate_preprocessor.transform(X_train).astype(np.float64) if X_test is not None: - X_test_processed = self._covariate_preprocessor.transform(X_test) + X_test_processed = self._covariate_preprocessor.transform(X_test).astype(np.float64) feature_types = np.asarray( self._covariate_preprocessor._processed_feature_types ) original_var_indices = ( self._covariate_preprocessor.fetch_original_feature_indices() ) - num_features = len(feature_types) # Determine whether a test set is provided self.has_test = X_test is not None @@ -946,32 +1101,16 @@ def sample( raise ValueError( "`previous_model_warmstart_sample_num` exceeds the number of samples in `previous_model_json`" ) - previous_model_decrement = True if num_chains > previous_model_warmstart_sample_num + 1: warnings.warn( "The number of chains being sampled exceeds the number of previous model samples available from the requested position in `previous_model_json`. All chains will be initialized from the same sample." ) - previous_model_decrement = False - previous_y_scale = previous_bart_model.y_std previous_model_num_samples = previous_bart_model.num_samples - if previous_bart_model.sample_sigma2_global: - previous_global_var_samples = previous_bart_model.global_var_samples / ( - previous_y_scale * previous_y_scale - ) - else: - previous_global_var_samples = None - if previous_bart_model.sample_sigma2_leaf: - previous_leaf_var_samples = previous_bart_model.leaf_scale_samples - else: - previous_leaf_var_samples = None if previous_model_warmstart_sample_num + 1 > previous_model_num_samples: raise ValueError( "`previous_model_warmstart_sample_num` exceeds the number of samples in `previous_model_json`" ) else: - previous_y_scale = None - previous_global_var_samples = None - previous_leaf_var_samples = None previous_model_num_samples = 0 # Update variable weights if the covariates have been resized (by e.g. one-hot encoding) @@ -1020,6 +1159,7 @@ def sample( # Preliminary runtime checks for probit link if not self.include_mean_forest: link_is_probit = False + sample_sigma2_leaf = False if link_is_probit: if np.unique(y_train).size != 2: raise ValueError( @@ -1086,203 +1226,6 @@ def sample( ) sample_sigma2_global = False - # Handle standardization, prior calibration, and initialization of forest - # differently for binary and continuous outcomes - if link_is_probit: - # Compute a probit-scale offset and fix scale to 1 - self.y_bar = norm.ppf(np.squeeze(np.mean(y_train))) - self.y_std = 1.0 - - # Set a pseudo outcome by subtracting mean(y_train) from y_train - resid_train = y_train - np.squeeze(np.mean(y_train)) - - # Set initial values of root nodes to 0.0 (in probit scale) - init_val_mean = 0.0 - - # Calibrate priors for sigma^2 and tau - # Set sigma2_init to 1, ignoring default provided - sigma2_init = 1.0 - current_sigma2 = sigma2_init - self.sigma2_init = sigma2_init - # Skip variance_forest_init, since variance forests are not supported with probit link - b_leaf = 1.0 / num_trees_mean if b_leaf is None else b_leaf - if self.has_basis: - if sigma2_leaf is None: - current_leaf_scale = np.zeros( - (self.num_basis, self.num_basis), dtype=float - ) - np.fill_diagonal( - current_leaf_scale, - 2.0 / num_trees_mean, - ) - elif isinstance(sigma2_leaf, (float, np.floating)): - current_leaf_scale = np.zeros( - (self.num_basis, self.num_basis), dtype=float - ) - np.fill_diagonal(current_leaf_scale, sigma2_leaf) - elif isinstance(sigma2_leaf, np.ndarray): - if sigma2_leaf.ndim != 2: - raise ValueError( - "sigma2_leaf must be a 2d symmetric numpy array if provided in matrix form" - ) - if sigma2_leaf.shape[0] != sigma2_leaf.shape[1]: - raise ValueError( - "sigma2_leaf must be a 2d symmetric numpy array if provided in matrix form" - ) - if sigma2_leaf.shape[0] != self.num_basis: - raise ValueError( - "sigma2_leaf must be a 2d symmetric numpy array with its dimensionality matching the basis dimension" - ) - current_leaf_scale = sigma2_leaf - else: - raise ValueError( - "sigma2_leaf must be either a scalar or a 2d symmetric numpy array" - ) - else: - if sigma2_leaf is None: - current_leaf_scale = np.array([[2.0 / num_trees_mean]]) - elif isinstance(sigma2_leaf, (float, np.floating)): - current_leaf_scale = np.array([[sigma2_leaf]]) - elif isinstance(sigma2_leaf, np.ndarray): - if sigma2_leaf.ndim != 2: - raise ValueError( - "sigma2_leaf must be a 2d symmetric numpy array if provided in matrix form" - ) - if sigma2_leaf.shape[0] != sigma2_leaf.shape[1]: - raise ValueError( - "sigma2_leaf must be a 2d symmetric numpy array if provided in matrix form" - ) - if sigma2_leaf.shape[0] != 1: - raise ValueError( - "sigma2_leaf must be a 1x1 numpy array for this leaf model" - ) - current_leaf_scale = sigma2_leaf - else: - raise ValueError( - "sigma2_leaf must be either a scalar or a 2d numpy array" - ) - elif link_is_cloglog: - # Fix offset to 0 and scale to 1 - self.y_bar = 0 - self.y_std = 1 - - # Remap outcomes to start from 0 - resid_train = y_train - np.min(unique_outcomes) - cloglog_num_categories = int(np.max(resid_train)) + 1 - - # Set initial values of root nodes to 0.0 (in linear scale) - init_val_mean = 0.0 - - # Calibrate priors for sigma^2 and tau - sigma2_init = 1.0 - current_sigma2 = sigma2_init - self.sigma2_init = sigma2_init - current_leaf_scale = np.array([[2.0 / num_trees_mean]]) - - # Set first cutpoint to 0 for identifiability - cloglog_cutpoint_0 = 0.0 - - # Set shape and rate parameters for conditional gamma model - cloglog_forest_shape = 2.0 - cloglog_forest_rate = 2.0 - else: - # Standardize if requested - if self.standardize: - self.y_bar = np.squeeze(np.mean(y_train)) - self.y_std = np.squeeze(np.std(y_train)) - else: - self.y_bar = 0 - self.y_std = 1 - - # Compute residual value - resid_train = (y_train - self.y_bar) / self.y_std - - # Compute initial value of root nodes in mean forest - init_val_mean = np.squeeze(np.mean(resid_train)) - - # Calibrate priors for global sigma^2 and sigma2_leaf - if not sigma2_init: - sigma2_init = 1.0 * np.var(resid_train) - if not variance_forest_leaf_init: - variance_forest_leaf_init = 0.6 * np.var(resid_train) - current_sigma2 = sigma2_init - self.sigma2_init = sigma2_init - if self.include_mean_forest: - b_leaf = ( - np.squeeze(np.var(resid_train)) / num_trees_mean - if b_leaf is None - else b_leaf - ) - if self.has_basis: - if sigma2_leaf is None: - current_leaf_scale = np.zeros( - (self.num_basis, self.num_basis), dtype=float - ) - np.fill_diagonal( - current_leaf_scale, - np.squeeze(np.var(resid_train)) / num_trees_mean, - ) - elif isinstance(sigma2_leaf, (float, np.floating)): - current_leaf_scale = np.zeros( - (self.num_basis, self.num_basis), dtype=float - ) - np.fill_diagonal(current_leaf_scale, sigma2_leaf) - elif isinstance(sigma2_leaf, np.ndarray): - if sigma2_leaf.ndim != 2: - raise ValueError( - "sigma2_leaf must be a 2d symmetric numpy array if provided in matrix form" - ) - if sigma2_leaf.shape[0] != sigma2_leaf.shape[1]: - raise ValueError( - "sigma2_leaf must be a 2d symmetric numpy array if provided in matrix form" - ) - if sigma2_leaf.shape[0] != self.num_basis: - raise ValueError( - "sigma2_leaf must be a 2d symmetric numpy array with its dimensionality matching the basis dimension" - ) - current_leaf_scale = sigma2_leaf - else: - raise ValueError( - "sigma2_leaf must be either a scalar or a 2d symmetric numpy array" - ) - else: - if sigma2_leaf is None: - current_leaf_scale = np.array([ - [np.squeeze(np.var(resid_train)) / num_trees_mean] - ]) - elif isinstance(sigma2_leaf, (float, np.floating)): - current_leaf_scale = np.array([[sigma2_leaf]]) - elif isinstance(sigma2_leaf, np.ndarray): - if sigma2_leaf.ndim != 2: - raise ValueError( - "sigma2_leaf must be a 2d symmetric numpy array if provided in matrix form" - ) - if sigma2_leaf.shape[0] != sigma2_leaf.shape[1]: - raise ValueError( - "sigma2_leaf must be a 2d symmetric numpy array if provided in matrix form" - ) - if sigma2_leaf.shape[0] != 1: - raise ValueError( - "sigma2_leaf must be a 1x1 numpy array for this leaf model" - ) - current_leaf_scale = sigma2_leaf - else: - raise ValueError( - "sigma2_leaf must be either a scalar or a 2d numpy array" - ) - else: - current_leaf_scale = np.array([[1.0]]) - if self.include_variance_forest: - if not a_forest: - a_forest = num_trees_variance / a_0**2 + 0.5 - if not b_forest: - b_forest = num_trees_variance / a_0**2 - else: - if not a_forest: - a_forest = 1.0 - if not b_forest: - b_forest = 1.0 - # Runtime checks on RFX group ids self.has_rfx = False has_rfx_test = False @@ -1324,892 +1267,494 @@ def sample( elif self.rfx_model_spec == "intercept_only": if rfx_basis_test is None: rfx_basis_test = np.ones((rfx_group_ids_test.shape[0], 1)) - # Set up random effects structures - if self.has_rfx: - # Prior parameters - if rfx_working_parameter_prior_mean is None: - if num_rfx_components == 1: - alpha_init = np.array([0.0], dtype=float) - elif num_rfx_components > 1: - alpha_init = np.zeros(num_rfx_components, dtype=float) - else: - raise ValueError("There must be at least 1 random effect component") - else: - alpha_init = _expand_dims_1d( - rfx_working_parameter_prior_mean, num_rfx_components - ) - - if rfx_group_parameter_prior_mean is None: - xi_init = np.tile(np.expand_dims(alpha_init, 1), (1, num_rfx_groups)) - else: - xi_init = _expand_dims_2d( - rfx_group_parameter_prior_mean, num_rfx_components, num_rfx_groups - ) - - if rfx_working_parameter_prior_cov is None: - sigma_alpha_init = np.identity(num_rfx_components) - else: - sigma_alpha_init = _expand_dims_2d_diag( - rfx_working_parameter_prior_cov, num_rfx_components - ) - - if rfx_group_parameter_prior_cov is None: - sigma_xi_init = np.identity(num_rfx_components) - else: - sigma_xi_init = _expand_dims_2d_diag( - rfx_group_parameter_prior_cov, num_rfx_components - ) - - sigma_xi_shape = rfx_variance_prior_shape - sigma_xi_scale = rfx_variance_prior_scale - - # Random effects sampling data structures - rfx_dataset_train = RandomEffectsDataset() - rfx_dataset_train.add_group_labels(rfx_group_ids_train) - rfx_dataset_train.add_basis(rfx_basis_train) - rfx_tracker = RandomEffectsTracker(rfx_group_ids_train) - rfx_model = RandomEffectsModel(num_rfx_components, num_rfx_groups) - rfx_model.set_working_parameter(alpha_init) - rfx_model.set_group_parameters(xi_init) - rfx_model.set_working_parameter_covariance(sigma_alpha_init) - rfx_model.set_group_parameter_covariance(sigma_xi_init) - rfx_model.set_variance_prior_shape(sigma_xi_shape) - rfx_model.set_variance_prior_scale(sigma_xi_scale) - self.rfx_container = RandomEffectsContainer() - self.rfx_container.load_new_container( - num_rfx_components, num_rfx_groups, rfx_tracker - ) - - # Container of variance parameter samples - self.num_gfr = num_gfr - self.num_burnin = num_burnin - self.num_mcmc = num_mcmc - self.num_chains = num_chains - self.keep_every = keep_every - num_temp_samples = num_gfr + num_burnin + num_mcmc * keep_every - num_retained_samples = num_mcmc * num_chains - # Delete GFR samples from these containers after the fact if desired - # if keep_gfr: - # num_retained_samples += num_gfr - num_retained_samples += num_gfr - if keep_burnin: - num_retained_samples += num_burnin * num_chains - self.num_samples = num_retained_samples - self.sample_sigma2_global = sample_sigma2_global - self.sample_sigma2_leaf = sample_sigma2_leaf - if sample_sigma2_global: - self.global_var_samples = np.empty(self.num_samples, dtype=np.float64) - if sample_sigma2_leaf: - self.leaf_scale_samples = np.empty(self.num_samples, dtype=np.float64) - if self.include_mean_forest: - yhat_train_raw = np.empty( - (self.n_train, self.num_samples), dtype=np.float64 - ) - if self.include_variance_forest: - sigma2_x_train_raw = np.empty( - (self.n_train, self.num_samples), dtype=np.float64 - ) - sample_counter = -1 - - # Forest Dataset (covariates and optional basis) - forest_dataset_train = Dataset() - forest_dataset_train.add_covariates(X_train_processed) - if self.has_basis: - forest_dataset_train.add_basis(leaf_basis_train) - if observation_weights is not None: - forest_dataset_train.add_variance_weights(observation_weights_) - if self.has_test: - forest_dataset_test = Dataset() - forest_dataset_test.add_covariates(X_test_processed) - if self.has_basis: - forest_dataset_test.add_basis(leaf_basis_test) - - # Residual - residual_train = Residual(resid_train) - - # C++ and Numpy random number generator - if random_seed is None: - cpp_rng = RNG(-1) - self.rng = np.random.default_rng() - else: - cpp_rng = RNG(random_seed) - self.rng = np.random.default_rng(random_seed) - - # Set variance leaf model type (currently only one option) - leaf_model_variance_forest = 3 - leaf_dimension_variance = 1 # Determine the mean forest leaf model type if link_is_cloglog and not self.has_basis: leaf_model_mean_forest = 4 - leaf_dimension_mean = 1 elif not self.has_basis: leaf_model_mean_forest = 0 - leaf_dimension_mean = 1 elif self.num_basis == 1: leaf_model_mean_forest = 1 - leaf_dimension_mean = 1 else: leaf_model_mean_forest = 2 - leaf_dimension_mean = self.num_basis - - # Sampling data structures - global_model_config = GlobalModelConfig(global_error_variance=current_sigma2) - if self.include_mean_forest: - forest_model_config_mean = ForestModelConfig( - num_trees=num_trees_mean, - num_features=num_features, - num_observations=self.n_train, - feature_types=feature_types, - variable_weights=variable_weights_mean, - leaf_dimension=leaf_dimension_mean, - alpha=alpha_mean, - beta=beta_mean, - min_samples_leaf=min_samples_leaf_mean, - max_depth=max_depth_mean, - leaf_model_type=leaf_model_mean_forest, - leaf_model_scale=current_leaf_scale, - cutpoint_grid_size=cutpoint_grid_size, - num_features_subsample=num_features_subsample_mean, - ) - if link_is_cloglog: - forest_model_config_mean.update_cloglog_forest_shape( - cloglog_forest_shape - ) - forest_model_config_mean.update_cloglog_forest_rate(cloglog_forest_rate) - forest_sampler_mean = ForestSampler( - forest_dataset_train, - global_model_config, - forest_model_config_mean, - ) - if self.include_variance_forest: - forest_model_config_variance = ForestModelConfig( - num_trees=num_trees_variance, - num_features=num_features, - num_observations=self.n_train, - feature_types=feature_types, - variable_weights=variable_weights_variance, - leaf_dimension=leaf_dimension_variance, - alpha=alpha_variance, - beta=beta_variance, - min_samples_leaf=min_samples_leaf_variance, - max_depth=max_depth_variance, - leaf_model_type=leaf_model_variance_forest, - cutpoint_grid_size=cutpoint_grid_size, - variance_forest_shape=a_forest, - variance_forest_scale=b_forest, - num_features_subsample=num_features_subsample_variance, - ) - forest_sampler_variance = ForestSampler( - forest_dataset_train, - global_model_config, - forest_model_config_variance, - ) - - # Container of forest samples - if self.include_mean_forest: - self.forest_container_mean = ( - ForestContainer(num_trees_mean, 1, True, False) - if not self.has_basis - else ForestContainer(num_trees_mean, self.num_basis, False, False) - ) - active_forest_mean = ( - Forest(num_trees_mean, 1, True, False) - if not self.has_basis - else Forest(num_trees_mean, self.num_basis, False, False) - ) - if self.include_variance_forest: - self.forest_container_variance = ForestContainer( - num_trees_variance, 1, True, True - ) - active_forest_variance = Forest(num_trees_variance, 1, True, True) - - # Variance samplers - if self.sample_sigma2_global: - global_var_model = GlobalVarianceModel() - if self.sample_sigma2_leaf: - leaf_var_model = LeafVarianceModel() - - # Initialize the leaves of each tree in the mean forest - if self.include_mean_forest: - if self.has_basis: - init_val_mean = np.repeat(0.0, leaf_basis_train.shape[1]) - else: - init_val_mean = np.array([0.0]) - forest_sampler_mean.prepare_for_sampler( - forest_dataset_train, - residual_train, - active_forest_mean, - leaf_model_mean_forest, - init_val_mean, - ) - - # Initialize the leaves of each tree in the variance forest - if self.include_variance_forest: - init_val_variance = np.array([variance_forest_leaf_init]) - forest_sampler_variance.prepare_for_sampler( - forest_dataset_train, - residual_train, - active_forest_variance, - leaf_model_variance_forest, - init_val_variance, - ) - - # Initialize auxiliary data and ordinal sampler for cloglog - if link_is_cloglog: - ordinal_sampler = OrdinalSampler() - train_size = self.n_train - - # Slot 0: Latent variable Z (size n_train) - forest_dataset_train.add_auxiliary_dimension(train_size) - # Slot 1: Forest predictions eta (size n_train) - forest_dataset_train.add_auxiliary_dimension(train_size) - # Slot 2: Log-scale cutpoints gamma (size cloglog_num_categories - 1) - forest_dataset_train.add_auxiliary_dimension(cloglog_num_categories - 1) - # Slot 3: Cumulative exp cutpoints seg (size cloglog_num_categories) - forest_dataset_train.add_auxiliary_dimension(cloglog_num_categories) - - # Initialize all slots to 0 - for j in range(train_size): - forest_dataset_train.set_auxiliary_data_value(0, j, 0.0) - forest_dataset_train.set_auxiliary_data_value(1, j, 0.0) - for j in range(cloglog_num_categories - 1): - forest_dataset_train.set_auxiliary_data_value(2, j, 0.0) - - # Compute initial cumulative exp sums - ordinal_sampler.update_cumulative_exp_sums(forest_dataset_train) - - # Allocate storage for cutpoint samples - cloglog_cutpoint_samples = np.full( - (cloglog_num_categories - 1, num_retained_samples), np.nan - ) - # Run GFR (warm start) if specified - if self.num_gfr > 0: - for i in range(self.num_gfr): - # Keep all GFR samples at this stage -- remove from ForestSamples after MCMC - # keep_sample = keep_gfr - keep_sample = True - if keep_sample: - sample_counter += 1 - if self.include_mean_forest: - if link_is_probit: - # Sample latent probit variable z | - - outcome_pred = active_forest_mean.predict(forest_dataset_train) - if self.has_rfx: - rfx_pred = rfx_model.predict(rfx_dataset_train, rfx_tracker) - outcome_pred = outcome_pred + rfx_pred - # Full probit-scale predictor: forest learns z - y_bar, so add y_bar back - eta_pred = outcome_pred + self.y_bar - mu0 = eta_pred[y_train[:, 0] == 0] - mu1 = eta_pred[y_train[:, 0] == 1] - n0 = np.sum(y_train[:, 0] == 0) - n1 = np.sum(y_train[:, 0] == 1) - u0 = self.rng.uniform( - low=0.0, - high=norm.cdf(0 - mu0), - size=n0, - ) - u1 = self.rng.uniform( - low=norm.cdf(0 - mu1), - high=1.0, - size=n1, - ) - resid_train[y_train[:, 0] == 0, 0] = mu0 + norm.ppf(u0) - resid_train[y_train[:, 0] == 1, 0] = mu1 + norm.ppf(u1) - - # Update outcome: center z by y_bar before passing to forest - new_outcome = ( - np.squeeze(resid_train) - self.y_bar - outcome_pred - ) - residual_train.update_data(new_outcome) - - # Sample the mean forest - forest_sampler_mean.sample_one_iteration( - self.forest_container_mean, - active_forest_mean, - forest_dataset_train, - residual_train, - cpp_rng, - global_model_config, - forest_model_config_mean, - keep_sample, - True, - num_threads, - ) - # Cache train set predictions since they are already computed during sampling - if keep_sample: - yhat_train_raw[:, sample_counter] = ( - forest_sampler_mean.get_cached_forest_predictions() - ) - - # Sample the variance forest - if self.include_variance_forest: - forest_sampler_variance.sample_one_iteration( - self.forest_container_variance, - active_forest_variance, - forest_dataset_train, - residual_train, - cpp_rng, - global_model_config, - forest_model_config_variance, - keep_sample, - True, - num_threads, - ) - - # Cache train set predictions since they are already computed during sampling - if keep_sample: - sigma2_x_train_raw[:, sample_counter] = ( - forest_sampler_variance.get_cached_forest_predictions() - ) + # Determine cloglog number of classes + cloglog_num_categories = ( + int(np.max(y_train - np.min(y_train))) + 1 if link_is_cloglog else 0 + ) - # Sample variance parameters (if requested) - if self.sample_sigma2_global: - current_sigma2 = global_var_model.sample_one_iteration( - residual_train, cpp_rng, a_global, b_global - ) - global_model_config.update_global_error_variance(current_sigma2) - if keep_sample: - self.global_var_samples[sample_counter] = current_sigma2 - if self.sample_sigma2_leaf: - current_leaf_scale[0, 0] = leaf_var_model.sample_one_iteration( - active_forest_mean, cpp_rng, a_leaf, b_leaf - ) - forest_model_config_mean.update_leaf_model_scale(current_leaf_scale) - if keep_sample: - self.leaf_scale_samples[sample_counter] = current_leaf_scale[ - 0, 0 - ] - - # Sample random effects - if self.has_rfx: - rfx_model.sample( - rfx_dataset_train, - residual_train, - rfx_tracker, - self.rfx_container, - keep_sample, - current_sigma2, - cpp_rng, - ) + # Expand dimensions on RFX prior parameters if provided + # Working parameter (should be expanded to a 1d array if provided as a scalar) + if rfx_working_parameter_prior_mean is not None: + rfx_working_parameter_prior_mean = _expand_dims_1d( + rfx_working_parameter_prior_mean, num_rfx_components + ) + + # Group parameter (should be expanded to a 2d array if provided as a scalar) + if rfx_group_parameter_prior_mean is not None: + rfx_group_parameter_prior_mean = _expand_dims_2d( + rfx_group_parameter_prior_mean, num_rfx_components, num_rfx_groups + ) + + # Working parameter (should be expanded to a diagonal matrix if provided as a scalar) + if rfx_working_parameter_prior_cov is not None: + rfx_working_parameter_prior_cov = _expand_dims_2d_diag( + rfx_working_parameter_prior_cov, num_rfx_components + ) + + # Group parameter (should be expanded to a diagonal matrix if provided as a scalar) + if rfx_group_parameter_prior_cov is not None: + rfx_group_parameter_prior_cov = _expand_dims_2d_diag( + rfx_group_parameter_prior_cov, num_rfx_components + ) + + # Arrange all config in a large python dictionary + bart_config = { + "standardize_outcome": self.standardize, + "num_threads": num_threads, + "verbose": verbose, + "cutpoint_grid_size": cutpoint_grid_size, + "link_function": 0 + if self.outcome_model.link == "identity" + else (1 if self.outcome_model.link == "probit" else 2), + "outcome_type": 0 + if self.outcome_model.outcome == "continuous" + else (1 if self.outcome_model.outcome == "binary" else 2), + "random_seed": random_seed, + "keep_gfr": keep_gfr, + "keep_burnin": keep_burnin, + "a_sigma2_global": a_global, + "b_sigma2_global": b_global, + "sigma2_global_init": sigma2_init, + "sample_sigma2_global": sample_sigma2_global, + "num_trees_mean": num_trees_mean, + "alpha_mean": alpha_mean, + "beta_mean": beta_mean, + "min_samples_leaf_mean": min_samples_leaf_mean, + "max_depth_mean": max_depth_mean, + "leaf_constant_mean": False if self.has_basis else True, + "leaf_dim_mean": self.num_basis if self.has_basis else 1, + "exponentiated_leaf_mean": False, + "num_features_subsample_mean": num_features_subsample_mean, + "a_sigma2_mean": a_leaf, + "b_sigma2_mean": b_leaf, + "sigma2_mean_init": sigma2_leaf + if isinstance(sigma2_leaf, float) + else -1.0, + "sample_sigma2_leaf_mean": sample_sigma2_leaf, + "mean_leaf_model_type": leaf_model_mean_forest, + "sigma2_leaf_mean_matrix": sigma2_leaf.flatten(order="F") + if isinstance(sigma2_leaf, np.ndarray) + else None, + "num_classes_cloglog": cloglog_num_categories, + "cloglog_leaf_prior_shape": cloglog_leaf_prior_shape, + "cloglog_leaf_prior_scale": cloglog_leaf_prior_scale, + "cloglog_cutpoint_0": 0, + "num_trees_variance": num_trees_variance, + "leaf_prior_calibration_param": a_0, + "shape_variance_forest": a_forest, + "scale_variance_forest": b_forest, + "variance_forest_leaf_init": variance_forest_leaf_init, + "alpha_variance": alpha_variance, + "beta_variance": beta_variance, + "min_samples_leaf_variance": min_samples_leaf_variance, + "max_depth_variance": max_depth_variance, + "leaf_constant_variance": True, + "leaf_dim_variance": 1, + "exponentiated_leaf_variance": True, + "num_features_subsample_variance": num_features_subsample_variance, + "feature_types": feature_types.astype(int), + "sweep_update_indices_mean": list(range(num_trees_mean)) + if num_trees_mean > 0 + else None, + "sweep_update_indices_variance": list(range(num_trees_variance)) + if num_trees_variance > 0 + else None, + "var_weights_mean": variable_weights_mean, + "var_weights_variance": variable_weights_variance, + "has_random_effects": self.has_rfx, + "rfx_model_spec": 0 + if self.rfx_model_spec == "custom" + else (1 if self.rfx_model_spec == "intercept_only" else None), + "rfx_working_parameter_mean_prior": rfx_working_parameter_prior_mean + if self.has_rfx + else None, + "rfx_group_parameter_mean_prior": rfx_group_parameter_prior_mean + if self.has_rfx + else None, + "rfx_working_parameter_cov_prior": rfx_working_parameter_prior_cov + if self.has_rfx + else None, + "rfx_group_parameter_cov_prior": rfx_group_parameter_prior_cov + if self.has_rfx + else None, + "rfx_variance_prior_shape": rfx_variance_prior_shape + if self.has_rfx + else None, + "rfx_variance_prior_scale": rfx_variance_prior_scale + if self.has_rfx + else None, + } - # Cloglog Gibbs updates - if link_is_cloglog: - # Update auxiliary data slot 1 with current forest predictions - forest_pred_current = ( - forest_sampler_mean.get_cached_forest_predictions() - ) - for j in range(train_size): - forest_dataset_train.set_auxiliary_data_value( - 1, j, forest_pred_current[j] - ) + # Remove None values from config (alternative is to check for Nones on the C++ side when unpacking into non-optional types) + bart_config = {k: v for k, v in bart_config.items() if v is not None} + + # Cache the config so that continue_sampling() can warm-start with the + # exact same priors/structure without re-deriving them. + self._cached_bart_config = bart_config + + # Convert arrays to F-contiguous (column-major) before calling C++. + # convert_numpy_to_bart_data stores raw pointers into these arrays; if + # pybind11 has to make a copy (wrong dtype or wrong order) that copy is + # destroyed when the helper returns, leaving a dangling pointer. + # Passing already-correct arrays causes pybind11 to return a view of + # the original, which remains alive in this Python scope. + X_train_cpp = np.asfortranarray(X_train_processed) + y_train_remapped = y_train - np.min(y_train) if link_is_cloglog else y_train + y_train_cpp = np.asfortranarray(y_train_remapped, dtype=np.float64) + X_test_cpp = np.asfortranarray(X_test_processed) if self.has_test else None + basis_train_cpp = ( + np.asfortranarray(leaf_basis_train.astype(np.float64)) if self.has_basis else None + ) + basis_test_cpp = ( + np.asfortranarray(leaf_basis_test.astype(np.float64)) + if self.has_basis and self.has_test + else None + ) + # rfx group IDs must be int32: pybind11 casts int64→int32 via a temporary + # inside convert_numpy_to_bart_data, making the returned raw pointer dangle. + rfx_group_ids_train_cpp = ( + rfx_group_ids_train.astype(np.int32) if rfx_group_ids_train is not None else None + ) + rfx_group_ids_test_cpp = ( + rfx_group_ids_test.astype(np.int32) if rfx_group_ids_test is not None else None + ) + rfx_basis_train_cpp = ( + np.asfortranarray(rfx_basis_train.astype(np.float64)) if rfx_basis_train is not None else None + ) + rfx_basis_test_cpp = ( + np.asfortranarray(rfx_basis_test.astype(np.float64)) if rfx_basis_test is not None else None + ) - # Sample latent z_i's using truncated exponential - ordinal_sampler.update_latent_variables( - forest_dataset_train, residual_train, cpp_rng - ) + # Run the BART sampler from C++ + bart_results = bart_sample_cpp( + X_train=X_train_cpp, + y_train=y_train_cpp, + X_test=X_test_cpp, + n_train=X_train_cpp.shape[0], + n_test=X_test_cpp.shape[0] if self.has_test else 0, + p=X_train_cpp.shape[1], + basis_train=basis_train_cpp, + basis_test=basis_test_cpp, + basis_dim=self.num_basis if self.has_basis else 0, + obs_weights_train=observation_weights_train + if observation_weights_train is not None + else None, + obs_weights_test=None, + rfx_group_ids_train=rfx_group_ids_train_cpp, + rfx_group_ids_test=rfx_group_ids_test_cpp, + rfx_basis_train=rfx_basis_train_cpp, + rfx_basis_test=rfx_basis_test_cpp, + rfx_num_groups=num_rfx_groups if self.has_rfx else 0, + rfx_basis_dim=self.num_rfx_basis if self.has_rfx else 0, + num_gfr=num_gfr, + num_burnin=num_burnin, + keep_every=keep_every, + num_mcmc=num_mcmc, + num_chains=num_chains, + config_input=bart_config, + ) - # Sample gamma parameters (cutpoints) - ordinal_sampler.update_gamma_params( - forest_dataset_train, - residual_train, - cloglog_forest_shape, - cloglog_forest_rate, - cloglog_cutpoint_0, - cpp_rng, - ) + # Store high level model metadata from C++ results + self.num_gfr = num_gfr + self.num_burnin = num_burnin + self.keep_every = keep_every + self.num_mcmc = num_mcmc + self.num_chains = num_chains + self.sample_sigma2_global = sample_sigma2_global + self.sample_sigma2_leaf = sample_sigma2_leaf - # Update cumulative sum of exp(gamma) values - ordinal_sampler.update_cumulative_exp_sums(forest_dataset_train) + # Persist the final RNG state so continue_sampling() can resume the random + # stream for bit-identical continuation. Only populated for single-chain runs + # (multi-chain runs use one RNG per chain, so there is no single stream to resume). + self.rng_state = bart_results.get("rng_state", None) + # Persist the leaf normal sampler's cached spare value (Marsaglia-polar) so a + # continued chain reproduces the one-shot leaf draws exactly. + self.leaf_normal_cache = bart_results.get("leaf_normal_cache", None) + + # Unpack standardization params computed by C++ sampler + self.y_bar = bart_results["y_bar"] + self.y_std = bart_results["y_std"] + self.sigma2_init = bart_results["sigma2_init"] + self.sigma2_leaf_init = ( + bart_results["sigma2_mean_init"] if self.include_mean_forest else None + ) + self.b_leaf = ( + bart_results["b_sigma2_mean"] if self.include_mean_forest else None + ) + self.shape_variance_forest = ( + bart_results["shape_variance_forest"] + if self.include_variance_forest + else None + ) + self.scale_variance_forest = ( + bart_results["scale_variance_forest"] + if self.include_variance_forest + else None + ) - # Retain cutpoint draw - if keep_sample: - cloglog_cutpoint_samples[:, sample_counter] = ( - forest_dataset_train.get_auxiliary_data_vector(2) - ) + # Unpack mean forest predictions. The forests themselves are owned by self._samples + # (assembled below); predictions are transient per-data outputs. + if self.include_mean_forest: + mean_forest_preds_train = bart_results[ + "mean_forest_predictions_train" + ].reshape(self.n_train, bart_results["num_samples"], order="F") + self.y_hat_train = mean_forest_preds_train * self.y_std + self.y_bar + if self.has_test: + mean_forest_preds_test = bart_results[ + "mean_forest_predictions_test" + ].reshape(self.n_test, bart_results["num_samples"], order="F") + self.y_hat_test = mean_forest_preds_test * self.y_std + self.y_bar - # Run MCMC - if self.num_burnin + self.num_mcmc > 0: - for chain_num in range(num_chains): - if num_gfr > 0: - forest_ind = num_gfr - chain_num - 1 - # Reset mean forest - if self.include_mean_forest: - active_forest_mean.reset(self.forest_container_mean, forest_ind) - forest_sampler_mean.reconstitute_from_forest( - active_forest_mean, - forest_dataset_train, - residual_train, - True, - ) - # Undo residual update done by reconstitute_from_forest for the cloglog model - if link_is_cloglog: - residual_train.update_data(resid_train[:, 0]) - # Reset leaf scale - if sample_sigma2_leaf: - leaf_scale_double = self.leaf_scale_samples[forest_ind] - current_leaf_scale[0, 0] = leaf_scale_double - forest_model_config_mean.update_leaf_model_scale( - leaf_scale_double - ) - # Reset variance forest - if self.include_variance_forest: - active_forest_variance.reset( - self.forest_container_variance, forest_ind - ) - forest_sampler_variance.reconstitute_from_forest( - active_forest_variance, - forest_dataset_train, - residual_train, - False, - ) - # Reset global error scale - if sample_sigma2_global: - current_sigma2 = self.global_var_samples[forest_ind] - global_model_config.update_global_error_variance(current_sigma2) - # Reset random effects - if self.has_rfx: - rfx_model.reset( - self.rfx_container, forest_ind, sigma_alpha_init - ) - rfx_tracker.reset( - rfx_model, - rfx_dataset_train, - residual_train, - self.rfx_container, - ) - # Reset cloglog auxiliary data - if link_is_cloglog: - # Reset cutpoints from saved GFR samples - current_cutpoints = cloglog_cutpoint_samples[:, forest_ind] - for j in range(len(current_cutpoints)): - forest_dataset_train.set_auxiliary_data_value( - 2, j, current_cutpoints[j] - ) - ordinal_sampler.update_cumulative_exp_sums(forest_dataset_train) - # Reset forest predictions by re-predicting from active forest - active_forest_preds = active_forest_mean.predict( - forest_dataset_train - ) - for j in range(train_size): - forest_dataset_train.set_auxiliary_data_value( - 1, j, active_forest_preds[j] - ) - # Latent variables must be reset to 0 and burnt in - forest_dataset_train.set_auxiliary_data_value(0, j, 0.0) - elif has_prev_model: - warmstart_index = ( - previous_model_warmstart_sample_num - chain_num - if previous_model_decrement - else previous_model_warmstart_sample_num + # Unpack RFX results + if self.has_rfx: + self.rfx_container = RandomEffectsContainer() + self.rfx_container.rfx_container_cpp = bart_results["rfx_container"] + self.rfx_container.rfx_label_mapper_cpp = bart_results["rfx_label_mapper"] + self.rfx_container.rfx_group_ids = bart_results["rfx_label_mapper"].GetUniqueGroupIds() + rfx_preds_train = ( + bart_results["rfx_predictions_train"].reshape( + self.n_train, bart_results["num_samples"], order="F" + ) + * self.y_std + ) + self.y_hat_train = ( + self.y_hat_train + rfx_preds_train + if self.include_mean_forest + else rfx_preds_train + ) + if self.has_test: + rfx_preds_test = ( + bart_results["rfx_predictions_test"].reshape( + self.n_test, bart_results["num_samples"], order="F" ) - # Reset mean forest - if self.include_mean_forest: - active_forest_mean.reset( - previous_bart_model.forest_container_mean, - warmstart_index, - ) - forest_sampler_mean.reconstitute_from_forest( - active_forest_mean, - forest_dataset_train, - residual_train, - True, - ) - # Undo residual update done by reconstitute_from_forest for the cloglog model - if link_is_cloglog: - residual_train.update_data(resid_train[:, 0]) - # Reset leaf scale - if sample_sigma2_leaf and previous_leaf_var_samples is not None: - leaf_scale_double = previous_leaf_var_samples[ - warmstart_index - ] - current_leaf_scale[0, 0] = leaf_scale_double - forest_model_config_mean.update_leaf_model_scale( - leaf_scale_double - ) - # Reset variance forest - if self.include_variance_forest: - active_forest_variance.reset( - previous_bart_model.forest_container_variance, - warmstart_index, - ) - forest_sampler_variance.reconstitute_from_forest( - active_forest_variance, - forest_dataset_train, - residual_train, - True, - ) - # Reset global error scale - if self.sample_sigma2_global: - current_sigma2 = previous_global_var_samples[warmstart_index] - global_model_config.update_global_error_variance(current_sigma2) - # Reset random effects - if self.has_rfx: - rfx_model.reset( - previous_bart_model.rfx_container, - warmstart_index, - sigma_alpha_init, - ) - rfx_tracker.reset( - rfx_model, - rfx_dataset_train, - residual_train, - previous_bart_model.rfx_container, - ) - # Reset cloglog auxiliary data from previous model - if link_is_cloglog: - previous_cloglog_cutpoint_samples = getattr( - previous_bart_model, "cloglog_cutpoint_samples", None - ) - if previous_cloglog_cutpoint_samples is not None: - current_cutpoints = previous_cloglog_cutpoint_samples[ - :, warmstart_index - ] - for j in range(len(current_cutpoints)): - forest_dataset_train.set_auxiliary_data_value( - 2, j, current_cutpoints[j] - ) - ordinal_sampler.update_cumulative_exp_sums( - forest_dataset_train - ) - active_forest_preds = active_forest_mean.predict( - forest_dataset_train - ) - for j in range(train_size): - forest_dataset_train.set_auxiliary_data_value( - 1, j, active_forest_preds[j] - ) - # Latent variables must be reset to 0 and burnt in - forest_dataset_train.set_auxiliary_data_value(0, j, 0.0) - else: - # Reset mean forest - if self.include_mean_forest: - active_forest_mean.reset_root() - if init_val_mean.shape[0] == 1: - active_forest_mean.set_root_leaves( - init_val_mean[0] / num_trees_mean - ) - else: - active_forest_mean.set_root_leaves( - init_val_mean / num_trees_mean - ) - forest_sampler_mean.reconstitute_from_forest( - active_forest_mean, - forest_dataset_train, - residual_train, - True, - ) - # Undo residual update done by reconstitute_from_forest for the cloglog model - if link_is_cloglog: - residual_train.update_data(resid_train[:, 0]) - # Reset mean forest leaf scale - if sample_sigma2_leaf and previous_leaf_var_samples is not None: - current_leaf_scale[0, 0] = sigma2_leaf - forest_model_config_mean.update_leaf_model_scale( - current_leaf_scale - ) - if link_is_cloglog: - # Reset all cloglog parameters to default values - for j in range(train_size): - forest_dataset_train.set_auxiliary_data_value(1, j, 0.0) - forest_dataset_train.set_auxiliary_data_value(0, j, 0.0) - # Initialize log-scale cutpoints to 0 - initial_gamma = np.zeros(cloglog_num_categories - 1) - for j in range(cloglog_num_categories - 1): - forest_dataset_train.set_auxiliary_data_value( - 2, j, initial_gamma[j] - ) - # Convert to cumulative exponentiated cutpoints - ordinal_sampler.update_cumulative_exp_sums( - forest_dataset_train - ) - # Reset variance forest - if self.include_variance_forest: - active_forest_variance.reset_root() - active_forest_variance.set_root_leaves( - log(variance_forest_leaf_init) / num_trees_variance - ) - forest_sampler_variance.reconstitute_from_forest( - active_forest_variance, - forest_dataset_train, - residual_train, - False, - ) - # Reset global error scale - if self.sample_sigma2_global: - current_sigma2 = sigma2_init - global_model_config.update_global_error_variance(current_sigma2) - # Reset random effects terms - if self.has_rfx: - rfx_model.root_reset( - alpha_init, - xi_init, - sigma_alpha_init, - sigma_xi_init, - sigma_xi_shape, - sigma_xi_scale, - ) - rfx_tracker.root_reset( - rfx_model, - rfx_dataset_train, - residual_train, - self.rfx_container, - ) - # Sample MCMC and burnin for each chain - for i in range(self.num_gfr, num_temp_samples): - is_mcmc = i + 1 > num_gfr + num_burnin - if is_mcmc: - mcmc_counter = i - num_gfr - num_burnin + 1 - if mcmc_counter % keep_every == 0: - keep_sample = True - else: - keep_sample = False - else: - if keep_burnin: - keep_sample = True - else: - keep_sample = False - if keep_sample: - sample_counter += 1 - - if self.include_mean_forest: - if link_is_probit: - # Sample latent probit variable z | - - outcome_pred = active_forest_mean.predict( - forest_dataset_train - ) - if self.has_rfx: - rfx_pred = rfx_model.predict( - rfx_dataset_train, rfx_tracker - ) - outcome_pred = outcome_pred + rfx_pred - # Full probit-scale predictor: forest learns z - y_bar, so add y_bar back - eta_pred = outcome_pred + self.y_bar - mu0 = eta_pred[y_train[:, 0] == 0] - mu1 = eta_pred[y_train[:, 0] == 1] - n0 = np.sum(y_train[:, 0] == 0) - n1 = np.sum(y_train[:, 0] == 1) - u0 = self.rng.uniform( - low=0.0, - high=norm.cdf(0 - mu0), - size=n0, - ) - u1 = self.rng.uniform( - low=norm.cdf(0 - mu1), - high=1.0, - size=n1, - ) - resid_train[y_train[:, 0] == 0, 0] = mu0 + norm.ppf(u0) - resid_train[y_train[:, 0] == 1, 0] = mu1 + norm.ppf(u1) - - # Update outcome: center z by y_bar before passing to forest - new_outcome = ( - np.squeeze(resid_train) - self.y_bar - outcome_pred - ) - residual_train.update_data(new_outcome) - - # Sample the mean forest - forest_sampler_mean.sample_one_iteration( - self.forest_container_mean, - active_forest_mean, - forest_dataset_train, - residual_train, - cpp_rng, - global_model_config, - forest_model_config_mean, - keep_sample, - False, - num_threads, - ) + * self.y_std + ) + self.y_hat_test = ( + self.y_hat_test + rfx_preds_test + if self.include_mean_forest + else rfx_preds_test + ) - if keep_sample: - yhat_train_raw[:, sample_counter] = ( - forest_sampler_mean.get_cached_forest_predictions() - ) - - # Sample the variance forest - if self.include_variance_forest: - forest_sampler_variance.sample_one_iteration( - self.forest_container_variance, - active_forest_variance, - forest_dataset_train, - residual_train, - cpp_rng, - global_model_config, - forest_model_config_variance, - keep_sample, - False, - num_threads, - ) + # Unpack variance forest predictions (forest owned by self._samples, assembled below) + if self.include_variance_forest: + self.sigma2_x_train = bart_results[ + "variance_forest_predictions_train" + ].reshape(self.n_train, bart_results["num_samples"], order="F") + if self.has_test: + self.sigma2_x_test = bart_results[ + "variance_forest_predictions_test" + ].reshape(self.n_test, bart_results["num_samples"], order="F") - if keep_sample: - sigma2_x_train_raw[:, sample_counter] = ( - forest_sampler_variance.get_cached_forest_predictions() - ) + # Parameter sampling flags (the param traces themselves are owned by self._samples). + self.sample_sigma2_global = sample_sigma2_global + self.sample_sigma2_leaf = sample_sigma2_leaf + # cloglog cutpoints are not routed through the samples wrapper yet (guarded there); keep + # them as a separate attribute. + if link_is_cloglog: + self.cloglog_num_categories = cloglog_num_categories + if not outcome_is_binary: + self.cloglog_cutpoint_samples = bart_results[ + "cloglog_cutpoint_samples" + ].reshape( + cloglog_num_categories - 1, bart_results["num_samples"], order="F" + ) + + # Assemble the single source of truth: forests + parameter traces in one BARTSamplesCpp. + # forest_container_*/global_var_samples/leaf_scale_samples/num_samples are properties off it. + self._set_samples( + BARTSamplesCpp.from_components( + bart_results["forest_container_mean"] if self.include_mean_forest else None, + bart_results["forest_container_variance"] if self.include_variance_forest else None, + bart_results["global_var_samples"] if self.sample_sigma2_global else None, + bart_results["leaf_scale_samples"] if self.sample_sigma2_leaf else None, + float(self.y_bar), + float(self.y_std), + int(bart_results["num_samples"]), + ) + ) + self.sampled = True - # Sample variance parameters (if requested) - if self.sample_sigma2_global: - current_sigma2 = global_var_model.sample_one_iteration( - residual_train, cpp_rng, a_global, b_global - ) - global_model_config.update_global_error_variance(current_sigma2) - if keep_sample: - self.global_var_samples[sample_counter] = current_sigma2 - if self.sample_sigma2_leaf: - current_leaf_scale[0, 0] = leaf_var_model.sample_one_iteration( - active_forest_mean, cpp_rng, a_leaf, b_leaf - ) - forest_model_config_mean.update_leaf_model_scale( - current_leaf_scale - ) - if keep_sample: - self.leaf_scale_samples[sample_counter] = ( - current_leaf_scale[0, 0] - ) - - # Sample random effects - if self.has_rfx: - rfx_model.sample( - rfx_dataset_train, - residual_train, - rfx_tracker, - self.rfx_container, - keep_sample, - current_sigma2, - cpp_rng, - ) + return self - # Cloglog Gibbs updates - if link_is_cloglog: - # Update auxiliary data slot 1 with current forest predictions - forest_pred_current = ( - forest_sampler_mean.get_cached_forest_predictions() - ) - for j in range(train_size): - forest_dataset_train.set_auxiliary_data_value( - 1, j, forest_pred_current[j] - ) - - # Sample latent z_i's using truncated exponential - ordinal_sampler.update_latent_variables( - forest_dataset_train, residual_train, cpp_rng - ) + def continue_sampling( + self, + X_train: Union[np.array, pd.DataFrame], + y_train: np.array, + num_mcmc: int = 100, + num_burnin: int = 0, + keep_every: int = 1, + leaf_basis_train: np.array = None, + random_seed: int = None, + ): + """Continue (warm-start) MCMC sampling from an already-fit BART model, appending + additional samples to the existing posterior draws. - # Sample gamma parameters (cutpoints) - ordinal_sampler.update_gamma_params( - forest_dataset_train, - residual_train, - cloglog_forest_shape, - cloglog_forest_rate, - cloglog_cutpoint_0, - cpp_rng, - ) + The training data must be re-supplied (it is not retained on the model). The + warm-start initializes the sampler from the last retained sample so that the new + draws form a continuous chain with the existing ones. - # Update cumulative sum of exp(gamma) values - ordinal_sampler.update_cumulative_exp_sums(forest_dataset_train) + .. note:: + This is currently supported only for Gaussian (identity-link) models with a + mean forest and no variance forest or random effects. - # Retain cutpoint draw - if keep_sample: - cloglog_cutpoint_samples[:, sample_counter] = ( - forest_dataset_train.get_auxiliary_data_vector(2) - ) + Parameters + ---------- + X_train : np.array or pd.DataFrame + Training covariates (re-supplied; must match the structure used to fit the model). + y_train : np.array + Training outcome (re-supplied). + num_mcmc : int, optional + Number of additional retained MCMC samples. Defaults to ``100``. + num_burnin : int, optional + Number of additional burn-in iterations to discard before retaining. Defaults to ``0``. + keep_every : int, optional + Thinning interval for the additional samples. Defaults to ``1``. + leaf_basis_train : np.array, optional + Training leaf basis (required if the model was fit with a leaf regression basis). + random_seed : int, optional + If supplied, re-seeds the sampler RNG for the continued draws. By default + (``None``), the RNG state from the prior run is resumed, so the continued chain + is bit-identical to a single run of the combined length. - # Mark the model as sampled - self.sampled = True + Returns + ------- + self + """ + if not getattr(self, "sampled", False): + raise RuntimeError("Cannot continue sampling: this model has not been sampled yet") + if not self.include_mean_forest: + raise NotImplementedError("Continued sampling currently requires a mean forest") + if self.include_variance_forest: + raise NotImplementedError( + "Continued sampling is not yet supported for models with a variance forest" + ) + if getattr(self, "has_rfx", False): + raise NotImplementedError( + "Continued sampling is not yet supported for models with random effects" + ) + cfg = getattr(self, "_cached_bart_config", None) + if cfg is None: + raise RuntimeError( + "Cannot continue sampling: cached sampler configuration is unavailable " + "(continuation is not supported for deserialized models yet)" + ) + if cfg.get("link_function", 0) != 0: + raise NotImplementedError( + "Continued sampling is not yet supported for probit or cloglog link functions" + ) - # Remove GFR samples if they are not to be retained - if not keep_gfr and num_gfr > 0: - for i in range(num_gfr): - if self.include_mean_forest: - self.forest_container_mean.delete_sample(0) - if self.include_variance_forest: - self.forest_container_variance.delete_sample(0) - if self.has_rfx: - self.rfx_container.delete_sample(0) - if self.sample_sigma2_global: - self.global_var_samples = self.global_var_samples[num_gfr:] - if self.sample_sigma2_leaf: - self.leaf_scale_samples = self.leaf_scale_samples[num_gfr:] - if self.include_mean_forest: - yhat_train_raw = yhat_train_raw[:, num_gfr:] - if self.include_variance_forest: - sigma2_x_train_raw = sigma2_x_train_raw[:, num_gfr:] - if link_is_cloglog: - cloglog_cutpoint_samples = cloglog_cutpoint_samples[:, num_gfr:] - self.num_samples -= num_gfr + # Preprocess the re-supplied covariates with the fitted preprocessor and validate structure + X_train_processed = self._covariate_preprocessor.transform(X_train).astype(np.float64) + if X_train_processed.shape[1] != self.num_covariates: + raise ValueError( + f"Re-supplied covariates have {X_train_processed.shape[1]} columns; " + f"model expects {self.num_covariates}" + ) + y_train = np.asarray(y_train).astype(np.float64).reshape(-1) + if X_train_processed.shape[0] != y_train.shape[0]: + raise ValueError("X_train and y_train have differing numbers of observations") - # Store cloglog results (cutpoints only for ordinal, num_categories always) - if link_is_cloglog: - self.cloglog_num_categories = cloglog_num_categories - if not outcome_is_binary: - self.cloglog_cutpoint_samples = cloglog_cutpoint_samples + if self.has_basis and leaf_basis_train is None: + raise ValueError( + "This model was fit with a leaf basis; leaf_basis_train must be supplied to continue sampling" + ) + if self.has_basis: + leaf_basis_train = np.atleast_2d(leaf_basis_train) + if leaf_basis_train.shape[1] != self.num_basis: + raise ValueError( + f"Re-supplied leaf basis has {leaf_basis_train.shape[1]} columns; " + f"model expects {self.num_basis}" + ) - # Store predictions - if self.sample_sigma2_global: - self.global_var_samples = self.global_var_samples * self.y_std * self.y_std + X_train_cpp = np.asfortranarray(X_train_processed) + y_train_cpp = np.asfortranarray(y_train, dtype=np.float64) + basis_train_cpp = ( + np.asfortranarray(leaf_basis_train.astype(np.float64)) if self.has_basis else None + ) - if self.sample_sigma2_leaf: - self.leaf_scale_samples = self.leaf_scale_samples + # RNG continuation: by default resume the saved stream for bit-identical results. + # If the user supplies a new seed, override the cached config seed and tell the + # binding to keep the fresh seed instead of restoring the saved state. + override_seed = random_seed is not None + if override_seed: + cfg = dict(cfg) + cfg["random_seed"] = random_seed + rng_state_in = self.rng_state if (not override_seed and self.rng_state is not None) else "" + leaf_normal_cache_in = ( + self.leaf_normal_cache + if (not override_seed and getattr(self, "leaf_normal_cache", None) is not None) + else "" + ) - if self.include_mean_forest: - self.y_hat_train = yhat_train_raw * self.y_std + self.y_bar - if self.has_test: - yhat_test_raw = self.forest_container_mean.forest_container_cpp.Predict( - forest_dataset_test.dataset_cpp - ) - self.y_hat_test = yhat_test_raw * self.y_std + self.y_bar + bart_results = bart_continue_sample_cpp( + X_train=X_train_cpp, + y_train=y_train_cpp, + X_test=None, + n_train=X_train_cpp.shape[0], + n_test=0, + p=X_train_cpp.shape[1], + basis_train=basis_train_cpp, + basis_test=None, + basis_dim=self.num_basis if self.has_basis else 0, + obs_weights_train=None, + obs_weights_test=None, + rfx_group_ids_train=None, + rfx_group_ids_test=None, + rfx_basis_train=None, + rfx_basis_test=None, + rfx_num_groups=0, + rfx_basis_dim=0, + num_burnin=num_burnin, + keep_every=keep_every, + num_mcmc=num_mcmc, + mean_forest_container=self._forest_container_mean.forest_container_cpp, + global_var_samples=self.global_var_samples if self.sample_sigma2_global else None, + leaf_scale_samples=self.leaf_scale_samples if self.sample_sigma2_leaf else None, + y_bar=float(self.y_bar), + y_std=float(self.y_std), + rng_state_in=rng_state_in, + override_seed=override_seed, + leaf_normal_cache_in=leaf_normal_cache_in, + config_input=cfg, + ) - # TODO: make rfx_preds_train and rfx_preds_test persistent properties - if self.has_rfx: - rfx_preds_train = ( - self.rfx_container.predict(rfx_group_ids_train, rfx_basis_train) - * self.y_std + # Rebuild the single source of truth from the (history + new) continuation results. + # Continuation supports a mean forest only (no variance forest), so variance is None. + self._set_samples( + BARTSamplesCpp.from_components( + bart_results["forest_container_mean"], + None, + bart_results["global_var_samples"] if self.sample_sigma2_global else None, + bart_results["leaf_scale_samples"] if self.sample_sigma2_leaf else None, + float(self.y_bar), + float(self.y_std), + int(bart_results["num_samples"]), ) - if has_rfx_test: - rfx_preds_test = ( - self.rfx_container.predict(rfx_group_ids_test, rfx_basis_test) - * self.y_std - ) - if self.include_mean_forest: - self.y_hat_train = self.y_hat_train + rfx_preds_train - if self.has_test: - self.y_hat_test = self.y_hat_test + rfx_preds_test - else: - self.y_hat_train = rfx_preds_train - if self.has_test: - self.y_hat_test = rfx_preds_test + ) + self.num_mcmc = (self.num_mcmc or 0) + num_mcmc + # Carry the new final RNG + leaf-cache state forward so further continuations stay bit-identical + self.rng_state = bart_results.get("rng_state", None) + self.leaf_normal_cache = bart_results.get("leaf_normal_cache", None) + + # Recompute training predictions over the full (history + new) set of samples + self.y_hat_train = self.predict( + X_train, + leaf_basis=leaf_basis_train if self.has_basis else None, + terms="y_hat", + ) - if self.include_variance_forest: - if self.sample_sigma2_global: - self.sigma2_x_train = np.empty_like(sigma2_x_train_raw) - for i in range(self.num_samples): - self.sigma2_x_train[:, i] = ( - np.exp(sigma2_x_train_raw[:, i]) * self.global_var_samples[i] - ) - else: - self.sigma2_x_train = ( - np.exp(sigma2_x_train_raw) - * self.sigma2_init - * self.y_std - * self.y_std - ) - if self.has_test: - sigma2_x_test_raw = ( - self.forest_container_variance.forest_container_cpp.Predict( - forest_dataset_test.dataset_cpp - ) - ) - if self.sample_sigma2_global: - self.sigma2_x_test = sigma2_x_test_raw - for i in range(self.num_samples): - self.sigma2_x_test[:, i] = ( - sigma2_x_test_raw[:, i] * self.global_var_samples[i] - ) - else: - self.sigma2_x_test = ( - sigma2_x_test_raw * self.sigma2_init * self.y_std * self.y_std - ) + return self def predict( self, @@ -2396,41 +1941,6 @@ def predict( else: X_processed = self._covariate_preprocessor.transform(X) - # Dataset construction - pred_dataset = Dataset() - pred_dataset.add_covariates(X_processed) - if leaf_basis is not None: - pred_dataset.add_basis(leaf_basis) - - # Variance forest predictions - if predict_variance_forest: - variance_pred_raw = ( - self.forest_container_variance.forest_container_cpp.Predict( - pred_dataset.dataset_cpp - ) - ) - if self.sample_sigma2_global: - variance_forest_predictions = np.empty_like(variance_pred_raw) - for i in range(self.num_samples): - variance_forest_predictions[:, i] = ( - variance_pred_raw[:, i] * self.global_var_samples[i] - ) - else: - variance_forest_predictions = ( - variance_pred_raw * self.sigma2_init * self.y_std * self.y_std - ) - if predict_mean: - variance_forest_predictions = np.mean( - variance_forest_predictions, axis=1 - ) - - # Forest predictions - if predict_mean_forest or predict_mean_forest_intermediate: - mean_pred_raw = self.forest_container_mean.forest_container_cpp.Predict( - pred_dataset.dataset_cpp - ) - mean_forest_predictions = mean_pred_raw * self.y_std + self.y_bar - # Random effects data checks if predict_rfx and rfx_group_ids is None: raise ValueError( @@ -2445,154 +1955,111 @@ def predict( raise ValueError( "Random effects basis has a different dimension than the basis used to train this model" ) + + # Convert prediction scale info to integer code for easy conversion to enum in C++ + scale_int = 0 if not probability_scale and not class_scale else (1 if probability_scale else 2) + + # # Convert cloglog cutpoint samples to fortran (column-major) array if present and not already aligned as such + # cloglog_cutpoints = getattr(self, "cloglog_cutpoint_samples", None) + # if cloglog_cutpoints is not None: + # cloglog_cutpoints = np.asfortranarray(cloglog_cutpoints) + + # Construct dictionary of model components to pass to C++ prediction function, with None for any components not present in the model + bart_model_dict = { + "mean_forests": self._forest_container_mean.forest_container_cpp if self.include_mean_forest else None, + "variance_forests": self._forest_container_variance.forest_container_cpp if self.include_variance_forest else None, + "rfx_container": self.rfx_container.rfx_container_cpp if has_rfx else None, + "rfx_label_mapper": self.rfx_container.rfx_label_mapper_cpp if has_rfx else None, + "sigma2_global_samples": getattr(self, 'global_var_samples', None), + "sigma2_leaf_samples": getattr(self, 'leaf_scale_samples', None), + "cloglog_cutpoint_samples": np.asfortranarray(self.cloglog_cutpoint_samples) if getattr(self, 'cloglog_cutpoint_samples', None) is not None else None, + "num_samples": int(self.num_samples), + "y_bar": float(self.y_bar), + "y_std": float(self.y_std), + "include_variance_forest": has_variance_forest, + "has_rfx": has_rfx, + "rfx_model_spec": self.rfx_model_spec if has_rfx else "", + "link_function": self.outcome_model.link, + "outcome_type": self.outcome_model.outcome, + } + if is_cloglog: + bart_model_dict["cloglog_num_classes"] = int(self.cloglog_num_categories) + + # Data dimensions + n, p = X_processed.shape + num_basis = int(self.num_basis) if self.has_basis else 0 + rfx_basis_dim = int(self.num_rfx_basis) if has_rfx else 0 + + # Call the C++ prediction function, returning results as a dictionary + output = bart_predict_cpp( + bart_model_dict=bart_model_dict, + X=np.asfortranarray(X_processed), + leaf_basis=np.asfortranarray(leaf_basis) if leaf_basis is not None else None, + n=n, + p=p, + num_basis=num_basis, + rfx_group_ids=rfx_group_ids.astype(np.int32) if rfx_group_ids is not None else None, + rfx_basis=np.asfortranarray(rfx_basis) if rfx_basis is not None else None, + rfx_num_groups=self.rfx_container.num_groups() if has_rfx else 0, + rfx_basis_dim=rfx_basis_dim, + posterior=(type == "posterior"), + scale=scale_int, + predict_y_hat=predict_y_hat, + predict_mean_forest=predict_mean_forest, + predict_variance_forest=predict_variance_forest, + predict_random_effects=predict_rfx, + ) - # Convert rfx_group_ids to their corresponding array position indices in the random effects parameter sample arrays - if rfx_group_ids is not None: - rfx_group_id_indices = self.rfx_container.map_group_ids_to_array_indices( - rfx_group_ids - ) - - # Random effects predictions - if predict_rfx or predict_rfx_intermediate: - if rfx_basis is not None: - rfx_predictions = ( - self.rfx_container.predict(rfx_group_ids, rfx_basis) * self.y_std - ) - else: - # Sanity check -- this branch should only occur if rfx_model_spec == "intercept_only" - if not rfx_intercept: - raise ValueError( - "A user-provided basis (`rfx_basis`) must be provided when the model was sampled with a random effects model spec set to 'custom'" - ) - - # Extract the raw RFX samples and scale by train set outcome standard deviation - rfx_samples_raw = self.rfx_container.extract_parameter_samples() - rfx_beta_draws = rfx_samples_raw["beta_samples"] * self.y_std - - # Construct an array with the appropriate group random effects arranged for each observation - n_train = X.shape[0] - if rfx_beta_draws.ndim != 2: - raise ValueError( - "BART models fit with random intercept models should only yield 2 dimensional random effect sample matrices" - ) - else: - rfx_predictions_raw = np.empty( - shape=(n_train, 1, rfx_beta_draws.shape[1]) - ) - for i in range(n_train): - rfx_predictions_raw[i, 0, :] = rfx_beta_draws[ - rfx_group_id_indices[i], : - ] - rfx_predictions = np.squeeze(rfx_predictions_raw[:, 0, :]) - - # Combine into y hat predictions - if probability_scale or class_scale: - if is_probit: - if predict_y_hat and has_mean_forest and has_rfx: - y_hat = norm.cdf(mean_forest_predictions + rfx_predictions) - elif predict_y_hat and has_mean_forest: - y_hat = norm.cdf(mean_forest_predictions) - elif predict_y_hat and has_rfx: - y_hat = norm.cdf(rfx_predictions) - if ( - predict_mean_forest or predict_mean_forest_intermediate - ) and has_mean_forest: - mean_forest_predictions = norm.cdf(mean_forest_predictions) - if (predict_rfx or predict_rfx_intermediate) and has_rfx: - rfx_predictions = norm.cdf(rfx_predictions) - elif is_binary_cloglog: - mean_forest_predictions = np.exp(-np.exp(mean_forest_predictions)) - if predict_y_hat: - y_hat = mean_forest_predictions - elif is_ordinal_cloglog: - cloglog_num_categories = self.cloglog_num_categories - cloglog_cutpoint_samples = self.cloglog_cutpoint_samples - n_obs = X.shape[0] - num_samples = self.num_samples - # Sequential ordinal cloglog: P(Y=k) = prod_{j str: # Initialize JSONSerializer object bart_json = JSONSerializer() - # Add the forests + # Add the forests under self-describing named keys if self.include_mean_forest: - bart_json.add_forest(self.forest_container_mean) + bart_json.add_forest(self._forest_container_mean, "mean_forest") if self.include_variance_forest: - bart_json.add_forest(self.forest_container_variance) + bart_json.add_forest(self._forest_container_variance, "variance_forest") # Add the rfx if self.has_rfx: bart_json.add_random_effects(self.rfx_container) + # Python rfx group ids are always integer-valued, so an rfx model is + # cross-platform compatible on the rfx axis. + bart_json.add_boolean( + "cross_platform_compatible", True, subfolder_name="random_effects" + ) # Add version stamp and global parameters bart_json.add_string("stochtree_version", _get_stochtree_version()) + bart_json.add_string("platform", "python") + bart_json.add_integer("schema_version", SCHEMA_VERSION) + # Covariate count: written so a cross-platform (R) loader can validate + # prediction-set dimension. For all-numeric (portable) models the + # processed count equals the original count R expects. + bart_json.add_scalar("num_covariates", self.num_covariates) bart_json.add_scalar("outcome_scale", self.y_std) bart_json.add_scalar("outcome_mean", self.y_bar) bart_json.add_boolean("standardize", self.standardize) @@ -3235,6 +2713,8 @@ def from_json(self, json_string: str) -> None: _ver = _infer_stochtree_version(json_string) bart_json = JSONSerializer() bart_json.load_from_json_string(json_string) + resolve_schema_version(bart_json, migrate=_migrate_bart_v0_to_v1) + cross_platform = enforce_cross_platform_gate(bart_json, "python") # Unpack forests self.include_mean_forest = bart_json.get_boolean("include_mean_forest") @@ -3253,23 +2733,19 @@ def from_json(self, json_string: str) -> None: f"Re-save your model to suppress this warning." ) + # v1 forests are stored under self-describing named keys. Load into temporary + # ForestContainer objects, then deep-copy into the single owned samples object below. + mean_forest_loaded = None + variance_forest_loaded = None if self.include_mean_forest: - # TODO: don't just make this a placeholder that we overwrite - self.forest_container_mean = ForestContainer(0, 0, False, False) - self.forest_container_mean.forest_container_cpp.LoadFromJson( - bart_json.json_cpp, "forest_0" + mean_forest_loaded = ForestContainer(0, 0, False, False) + mean_forest_loaded.forest_container_cpp.LoadFromJson( + bart_json.json_cpp, "mean_forest" ) - if self.include_variance_forest: - # TODO: don't just make this a placeholder that we overwrite - self.forest_container_variance = ForestContainer(0, 0, False, False) - self.forest_container_variance.forest_container_cpp.LoadFromJson( - bart_json.json_cpp, "forest_1" - ) - else: - # TODO: don't just make this a placeholder that we overwrite - self.forest_container_variance = ForestContainer(0, 0, False, False) - self.forest_container_variance.forest_container_cpp.LoadFromJson( - bart_json.json_cpp, "forest_0" + if self.include_variance_forest: + variance_forest_loaded = ForestContainer(0, 0, False, False) + variance_forest_loaded.forest_container_cpp.LoadFromJson( + bart_json.json_cpp, "variance_forest" ) # Unpack random effects @@ -3287,7 +2763,7 @@ def from_json(self, json_string: str) -> None: self.num_gfr = bart_json.get_integer("num_gfr") self.num_burnin = bart_json.get_integer("num_burnin") self.num_mcmc = bart_json.get_integer("num_mcmc") - self.num_samples = bart_json.get_integer("num_samples") + num_samples_loaded = bart_json.get_integer("num_samples") self.num_basis = bart_json.get_integer("num_basis") self.has_basis = bart_json.get_boolean("requires_basis") @@ -3344,32 +2820,53 @@ def from_json(self, json_string: str) -> None: f"Re-save your model to suppress this warning." ) - # Unpack parameter samples - if self.sample_sigma2_global: - self.global_var_samples = bart_json.get_numeric_vector( - "sigma2_global_samples", "parameters" - ) - if self.sample_sigma2_leaf: - self.leaf_scale_samples = bart_json.get_numeric_vector( - "sigma2_leaf_samples", "parameters" - ) + # Unpack parameter samples (owned by self._samples, built below) + global_var_loaded = ( + bart_json.get_numeric_vector("sigma2_global_samples", "parameters") + if self.sample_sigma2_global + else None + ) + leaf_scale_loaded = ( + bart_json.get_numeric_vector("sigma2_leaf_samples", "parameters") + if self.sample_sigma2_leaf + else None + ) - # Unpack cloglog parameters (num_categories always, cutpoints only for ordinal) + # Unpack cloglog parameters (num_categories always, cutpoints only for ordinal); these stay + # a separate attribute (not routed through the samples wrapper yet). if self.outcome_model.link == "cloglog": self.cloglog_num_categories = bart_json.get_integer( "cloglog_num_categories" ) if self.outcome_model.outcome == "ordinal": self.cloglog_cutpoint_samples = np.full( - (self.cloglog_num_categories - 1, self.num_samples), np.nan + (self.cloglog_num_categories - 1, num_samples_loaded), np.nan ) for i in range(self.cloglog_num_categories - 1): self.cloglog_cutpoint_samples[i, :] = bart_json.get_numeric_vector( f"cloglog_cutpoint_samples_{i + 1}", "parameters" ) + # Assemble the single source of truth (forests + parameter traces) from the loaded parts. + self._set_samples( + BARTSamplesCpp.from_components( + mean_forest_loaded.forest_container_cpp if self.include_mean_forest else None, + variance_forest_loaded.forest_container_cpp if self.include_variance_forest else None, + global_var_loaded, + leaf_scale_loaded, + float(self.y_bar), + float(self.y_std), + int(num_samples_loaded), + ) + ) + # Unpack covariate preprocessor - if "covariate_preprocessor" in _raw: + if cross_platform: + # The foreign native preprocessor can't be reconstructed; the gate + # guarantees the model is all-numeric, so an unfitted (identity) + # preprocessor passes numeric covariates through at predict time. + self._covariate_preprocessor = CovariatePreprocessor() + elif "covariate_preprocessor" in _raw: covariate_preprocessor_string = bart_json.get_string( "covariate_preprocessor" ) @@ -3405,51 +2902,44 @@ def from_json_string_list(self, json_string_list: list[str]) -> None: # For scalar / preprocessing details which aren't sample-dependent, defer to the first json json_object_default = json_object_list[0] + for json_object in json_object_list: + resolve_schema_version(json_object, migrate=_migrate_bart_v0_to_v1) + cross_platform = enforce_cross_platform_gate(json_object_default, "python") _raw = json.loads(json_string_list[0]) _ver = _infer_stochtree_version(json_string_list[0]) - # Unpack forests + # Unpack forests (v1 named keys) self.include_mean_forest = json_object_default.get_boolean( "include_mean_forest" ) self.include_variance_forest = json_object_default.get_boolean( "include_variance_forest" ) + # Load each chain's forests into temporary containers (append across chains), then deep-copy + # the combined result into the single owned samples object below. + mean_forest_loaded = None + variance_forest_loaded = None if self.include_mean_forest: - # TODO: don't just make this a placeholder that we overwrite - self.forest_container_mean = ForestContainer(0, 0, False, False) + mean_forest_loaded = ForestContainer(0, 0, False, False) for i in range(len(json_object_list)): if i == 0: - self.forest_container_mean.forest_container_cpp.LoadFromJson( - json_object_list[i].json_cpp, "forest_0" + mean_forest_loaded.forest_container_cpp.LoadFromJson( + json_object_list[i].json_cpp, "mean_forest" ) else: - self.forest_container_mean.forest_container_cpp.AppendFromJson( - json_object_list[i].json_cpp, "forest_0" + mean_forest_loaded.forest_container_cpp.AppendFromJson( + json_object_list[i].json_cpp, "mean_forest" ) - if self.include_variance_forest: - # TODO: don't just make this a placeholder that we overwrite - self.forest_container_variance = ForestContainer(0, 0, False, False) - for i in range(len(json_object_list)): - if i == 0: - self.forest_container_variance.forest_container_cpp.LoadFromJson( - json_object_list[i].json_cpp, "forest_1" - ) - else: - self.forest_container_variance.forest_container_cpp.AppendFromJson( - json_object_list[i].json_cpp, "forest_1" - ) - else: - # TODO: don't just make this a placeholder that we overwrite - self.forest_container_variance = ForestContainer(0, 0, False, False) + if self.include_variance_forest: + variance_forest_loaded = ForestContainer(0, 0, False, False) for i in range(len(json_object_list)): if i == 0: - self.forest_container_variance.forest_container_cpp.LoadFromJson( - json_object_list[i].json_cpp, "forest_0" + variance_forest_loaded.forest_container_cpp.LoadFromJson( + json_object_list[i].json_cpp, "variance_forest" ) else: - self.forest_container_variance.forest_container_cpp.AppendFromJson( - json_object_list[i].json_cpp, "forest_0" + variance_forest_loaded.forest_container_cpp.AppendFromJson( + json_object_list[i].json_cpp, "variance_forest" ) # Unpack random effects @@ -3547,43 +3037,31 @@ def from_json_string_list(self, json_string_list: list[str]) -> None: f"Re-save your model to suppress this warning." ) - # Unpack number of samples + # Combined number of samples across chains (owned by self._samples, built below) + num_samples_loaded = 0 for i in range(len(json_object_list)): - if i == 0: - self.num_samples = json_object_list[i].get_integer("num_samples") - else: - self.num_samples += json_object_list[i].get_integer("num_samples") + num_samples_loaded += json_object_list[i].get_integer("num_samples") - # Unpack parameter samples + # Unpack parameter samples (concatenated across chains; owned by self._samples below) + global_var_loaded = None if self.sample_sigma2_global: for i in range(len(json_object_list)): - if i == 0: - self.global_var_samples = json_object_list[i].get_numeric_vector( - "sigma2_global_samples", "parameters" - ) - else: - global_var_samples = json_object_list[i].get_numeric_vector( - "sigma2_global_samples", "parameters" - ) - self.global_var_samples = np.concatenate(( - self.global_var_samples, - global_var_samples, - )) + gv_i = json_object_list[i].get_numeric_vector( + "sigma2_global_samples", "parameters" + ) + global_var_loaded = ( + gv_i if i == 0 else np.concatenate((global_var_loaded, gv_i)) + ) + leaf_scale_loaded = None if self.sample_sigma2_leaf: for i in range(len(json_object_list)): - if i == 0: - self.leaf_scale_samples = json_object_list[i].get_numeric_vector( - "sigma2_leaf_samples", "parameters" - ) - else: - leaf_scale_samples = json_object_list[i].get_numeric_vector( - "sigma2_leaf_samples", "parameters" - ) - self.leaf_scale_samples = np.concatenate(( - self.leaf_scale_samples, - leaf_scale_samples, - )) + ls_i = json_object_list[i].get_numeric_vector( + "sigma2_leaf_samples", "parameters" + ) + leaf_scale_loaded = ( + ls_i if i == 0 else np.concatenate((leaf_scale_loaded, ls_i)) + ) # Unpack cloglog parameters (num_categories always, cutpoints only for ordinal) if self.outcome_model.link == "cloglog": @@ -3607,8 +3085,25 @@ def from_json_string_list(self, json_string_list: list[str]) -> None: (self.cloglog_cutpoint_samples, cutpoints_i), axis=1 ) + # Assemble the single source of truth from the combined (across-chain) parts. + self._set_samples( + BARTSamplesCpp.from_components( + mean_forest_loaded.forest_container_cpp if self.include_mean_forest else None, + variance_forest_loaded.forest_container_cpp if self.include_variance_forest else None, + global_var_loaded, + leaf_scale_loaded, + float(self.y_bar), + float(self.y_std), + int(num_samples_loaded), + ) + ) + # Unpack covariate preprocessor - if "covariate_preprocessor" in _raw: + if cross_platform: + # Identity preprocessor for the cross-platform all-numeric path (gate + # enforced); the foreign native preprocessor is not reconstructed. + self._covariate_preprocessor = CovariatePreprocessor() + elif "covariate_preprocessor" in _raw: covariate_preprocessor_string = json_object_default.get_string( "covariate_preprocessor" ) diff --git a/stochtree/bcf.py b/stochtree/bcf.py index bfcd62e8..caae3542 100644 --- a/stochtree/bcf.py +++ b/stochtree/bcf.py @@ -2,7 +2,7 @@ import warnings from numbers import Integral from typing import Any, Dict, Optional, Union -from math import floor, log +from math import floor import numpy as np import pandas as pd @@ -10,18 +10,16 @@ from scipy.stats import norm from .bart import BARTModel -from .config import ForestModelConfig, GlobalModelConfig -from .data import Dataset, Residual -from .forest import Forest, ForestContainer +from .forest import ForestContainer from .preprocessing import CovariatePreprocessor, _preprocess_params -from .random_effects import ( - RandomEffectsContainer, - RandomEffectsDataset, - RandomEffectsModel, - RandomEffectsTracker, +from .random_effects import RandomEffectsContainer +from .serialization import ( + JSONSerializer, + SCHEMA_VERSION, + resolve_schema_version, + infer_platform_v0, + enforce_cross_platform_gate, ) -from .sampler import RNG, ForestSampler, GlobalVarianceModel, LeafVarianceModel -from .serialization import JSONSerializer from .utils import ( OutcomeModel, NotSampledError, @@ -33,6 +31,23 @@ _posterior_predictive_heuristic_multiplier, _summarize_interval, ) +from stochtree_cpp import bcf_sample_cpp, bcf_continue_sample_cpp, bcf_predict_cpp, BCFSamplesCpp + + +def _migrate_bcf_v0_to_v1(serializer: JSONSerializer, loaded_version: int) -> None: + """In-place v0 -> v1 migration for a BCF model envelope. + + Positional forest keys -> named keys: ``forest_0`` -> ``prognostic_forest``, + ``forest_1`` -> ``treatment_forest``, and (when present) ``forest_2`` -> + ``variance_forest``. The prognostic (mu) and treatment (tau) forests are always + present in a BCF model; the variance forest is optional. Also stamps the + writer ``platform`` (inferred from structural fingerprints). + """ + serializer.add_string("platform", infer_platform_v0(serializer, "python")) + serializer.rename_field("forest_0", "prognostic_forest", subfolder_name="forests") + serializer.rename_field("forest_1", "treatment_forest", subfolder_name="forests") + if serializer.get_boolean_or_default("include_variance_forest", False): + serializer.rename_field("forest_2", "variance_forest", subfolder_name="forests") class BCFModel: @@ -82,6 +97,148 @@ class BCFModel: def __init__(self) -> None: # Internal flag for whether the sample() method has been run self.sampled = False + # Single source of truth for the sampled forests + parameter traces (BCFSamplesCpp). + # forest_container_mu/tau/variance, global_var_samples, leaf_scale_mu/tau_samples, and + # num_samples are properties backed by this; _fc_*_cache hold the materialized forest views. + # tau_0_samples / b0_samples / b1_samples remain separate attributes (specialized shape/scale; + # the wrapper guards/omits them for now). + self._samples = None + self._fc_mu_cache = None + self._fc_tau_cache = None + self._fc_variance_cache = None + + def _set_samples(self, samples) -> None: + """Install a new BCFSamplesCpp as the single source of truth and invalidate the + internal materialized forest-container caches.""" + self._samples = samples + self._fc_mu_cache = None + self._fc_tau_cache = None + self._fc_variance_cache = None + + @property + def samples(self): + """The single-owner ``BCFSamplesCpp`` holding the sampled forests and parameter traces.""" + return self._samples + + def extract_forest(self, forest: str = "prognostic"): + """Return a standalone deep copy of a sampled forest as a ``ForestContainer``. + + Parameters + ---------- + forest : str + Which forest to extract: ``"prognostic"`` (mu), ``"treatment"`` (tau), or ``"variance"``. + + Returns + ------- + ForestContainer or None + A deep copy independent of the model, or ``None`` if that forest was not sampled. + """ + if self._samples is None: + raise RuntimeError("Model has not been sampled; no forests to extract.") + if forest in ("prognostic", "mu"): + if not self._samples.has_mu_forest(): + return None + cpp = self._samples.materialize_mu_forest() + fc = ForestContainer(cpp.NumTrees(), 1, True, False) + fc.forest_container_cpp = cpp + return fc + elif forest in ("treatment", "tau"): + if not self._samples.has_tau_forest(): + return None + cpp = self._samples.materialize_tau_forest() + fc = ForestContainer(cpp.NumTrees(), self.treatment_dim, False, False) + fc.forest_container_cpp = cpp + return fc + elif forest == "variance": + if not self._samples.has_variance_forest(): + return None + cpp = self._samples.materialize_variance_forest() + fc = ForestContainer(cpp.NumTrees(), 1, True, True) + fc.forest_container_cpp = cpp + return fc + raise ValueError( + f"Unknown forest '{forest}'; expected 'prognostic', 'treatment', or 'variance'." + ) + + @property + def _forest_container_mu(self): + # Internal cached deep copy for prediction / serialization (no borrowed-ptr accessor). + if self._samples is None or not self._samples.has_mu_forest(): + return None + if self._fc_mu_cache is None: + cpp = self._samples.materialize_mu_forest() + fc = ForestContainer(cpp.NumTrees(), 1, True, False) + fc.forest_container_cpp = cpp + self._fc_mu_cache = fc + return self._fc_mu_cache + + @property + def _forest_container_tau(self): + if self._samples is None or not self._samples.has_tau_forest(): + return None + if self._fc_tau_cache is None: + cpp = self._samples.materialize_tau_forest() + fc = ForestContainer(cpp.NumTrees(), self.treatment_dim, False, False) + fc.forest_container_cpp = cpp + self._fc_tau_cache = fc + return self._fc_tau_cache + + @property + def _forest_container_variance(self): + if self._samples is None or not self._samples.has_variance_forest(): + return None + if self._fc_variance_cache is None: + cpp = self._samples.materialize_variance_forest() + fc = ForestContainer(cpp.NumTrees(), 1, True, True) + fc.forest_container_cpp = cpp + self._fc_variance_cache = fc + return self._fc_variance_cache + + @property + def forest_container_mu(self): + raise AttributeError( + "`BCFModel.forest_container_mu` has been removed. The sampled forests are owned by " + "`model.samples`; extract a standalone copy with `model.extract_forest('prognostic')`." + ) + + @property + def forest_container_tau(self): + raise AttributeError( + "`BCFModel.forest_container_tau` has been removed. The sampled forests are owned by " + "`model.samples`; extract a standalone copy with `model.extract_forest('treatment')`." + ) + + @property + def forest_container_variance(self): + raise AttributeError( + "`BCFModel.forest_container_variance` has been removed. The sampled forests are owned " + "by `model.samples`; extract a standalone copy with `model.extract_forest('variance')`." + ) + + @property + def global_var_samples(self): + if self._samples is None: + return None + arr = self._samples.global_var_samples() + return arr if arr.size else None + + @property + def leaf_scale_mu_samples(self): + if self._samples is None: + return None + arr = self._samples.leaf_scale_mu_samples() + return arr if arr.size else None + + @property + def leaf_scale_tau_samples(self): + if self._samples is None: + return None + arr = self._samples.leaf_scale_tau_samples() + return arr if arr.size else None + + @property + def num_samples(self): + return self._samples.num_samples() if self._samples is not None else None def sample( self, @@ -96,6 +253,7 @@ def sample( propensity_test: np.array = None, rfx_group_ids_test: np.array = None, rfx_basis_test: np.array = None, + observation_weights_train: Optional[np.ndarray] = None, observation_weights: Optional[np.ndarray] = None, num_gfr: int = 5, num_burnin: int = 0, @@ -139,7 +297,7 @@ def sample( test set evaluation for group labels that were not in the training set. rfx_basis_test : np.array, optional Optional test set basis for "random-slope" regression in additive random effects model. - observation_weights : np.array, optional + observation_weights_train : np.array, optional Optional vector of observation weights of length ``n_train``. Weights are applied as ``y_i | - ~ N(mu(X_i), sigma^2 / w_i)``, so larger weights increase an observation's influence on the fit. All weights must be non-negative. Defaults to ``None`` (all @@ -148,6 +306,8 @@ def sample( not compatible with cloglog link functions. Note: these are referred to internally in the C++ layer as "variance weights" (``var_weights``), since they scale the residual variance. + observation_weights : np.array, optional + Deprecated alias for ``observation_weights_train``; will be removed in a future release. num_gfr : int, optional Number of "warm-start" iterations run using the grow-from-root algorithm (He and Hahn, 2021). Defaults to `5`. num_burnin : int, optional @@ -226,6 +386,7 @@ def sample( - **outcome_model** (*stochtree.OutcomeModel*): An object of class ``OutcomeModel`` specifying the outcome model. Default: ``OutcomeModel(outcome="continuous", link="identity")``. Pre-empts the deprecated ``probit_outcome_model`` parameter if specified. - **probit_outcome_model** (*bool*): Deprecated in favor of ``outcome_model``. Whether or not the outcome should be modeled as explicitly binary via a probit link. If ``True``, ``y`` must only contain the values ``0`` and ``1``. Default: ``False``. - **num_threads** (*int*): Number of threads to use in the GFR and MCMC algorithms, as well as prediction. Defaults to ``1`` (single-threaded). Set to ``-1`` to use the maximum number of available threads, or a positive integer for a specific count. OpenMP must be available for values other than ``1``. + - **verbose** (*bool*): Whether to print sampler progress (GFR / MCMC iteration updates) to the console during sampling. Defaults to ``False``. **prognostic_forest_params keys** @@ -307,6 +468,7 @@ def sample( "outcome_model": OutcomeModel(outcome="continuous", link="identity"), "probit_outcome_model": False, "num_threads": 1, + "verbose": False, } general_params_updated = _preprocess_params( general_params_default, general_params @@ -407,6 +569,7 @@ def sample( self.probit_outcome_model = general_params_updated["probit_outcome_model"] self.outcome_model = general_params_updated["outcome_model"] num_threads = general_params_updated["num_threads"] + verbose = general_params_updated["verbose"] # 2. Mu forest parameters num_trees_mu = prognostic_forest_params_updated["num_trees"] @@ -595,77 +758,58 @@ def sample( raise ValueError( "`previous_model_warmstart_sample_num` exceeds the number of samples in `previous_model_json`" ) - previous_model_decrement = True if num_chains > previous_model_warmstart_sample_num + 1: warnings.warn( "The number of chains being sampled exceeds the number of previous model samples available from the requested position in `previous_model_json`. All chains will be initialized from the same sample." ) - previous_model_decrement = False - previous_y_scale = previous_bcf_model.y_std previous_model_num_samples = previous_bcf_model.num_samples - if previous_bcf_model.sample_sigma2_global: - previous_global_var_samples = previous_bcf_model.global_var_samples / ( - previous_y_scale * previous_y_scale - ) - else: - previous_global_var_samples = None - if previous_bcf_model.sample_sigma2_leaf_mu: - previous_leaf_var_mu_samples = previous_bcf_model.leaf_scale_mu_samples - else: - previous_leaf_var_mu_samples = None - if previous_bcf_model.sample_sigma2_leaf_tau: - previous_leaf_var_tau_samples = ( - previous_bcf_model.leaf_scale_tau_samples - ) - else: - previous_leaf_var_tau_samples = None - if previous_bcf_model.adaptive_coding: - previous_b0_samples = previous_bcf_model.b0_samples - previous_b1_samples = previous_bcf_model.b1_samples - else: - previous_b0_samples = None - previous_b1_samples = None if previous_model_warmstart_sample_num + 1 > previous_model_num_samples: raise ValueError( "`previous_model_warmstart_sample_num` exceeds the number of samples in `previous_model_json`" ) else: - previous_y_scale = None - previous_global_var_samples = None - previous_leaf_var_mu_samples = None - previous_leaf_var_tau_samples = None previous_model_num_samples = 0 - previous_b0_samples = None - previous_b1_samples = None # Determine whether conditional variance model will be fit self.include_variance_forest = True if num_trees_variance > 0 else False - # observation_weights validation and compatibility checks + # `observation_weights` was renamed to `observation_weights_train`; honor the + # deprecated argument for one release cycle. if observation_weights is not None: - if not isinstance(observation_weights, np.ndarray): - raise ValueError("observation_weights must be a numpy array") - observation_weights_ = np.squeeze(observation_weights) - if observation_weights_.ndim != 1: + warnings.warn( + "`observation_weights` is deprecated and will be removed in a future " + "release; use `observation_weights_train` instead.", + DeprecationWarning, + stacklevel=2, + ) + if observation_weights_train is None: + observation_weights_train = observation_weights + + # observation_weights_train validation and compatibility checks + if observation_weights_train is not None: + if not isinstance(observation_weights_train, np.ndarray): + raise ValueError("observation_weights_train must be a numpy array") + observation_weights_train_ = np.squeeze(observation_weights_train) + if observation_weights_train_.ndim != 1: raise ValueError( - "observation_weights must be a 1-dimensional numpy array" + "observation_weights_train must be a 1-dimensional numpy array" ) - if np.any(observation_weights_ < 0): - raise ValueError("observation_weights cannot have any negative values") - if np.all(observation_weights_ == 0) and num_gfr > 0: + if np.any(observation_weights_train_ < 0): + raise ValueError("observation_weights_train cannot have any negative values") + if np.all(observation_weights_train_ == 0) and num_gfr > 0: raise ValueError( - "observation_weights are all zero (prior sampling mode) but num_gfr > 0. " + "observation_weights_train are all zero (prior sampling mode) but num_gfr > 0. " "GFR warm-start is data-dependent and ill-defined with zero weights. " - "Set num_gfr=0 when using all-zero observation_weights." + "Set num_gfr=0 when using all-zero observation_weights_train." ) if link_is_cloglog: raise ValueError( - "observation_weights are not compatible with cloglog link functions." + "observation_weights_train are not compatible with cloglog link functions." ) if self.include_variance_forest: raise ValueError( - "observation_weights are not compatible with a variance forest model. " - "Use either observation_weights or a variance forest, not both." + "observation_weights_train are not compatible with a variance forest model. " + "Use either observation_weights_train or a variance forest, not both." ) # Check data inputs @@ -719,9 +863,13 @@ def sample( if Z_train is not None: if Z_train.ndim == 1: Z_train = np.expand_dims(Z_train, 1) + if not np.issubdtype(Z_train.dtype, np.float64): + Z_train = Z_train.astype(np.float64) if propensity_train is not None: if propensity_train.ndim == 1: propensity_train = np.expand_dims(propensity_train, 1) + if not np.issubdtype(propensity_train.dtype, np.float64): + propensity_train = propensity_train.astype(np.float64) if y_train.ndim == 1: y_train = np.expand_dims(y_train, 1) if not np.issubdtype(y_train.dtype, np.float64): @@ -733,9 +881,13 @@ def sample( if Z_test is not None: if Z_test.ndim == 1: Z_test = np.expand_dims(Z_test, 1) + if not np.issubdtype(Z_test.dtype, np.float64): + Z_test = Z_test.astype(np.float64) if propensity_test is not None: if propensity_test.ndim == 1: propensity_test = np.expand_dims(propensity_test, 1) + if not np.issubdtype(propensity_test.dtype, np.float64): + propensity_test = propensity_test.astype(np.float64) if rfx_group_ids_train is not None: if rfx_group_ids_train.ndim != 1: rfx_group_ids_train = np.squeeze(rfx_group_ids_train) @@ -858,20 +1010,12 @@ def sample( "Consider converting binary variables to ordered categorical (i.e. `pd.Categorical(..., ordered = True)`." ) - # Prognostic model details - leaf_dimension_mu = 1 - leaf_model_mu = 0 - # Treatment details self.treatment_dim = Z_train.shape[1] self.multivariate_treatment = True if self.treatment_dim > 1 else False - leaf_dimension_tau = self.treatment_dim leaf_model_tau = 2 if self.multivariate_treatment else 1 # treatment_leaf_model = 2 if self.multivariate_treatment else 1 - # Set variance leaf model type (currently only one option) - leaf_dimension_variance = 1 - leaf_model_variance = 3 # Check parameters if sigma2_leaf_tau is not None: @@ -1131,9 +1275,9 @@ def sample( # Covariate preprocessing self._covariate_preprocessor = CovariatePreprocessor() self._covariate_preprocessor.fit(X_train) - X_train_processed = self._covariate_preprocessor.transform(X_train) + X_train_processed = self._covariate_preprocessor.transform(X_train).astype(np.float64) if X_test is not None: - X_test_processed = self._covariate_preprocessor.transform(X_test) + X_test_processed = self._covariate_preprocessor.transform(X_test).astype(np.float64) feature_types = np.asarray( self._covariate_preprocessor._processed_feature_types ) @@ -1560,7 +1704,7 @@ def sample( else: self.internal_propensity_model = False - # Runtime checks on RFX group ids + # Runtime checks on RFX group ids self.has_rfx = False has_rfx_test = False if rfx_group_ids_train is not None: @@ -1646,219 +1790,9 @@ def sample( "Sampling global error variance not yet supported for models with variance forests, so the global error variance parameter will not be sampled in this model." ) sample_sigma2_global = False - - # Handle standardization, prior calibration, and initialization of forest - # differently for binary and continuous outcomes - if link_is_probit: - # Compute a probit-scale offset and fix scale to 1 - self.y_bar = norm.ppf(np.squeeze(np.mean(y_train))) - self.y_std = 1.0 - - # Set a pseudo outcome by subtracting mean(y_train) from y_train - resid_train = y_train - np.squeeze(np.mean(y_train)) - - # Set initial value for the mu forest - init_mu = 0.0 - - # Calibrate priors for sigma^2 and tau - # Set sigma2_init to 1, ignoring default provided - sigma2_init = 1.0 - current_sigma2 = sigma2_init - self.sigma2_init = sigma2_init - # Skip variance_forest_init, since variance forests are not supported with probit link - b_leaf_mu = 1.0 / num_trees_mu if b_leaf_mu is None else b_leaf_mu - b_leaf_tau = 1.0 / (2 * num_trees_tau) if b_leaf_tau is None else b_leaf_tau - sigma2_leaf_mu = ( - 1 / num_trees_mu if sigma2_leaf_mu is None else sigma2_leaf_mu - ) - if isinstance(sigma2_leaf_mu, (float, np.floating)): - current_leaf_scale_mu = np.array([[sigma2_leaf_mu]]) - else: - raise ValueError("sigma2_leaf_mu must be a scalar") - # Calibrate prior so that P(abs(tau(X)) < delta_max / dnorm(0)) = p - # Use p = 0.9 as an internal default rather than adding another - # user-facing "parameter" of the binary outcome BCF prior. - # Can be overriden by specifying `sigma2_leaf_init` in - # treatment_effect_forest_params. - p = 0.6827 - q_quantile = norm.ppf((p + 1) / 2.0) - sigma2_leaf_tau = ( - ((delta_max / (q_quantile * norm.pdf(0))) ** 2) / num_trees_tau - if sigma2_leaf_tau is None - else sigma2_leaf_tau - ) - if self.multivariate_treatment: - if not isinstance(sigma2_leaf_tau, np.ndarray): - sigma2_leaf_tau = np.diagflat( - np.repeat(sigma2_leaf_tau, self.treatment_dim) - ) - if isinstance(sigma2_leaf_tau, (float, np.floating)): - if Z_train.shape[1] > 1: - current_leaf_scale_tau = np.zeros( - (Z_train.shape[1], Z_train.shape[1]), dtype=float - ) - np.fill_diagonal(current_leaf_scale_tau, sigma2_leaf_tau) - else: - current_leaf_scale_tau = np.array([[sigma2_leaf_tau]]) - elif isinstance(sigma2_leaf_tau, np.ndarray): - if sigma2_leaf_tau.ndim != 2: - raise ValueError( - "sigma2_leaf_tau must be a 2d symmetric numpy array if provided in matrix form" - ) - if sigma2_leaf_tau.shape[0] != sigma2_leaf_tau.shape[1]: - raise ValueError( - "sigma2_leaf_tau must be a 2d symmetric numpy array if provided in matrix form" - ) - if sigma2_leaf_tau.shape[0] != Z_train.shape[1]: - raise ValueError( - "sigma2_leaf_tau must be a 2d numpy array with dimension matching that of the treatment vector" - ) - current_leaf_scale_tau = sigma2_leaf_tau - else: - raise ValueError("sigma2_leaf_tau must be a scalar or a 2d numpy array") - else: - # Standardize if requested - if self.standardize: - self.y_bar = np.squeeze(np.mean(y_train)) - self.y_std = np.squeeze(np.std(y_train)) - else: - self.y_bar = 0 - self.y_std = 1 - - # Compute residual value - resid_train = (y_train - self.y_bar) / self.y_std - - # Compute initial value of root nodes in mean forest - init_mu = np.squeeze(np.mean(resid_train)) - - # Calibrate priors for global sigma^2 and sigma2_leaf - if not sigma2_init: - sigma2_init = 1.0 * np.var(resid_train) - if not variance_forest_leaf_init: - variance_forest_leaf_init = 0.6 * np.var(resid_train) - current_sigma2 = sigma2_init - self.sigma2_init = sigma2_init - b_leaf_mu = ( - np.squeeze(np.var(resid_train)) / num_trees_mu - if b_leaf_mu is None - else b_leaf_mu - ) - b_leaf_tau = ( - np.squeeze(np.var(resid_train)) / (2 * num_trees_tau) - if b_leaf_tau is None - else b_leaf_tau - ) - sigma2_leaf_mu = ( - np.squeeze(2 * np.var(resid_train)) / num_trees_mu - if sigma2_leaf_mu is None - else sigma2_leaf_mu - ) - if isinstance(sigma2_leaf_mu, (float, np.floating)): - current_leaf_scale_mu = np.array([[sigma2_leaf_mu]]) - else: - raise ValueError("sigma2_leaf_mu must be a scalar") - sigma2_leaf_tau = ( - np.squeeze(0.5 * np.var(resid_train)) / (num_trees_tau) - if sigma2_leaf_tau is None - else sigma2_leaf_tau - ) - if self.multivariate_treatment: - if not isinstance(sigma2_leaf_tau, np.ndarray): - sigma2_leaf_tau = np.diagflat( - np.repeat(sigma2_leaf_tau, self.treatment_dim) - ) - if isinstance(sigma2_leaf_tau, (float, np.floating)): - if Z_train.shape[1] > 1: - current_leaf_scale_tau = np.zeros( - (Z_train.shape[1], Z_train.shape[1]), dtype=float - ) - np.fill_diagonal(current_leaf_scale_tau, sigma2_leaf_tau) - else: - current_leaf_scale_tau = np.array([[sigma2_leaf_tau]]) - elif isinstance(sigma2_leaf_tau, np.ndarray): - if sigma2_leaf_tau.ndim != 2: - raise ValueError( - "sigma2_leaf_tau must be a 2d symmetric numpy array if provided in matrix form" - ) - if sigma2_leaf_tau.shape[0] != sigma2_leaf_tau.shape[1]: - raise ValueError( - "sigma2_leaf_tau must be a 2d symmetric numpy array if provided in matrix form" - ) - if sigma2_leaf_tau.shape[0] != Z_train.shape[1]: - raise ValueError( - "sigma2_leaf_tau must be a 2d numpy array with dimension matching that of the treatment vector" - ) - current_leaf_scale_tau = sigma2_leaf_tau - else: - raise ValueError("sigma2_leaf_tau must be a scalar or a 2d numpy array") - if self.include_variance_forest: - if not a_forest: - a_forest = num_trees_variance / a_0**2 + 0.5 - if not b_forest: - b_forest = num_trees_variance / a_0**2 - else: - if not a_forest: - a_forest = 1.0 - if not b_forest: - b_forest = 1.0 - - # Set up random effects structures - if self.has_rfx: - # Prior parameters - if rfx_working_parameter_prior_mean is None: - if num_rfx_components == 1: - alpha_init = np.array([0.0], dtype=float) - elif num_rfx_components > 1: - alpha_init = np.zeros(num_rfx_components, dtype=float) - else: - raise ValueError("There must be at least 1 random effect component") - else: - alpha_init = _expand_dims_1d( - rfx_working_parameter_prior_mean, num_rfx_components - ) - - if rfx_group_parameter_prior_mean is None: - xi_init = np.tile(np.expand_dims(alpha_init, 1), (1, num_rfx_groups)) - else: - xi_init = _expand_dims_2d( - rfx_group_parameter_prior_mean, num_rfx_components, num_rfx_groups - ) - - if rfx_working_parameter_prior_cov is None: - sigma_alpha_init = np.identity(num_rfx_components) - else: - sigma_alpha_init = _expand_dims_2d_diag( - rfx_working_parameter_prior_cov, num_rfx_components - ) - - if rfx_group_parameter_prior_cov is None: - sigma_xi_init = np.identity(num_rfx_components) - else: - sigma_xi_init = _expand_dims_2d_diag( - rfx_group_parameter_prior_cov, num_rfx_components - ) - - sigma_xi_shape = rfx_variance_prior_shape - sigma_xi_scale = rfx_variance_prior_scale - - # Random effects sampling data structures - rfx_dataset_train = RandomEffectsDataset() - rfx_dataset_train.add_group_labels(rfx_group_ids_train) - rfx_dataset_train.add_basis(rfx_basis_train) - rfx_tracker = RandomEffectsTracker(rfx_group_ids_train) - rfx_model = RandomEffectsModel(num_rfx_components, num_rfx_groups) - rfx_model.set_working_parameter(alpha_init) - rfx_model.set_group_parameters(xi_init) - rfx_model.set_working_parameter_covariance(sigma_alpha_init) - rfx_model.set_group_parameter_covariance(sigma_xi_init) - rfx_model.set_variance_prior_shape(sigma_xi_shape) - rfx_model.set_variance_prior_scale(sigma_xi_scale) - self.rfx_container = RandomEffectsContainer() - self.rfx_container.load_new_container( - num_rfx_components, num_rfx_groups, rfx_tracker - ) - - # Update variable weights + + # Update variable weights for one-hot-expanded covariates and zero out + # excluded variables. Runs for both the C++ and Python sampling paths. variable_counts = [original_var_indices.count(i) for i in original_var_indices] variable_weights_mu_adj = [1 / i for i in variable_counts] variable_weights_tau_adj = [1 / i for i in variable_counts] @@ -1885,7 +1819,8 @@ def sample( [variable_subset_variance.count(i) == 0 for i in original_var_indices] ] = 0 - # Update covariates to include propensities if requested + # Append propensity score to X and update feature_types / variable weights. + # Runs for both the C++ and Python sampling paths. if propensity_covariate != "none": feature_types = np.append( feature_types, np.repeat(0, propensity_train.shape[1]) @@ -1933,1298 +1868,592 @@ def sample( # Store propensity score requirements of the BCF forests self.propensity_covariate = propensity_covariate - # Set num_features_subsample to default, ncol(X_train), if not already set + # Resolve num_features_subsample defaults now that X_train_processed has its + # final shape (propensity appended, one-hot expanded, etc.). Both the C++ and + # Python paths rely on these; if left as None the C++ path would receive 0 from + # get_config_scalar_default, which disables all GFR splits. + num_features_total = X_train_processed.shape[1] if num_features_subsample_mu is None: - num_features_subsample_mu = X_train_processed.shape[1] + num_features_subsample_mu = num_features_total if num_features_subsample_tau is None: - num_features_subsample_tau = X_train_processed.shape[1] + num_features_subsample_tau = num_features_total if num_features_subsample_variance is None: - num_features_subsample_variance = X_train_processed.shape[1] + num_features_subsample_variance = num_features_total - # Container of variance parameter samples - self.num_gfr = num_gfr - self.num_burnin = num_burnin - self.num_mcmc = num_mcmc - self.num_chains = num_chains - self.keep_every = keep_every - num_temp_samples = num_gfr + num_burnin + num_mcmc * keep_every - num_retained_samples = num_mcmc * num_chains - # Delete GFR samples from these containers after the fact if desired - # if keep_gfr: - # num_retained_samples += num_gfr - num_retained_samples += num_gfr - if keep_burnin: - num_retained_samples += num_burnin - self.num_samples = num_retained_samples - self.sample_sigma2_global = sample_sigma2_global - self.sample_sigma2_leaf_mu = sample_sigma2_leaf_mu - self.sample_sigma2_leaf_tau = sample_sigma2_leaf_tau - if sample_sigma2_global: - self.global_var_samples = np.empty(self.num_samples, dtype=np.float64) - if sample_sigma2_leaf_mu: - self.leaf_scale_mu_samples = np.empty(self.num_samples, dtype=np.float64) - if sample_sigma2_leaf_tau: - self.leaf_scale_tau_samples = np.empty(self.num_samples, dtype=np.float64) - if self.sample_tau_0: - p_tau0 = Z_train.shape[1] if Z_train.ndim > 1 else 1 - self.tau_0_samples = np.empty((p_tau0, self.num_samples), dtype=np.float64) - muhat_train_raw = np.empty((self.n_train, self.num_samples), dtype=np.float64) - if self.include_variance_forest: - sigma2_x_train_raw = np.empty( - (self.n_train, self.num_samples), dtype=np.float64 + # Expand dimensions on RFX prior parameters if provided + # Working parameter (should be expanded to a 1d array if provided as a scalar) + if rfx_working_parameter_prior_mean is not None: + rfx_working_parameter_prior_mean = _expand_dims_1d( + rfx_working_parameter_prior_mean, num_rfx_components ) - sample_counter = -1 - # Prepare adaptive coding structure - if self.adaptive_coding: - if np.size(b_0) > 1 or np.size(b_1) > 1: - raise ValueError("b_0 and b_1 must be single numeric values") - if not ( - isinstance(b_0, (int, float, np.floating)) - or isinstance(b_1, (int, float, np.floating)) - ): - raise ValueError("b_0 and b_1 must be numeric values") - self.b0_samples = np.empty(self.num_samples, dtype=np.float64) - self.b1_samples = np.empty(self.num_samples, dtype=np.float64) - current_b_0 = b_0 - current_b_1 = b_1 - tau_basis_train = (1 - Z_train) * current_b_0 + Z_train * current_b_1 - if self.has_test: - tau_basis_test = (1 - Z_test) * current_b_0 + Z_test * current_b_1 - else: - tau_basis_train = Z_train - if self.has_test: - tau_basis_test = Z_test + # Group parameter (should be expanded to a 2d array if provided as a scalar) + if rfx_group_parameter_prior_mean is not None: + rfx_group_parameter_prior_mean = _expand_dims_2d( + rfx_group_parameter_prior_mean, num_rfx_components, num_rfx_groups + ) - # Prepare tau_0 (global treatment effect intercept) structure - if self.sample_tau_0: - tau_0 = np.zeros(p_tau0) - # Auto-calibrate prior variance if not provided - if tau_0_prior_var is None: - tau_0_prior_var = np.var(resid_train) - - # Prognostic Forest Dataset (covariates) - forest_dataset_train = Dataset() - forest_dataset_train.add_covariates(X_train_processed) - forest_dataset_train.add_basis(tau_basis_train) - if observation_weights is not None: - forest_dataset_train.add_variance_weights(observation_weights_) - if self.has_test: - forest_dataset_test = Dataset() - forest_dataset_test.add_covariates(X_test_processed) - forest_dataset_test.add_basis(tau_basis_test) + # Working parameter (should be expanded to a diagonal matrix if provided as a scalar) + if rfx_working_parameter_prior_cov is not None: + rfx_working_parameter_prior_cov = _expand_dims_2d_diag( + rfx_working_parameter_prior_cov, num_rfx_components + ) - # Residual - residual_train = Residual(resid_train) + # Group parameter (should be expanded to a diagonal matrix if provided as a scalar) + if rfx_group_parameter_prior_cov is not None: + rfx_group_parameter_prior_cov = _expand_dims_2d_diag( + rfx_group_parameter_prior_cov, num_rfx_components + ) - # C++ and numpy random number generator - if random_seed is None: - cpp_rng = RNG(-1) - self.rng = np.random.default_rng() - else: - cpp_rng = RNG(random_seed) - self.rng = np.random.default_rng(random_seed) - - # Sampling data structures - global_model_config = GlobalModelConfig(global_error_variance=current_sigma2) - forest_model_config_mu = ForestModelConfig( - num_trees=num_trees_mu, - num_features=forest_dataset_train.num_covariates(), - num_observations=self.n_train, - feature_types=feature_types, - variable_weights=variable_weights_mu, - leaf_dimension=leaf_dimension_mu, - alpha=alpha_mu, - beta=beta_mu, - min_samples_leaf=min_samples_leaf_mu, - max_depth=max_depth_mu, - leaf_model_type=leaf_model_mu, - leaf_model_scale=current_leaf_scale_mu, - cutpoint_grid_size=cutpoint_grid_size, - num_features_subsample=num_features_subsample_mu, + # Validate delta_max and, for binary (probit) outcomes, calibrate the treatment + # effect leaf scale from it when sigma2_leaf_init was not set directly for tau. + if (delta_max <= 0) or (delta_max >= 1): + raise ValueError("delta_max must be > 0 and < 1") + if link_is_probit and sigma2_leaf_tau is None: + # Calibrate the prior so that P(|tau(X)| < delta_max / dnorm(0)) = p, + # using p = 0.6827 as an internal default (overridden by setting + # sigma2_leaf_init in treatment_effect_forest_params). + p = 0.6827 + q_quantile = norm.ppf((p + 1) / 2) + sigma2_leaf_tau_scalar = ( + (delta_max / (q_quantile * norm.pdf(0))) ** 2 + ) / num_trees_tau + if self.multivariate_treatment: + sigma2_leaf_tau = np.diag( + np.full(self.treatment_dim, sigma2_leaf_tau_scalar) + ) + else: + sigma2_leaf_tau = float(sigma2_leaf_tau_scalar) + + # Arrange all config in a large python dictionary + bcf_config = { + "standardize_outcome": self.standardize, + "num_threads": num_threads, + "verbose": verbose, + "cutpoint_grid_size": cutpoint_grid_size, + "link_function": 0 + if self.outcome_model.link == "identity" + else (1 if self.outcome_model.link == "probit" else 2), + "outcome_type": 0 + if self.outcome_model.outcome == "continuous" + else (1 if self.outcome_model.outcome == "binary" else 2), + "random_seed": random_seed, + "keep_gfr": keep_gfr, + "keep_burnin": keep_burnin, + "adaptive_coding": self.adaptive_coding, + "b_0_init": b_0, + "b_1_init": b_1, + "a_sigma2_global": a_global, + "b_sigma2_global": b_global, + "sigma2_global_init": sigma2_init, + "sample_sigma2_global": sample_sigma2_global, + "num_trees_mu": num_trees_mu, + "alpha_mu": alpha_mu, + "beta_mu": beta_mu, + "min_samples_leaf_mu": min_samples_leaf_mu, + "max_depth_mu": max_depth_mu, + "leaf_constant_mu": True, + "leaf_dim_mu": 1, + "exponentiated_leaf_mu": False, + "num_features_subsample_mu": num_features_subsample_mu, + "a_sigma2_mu": a_leaf_mu, + "b_sigma2_mu": b_leaf_mu, + "sigma2_mu_init": sigma2_leaf_mu if isinstance(sigma2_leaf_mu, float) else -1.0, + "sample_sigma2_leaf_mu": sample_sigma2_leaf_mu, + "num_trees_tau": num_trees_tau, + "alpha_tau": alpha_tau, + "beta_tau": beta_tau, + "min_samples_leaf_tau": min_samples_leaf_tau, + "max_depth_tau": max_depth_tau, + "leaf_constant_tau": False, + "leaf_dim_tau": self.treatment_dim, + "exponentiated_leaf_tau": False, + "num_features_subsample_tau": num_features_subsample_tau, + "a_sigma2_tau": a_leaf_tau, + "b_sigma2_tau": b_leaf_tau, + "sigma2_tau_init": sigma2_leaf_tau if isinstance(sigma2_leaf_tau, float) else -1.0, + "sample_sigma2_leaf_tau": sample_sigma2_leaf_tau, + "sample_tau_0": self.sample_tau_0, + "tau_0_prior_var_scalar": tau_0_prior_var + if (self.sample_tau_0 and not self.multivariate_treatment and tau_0_prior_var is not None) + else None, + "tau_leaf_model_type": leaf_model_tau, + "sigma2_leaf_tau_matrix": sigma2_leaf_tau.flatten(order="F") + if isinstance(sigma2_leaf_tau, np.ndarray) + else None, + "num_trees_variance": num_trees_variance, + "leaf_prior_calibration_param": a_0, + "shape_variance_forest": a_forest, + "scale_variance_forest": b_forest, + "variance_forest_leaf_init": variance_forest_leaf_init, + "alpha_variance": alpha_variance, + "beta_variance": beta_variance, + "min_samples_leaf_variance": min_samples_leaf_variance, + "max_depth_variance": max_depth_variance, + "leaf_constant_variance": True, + "leaf_dim_variance": 1, + "exponentiated_leaf_variance": True, + "num_features_subsample_variance": num_features_subsample_variance, + "feature_types": feature_types.astype(int), + "sweep_update_indices_mu": list(range(num_trees_mu)), + "sweep_update_indices_tau": list(range(num_trees_tau)), + "sweep_update_indices_variance": list(range(num_trees_variance)) + if num_trees_variance > 0 + else None, + "var_weights_mu": variable_weights_mu, + "var_weights_tau": variable_weights_tau, + "var_weights_variance": variable_weights_variance, + "has_random_effects": self.has_rfx, + "rfx_model_spec": 0 + if self.rfx_model_spec == "custom" + else (1 if self.rfx_model_spec == "intercept_only" else None), + "rfx_working_parameter_mean_prior": rfx_working_parameter_prior_mean + if self.has_rfx + else None, + "rfx_group_parameter_mean_prior": rfx_group_parameter_prior_mean + if self.has_rfx + else None, + "rfx_working_parameter_cov_prior": rfx_working_parameter_prior_cov + if self.has_rfx + else None, + "rfx_group_parameter_cov_prior": rfx_group_parameter_prior_cov + if self.has_rfx + else None, + "rfx_variance_prior_shape": rfx_variance_prior_shape + if self.has_rfx + else None, + "rfx_variance_prior_scale": rfx_variance_prior_scale + if self.has_rfx + else None, + } + + # Remove None values from config (alternative is to check for Nones on the C++ side when unpacking into non-optional types) + bcf_config = {k: v for k, v in bcf_config.items() if v is not None} + + # Cache the config so that continue_sampling() can warm-start with the same sampler + # configuration (the C++ BCFConfig is reconstructed from this dict at continue-time). + self._cached_bcf_config = bcf_config + + # Convert arrays to F-contiguous (column-major) before calling C++. + # convert_numpy_to_bart_data stores raw pointers into these arrays; if + # pybind11 has to make a copy (wrong dtype or order) that copy is destroyed + # when the helper returns, leaving a dangling pointer. + # Passing already-correct arrays causes pybind11 to return a view of + # the original, which remains alive in this Python scope. + X_train_cpp = np.asfortranarray(X_train_processed.astype(np.float64)) + # y_train_remapped = y_train - np.min(y_train) if link_is_cloglog else y_train + y_train_remapped = y_train + y_train_cpp = np.asfortranarray(y_train_remapped) + X_test_cpp = np.asfortranarray(X_test_processed.astype(np.float64)) if self.has_test else None + Z_train_cpp = np.asfortranarray(Z_train.astype(np.float64)) + Z_test_cpp = np.asfortranarray(Z_test.astype(np.float64)) if self.has_test else None + # rfx group IDs must be int32: pybind11 casts int64→int32 via a temporary + # inside convert_numpy_to_bart_data, making the returned raw pointer dangle. + rfx_group_ids_train_cpp = ( + rfx_group_ids_train.astype(np.int32) if rfx_group_ids_train is not None else None ) - forest_sampler_mu = ForestSampler( - forest_dataset_train, - global_model_config, - forest_model_config_mu, + rfx_group_ids_test_cpp = ( + rfx_group_ids_test.astype(np.int32) if rfx_group_ids_test is not None else None ) - forest_model_config_tau = ForestModelConfig( - num_trees=num_trees_tau, - num_features=forest_dataset_train.num_covariates(), - num_observations=self.n_train, - feature_types=feature_types, - variable_weights=variable_weights_tau, - leaf_dimension=leaf_dimension_tau, - alpha=alpha_tau, - beta=beta_tau, - min_samples_leaf=min_samples_leaf_tau, - max_depth=max_depth_tau, - leaf_model_type=leaf_model_tau, - leaf_model_scale=current_leaf_scale_tau, - cutpoint_grid_size=cutpoint_grid_size, - num_features_subsample=num_features_subsample_tau, + rfx_basis_train_cpp = ( + np.asfortranarray(rfx_basis_train.astype(np.float64)) if rfx_basis_train is not None else None ) - forest_sampler_tau = ForestSampler( - forest_dataset_train, - global_model_config, - forest_model_config_tau, + rfx_basis_test_cpp = ( + np.asfortranarray(rfx_basis_test.astype(np.float64)) if rfx_basis_test is not None else None ) - if self.include_variance_forest: - forest_model_config_variance = ForestModelConfig( - num_trees=num_trees_variance, - num_features=forest_dataset_train.num_covariates(), - num_observations=self.n_train, - feature_types=feature_types, - variable_weights=variable_weights_variance, - leaf_dimension=leaf_dimension_variance, - alpha=alpha_variance, - beta=beta_variance, - min_samples_leaf=min_samples_leaf_variance, - max_depth=max_depth_variance, - leaf_model_type=leaf_model_variance, - cutpoint_grid_size=cutpoint_grid_size, - variance_forest_shape=a_forest, - variance_forest_scale=b_forest, - num_features_subsample=num_features_subsample_variance, - ) - forest_sampler_variance = ForestSampler( - forest_dataset_train, global_model_config, forest_model_config_variance - ) - # Container of forest samples - self.forest_container_mu = ForestContainer( - num_trees_mu, leaf_dimension_mu, True, False + # Run the BCF sampler from C++ + bcf_results = bcf_sample_cpp( + X_train=X_train_cpp, + Z_train=Z_train_cpp, + y_train=y_train_cpp, + X_test=X_test_cpp, + Z_test=Z_test_cpp, + n_train=X_train_cpp.shape[0], + n_test=X_test_cpp.shape[0] if self.has_test else 0, + p=X_train_cpp.shape[1], + treatment_dim=self.treatment_dim, + obs_weights_train=observation_weights_train + if observation_weights_train is not None + else None, + obs_weights_test=None, + rfx_group_ids_train=rfx_group_ids_train_cpp, + rfx_group_ids_test=rfx_group_ids_test_cpp, + rfx_basis_train=rfx_basis_train_cpp, + rfx_basis_test=rfx_basis_test_cpp, + rfx_num_groups=num_rfx_groups if self.has_rfx else 0, + rfx_basis_dim=self.num_rfx_basis if self.has_rfx else 0, + num_gfr=num_gfr, + num_burnin=num_burnin, + keep_every=keep_every, + num_mcmc=num_mcmc, + num_chains=num_chains, + adaptive_coding=self.adaptive_coding, + config_input=bcf_config, ) - self.forest_container_tau = ForestContainer( - num_trees_tau, leaf_dimension_tau, False, False - ) - active_forest_mu = Forest(num_trees_mu, 1, True, False) - active_forest_tau = Forest(num_trees_tau, Z_train.shape[1], False, False) - if self.include_variance_forest: - self.forest_container_variance = ForestContainer( - num_trees_variance, 1, True, True - ) - active_forest_variance = Forest(num_trees_variance, 1, True, True) - # Variance samplers - if self.sample_sigma2_global: - global_var_model = GlobalVarianceModel() - if self.sample_sigma2_leaf_mu: - leaf_var_model_mu = LeafVarianceModel() - if self.sample_sigma2_leaf_tau: - leaf_var_model_tau = LeafVarianceModel() - - # Initialize the leaves of each tree in the prognostic forest - if not isinstance(init_mu, np.ndarray): - init_mu = np.array([init_mu]) - forest_sampler_mu.prepare_for_sampler( - forest_dataset_train, - residual_train, - active_forest_mu, - leaf_model_mu, - init_mu, - ) + # Store high level model metadata from C++ results + self.num_gfr = num_gfr + self.num_burnin = num_burnin + self.keep_every = keep_every + self.num_mcmc = num_mcmc + self.num_chains = num_chains + self.sample_sigma2_global = sample_sigma2_global + self.sample_sigma2_leaf_mu = sample_sigma2_leaf_mu + self.sample_sigma2_leaf_tau = sample_sigma2_leaf_tau - # Initialize the leaves of each tree in the treatment forest - if self.multivariate_treatment: - init_tau = np.zeros(Z_train.shape[1], dtype=float) - else: - init_tau = np.array([0.0]) - forest_sampler_tau.prepare_for_sampler( - forest_dataset_train, - residual_train, - active_forest_tau, - leaf_model_tau, - init_tau, + # Unpack standardization params computed by C++ sampler + self.y_bar = bcf_results["y_bar"] + self.y_std = bcf_results["y_std"] + self.sigma2_init = bcf_results["sigma2_init"] + self.sigma2_leaf_mu_init = bcf_results["sigma2_mu_init"] + self.sigma2_leaf_tau_init = bcf_results["sigma2_tau_init"] + self.b_leaf_mu = bcf_results["b_sigma2_mu"] + self.b_leaf_tau = bcf_results["b_sigma2_tau"] + self.shape_variance_forest = ( + bcf_results["shape_variance_forest"] + if self.include_variance_forest + else None + ) + self.scale_variance_forest = ( + bcf_results["scale_variance_forest"] + if self.include_variance_forest + else None ) - # Initialize the leaves of each tree in the variance forest - if self.include_variance_forest: - init_val_variance = np.array([variance_forest_leaf_init]) - forest_sampler_variance.prepare_for_sampler( - forest_dataset_train, - residual_train, - active_forest_variance, - leaf_model_variance, - init_val_variance, - ) + # Unpack mu forest predictions (forest owned by self._samples, assembled below) + mu_forest_preds_train = bcf_results[ + "mu_forest_predictions_train" + ].reshape(self.n_train, bcf_results["num_samples"], order="F") + self.mu_hat_train = mu_forest_preds_train * self.y_std + self.y_bar + if self.has_test: + mu_forest_preds_test = bcf_results[ + "mu_forest_predictions_test" + ].reshape(self.n_test, bcf_results["num_samples"], order="F") + self.mu_hat_test = mu_forest_preds_test * self.y_std + self.y_bar - # Run GFR (warm start) if specified - if num_gfr > 0: - for i in range(num_gfr): - # Keep all GFR samples at this stage -- remove from ForestSamples after MCMC - # keep_sample = keep_gfr - keep_sample = True - if keep_sample: - sample_counter += 1 - - if link_is_probit: - # Sample latent probit variable z | - - # outcome_pred is the centered forest prediction (not including y_bar_train). - # The truncated normal mean is outcome_pred + y_bar_train (the full eta on the probit scale). - # The residual stored is z - y_bar_train - outcome_pred so the forest sees a - # zero-centered signal and the prior shrinkage toward 0 is well-calibrated. - forest_pred_mu = active_forest_mu.predict(forest_dataset_train) - forest_pred_tau = active_forest_tau.predict(forest_dataset_train) - outcome_pred = forest_pred_mu + forest_pred_tau - if self.has_rfx: - rfx_pred = rfx_model.predict(rfx_dataset_train, rfx_tracker) - outcome_pred = outcome_pred + rfx_pred - eta_pred = outcome_pred + self.y_bar - mu0 = eta_pred[y_train[:, 0] == 0] - mu1 = eta_pred[y_train[:, 0] == 1] - n0 = np.sum(y_train[:, 0] == 0) - n1 = np.sum(y_train[:, 0] == 1) - u0 = self.rng.uniform( - low=0.0, - high=norm.cdf(0 - mu0), - size=n0, - ) - u1 = self.rng.uniform( - low=norm.cdf(0 - mu1), - high=1.0, - size=n1, - ) - resid_train[y_train[:, 0] == 0, 0] = mu0 + norm.ppf(u0) - resid_train[y_train[:, 0] == 1, 0] = mu1 + norm.ppf(u1) - - # Update outcome: center z by y_bar before passing to forests - new_outcome = (np.squeeze(resid_train) - self.y_bar) - outcome_pred - residual_train.update_data(new_outcome) - - # Sample the prognostic forest - forest_sampler_mu.sample_one_iteration( - self.forest_container_mu, - active_forest_mu, - forest_dataset_train, - residual_train, - cpp_rng, - global_model_config, - forest_model_config_mu, - keep_sample, - True, - num_threads, + # Unpack tau forest predictions (forest owned by self._samples, assembled below) + if self.multivariate_treatment: + tau_forest_preds_train = bcf_results[ + "tau_forest_predictions_train" + ].reshape( + self.n_train, self.treatment_dim, bcf_results["num_samples"], order="F" + ) + self.tau_hat_train = tau_forest_preds_train * self.y_std + if self.has_test: + tau_forest_preds_test = bcf_results[ + "tau_forest_predictions_test" + ].reshape( + self.n_test, self.treatment_dim, bcf_results["num_samples"], order="F" ) + self.tau_hat_test = tau_forest_preds_test * self.y_std + else: + tau_forest_preds_train = bcf_results[ + "tau_forest_predictions_train" + ].reshape(self.n_train, bcf_results["num_samples"], order="F") + self.tau_hat_train = tau_forest_preds_train * self.y_std + if self.has_test: + tau_forest_preds_test = bcf_results[ + "tau_forest_predictions_test" + ].reshape(self.n_test, bcf_results["num_samples"], order="F") + self.tau_hat_test = tau_forest_preds_test * self.y_std + + # Unpack y_hat results + self.y_hat_train = bcf_results[ + "y_hat_train" + ].reshape(self.n_train, bcf_results["num_samples"], order="F") + if self.has_test: + self.y_hat_test = bcf_results[ + "y_hat_test" + ].reshape(self.n_test, bcf_results["num_samples"], order="F") - # Cache train set predictions since they are already computed during sampling - if keep_sample: - muhat_train_raw[:, sample_counter] = ( - forest_sampler_mu.get_cached_forest_predictions() - ) - - # Sample variance parameters (if requested) - if self.sample_sigma2_global: - current_sigma2 = global_var_model.sample_one_iteration( - residual_train, cpp_rng, a_global, b_global - ) - global_model_config.update_global_error_variance(current_sigma2) - if self.sample_sigma2_leaf_mu: - current_leaf_scale_mu[0, 0] = ( - leaf_var_model_mu.sample_one_iteration( - active_forest_mu, cpp_rng, a_leaf_mu, b_leaf_mu - ) - ) - forest_model_config_mu.update_leaf_model_scale( - current_leaf_scale_mu - ) - if keep_sample: - self.leaf_scale_mu_samples[sample_counter] = ( - current_leaf_scale_mu[0, 0] - ) - - # Sample tau_0 (global treatment effect intercept, if requested) - if self.sample_tau_0: - mu_x_tau0 = np.squeeze( - active_forest_mu.predict_raw(forest_dataset_train) - ) - tau_x_raw_tau0 = active_forest_tau.predict_raw(forest_dataset_train) - Z_basis = ( - tau_basis_train.reshape(-1, 1) - if tau_basis_train.ndim == 1 - else tau_basis_train - ) - tau_x_raw_2d = tau_x_raw_tau0.reshape(self.n_train, -1) - tau_x_full = np.sum(Z_basis * tau_x_raw_2d, axis=1) - # Center z by y_bar so tau_0 does not absorb the probit intercept - resid_for_tau0 = ( - (np.squeeze(resid_train) - self.y_bar) - if link_is_probit - else np.squeeze(resid_train) - ) - partial_resid_tau0 = resid_for_tau0 - mu_x_tau0 - tau_x_full - if self.has_rfx: - partial_resid_tau0 = partial_resid_tau0 - np.squeeze( - rfx_model.predict(rfx_dataset_train, rfx_tracker) - ) - Ztr = Z_basis.T @ partial_resid_tau0 - ZtZ_current = Z_basis.T @ Z_basis - Sigma_post = np.linalg.inv( - ZtZ_current / current_sigma2 + np.eye(p_tau0) / tau_0_prior_var - ) - mu_post = Sigma_post @ Ztr / current_sigma2 - tau_0_new = self.rng.multivariate_normal( - mean=mu_post, cov=Sigma_post - ) - residual_train.add_vector( - -np.squeeze(Z_basis @ (tau_0_new - tau_0)) - ) - tau_0 = tau_0_new - if keep_sample: - self.tau_0_samples[:, sample_counter] = tau_0 - - # Sample the treatment forest - forest_sampler_tau.sample_one_iteration( - self.forest_container_tau, - active_forest_tau, - forest_dataset_train, - residual_train, - cpp_rng, - global_model_config, - forest_model_config_tau, - keep_sample, - True, - num_threads, + # Unpack RFX results + if self.has_rfx: + rfx_wrapper = RandomEffectsContainer() + rfx_wrapper.rfx_container_cpp = bcf_results["rfx_container"] + rfx_wrapper.rfx_label_mapper_cpp = bcf_results["rfx_label_mapper"] + rfx_wrapper.rfx_group_ids = bcf_results["rfx_label_mapper"].GetUniqueGroupIds() + self.rfx_container = rfx_wrapper + rfx_preds_train = ( + bcf_results["rfx_predictions_train"].reshape( + self.n_train, bcf_results["num_samples"], order="F" ) - - # Cannot cache train set predictions for tau because the cached predictions in the - # tracking data structures are pre-multiplied by the basis (treatment) - # ... - - # Sample coding parameters (if requested) - if self.adaptive_coding: - mu_x = active_forest_mu.predict_raw(forest_dataset_train) - tau_x = np.squeeze( - active_forest_tau.predict_raw(forest_dataset_train) - ) - # Center z by y_bar so coding regression does not absorb the probit intercept - resid_for_coding = ( - (resid_train - self.y_bar) if link_is_probit else resid_train - ) - partial_resid_train = np.squeeze(resid_for_coding - mu_x) - if self.has_rfx: - rfx_pred = np.squeeze( - rfx_model.predict(rfx_dataset_train, rfx_tracker) - ) - partial_resid_train = partial_resid_train - rfx_pred - # Use tau_total = tau_0 + tau(X) for sufficient stats when sample_tau_0 - tau_x_for_coding = ( - (tau_x + tau_0[0]) if self.sample_tau_0 else tau_x - ) - s_tt0 = np.sum( - tau_x_for_coding * tau_x_for_coding * (np.squeeze(Z_train) == 0) - ) - s_tt1 = np.sum( - tau_x_for_coding * tau_x_for_coding * (np.squeeze(Z_train) == 1) - ) - s_ty0 = np.sum( - tau_x_for_coding - * partial_resid_train - * (np.squeeze(Z_train) == 0) - ) - s_ty1 = np.sum( - tau_x_for_coding - * partial_resid_train - * (np.squeeze(Z_train) == 1) - ) - current_b_0 = self.rng.normal( - loc=(s_ty0 / (s_tt0 + 2 * current_sigma2)), - scale=np.sqrt(current_sigma2 / (s_tt0 + 2 * current_sigma2)), - size=1, - )[0] - current_b_1 = self.rng.normal( - loc=(s_ty1 / (s_tt1 + 2 * current_sigma2)), - scale=np.sqrt(current_sigma2 / (s_tt1 + 2 * current_sigma2)), - size=1, - )[0] - if self.sample_tau_0: - tau_basis_old = np.squeeze(tau_basis_train).copy() - tau_basis_train = ( - 1 - np.squeeze(Z_train) - ) * current_b_0 + np.squeeze(Z_train) * current_b_1 - forest_dataset_train.update_basis(tau_basis_train) - if self.has_test: - tau_basis_test = ( - 1 - np.squeeze(Z_test) - ) * current_b_0 + np.squeeze(Z_test) * current_b_1 - forest_dataset_test.update_basis(tau_basis_test) - if keep_sample: - self.b0_samples[sample_counter] = current_b_0 - self.b1_samples[sample_counter] = current_b_1 - - # Update residual to reflect adjusted basis - forest_sampler_tau.propagate_basis_update( - forest_dataset_train, residual_train, active_forest_tau - ) - - # Fix tau_0 component of residual after basis change - if self.sample_tau_0: - residual_train.add_vector( - -(np.squeeze(tau_basis_train) - tau_basis_old) * tau_0[0] - ) - - # Sample the variance forest - if self.include_variance_forest: - forest_sampler_variance.sample_one_iteration( - self.forest_container_variance, - active_forest_variance, - forest_dataset_train, - residual_train, - cpp_rng, - global_model_config, - forest_model_config_variance, - keep_sample, - True, - num_threads, - ) - - # Cache train set predictions since they are already computed during sampling - if keep_sample: - sigma2_x_train_raw[:, sample_counter] = ( - forest_sampler_variance.get_cached_forest_predictions() - ) - - # Sample variance parameters (if requested) - if self.sample_sigma2_global: - current_sigma2 = global_var_model.sample_one_iteration( - residual_train, cpp_rng, a_global, b_global - ) - global_model_config.update_global_error_variance(current_sigma2) - if keep_sample: - self.global_var_samples[sample_counter] = current_sigma2 - if self.sample_sigma2_leaf_tau: - current_leaf_scale_tau[0, 0] = ( - leaf_var_model_tau.sample_one_iteration( - active_forest_tau, cpp_rng, a_leaf_tau, b_leaf_tau - ) - ) - forest_model_config_tau.update_leaf_model_scale( - current_leaf_scale_tau - ) - if keep_sample: - self.leaf_scale_tau_samples[sample_counter] = ( - current_leaf_scale_tau[0, 0] - ) - - # Sample random effects - if self.has_rfx: - rfx_model.sample( - rfx_dataset_train, - residual_train, - rfx_tracker, - self.rfx_container, - keep_sample, - current_sigma2, - cpp_rng, - ) - - # Run MCMC - if num_burnin + num_mcmc > 0: - for chain_num in range(num_chains): - if num_gfr > 0: - forest_ind = num_gfr - chain_num - 1 - # Reset prognostic forest - active_forest_mu.reset(self.forest_container_mu, forest_ind) - forest_sampler_mu.reconstitute_from_forest( - active_forest_mu, - forest_dataset_train, - residual_train, - True, - ) - # Reset CATE forest - active_forest_tau.reset(self.forest_container_tau, forest_ind) - forest_sampler_tau.reconstitute_from_forest( - active_forest_tau, - forest_dataset_train, - residual_train, - True, - ) - # Reset variance forest - if self.include_variance_forest: - active_forest_variance.reset( - self.forest_container_variance, forest_ind - ) - forest_sampler_variance.reconstitute_from_forest( - active_forest_variance, - forest_dataset_train, - residual_train, - False, - ) - # Reset global error scale - if sample_sigma2_global: - current_sigma2 = self.global_var_samples[forest_ind] - global_model_config.update_global_error_variance(current_sigma2) - # Reset mu forest leaf scale - if sample_sigma2_leaf_mu: - leaf_scale_double_mu = self.leaf_scale_mu_samples[forest_ind] - current_leaf_scale_mu[0, 0] = leaf_scale_double_mu - forest_model_config_mu.update_leaf_model_scale( - current_leaf_scale_mu - ) - # Reset tau forest leaf scale - if sample_sigma2_leaf_tau: - leaf_scale_double_tau = self.leaf_scale_tau_samples[forest_ind] - current_leaf_scale_tau[0, 0] = leaf_scale_double_tau - forest_model_config_tau.update_leaf_model_scale( - current_leaf_scale_tau - ) - # Reset adaptive coding parameters - if self.adaptive_coding: - tau_basis_train_old = tau_basis_train.copy() - if self.b0_samples is not None: - current_b_0 = self.b0_samples[forest_ind] - else: - current_b_0 = b_0 - if self.b1_samples is not None: - current_b_1 = self.b1_samples[forest_ind] - else: - current_b_1 = b_1 - tau_basis_train = ( - 1 - np.squeeze(Z_train) - ) * current_b_0 + np.squeeze(Z_train) * current_b_1 - forest_dataset_train.update_basis(tau_basis_train) - if self.has_test: - tau_basis_test = ( - 1 - np.squeeze(Z_test) - ) * current_b_0 + np.squeeze(Z_test) * current_b_1 - forest_dataset_test.update_basis(tau_basis_test) - forest_sampler_tau.propagate_basis_update( - forest_dataset_train, residual_train, active_forest_tau - ) - # Correct residual for tau_0 component of the basis change - if self.sample_tau_0: - residual_train.add_vector( - -( - np.squeeze(tau_basis_train) - - np.squeeze(tau_basis_train_old) - ) - * tau_0[0] - ) - # Reset tau_0 intercept and correct the running residual - if self.sample_tau_0: - tau_0_old = tau_0.copy() - tau_0 = self.tau_0_samples[:, forest_ind].copy() - Z_basis_gfr = ( - tau_basis_train.reshape(-1, 1) - if tau_basis_train.ndim == 1 - else tau_basis_train - ) - residual_train.add_vector( - -np.squeeze(Z_basis_gfr @ (tau_0 - tau_0_old)) - ) - # Reset random effects terms - if self.has_rfx: - rfx_model.reset( - self.rfx_container, forest_ind, sigma_alpha_init - ) - rfx_tracker.reset( - rfx_model, - rfx_dataset_train, - residual_train, - self.rfx_container, - ) - elif has_prev_model: - warmstart_index = ( - previous_model_warmstart_sample_num - chain_num - if previous_model_decrement - else previous_model_warmstart_sample_num - ) - # Reset prognostic forest - active_forest_mu.reset( - previous_bcf_model.forest_container_mu, warmstart_index - ) - forest_sampler_mu.reconstitute_from_forest( - active_forest_mu, - forest_dataset_train, - residual_train, - True, - ) - # Reset CATE forest - active_forest_tau.reset( - previous_bcf_model.forest_container_tau, warmstart_index - ) - forest_sampler_tau.reconstitute_from_forest( - active_forest_tau, - forest_dataset_train, - residual_train, - True, - ) - # Reset variance forest - if self.include_variance_forest: - active_forest_variance.reset( - previous_bcf_model.forest_container_variance, - warmstart_index, - ) - forest_sampler_variance.reconstitute_from_forest( - active_forest_variance, - forest_dataset_train, - residual_train, - True, - ) - # Reset global error scale - if self.sample_sigma2_global: - current_sigma2 = previous_global_var_samples[warmstart_index] - global_model_config.update_global_error_variance(current_sigma2) - # Reset mu forest leaf scale - if ( - sample_sigma2_leaf_mu - and previous_leaf_var_mu_samples is not None - ): - leaf_scale_double_mu = previous_leaf_var_mu_samples[ - warmstart_index - ] - current_leaf_scale_mu[0, 0] = leaf_scale_double_mu - forest_model_config_mu.update_leaf_model_scale( - current_leaf_scale_mu - ) - # Reset mu forest leaf scale - if ( - sample_sigma2_leaf_tau - and previous_leaf_var_tau_samples is not None - ): - leaf_scale_double_tau = previous_leaf_var_tau_samples[ - warmstart_index - ] - current_leaf_scale_tau[0, 0] = leaf_scale_double_tau - forest_model_config_tau.update_leaf_model_scale( - current_leaf_scale_tau - ) - # Reset adaptive coding parameters - if self.adaptive_coding: - tau_basis_train_old = tau_basis_train.copy() - if previous_b0_samples is not None: - current_b_0 = previous_b0_samples[warmstart_index] - if previous_b1_samples is not None: - current_b_1 = previous_b1_samples[warmstart_index] - tau_basis_train = ( - 1 - np.squeeze(Z_train) - ) * current_b_0 + np.squeeze(Z_train) * current_b_1 - forest_dataset_train.update_basis(tau_basis_train) - if self.has_test: - tau_basis_test = ( - 1 - np.squeeze(Z_test) - ) * current_b_0 + np.squeeze(Z_test) * current_b_1 - forest_dataset_test.update_basis(tau_basis_test) - forest_sampler_tau.propagate_basis_update( - forest_dataset_train, residual_train, active_forest_tau - ) - # Correct residual for tau_0 component of the basis change - if self.sample_tau_0: - residual_train.add_vector( - -( - np.squeeze(tau_basis_train) - - np.squeeze(tau_basis_train_old) - ) - * tau_0[0] - ) - # Reset tau_0 intercept and correct the running residual - if self.sample_tau_0: - prev_tau_0_samples = getattr( - previous_bcf_model, "tau_0_samples", None - ) - if prev_tau_0_samples is not None: - tau_0_old = tau_0.copy() - # tau_0_samples in previous model are in original scale; convert back - tau_0 = ( - prev_tau_0_samples[:, warmstart_index] - / previous_bcf_model.y_std - ).copy() - Z_basis_ws = ( - tau_basis_train.reshape(-1, 1) - if tau_basis_train.ndim == 1 - else tau_basis_train - ) - residual_train.add_vector( - -np.squeeze(Z_basis_ws @ (tau_0 - tau_0_old)) - ) - # Reset random effects terms - if self.has_rfx: - rfx_model.reset( - previous_bcf_model.rfx_container, - warmstart_index, - sigma_alpha_init, - ) - rfx_tracker.reset( - rfx_model, - rfx_dataset_train, - residual_train, - previous_bcf_model.rfx_container, - ) - else: - # Reset prognostic forest - active_forest_mu.reset_root() - if init_mu.shape[0] == 1: - active_forest_mu.set_root_leaves(init_mu[0] / num_trees_mu) - else: - active_forest_mu.set_root_leaves(init_mu / num_trees_mu) - forest_sampler_mu.reconstitute_from_forest( - active_forest_mu, - forest_dataset_train, - residual_train, - True, - ) - # Reset CATE forest - active_forest_tau.reset_root() - if init_tau.shape[0] == 1: - active_forest_tau.set_root_leaves(init_tau[0] / num_trees_tau) - else: - active_forest_tau.set_root_leaves(init_tau / num_trees_tau) - forest_sampler_tau.reconstitute_from_forest( - active_forest_tau, - forest_dataset_train, - residual_train, - True, - ) - # Reset variance forest - if self.include_variance_forest: - active_forest_variance.reset_root() - active_forest_variance.set_root_leaves( - log(variance_forest_leaf_init) / num_trees_variance - ) - forest_sampler_variance.reconstitute_from_forest( - active_forest_variance, - forest_dataset_train, - residual_train, - False, - ) - # Reset global error scale - if self.sample_sigma2_global: - current_sigma2 = sigma2_init - global_model_config.update_global_error_variance(current_sigma2) - # Reset mu forest leaf scale - if ( - sample_sigma2_leaf_mu - and previous_leaf_var_mu_samples is not None - ): - current_leaf_scale_mu[0, 0] = sigma2_leaf_mu - forest_model_config_mu.update_leaf_model_scale( - current_leaf_scale_mu - ) - # Reset mu forest leaf scale - if ( - sample_sigma2_leaf_tau - and previous_leaf_var_tau_samples is not None - ): - current_leaf_scale_tau[0, 0] = sigma2_leaf_tau - forest_model_config_tau.update_leaf_model_scale( - current_leaf_scale_tau - ) - # Reset adaptive coding parameters - if self.adaptive_coding: - tau_basis_train_old = tau_basis_train.copy() - current_b_0 = b_0 - current_b_1 = b_1 - tau_basis_train = ( - 1 - np.squeeze(Z_train) - ) * current_b_0 + np.squeeze(Z_train) * current_b_1 - forest_dataset_train.update_basis(tau_basis_train) - if self.has_test: - tau_basis_test = ( - 1 - np.squeeze(Z_test) - ) * current_b_0 + np.squeeze(Z_test) * current_b_1 - forest_dataset_test.update_basis(tau_basis_test) - forest_sampler_tau.propagate_basis_update( - forest_dataset_train, residual_train, active_forest_tau - ) - # Correct residual for tau_0 component of the basis change - if self.sample_tau_0: - residual_train.add_vector( - -( - np.squeeze(tau_basis_train) - - np.squeeze(tau_basis_train_old) - ) - * tau_0[0] - ) - # Reset tau_0 to initial value (0) and correct the running residual - if self.sample_tau_0: - tau_0_old = tau_0.copy() - tau_0 = np.zeros_like(tau_0) - Z_basis_reset = ( - tau_basis_train.reshape(-1, 1) - if tau_basis_train.ndim == 1 - else tau_basis_train - ) - residual_train.add_vector( - -np.squeeze(Z_basis_reset @ (tau_0 - tau_0_old)) - ) - # Reset random effects terms - if self.has_rfx: - rfx_model.root_reset( - alpha_init, - xi_init, - sigma_alpha_init, - sigma_xi_init, - sigma_xi_shape, - sigma_xi_scale, - ) - rfx_tracker.root_reset( - rfx_model, - rfx_dataset_train, - residual_train, - self.rfx_container, - ) - # Sample MCMC and burnin for each chain - for i in range(num_gfr, num_temp_samples): - is_mcmc = i + 1 > num_gfr + num_burnin - if is_mcmc: - mcmc_counter = i - num_gfr - num_burnin + 1 - if mcmc_counter % keep_every == 0: - keep_sample = True - else: - keep_sample = False - else: - if keep_burnin: - keep_sample = True - else: - keep_sample = False - if keep_sample: - sample_counter += 1 - - if link_is_probit: - # Sample latent probit variable z | - - forest_pred_mu = active_forest_mu.predict(forest_dataset_train) - forest_pred_tau = active_forest_tau.predict( - forest_dataset_train - ) - outcome_pred = forest_pred_mu + forest_pred_tau - if self.has_rfx: - rfx_pred = rfx_model.predict(rfx_dataset_train, rfx_tracker) - outcome_pred = outcome_pred + rfx_pred - # Full probit-scale predictor: forests learn z - y_bar, so add y_bar back - eta_pred = outcome_pred + self.y_bar - mu0 = eta_pred[y_train[:, 0] == 0] - mu1 = eta_pred[y_train[:, 0] == 1] - n0 = np.sum(y_train[:, 0] == 0) - n1 = np.sum(y_train[:, 0] == 1) - u0 = self.rng.uniform( - low=0.0, - high=norm.cdf(0 - mu0), - size=n0, - ) - u1 = self.rng.uniform( - low=norm.cdf(0 - mu1), - high=1.0, - size=n1, - ) - resid_train[y_train[:, 0] == 0, 0] = mu0 + norm.ppf(u0) - resid_train[y_train[:, 0] == 1, 0] = mu1 + norm.ppf(u1) - - # Update outcome: center z by y_bar before passing to forests - new_outcome = ( - np.squeeze(resid_train) - self.y_bar - outcome_pred - ) - residual_train.update_data(new_outcome) - - # Sample the prognostic forest - forest_sampler_mu.sample_one_iteration( - self.forest_container_mu, - active_forest_mu, - forest_dataset_train, - residual_train, - cpp_rng, - global_model_config, - forest_model_config_mu, - keep_sample, - False, - num_threads, - ) - - # Cache train set predictions since they are already computed during sampling - if keep_sample: - muhat_train_raw[:, sample_counter] = ( - forest_sampler_mu.get_cached_forest_predictions() - ) - - # Sample variance parameters (if requested) - if self.sample_sigma2_global: - current_sigma2 = global_var_model.sample_one_iteration( - residual_train, cpp_rng, a_global, b_global - ) - global_model_config.update_global_error_variance(current_sigma2) - if self.sample_sigma2_leaf_mu: - current_leaf_scale_mu[0, 0] = ( - leaf_var_model_mu.sample_one_iteration( - active_forest_mu, cpp_rng, a_leaf_mu, b_leaf_mu - ) - ) - forest_model_config_mu.update_leaf_model_scale( - current_leaf_scale_mu - ) - if keep_sample: - self.leaf_scale_mu_samples[sample_counter] = ( - current_leaf_scale_mu[0, 0] - ) - - # Sample tau_0 (global treatment effect intercept, if requested) - if self.sample_tau_0: - mu_x_tau0 = np.squeeze( - active_forest_mu.predict_raw(forest_dataset_train) - ) - tau_x_raw_tau0 = active_forest_tau.predict_raw( - forest_dataset_train - ) - Z_basis = ( - tau_basis_train.reshape(-1, 1) - if tau_basis_train.ndim == 1 - else tau_basis_train - ) - tau_x_raw_2d = tau_x_raw_tau0.reshape(self.n_train, -1) - tau_x_full = np.sum(Z_basis * tau_x_raw_2d, axis=1) - # Center z by y_bar so tau_0 does not absorb the probit intercept - resid_for_tau0 = ( - (np.squeeze(resid_train) - self.y_bar) - if link_is_probit - else np.squeeze(resid_train) - ) - partial_resid_tau0 = resid_for_tau0 - mu_x_tau0 - tau_x_full - if self.has_rfx: - partial_resid_tau0 = partial_resid_tau0 - np.squeeze( - rfx_model.predict(rfx_dataset_train, rfx_tracker) - ) - Ztr = Z_basis.T @ partial_resid_tau0 - ZtZ_current = Z_basis.T @ Z_basis - Sigma_post = np.linalg.inv( - ZtZ_current / current_sigma2 - + np.eye(p_tau0) / tau_0_prior_var - ) - mu_post = Sigma_post @ Ztr / current_sigma2 - tau_0_new = self.rng.multivariate_normal( - mean=mu_post, cov=Sigma_post - ) - residual_train.add_vector( - -np.squeeze(Z_basis @ (tau_0_new - tau_0)) - ) - tau_0 = tau_0_new - if keep_sample: - self.tau_0_samples[:, sample_counter] = tau_0 - - # Sample the treatment forest - forest_sampler_tau.sample_one_iteration( - self.forest_container_tau, - active_forest_tau, - forest_dataset_train, - residual_train, - cpp_rng, - global_model_config, - forest_model_config_tau, - keep_sample, - False, - num_threads, + * self.y_std + ) + self.y_hat_train = self.y_hat_train + rfx_preds_train + if self.has_test: + rfx_preds_test = ( + bcf_results["rfx_predictions_test"].reshape( + self.n_test, bcf_results["num_samples"], order="F" ) + * self.y_std + ) + self.y_hat_test = self.y_hat_test + rfx_preds_test - # Cannot cache train set predictions for tau because the cached predictions in the - # tracking data structures are pre-multiplied by the basis (treatment) - # ... - - # Sample coding parameters (if requested) - if self.adaptive_coding: - mu_x = active_forest_mu.predict_raw(forest_dataset_train) - tau_x = np.squeeze( - active_forest_tau.predict_raw(forest_dataset_train) - ) - # Center z by y_bar so coding regression does not absorb the probit intercept - resid_for_coding = ( - (resid_train - self.y_bar) - if link_is_probit - else resid_train - ) - partial_resid_train = np.squeeze(resid_for_coding - mu_x) - if self.has_rfx: - rfx_pred = np.squeeze( - rfx_model.predict(rfx_dataset_train, rfx_tracker) - ) - partial_resid_train = partial_resid_train - rfx_pred - # Use tau_total = tau_0 + tau(X) for sufficient stats when sample_tau_0 - tau_x_for_coding = ( - (tau_x + tau_0[0]) if self.sample_tau_0 else tau_x - ) - s_tt0 = np.sum( - tau_x_for_coding - * tau_x_for_coding - * (np.squeeze(Z_train) == 0) - ) - s_tt1 = np.sum( - tau_x_for_coding - * tau_x_for_coding - * (np.squeeze(Z_train) == 1) - ) - s_ty0 = np.sum( - tau_x_for_coding - * partial_resid_train - * (np.squeeze(Z_train) == 0) - ) - s_ty1 = np.sum( - tau_x_for_coding - * partial_resid_train - * (np.squeeze(Z_train) == 1) - ) - current_b_0 = self.rng.normal( - loc=(s_ty0 / (s_tt0 + 2 * current_sigma2)), - scale=np.sqrt( - current_sigma2 / (s_tt0 + 2 * current_sigma2) - ), - size=1, - )[0] - current_b_1 = self.rng.normal( - loc=(s_ty1 / (s_tt1 + 2 * current_sigma2)), - scale=np.sqrt( - current_sigma2 / (s_tt1 + 2 * current_sigma2) - ), - size=1, - )[0] - if self.sample_tau_0: - tau_basis_old = np.squeeze(tau_basis_train).copy() - tau_basis_train = ( - 1 - np.squeeze(Z_train) - ) * current_b_0 + np.squeeze(Z_train) * current_b_1 - forest_dataset_train.update_basis(tau_basis_train) - if self.has_test: - tau_basis_test = ( - 1 - np.squeeze(Z_test) - ) * current_b_0 + np.squeeze(Z_test) * current_b_1 - forest_dataset_test.update_basis(tau_basis_test) - if keep_sample: - self.b0_samples[sample_counter] = current_b_0 - self.b1_samples[sample_counter] = current_b_1 - - # Update residual to reflect adjusted basis - forest_sampler_tau.propagate_basis_update( - forest_dataset_train, residual_train, active_forest_tau - ) - - # Fix tau_0 component of residual after basis change - if self.sample_tau_0: - residual_train.add_vector( - -(tau_basis_train - tau_basis_old) * tau_0[0] - ) + # Unpack variance forest predictions (forest owned by self._samples, assembled below) + if self.include_variance_forest: + variance_forest_preds_train = bcf_results[ + "variance_forest_predictions_train" + ].reshape(self.n_train, bcf_results["num_samples"], order="F") + self.sigma2_x_train = variance_forest_preds_train + if self.has_test: + variance_forest_preds_test = bcf_results[ + "variance_forest_predictions_test" + ].reshape(self.n_test, bcf_results["num_samples"], order="F") + self.sigma2_x_test = variance_forest_preds_test - # Sample the variance forest - if self.include_variance_forest: - forest_sampler_variance.sample_one_iteration( - self.forest_container_variance, - active_forest_variance, - forest_dataset_train, - residual_train, - cpp_rng, - global_model_config, - forest_model_config_variance, - keep_sample, - True, - num_threads, - ) + # Parameter sampling flags (global var + leaf scales are owned by self._samples, built below) + self.sample_sigma2_global = sample_sigma2_global + self.sample_sigma2_leaf_mu = sample_sigma2_leaf_mu + self.sample_sigma2_leaf_tau = sample_sigma2_leaf_tau + # tau_0 / b0 / b1 stay separate attributes (specialized shape/scale; not routed through the + # samples wrapper yet). + if self.sample_tau_0: + tau_0_raw = bcf_results["tau_0_samples"] + if tau_0_raw is not None: + self.tau_0_samples = tau_0_raw.reshape( + self.treatment_dim, bcf_results["num_samples"], order="F" + ) * self.y_std + if self.adaptive_coding: + self.b0_samples = bcf_results["b0_samples"] + self.b1_samples = bcf_results["b1_samples"] + + # Persist the final RNG + leaf-cache state so continue_sampling() can resume the random + # stream for bit-identical continuation (single-chain runs only; None otherwise). + self.rng_state = bcf_results.get("rng_state", None) + self.leaf_normal_cache = bcf_results.get("leaf_normal_cache", None) + + # Assemble the single source of truth: forests + global var + leaf scales in one BCFSamplesCpp. + # forest_container_*/global_var_samples/leaf_scale_*_samples/num_samples are properties off it. + self._set_samples( + BCFSamplesCpp.from_components( + bcf_results["forest_container_mu"], + bcf_results["forest_container_tau"], + bcf_results["forest_container_variance"] if self.include_variance_forest else None, + bcf_results["global_var_samples"] if self.sample_sigma2_global else None, + bcf_results["leaf_scale_mu_samples"] if self.sample_sigma2_leaf_mu else None, + bcf_results["leaf_scale_tau_samples"] if self.sample_sigma2_leaf_tau else None, + None, # tau_0 stays a separate attribute + None, # b0 stays a separate attribute + None, # b1 stays a separate attribute + float(self.y_bar), + float(self.y_std), + int(bcf_results["num_samples"]), + int(self.treatment_dim), + ) + ) + self.sampled = True - # Cache train set predictions since they are already computed during sampling - if keep_sample: - sigma2_x_train_raw[:, sample_counter] = ( - forest_sampler_variance.get_cached_forest_predictions() - ) + return self - # Sample variance parameters (if requested) - if self.sample_sigma2_global: - current_sigma2 = global_var_model.sample_one_iteration( - residual_train, cpp_rng, a_global, b_global - ) - global_model_config.update_global_error_variance(current_sigma2) - if keep_sample: - self.global_var_samples[sample_counter] = current_sigma2 - if self.sample_sigma2_leaf_tau: - current_leaf_scale_tau[0, 0] = ( - leaf_var_model_tau.sample_one_iteration( - active_forest_tau, cpp_rng, a_leaf_tau, b_leaf_tau - ) - ) - forest_model_config_tau.update_leaf_model_scale( - current_leaf_scale_tau - ) - if keep_sample: - self.leaf_scale_tau_samples[sample_counter] = ( - current_leaf_scale_tau[0, 0] - ) + def continue_sampling( + self, + X_train: Union[np.array, pd.DataFrame], + Z_train: np.array, + y_train: np.array, + propensity_train: np.array = None, + num_mcmc: int = 100, + num_burnin: int = 0, + keep_every: int = 1, + observation_weights_train: Optional[np.ndarray] = None, + random_seed: int = None, + ): + """Continue (warm-start) MCMC sampling from an already-fit BCF model, appending + additional samples to the existing posterior draws. - # Sample random effects - if self.has_rfx: - rfx_model.sample( - rfx_dataset_train, - residual_train, - rfx_tracker, - self.rfx_container, - keep_sample, - current_sigma2, - cpp_rng, - ) + The training data must be re-supplied (it is not retained on the model). The + warm-start initializes the sampler from the last retained sample so that the new + draws form a continuous chain with the existing ones. - # Mark the model as sampled - self.sampled = True + .. note:: + This is currently supported only for Gaussian (identity-link), univariate-treatment + models without a variance forest, random effects, or adaptive coding. - # Remove GFR samples if they are not to be retained - if not keep_gfr and num_gfr > 0: - for i in range(num_gfr): - self.forest_container_mu.delete_sample(0) - self.forest_container_tau.delete_sample(0) - if self.include_variance_forest: - self.forest_container_variance.delete_sample(0) - if self.has_rfx: - self.rfx_container.delete_sample(0) - if self.adaptive_coding: - self.b1_samples = self.b1_samples[num_gfr:] - self.b0_samples = self.b0_samples[num_gfr:] - if self.sample_tau_0: - self.tau_0_samples = self.tau_0_samples[:, num_gfr:] - if self.sample_sigma2_global: - self.global_var_samples = self.global_var_samples[num_gfr:] - if self.sample_sigma2_leaf_mu: - self.leaf_scale_mu_samples = self.leaf_scale_mu_samples[num_gfr:] - if self.sample_sigma2_leaf_tau: - self.leaf_scale_tau_samples = self.leaf_scale_tau_samples[num_gfr:] - muhat_train_raw = muhat_train_raw[:, num_gfr:] - if self.include_variance_forest: - sigma2_x_train_raw = sigma2_x_train_raw[:, num_gfr:] - self.num_samples -= num_gfr + Parameters + ---------- + X_train : np.array or pd.DataFrame + Training covariates (re-supplied; must match the structure used to fit the model). + Z_train : np.array + Training treatment assignments (re-supplied). + y_train : np.array + Training outcome (re-supplied). + propensity_train : np.array, optional + Training propensity scores. Required if the model was fit with a propensity covariate + and an internal propensity model is not available to re-derive them. + num_mcmc : int, optional + Number of additional retained MCMC samples. Defaults to ``100``. + num_burnin : int, optional + Number of additional burn-in iterations to discard before retaining. Defaults to ``0``. + keep_every : int, optional + Thinning interval for the additional samples. Defaults to ``1``. + observation_weights_train : np.array, optional + Optional training observation weights (must match those used to fit the model). + random_seed : int, optional + If supplied, re-seeds the sampler RNG for the continued draws. By default + (``None``), the RNG state from the prior run is resumed, so the continued chain + is bit-identical to a single run of the combined length. - # Store predictions - self.mu_hat_train = muhat_train_raw * self.y_std + self.y_bar - tau_raw_train = self.forest_container_tau.forest_container_cpp.PredictRaw( - forest_dataset_train.dataset_cpp - ) - self.tau_hat_train = tau_raw_train - if self.adaptive_coding: - adaptive_coding_weights = np.expand_dims( - self.b1_samples - self.b0_samples, axis=(0, 2) + Returns + ------- + self + """ + if not getattr(self, "sampled", False): + raise RuntimeError("Cannot continue sampling: this model has not been sampled yet") + if self.include_variance_forest: + raise NotImplementedError( + "Continued sampling is not yet supported for models with a variance forest" + ) + if getattr(self, "has_rfx", False): + raise NotImplementedError( + "Continued sampling is not yet supported for models with random effects" ) - b0_weights = np.expand_dims(self.b0_samples, axis=(0, 2)) - control_adj_train = self.tau_hat_train * b0_weights * self.y_std - self.tau_hat_train = self.tau_hat_train * adaptive_coding_weights - self.mu_hat_train = self.mu_hat_train + np.squeeze(control_adj_train) - self.tau_hat_train = np.squeeze(self.tau_hat_train * self.y_std) - # Fold tau_0 into tau_hat_train so the attribute holds the full CATE (tau_0 + tau(X)) - if self.sample_tau_0: - tau_0_vec = self.tau_0_samples[ - 0, : - ] # num_samples vector (scalar treatment) - if self.adaptive_coding: - # CATE = (b_1 - b_0) * (tau_0 + tau(X)); control adj to mu = b_0 * (tau_0 + tau(X)) - self.tau_hat_train = self.tau_hat_train + ( - (self.b1_samples - self.b0_samples) * tau_0_vec * self.y_std - ) - self.mu_hat_train = self.mu_hat_train + ( - self.b0_samples * tau_0_vec * self.y_std - ) - elif self.multivariate_treatment: - for j in range(p_tau0): - self.tau_hat_train[:, :, j] = self.tau_hat_train[:, :, j] + ( - self.tau_0_samples[j, :] * self.y_std - ) - else: - self.tau_hat_train = self.tau_hat_train + tau_0_vec * self.y_std if self.multivariate_treatment: - treatment_term_train = np.multiply( - np.atleast_3d(Z_train).swapaxes(1, 2), self.tau_hat_train - ).sum(axis=2) - else: - treatment_term_train = Z_train * np.squeeze(self.tau_hat_train) - self.y_hat_train = self.mu_hat_train + treatment_term_train - if self.has_test: - mu_raw_test = self.forest_container_mu.forest_container_cpp.Predict( - forest_dataset_test.dataset_cpp + raise NotImplementedError( + "Continued sampling is not yet supported for multivariate treatment effects" + ) + if self.adaptive_coding: + raise NotImplementedError( + "Continued sampling is not yet supported for adaptive coding" ) - self.mu_hat_test = mu_raw_test * self.y_std + self.y_bar - tau_raw_test = self.forest_container_tau.forest_container_cpp.PredictRaw( - forest_dataset_test.dataset_cpp + cfg = getattr(self, "_cached_bcf_config", None) + if cfg is None: + raise RuntimeError( + "Cannot continue sampling: cached sampler configuration is unavailable " + "(continuation is not supported for deserialized models yet)" ) - self.tau_hat_test = tau_raw_test - if self.adaptive_coding: - adaptive_coding_weights_test = np.expand_dims( - self.b1_samples - self.b0_samples, axis=(0, 2) - ) - b0_weights = np.expand_dims(self.b0_samples, axis=(0, 2)) - control_adj_test = self.tau_hat_test * b0_weights * self.y_std - self.tau_hat_test = self.tau_hat_test * adaptive_coding_weights_test - self.mu_hat_test = self.mu_hat_test + np.squeeze(control_adj_test) - self.tau_hat_test = np.squeeze(self.tau_hat_test * self.y_std) - # Fold tau_0 into tau_hat_test so the attribute holds the full CATE (tau_0 + tau(X)) - if self.sample_tau_0: - if self.adaptive_coding: - self.tau_hat_test = self.tau_hat_test + ( - (self.b1_samples - self.b0_samples) * tau_0_vec * self.y_std - ) - self.mu_hat_test = self.mu_hat_test + ( - self.b0_samples * tau_0_vec * self.y_std + if cfg.get("link_function", 0) != 0: + raise NotImplementedError( + "Continued sampling is not yet supported for probit or cloglog link functions" + ) + + # Reconstruct the covariate matrix exactly as sample() did: preprocess, then append the + # propensity score column(s) if the model uses a propensity covariate. + X_train_processed = self._covariate_preprocessor.transform(X_train).astype(np.float64) + if getattr(self, "propensity_covariate", "none") != "none": + if propensity_train is None: + if getattr(self, "internal_propensity_model", False) and hasattr( + self, "bart_propensity_model" + ): + propensity_train = np.expand_dims( + self.bart_propensity_model.predict( + X=X_train_processed, terms="y_hat", type="mean" + ), + 1, ) - elif self.multivariate_treatment: - for j in range(p_tau0): - self.tau_hat_test[:, :, j] = self.tau_hat_test[:, :, j] + ( - self.tau_0_samples[j, :] * self.y_std - ) else: - self.tau_hat_test = self.tau_hat_test + tau_0_vec * self.y_std - if self.multivariate_treatment: - treatment_term_test = np.multiply( - np.atleast_3d(Z_test).swapaxes(1, 2), self.tau_hat_test - ).sum(axis=2) + raise ValueError( + "propensity_train must be supplied to continue sampling this model" + ) else: - treatment_term_test = Z_test * np.squeeze(self.tau_hat_test) - self.y_hat_test = self.mu_hat_test + treatment_term_test + propensity_train = np.atleast_2d(propensity_train) + if propensity_train.shape[0] == 1 and propensity_train.shape[1] != 1: + propensity_train = propensity_train.T + X_train_processed = np.c_[X_train_processed, propensity_train] - # TODO: make rfx_preds_train and rfx_preds_test persistent properties - if self.has_rfx: - rfx_preds_train = ( - self.rfx_container.predict(rfx_group_ids_train, rfx_basis_train) - * self.y_std + y_train = np.asarray(y_train).astype(np.float64).reshape(-1) + if X_train_processed.shape[0] != y_train.shape[0]: + raise ValueError("X_train and y_train have differing numbers of observations") + Z_train = np.atleast_2d(Z_train) + if Z_train.shape[0] == 1 and Z_train.shape[1] != X_train_processed.shape[0]: + Z_train = Z_train.T + if Z_train.shape[1] != self.treatment_dim: + raise ValueError( + f"Re-supplied treatment has {Z_train.shape[1]} columns; model expects {self.treatment_dim}" ) - if has_rfx_test: - rfx_preds_test = ( - self.rfx_container.predict(rfx_group_ids_test, rfx_basis_test) - * self.y_std - ) - self.y_hat_train = self.y_hat_train + rfx_preds_train - if self.has_test: - self.y_hat_test = self.y_hat_test + rfx_preds_test - - if self.sample_sigma2_global: - self.global_var_samples = self.global_var_samples * self.y_std * self.y_std - if self.sample_sigma2_leaf_mu: - self.leaf_scale_mu_samples = self.leaf_scale_mu_samples + X_train_cpp = np.asfortranarray(X_train_processed) + Z_train_cpp = np.asfortranarray(Z_train.astype(np.float64)) + y_train_cpp = np.asfortranarray(y_train) + + # RNG continuation: by default resume the saved stream for bit-identical results. + # If the user supplies a new seed, override the cached config seed and tell the + # binding to keep the fresh seed instead of restoring the saved state. + override_seed = random_seed is not None + if override_seed: + cfg = dict(cfg) + cfg["random_seed"] = random_seed + rng_state_in = ( + self.rng_state + if (not override_seed and getattr(self, "rng_state", None) is not None) + else "" + ) + leaf_normal_cache_in = ( + self.leaf_normal_cache + if (not override_seed and getattr(self, "leaf_normal_cache", None) is not None) + else "" + ) - if self.sample_sigma2_leaf_tau: - self.leaf_scale_tau_samples = self.leaf_scale_tau_samples + # tau_0 samples are stored as (treatment_dim, num_samples) on the original scale; flatten + # column-major to match the C++ col-major layout (the binding divides out y_std). + tau_0_in = None + if self.sample_tau_0 and getattr(self, "tau_0_samples", None) is not None: + tau_0_in = np.asarray(self.tau_0_samples).flatten(order="F") + + bcf_results = bcf_continue_sample_cpp( + X_train=X_train_cpp, + Z_train=Z_train_cpp, + y_train=y_train_cpp, + n_train=X_train_cpp.shape[0], + p=X_train_cpp.shape[1], + treatment_dim=self.treatment_dim, + obs_weights_train=observation_weights_train + if observation_weights_train is not None + else None, + num_burnin=num_burnin, + keep_every=keep_every, + num_mcmc=num_mcmc, + mu_forest_container=self._forest_container_mu.forest_container_cpp, + tau_forest_container=self._forest_container_tau.forest_container_cpp, + global_var_samples=self.global_var_samples if self.sample_sigma2_global else None, + leaf_scale_mu_samples=self.leaf_scale_mu_samples if self.sample_sigma2_leaf_mu else None, + leaf_scale_tau_samples=self.leaf_scale_tau_samples if self.sample_sigma2_leaf_tau else None, + tau_0_samples=tau_0_in, + y_bar=float(self.y_bar), + y_std=float(self.y_std), + rng_state_in=rng_state_in, + override_seed=override_seed, + leaf_normal_cache_in=leaf_normal_cache_in, + config_input=cfg, + ) - if self.adaptive_coding: - self.b0_samples = self.b0_samples - self.b1_samples = self.b1_samples + # Rebuild the single source of truth from the (history + new) continuation results. + # Continuation supports no variance forest, so variance is None. + self._set_samples( + BCFSamplesCpp.from_components( + bcf_results["forest_container_mu"], + bcf_results["forest_container_tau"], + None, + bcf_results["global_var_samples"] if self.sample_sigma2_global else None, + bcf_results["leaf_scale_mu_samples"] if self.sample_sigma2_leaf_mu else None, + bcf_results["leaf_scale_tau_samples"] if self.sample_sigma2_leaf_tau else None, + None, # tau_0 stays a separate attribute + None, # b0 + None, # b1 + float(self.y_bar), + float(self.y_std), + int(bcf_results["num_samples"]), + int(self.treatment_dim), + ) + ) + if self.sample_tau_0 and bcf_results["tau_0_samples"] is not None: + self.tau_0_samples = bcf_results["tau_0_samples"].reshape( + self.treatment_dim, self.num_samples, order="F" + ) * self.y_std + self.num_mcmc = (self.num_mcmc or 0) + num_mcmc + # Carry the new final RNG + leaf-cache state forward so further continuations stay bit-identical. + self.rng_state = bcf_results.get("rng_state", None) + self.leaf_normal_cache = bcf_results.get("leaf_normal_cache", None) + + # Recompute training predictions over the full (history + new) set of samples. + preds = self.predict( + X_train, Z_train, propensity=propensity_train, type="posterior", + terms=["y_hat", "mu", "tau"], + ) + self.y_hat_train = preds["y_hat"] + self.mu_hat_train = preds["mu_hat"] + self.tau_hat_train = preds["tau_hat"] - if self.sample_tau_0: - self.tau_0_samples = self.tau_0_samples * self.y_std - - if self.include_variance_forest: - if self.sample_sigma2_global: - self.sigma2_x_train = np.empty_like(sigma2_x_train_raw) - for i in range(self.num_samples): - self.sigma2_x_train[:, i] = ( - np.exp(sigma2_x_train_raw[:, i]) * self.global_var_samples[i] - ) - else: - self.sigma2_x_train = ( - np.exp(sigma2_x_train_raw) - * self.sigma2_init - * self.y_std - * self.y_std - ) - if self.has_test: - sigma2_x_test_raw = ( - self.forest_container_variance.forest_container_cpp.Predict( - forest_dataset_test.dataset_cpp - ) - ) - if self.sample_sigma2_global: - self.sigma2_x_test = np.empty_like(sigma2_x_test_raw) - for i in range(self.num_samples): - self.sigma2_x_test[:, i] = ( - sigma2_x_test_raw[:, i] * self.global_var_samples[i] - ) - else: - self.sigma2_x_test = ( - sigma2_x_test_raw * self.sigma2_init * self.y_std * self.y_std - ) + return self def predict( self, @@ -3322,30 +2551,30 @@ def predict( raise ValueError( f"term '{term}' was requested. Valid terms are 'y_hat', 'prognostic_function', 'mu', 'cate', 'tau', 'rfx', 'variance_forest', and 'all'" ) - has_mu_forest = True - has_tau_forest = True has_variance_forest = self.include_variance_forest has_rfx = self.has_rfx - has_y_hat = has_mu_forest or has_tau_forest or has_rfx - predict_y_hat = (has_y_hat and ("y_hat" in terms)) or ( - has_y_hat and ("all" in terms) - ) - predict_mu_forest = (has_mu_forest and ("mu" in terms)) or ( - has_mu_forest and ("all" in terms) - ) - predict_tau_forest = (has_tau_forest and ("tau" in terms)) or ( - has_tau_forest and ("all" in terms) - ) - predict_prog_function = ( - has_mu_forest and ("prognostic_function" in terms) - ) or (has_mu_forest and ("all" in terms)) - predict_cate_function = (has_tau_forest and ("cate" in terms)) or ( - has_tau_forest and ("all" in terms) - ) - predict_rfx = (has_rfx and ("rfx" in terms)) or (has_rfx and ("all" in terms)) - predict_variance_forest = ( - has_variance_forest and ("variance_forest" in terms) - ) or (has_variance_forest and ("all" in terms)) + has_mu_forest = self._forest_container_mu is not None + has_tau_forest = self._forest_container_tau is not None + predict_y_hat = ("y_hat" in terms) or ("all" in terms) + predict_mu_forest = ("mu" in terms) or ("all" in terms) + predict_tau_forest = ("tau" in terms) or ("all" in terms) + predict_prog_function = ("prognostic_function" in terms) or ("all" in terms) + predict_cate_function = ("cate" in terms) or ("all" in terms) + predict_rfx = ("rfx" in terms) or ("all" in terms) + predict_variance_forest = ("variance_forest" in terms) or ("all" in terms) + # Warn for individually requested terms that weren't fit + unavailable_terms = [] + if predict_rfx and not has_rfx: + unavailable_terms.append("rfx") + predict_rfx = False + if predict_variance_forest and not has_variance_forest: + unavailable_terms.append("variance_forest") + predict_variance_forest = False + if unavailable_terms: + term_str = ", ".join(unavailable_terms) + warnings.warn( + f"Requested term(s) '{term_str}' were not fit in this model and will be excluded from predictions" + ) predict_count = ( predict_y_hat + predict_mu_forest @@ -3442,263 +2671,118 @@ def predict( else: X_combined = np.c_[covariates_processed, propensity] - # Forest dataset - forest_dataset_test = Dataset() - forest_dataset_test.add_covariates(X_combined) - forest_dataset_test.add_basis(Z) - - # Compute predictions from the variance forest (if included) - if predict_variance_forest: - sigma2_x_raw = self.forest_container_variance.forest_container_cpp.Predict( - forest_dataset_test.dataset_cpp - ) - if self.sample_sigma2_global: - sigma2_x = np.empty_like(sigma2_x_raw) - for i in range(self.num_samples): - sigma2_x[:, i] = sigma2_x_raw[:, i] * self.global_var_samples[i] - else: - sigma2_x = sigma2_x_raw * self.sigma2_init * self.y_std * self.y_std - if predict_mean: - sigma2_x = np.mean(sigma2_x, axis=1) - - # Prognostic forest predictions - if predict_mu_forest or predict_mu_forest_intermediate: - mu_raw = self.forest_container_mu.forest_container_cpp.Predict( - forest_dataset_test.dataset_cpp - ) - mu_x_forest = mu_raw * self.y_std + self.y_bar - - # Treatment effect forest predictions - if predict_tau_forest or predict_tau_forest_intermediate: - tau_raw = self.forest_container_tau.forest_container_cpp.PredictRaw( - forest_dataset_test.dataset_cpp - ) - if self.adaptive_coding: - adaptive_coding_weights = np.expand_dims( - self.b1_samples - self.b0_samples, axis=(0, 2) - ) - if predict_mu_forest or predict_mu_forest_intermediate: - b0_weights = np.expand_dims(self.b0_samples, axis=(0, 2)) - control_adj = tau_raw * b0_weights * self.y_std - mu_x_forest = mu_x_forest + np.squeeze(control_adj) - tau_raw = tau_raw * adaptive_coding_weights - tau_x_forest = np.squeeze(tau_raw * self.y_std) - # tau_x_forest is the forest-only component tau(X); compute cate_x_forest - # (tau_0 + tau(X)) for the "cate" term and treatment_term used in y_hat - if getattr(self, "sample_tau_0", False) and hasattr(self, "tau_0_samples"): - tau_0_vec = self.tau_0_samples[0, :] - if self.adaptive_coding: - cate_x_forest = tau_x_forest + ( - (self.b1_samples - self.b0_samples) * tau_0_vec - ) - if predict_mu_forest or predict_mu_forest_intermediate: - mu_x_forest = mu_x_forest + (self.b0_samples * tau_0_vec) - elif Z.shape[1] > 1: - p_tau0 = Z.shape[1] - cate_x_forest = tau_x_forest.copy() - for j in range(p_tau0): - cate_x_forest[:, :, j] = ( - cate_x_forest[:, :, j] + (self.tau_0_samples[j, :]) - ) - else: - cate_x_forest = tau_x_forest + tau_0_vec - else: - cate_x_forest = tau_x_forest - if Z.shape[1] > 1: - treatment_term = np.multiply( - np.atleast_3d(Z).swapaxes(1, 2), cate_x_forest - ).sum(axis=2) - else: - treatment_term = Z * np.squeeze(cate_x_forest) - - # Random effects data checks - if has_rfx: - if rfx_group_ids is None: - raise ValueError( - "rfx_group_ids must be provided if rfx_basis is provided" - ) - - if self.rfx_model_spec == "custom": - if rfx_basis is None: - raise ValueError( - "A user-provided basis (`rfx_basis`) must be provided when the model was sampled with a random effects model spec set to 'custom'" - ) - elif self.rfx_model_spec == "intercept_only": - if rfx_basis is None: - rfx_basis = np.ones(shape=(X.shape[0], 1)) - elif self.rfx_model_spec == "intercept_plus_treatment": - if rfx_basis is None: - rfx_basis = np.concatenate( - (np.ones(shape=(X.shape[0], 1)), Z), axis=1 - ) - - if rfx_basis.ndim == 1: - rfx_basis = np.expand_dims(rfx_basis, 1) - if rfx_basis.shape[0] != X.shape[0]: - raise ValueError("X and rfx_basis must have the same number of rows") - if rfx_basis.shape[1] != self.num_rfx_basis: - raise ValueError( - "rfx_basis must have the same number of columns as the random effects basis used to sample this model" - ) - - # Convert rfx_group_ids to their corresponding array position indices in the random effects parameter sample arrays - if rfx_group_ids is not None: - rfx_group_id_indices = self.rfx_container.map_group_ids_to_array_indices( - rfx_group_ids - ) - - # Random effects predictions - if predict_rfx or predict_rfx_intermediate: - rfx_preds = ( - self.rfx_container.predict(rfx_group_ids, rfx_basis) * self.y_std - ) - - # Extract "raw" rfx predictions for each rfx basis term if needed - if predict_rfx_raw: - # Extract the raw RFX samples and scale by train set outcome standard deviation - rfx_samples_raw = self.rfx_container.extract_parameter_samples() - rfx_beta_draws = rfx_samples_raw["beta_samples"] * self.y_std - - # Construct an array with the appropriate group random effects arranged for each observation - if rfx_beta_draws.ndim == 3: - rfx_predictions_raw = np.empty( - shape=(X.shape[0], rfx_beta_draws.shape[0], rfx_beta_draws.shape[2]) - ) - for i in range(X.shape[0]): - rfx_predictions_raw[i, :, :] = rfx_beta_draws[ - :, rfx_group_id_indices[i], : - ] - elif rfx_beta_draws.ndim == 2: - rfx_predictions_raw = np.empty( - shape=(X.shape[0], 1, rfx_beta_draws.shape[1]) - ) - for i in range(X.shape[0]): - rfx_predictions_raw[i, 0, :] = rfx_beta_draws[ - rfx_group_id_indices[i], : - ] - else: - raise ValueError( - "Unexpected number of dimensions in extracted random effects samples" - ) - - # Add raw RFX predictions to mu and tau if warranted by the RFX model spec - if predict_prog_function: - if mu_prog_separate: - prognostic_function = mu_x_forest + np.squeeze( - rfx_predictions_raw[:, 0, :] - ) - else: - prognostic_function = mu_x_forest - if predict_cate_function: - if tau_cate_separate: - cate = cate_x_forest + np.squeeze(rfx_predictions_raw[:, 1:, :]) - else: - cate = cate_x_forest - - # Combine into y hat predictions - needs_mean_term_preds = ( - predict_y_hat - or predict_mu_forest - or predict_prog_function - or predict_tau_forest - or predict_cate_function - or predict_rfx + # Dimensions needed by both the C++ and R predict paths + n, p = X_combined.shape + treatment_dim = Z.shape[1] + obs_weights = None + rfx_num_groups = len(np.unique(rfx_group_ids)) if rfx_group_ids is not None else 0 + rfx_basis_dim = rfx_basis.shape[1] if rfx_basis is not None else 0 + + scale_int = { + "linear": 0, + "probability": 1, + "class": 2 + }.get(scale, 0) + + # Build a dictionary of model components that can be ingested and unpacked by bcf_predict_cpp + variance_forest_ptr = None + if has_variance_forest: + if self._forest_container_variance is not None: + variance_forest_ptr = self._forest_container_variance.forest_container_cpp + bcf_model_dict = { + "mu_forests": self._forest_container_mu.forest_container_cpp if self._forest_container_mu is not None else None, + "tau_forests": self._forest_container_tau.forest_container_cpp if self._forest_container_tau is not None else None, + "variance_forests": variance_forest_ptr, + "rfx_container": self.rfx_container.rfx_container_cpp if has_rfx else None, + "rfx_label_mapper": self.rfx_container.rfx_label_mapper_cpp if has_rfx else None, + "sigma2_global_samples": getattr(self, "global_var_samples", None), + "sigma2_leaf_mu_samples": getattr(self, "leaf_scale_mu_samples", None), + "sigma2_leaf_tau_samples": getattr(self, "leaf_scale_tau_samples", None), + "b0_samples": getattr(self, "b0_samples", None), + "b1_samples": getattr(self, "b1_samples", None), + "tau_0_samples": getattr(self, "tau_0_samples", None), + "num_samples": int(self.num_samples), + "y_bar": float(self.y_bar), + "y_std": float(self.y_std), + "include_variance_forest": has_variance_forest, + "has_rfx": has_rfx, + "rfx_model_spec": self.rfx_model_spec if has_rfx else "", + "adaptive_coding": self.adaptive_coding, + "sample_tau_0": self.sample_tau_0 + } + n, p = X_combined.shape + treatment_dim = Z.shape[1] + + output = bcf_predict_cpp( + bcf_model_dict = bcf_model_dict, + X = X_combined, + Z = Z, + n = n, + p = p, + treatment_dim = treatment_dim, + obs_weights = obs_weights, + rfx_group_ids = rfx_group_ids, + rfx_basis = rfx_basis, + rfx_num_groups = rfx_num_groups, + rfx_basis_dim = rfx_basis_dim, + posterior = type == "posterior", + scale = scale_int, + predict_y_hat = predict_y_hat, + predict_mu_x = predict_mu_forest, + predict_tau_x = predict_tau_forest, + predict_prognostic_function = predict_prog_function, + predict_cate = predict_cate_function, + predict_conditional_variance = predict_variance_forest, + predict_random_effects = predict_rfx ) - if needs_mean_term_preds: - if probability_scale: - if has_rfx: - if predict_y_hat: - y_hat = norm.cdf(mu_x_forest + treatment_term + rfx_preds) - if predict_rfx: - rfx_preds = norm.cdf(rfx_preds) - else: - if predict_y_hat: - y_hat = norm.cdf(mu_x_forest + treatment_term) - if predict_mu_forest: - mu_x = norm.cdf(mu_x_forest) - if predict_tau_forest: - tau_x = norm.cdf(cate_x_forest) - if predict_prog_function: - prognostic_function = norm.cdf(prognostic_function) - if predict_cate_function: - cate = norm.cdf(cate) - else: - if has_rfx: - if predict_y_hat: - y_hat = mu_x_forest + treatment_term + rfx_preds - else: - if predict_y_hat: - y_hat = mu_x_forest + treatment_term - if predict_mu_forest: - mu_x = mu_x_forest - if predict_tau_forest: - tau_x = cate_x_forest - if predict_prog_function: - prognostic_function = prognostic_function - if predict_cate_function: - cate = cate - - # Collapse to posterior mean predictions if requested - if predict_mean: - if predict_mu_forest: - mu_x = np.mean(mu_x, axis=1) - if predict_tau_forest: - tau_x = np.mean(tau_x, axis=1) - if predict_prog_function: - prognostic_function = np.mean(prognostic_function, axis=1) - if predict_cate_function: - cate = np.mean(cate, axis=1) - if predict_rfx: - rfx_preds = np.mean(rfx_preds, axis=1) - if predict_y_hat: - y_hat = np.mean(y_hat, axis=1) + # Reshape flat C++ output vectors to 2d or 3d arrays (n x num_samples) or (n x treatment_dim x num_samples) + # and rename fields to match the Python predict path. For type="mean", num_samples_output=1 return a 1d or 2d array. + num_samples_raw = self.num_samples + num_samples_output = num_samples_raw if type == "posterior" else 1 + def reshape_cpp_pred_2d(v, dim1, dim2): + if v is None: + return None + if dim2 == 1: + return v.flatten() + return np.reshape(v, (dim1, dim2), order='F') + def reshape_cpp_pred_3d(v, dim1, dim2, dim3): + if v is None: + return None + if dim2 == 1 and dim3 == 1: + return v.flatten() + if dim3 == 1: + return np.reshape(v, (dim1, dim2), order='F') + if dim2 == 1: + # Univariate treatment: squeeze to (n, num_samples) to match the Python path + return np.reshape(v, (dim1, dim3), order='F') + return np.reshape(v, (dim1, dim2, dim3), order='F') + + result = { + "y_hat": reshape_cpp_pred_2d(output["y_hat"], n, num_samples_output), + "mu_hat": reshape_cpp_pred_2d(output["mu_x"], n, num_samples_output), + "tau_hat": reshape_cpp_pred_3d(output["tau_x"], n, treatment_dim, num_samples_output), + "prognostic_function": reshape_cpp_pred_2d(output["prognostic_function"], n, num_samples_output), + "cate": reshape_cpp_pred_3d(output["cate"], n, treatment_dim, num_samples_output), + "rfx_predictions": reshape_cpp_pred_2d(output["random_effects"], n, num_samples_output), + "variance_forest_predictions": reshape_cpp_pred_2d(output["conditional_variance"], n, num_samples_output) + } if predict_count == 1: if predict_y_hat: - return y_hat + return result["y_hat"] elif predict_mu_forest: - return mu_x + return result["mu_hat"] elif predict_prog_function: - return prognostic_function + return result["prognostic_function"] elif predict_tau_forest: - return tau_x + return result["tau_hat"] elif predict_cate_function: - return cate + return result["cate"] elif predict_rfx: - return rfx_preds + return result["rfx_predictions"] elif predict_variance_forest: - return sigma2_x - else: - result = dict() - if predict_y_hat: - result["y_hat"] = y_hat - else: - result["y_hat"] = None - if predict_mu_forest: - result["mu_hat"] = mu_x - else: - result["mu_hat"] = None - if predict_tau_forest: - result["tau_hat"] = tau_x - else: - result["tau_hat"] = None - if predict_prog_function: - result["prognostic_function"] = prognostic_function - else: - result["prognostic_function"] = None - if predict_cate_function: - result["cate"] = cate - else: - result["cate"] = None - if predict_rfx: - result["rfx_predictions"] = rfx_preds - else: - result["rfx_predictions"] = None - if predict_variance_forest: - result["variance_forest_predictions"] = sigma2_x - else: - result["variance_forest_predictions"] = None - return result + return result["variance_forest_predictions"] + return result def compute_contrast( self, @@ -4201,18 +3285,29 @@ def to_json(self) -> str: # Initialize JSONSerializer object bcf_json = JSONSerializer() - # Add the forests - bcf_json.add_forest(self.forest_container_mu) - bcf_json.add_forest(self.forest_container_tau) + # Add the forests under self-describing named keys + bcf_json.add_forest(self._forest_container_mu, "prognostic_forest") + bcf_json.add_forest(self._forest_container_tau, "treatment_forest") if self.include_variance_forest: - bcf_json.add_forest(self.forest_container_variance) + bcf_json.add_forest(self._forest_container_variance, "variance_forest") # Add the rfx if self.has_rfx: bcf_json.add_random_effects(self.rfx_container) + # Python rfx group ids are always integer-valued, so an rfx model is + # cross-platform compatible on the rfx axis. + bcf_json.add_boolean( + "cross_platform_compatible", True, subfolder_name="random_effects" + ) # Add version stamp and global parameters bcf_json.add_string("stochtree_version", _get_stochtree_version()) + bcf_json.add_string("platform", "python") + bcf_json.add_integer("schema_version", SCHEMA_VERSION) + # Covariate count: written so a cross-platform (R) loader can validate + # prediction-set dimension. For all-numeric (portable) models the + # processed count equals the original count R expects. + bcf_json.add_scalar("num_covariates", self.p_x) bcf_json.add_scalar("outcome_scale", self.y_std) bcf_json.add_scalar("outcome_mean", self.y_bar) bcf_json.add_boolean("standardize", self.standardize) @@ -4289,6 +3384,8 @@ def from_json(self, json_string: str) -> None: # Parse string to a JSON object in C++ bcf_json = JSONSerializer() bcf_json.load_from_json_string(json_string) + resolve_schema_version(bcf_json, migrate=_migrate_bcf_v0_to_v1) + cross_platform = enforce_cross_platform_gate(bcf_json, "python") _raw = json.loads(json_string) _ver = _infer_stochtree_version(json_string) @@ -4313,21 +3410,21 @@ def from_json(self, json_string: str) -> None: f"Field 'multivariate_treatment' not found in BCF JSON " f"(inferred version: {_ver}). Defaulting to False." ) - # TODO: don't just make this a placeholder that we overwrite - self.forest_container_mu = ForestContainer(0, 0, False, False) - self.forest_container_mu.forest_container_cpp.LoadFromJson( - bcf_json.json_cpp, "forest_0" + # v1 forests are stored under self-describing named keys. Load into temporary containers, + # then deep-copy into the single owned samples object below. + mu_forest_loaded = ForestContainer(0, 0, False, False) + mu_forest_loaded.forest_container_cpp.LoadFromJson( + bcf_json.json_cpp, "prognostic_forest" ) - # TODO: don't just make this a placeholder that we overwrite - self.forest_container_tau = ForestContainer(0, 0, False, False) - self.forest_container_tau.forest_container_cpp.LoadFromJson( - bcf_json.json_cpp, "forest_1" + tau_forest_loaded = ForestContainer(0, 0, False, False) + tau_forest_loaded.forest_container_cpp.LoadFromJson( + bcf_json.json_cpp, "treatment_forest" ) + variance_forest_loaded = None if self.include_variance_forest: - # TODO: don't just make this a placeholder that we overwrite - self.forest_container_variance = ForestContainer(0, 0, False, False) - self.forest_container_variance.forest_container_cpp.LoadFromJson( - bcf_json.json_cpp, "forest_2" + variance_forest_loaded = ForestContainer(0, 0, False, False) + variance_forest_loaded.forest_container_cpp.LoadFromJson( + bcf_json.json_cpp, "variance_forest" ) # Unpack random effects @@ -4362,7 +3459,7 @@ def from_json(self, json_string: str) -> None: f"Field 'keep_every' not found in BCF JSON " f"(inferred version: {_ver}). Defaulting to 1." ) - self.num_samples = int(bcf_json.get_scalar("num_samples")) + num_samples_loaded = int(bcf_json.get_scalar("num_samples")) self.adaptive_coding = bcf_json.get_boolean("adaptive_coding") if "binary_treatment" in _raw: self.binary_treatment = bcf_json.get_boolean("binary_treatment") @@ -4442,19 +3539,23 @@ def from_json(self, json_string: str) -> None: f"(inferred version: {_ver}) but has_rfx=True." ) - # Unpack parameter samples - if self.sample_sigma2_global: - self.global_var_samples = bcf_json.get_numeric_vector( - "sigma2_global_samples", "parameters" - ) - if self.sample_sigma2_leaf_mu: - self.leaf_scale_mu_samples = bcf_json.get_numeric_vector( - "sigma2_leaf_mu_samples", "parameters" - ) - if self.sample_sigma2_leaf_tau: - self.leaf_scale_tau_samples = bcf_json.get_numeric_vector( - "sigma2_leaf_tau_samples", "parameters" - ) + # Unpack parameter samples (global var + leaf scales owned by self._samples, built below) + global_var_loaded = ( + bcf_json.get_numeric_vector("sigma2_global_samples", "parameters") + if self.sample_sigma2_global + else None + ) + leaf_scale_mu_loaded = ( + bcf_json.get_numeric_vector("sigma2_leaf_mu_samples", "parameters") + if self.sample_sigma2_leaf_mu + else None + ) + leaf_scale_tau_loaded = ( + bcf_json.get_numeric_vector("sigma2_leaf_tau_samples", "parameters") + if self.sample_sigma2_leaf_tau + else None + ) + # b0/b1/tau_0 stay separate attributes. if self.adaptive_coding: self.b1_samples = bcf_json.get_numeric_vector("b1_samples", "parameters") self.b0_samples = bcf_json.get_numeric_vector("b0_samples", "parameters") @@ -4463,6 +3564,25 @@ def from_json(self, json_string: str) -> None: tau_0_vec = bcf_json.get_numeric_vector("tau_0_samples", "parameters") self.tau_0_samples = tau_0_vec.reshape(tau_0_dim, -1) + # Assemble the single source of truth (forests + global var + leaf scales). + self._set_samples( + BCFSamplesCpp.from_components( + mu_forest_loaded.forest_container_cpp, + tau_forest_loaded.forest_container_cpp, + variance_forest_loaded.forest_container_cpp if self.include_variance_forest else None, + global_var_loaded, + leaf_scale_mu_loaded, + leaf_scale_tau_loaded, + None, + None, + None, + float(self.y_bar), + float(self.y_std), + int(num_samples_loaded), + int(self.treatment_dim), + ) + ) + # Unpack internal propensity model if self.internal_propensity_model: bart_propensity_string = bcf_json.get_string("bart_propensity_model") @@ -4470,7 +3590,11 @@ def from_json(self, json_string: str) -> None: self.bart_propensity_model.from_json(bart_propensity_string) # Unpack covariate preprocessor - if "covariate_preprocessor" in _raw: + if cross_platform: + # Identity preprocessor for the cross-platform all-numeric path (gate + # enforced); the foreign native preprocessor is not reconstructed. + self._covariate_preprocessor = CovariatePreprocessor() + elif "covariate_preprocessor" in _raw: covariate_preprocessor_string = bcf_json.get_string( "covariate_preprocessor" ) @@ -4505,46 +3629,50 @@ def from_json_string_list(self, json_string_list: list[str]) -> None: # For scalar / preprocessing details which aren't sample-dependent, defer to the first json json_object_default = json_object_list[0] + for json_object in json_object_list: + resolve_schema_version(json_object, migrate=_migrate_bcf_v0_to_v1) + cross_platform = enforce_cross_platform_gate(json_object_default, "python") _raw_default = json.loads(json_string_list[0]) _ver = _infer_stochtree_version(json_string_list[0]) - # Unpack forests - # Mu forest - self.forest_container_mu = ForestContainer(0, 0, False, False) + # Unpack forests (v1 named keys) into temporary containers (append across chains), then + # deep-copy the combined result into the single owned samples object below. + # Prognostic (mu) forest + mu_forest_loaded = ForestContainer(0, 0, False, False) for i in range(len(json_object_list)): if i == 0: - self.forest_container_mu.forest_container_cpp.LoadFromJson( - json_object_list[i].json_cpp, "forest_0" + mu_forest_loaded.forest_container_cpp.LoadFromJson( + json_object_list[i].json_cpp, "prognostic_forest" ) else: - self.forest_container_mu.forest_container_cpp.AppendFromJson( - json_object_list[i].json_cpp, "forest_0" + mu_forest_loaded.forest_container_cpp.AppendFromJson( + json_object_list[i].json_cpp, "prognostic_forest" ) - # Tau forest - self.forest_container_tau = ForestContainer(0, 0, False, False) + # Treatment (tau) forest + tau_forest_loaded = ForestContainer(0, 0, False, False) for i in range(len(json_object_list)): if i == 0: - self.forest_container_tau.forest_container_cpp.LoadFromJson( - json_object_list[i].json_cpp, "forest_1" + tau_forest_loaded.forest_container_cpp.LoadFromJson( + json_object_list[i].json_cpp, "treatment_forest" ) else: - self.forest_container_tau.forest_container_cpp.AppendFromJson( - json_object_list[i].json_cpp, "forest_1" + tau_forest_loaded.forest_container_cpp.AppendFromJson( + json_object_list[i].json_cpp, "treatment_forest" ) self.include_variance_forest = json_object_default.get_boolean( "include_variance_forest" ) + variance_forest_loaded = None if self.include_variance_forest: - # TODO: don't just make this a placeholder that we overwrite - self.forest_container_variance = ForestContainer(0, 0, False, False) + variance_forest_loaded = ForestContainer(0, 0, False, False) for i in range(len(json_object_list)): if i == 0: - self.forest_container_variance.forest_container_cpp.LoadFromJson( - json_object_list[i].json_cpp, "forest_2" + variance_forest_loaded.forest_container_cpp.LoadFromJson( + json_object_list[i].json_cpp, "variance_forest" ) else: - self.forest_container_variance.forest_container_cpp.AppendFromJson( - json_object_list[i].json_cpp, "forest_2" + variance_forest_loaded.forest_container_cpp.AppendFromJson( + json_object_list[i].json_cpp, "variance_forest" ) # Unpack random effects @@ -4693,58 +3821,43 @@ def from_json_string_list(self, json_string_list: list[str]) -> None: f"(inferred version: {_ver}) but has_rfx=True." ) - # Unpack number of samples + # Combined number of samples across chains (owned by self._samples, built below) + num_samples_loaded = 0 for i in range(len(json_object_list)): - if i == 0: - self.num_samples = json_object_list[i].get_integer("num_samples") - else: - self.num_samples += json_object_list[i].get_integer("num_samples") + num_samples_loaded += json_object_list[i].get_integer("num_samples") - # Unpack parameter samples + # Unpack parameter samples (concatenated across chains; owned by self._samples below) + global_var_loaded = None if self.sample_sigma2_global: for i in range(len(json_object_list)): - if i == 0: - self.global_var_samples = json_object_list[i].get_numeric_vector( - "sigma2_global_samples", "parameters" - ) - else: - global_var_samples = json_object_list[i].get_numeric_vector( - "sigma2_global_samples", "parameters" - ) - self.global_var_samples = np.concatenate(( - self.global_var_samples, - global_var_samples, - )) + gv_i = json_object_list[i].get_numeric_vector( + "sigma2_global_samples", "parameters" + ) + global_var_loaded = ( + gv_i if i == 0 else np.concatenate((global_var_loaded, gv_i)) + ) + leaf_scale_mu_loaded = None if self.sample_sigma2_leaf_mu: for i in range(len(json_object_list)): - if i == 0: - self.leaf_scale_mu_samples = json_object_list[i].get_numeric_vector( - "sigma2_leaf_mu_samples", "parameters" - ) - else: - leaf_scale_mu_samples = json_object_list[i].get_numeric_vector( - "sigma2_leaf_mu_samples", "parameters" - ) - self.leaf_scale_mu_samples = np.concatenate(( - self.leaf_scale_mu_samples, - leaf_scale_mu_samples, - )) + ls_i = json_object_list[i].get_numeric_vector( + "sigma2_leaf_mu_samples", "parameters" + ) + leaf_scale_mu_loaded = ( + ls_i if i == 0 else np.concatenate((leaf_scale_mu_loaded, ls_i)) + ) + # NOTE: fixes a pre-existing bug where the concatenated leaf-tau scale was assigned to the + # boolean flag self.sample_sigma2_leaf_tau instead of the samples array. + leaf_scale_tau_loaded = None if self.sample_sigma2_leaf_tau: for i in range(len(json_object_list)): - if i == 0: - self.sample_sigma2_leaf_tau = json_object_list[ - i - ].get_numeric_vector("sigma2_leaf_tau_samples", "parameters") - else: - sample_sigma2_leaf_tau = json_object_list[i].get_numeric_vector( - "sigma2_leaf_tau_samples", "parameters" - ) - self.sample_sigma2_leaf_tau = np.concatenate(( - self.sample_sigma2_leaf_tau, - sample_sigma2_leaf_tau, - )) + ls_i = json_object_list[i].get_numeric_vector( + "sigma2_leaf_tau_samples", "parameters" + ) + leaf_scale_tau_loaded = ( + ls_i if i == 0 else np.concatenate((leaf_scale_tau_loaded, ls_i)) + ) if self.sample_tau_0: tau_0_dim = int(json_object_default.get_scalar("tau_0_dim")) @@ -4758,6 +3871,25 @@ def from_json_string_list(self, json_string_list: list[str]) -> None: else: self.tau_0_samples = np.hstack((self.tau_0_samples, tau_0_mat_i)) + # Assemble the single source of truth from the combined (across-chain) parts. + self._set_samples( + BCFSamplesCpp.from_components( + mu_forest_loaded.forest_container_cpp, + tau_forest_loaded.forest_container_cpp, + variance_forest_loaded.forest_container_cpp if self.include_variance_forest else None, + global_var_loaded, + leaf_scale_mu_loaded, + leaf_scale_tau_loaded, + None, + None, + None, + float(self.y_bar), + float(self.y_std), + int(num_samples_loaded), + int(self.treatment_dim), + ) + ) + # Unpack internal propensity model if self.internal_propensity_model: bart_propensity_string = json_object_default.get_string( @@ -4767,7 +3899,11 @@ def from_json_string_list(self, json_string_list: list[str]) -> None: self.bart_propensity_model.from_json(bart_propensity_string) # Unpack covariate preprocessor - if "covariate_preprocessor" in _raw_default: + if cross_platform: + # Identity preprocessor for the cross-platform all-numeric path (gate + # enforced); the foreign native preprocessor is not reconstructed. + self._covariate_preprocessor = CovariatePreprocessor() + elif "covariate_preprocessor" in _raw_default: covariate_preprocessor_string = json_object_default.get_string( "covariate_preprocessor" ) diff --git a/stochtree/kernel.py b/stochtree/kernel.py index b76673e6..c4aea04c 100644 --- a/stochtree/kernel.py +++ b/stochtree/kernel.py @@ -88,24 +88,24 @@ def compute_forest_leaf_indices( raise ValueError( "Mean forest was not sampled for model_object, but requested by forest_type" ) - forest_container = model_object.forest_container_mean + forest_container = model_object.extract_forest("mean") else: if not model_object.include_variance_forest: raise ValueError( "Variance forest was not sampled for model_object, but requested by forest_type" ) - forest_container = model_object.forest_container_variance + forest_container = model_object.extract_forest("variance") elif model_type == "bcf": if forest_type == "prognostic": - forest_container = model_object.forest_container_mu + forest_container = model_object.extract_forest("prognostic") elif forest_type == "treatment": - forest_container = model_object.forest_container_tau + forest_container = model_object.extract_forest("treatment") else: if not model_object.include_variance_forest: raise ValueError( "Variance forest was not sampled for model_object, but requested by forest_type" ) - forest_container = model_object.forest_container_variance + forest_container = model_object.extract_forest("variance") else: forest_container = model_object @@ -230,24 +230,24 @@ def compute_forest_max_leaf_index( raise ValueError( "Mean forest was not sampled for model_object, but requested by forest_type" ) - forest_container = model_object.forest_container_mean + forest_container = model_object.extract_forest("mean") else: if not model_object.include_variance_forest: raise ValueError( "Variance forest was not sampled for model_object, but requested by forest_type" ) - forest_container = model_object.forest_container_variance + forest_container = model_object.extract_forest("variance") elif model_type == "bcf": if forest_type == "prognostic": - forest_container = model_object.forest_container_mu + forest_container = model_object.extract_forest("prognostic") elif forest_type == "treatment": - forest_container = model_object.forest_container_tau + forest_container = model_object.extract_forest("treatment") else: if not model_object.include_variance_forest: raise ValueError( "Variance forest was not sampled for model_object, but requested by forest_type" ) - forest_container = model_object.forest_container_variance + forest_container = model_object.extract_forest("variance") else: forest_container = model_object diff --git a/stochtree/preprocessing.py b/stochtree/preprocessing.py index e55645f2..74106184 100644 --- a/stochtree/preprocessing.py +++ b/stochtree/preprocessing.py @@ -635,6 +635,27 @@ def to_json(self) -> str: "original_feature_indices", self._original_feature_indices ) + # Cross-platform portability (RFC 0005): a model is portable across R / + # Python iff every covariate is numeric (no categorical encoding). The + # cross-platform loader peeks at this flag (a plain JSON parse) without + # reconstructing the native preprocessor. + portable = (self._num_ordinal_features == 0) and ( + self._num_onehot_features == 0 + ) + preprocessor_json.add_boolean("cross_platform_portable", portable) + if not portable: + # *_feature_index are indexed by original column; -1 means "not + # categorical", so a non-negative entry marks an offending column. + # Stored as a human-readable string for the cross-platform refusal + # message (the gate only branches on the boolean above). + offending = sorted( + {int(i) for i, v in enumerate(self._ordinal_feature_index) if v != -1} + | {int(i) for i, v in enumerate(self._onehot_feature_index) if v != -1} + ) + preprocessor_json.add_string( + "non_portable_columns", ", ".join(str(i) for i in offending) + ) + return preprocessor_json.return_json_string() def from_json(self, json_string: str) -> None: diff --git a/stochtree/serialization.py b/stochtree/serialization.py index 64f0d842..f4b8e0d8 100644 --- a/stochtree/serialization.py +++ b/stochtree/serialization.py @@ -1,3 +1,4 @@ +import json import warnings import numpy as np @@ -6,6 +7,126 @@ from .forest import ForestContainer from .random_effects import RandomEffectsContainer +# ----------------------------------------------------------------------------- +# Serialized model envelope schema version (RFC 0005). +# +# Integer identifying the *structure* of the serialized BART/BCF JSON envelope. +# Bumped ONLY on a breaking change (rename / remove / re-type a field, change a +# field's meaning, or change a structural convention). Additive, safely-defaulted +# fields do NOT bump it -- they are handled by "augmentation" on read (see the +# defaults registry below and the ``get_*_or_default`` helpers on JSONSerializer). +# +# Kept in sync with the R ``STOCHTREE_SCHEMA_VERSION`` (R/serialization.R). The two +# are independent constants by design (each language owns its serde); their +# agreement is enforced by the cross-platform golden fixtures, not by sharing a value. +SCHEMA_VERSION = 1 + +# ----------------------------------------------------------------------------- +# Augmentation defaults registry (schema_version = 1) +# +# Every OPTIONAL envelope field that may be ABSENT from a model written by an +# earlier release at the SAME schema_version must be read with a default that +# reproduces that model's pre-field behavior. This list is the single source of +# truth for "which fields need defaulting"; read those fields via +# ``get_scalar_or_default`` / ``get_boolean_or_default`` / ``get_string_or_default`` +# (never a bare required getter). +# +# field default (behavior when absent) +# "outcome" "continuous" +# "link" "identity" +# "multivariate_treatment" False +# "internal_propensity_model" False +# "has_rfx_basis" False +# ... (add a row whenever you add an additive field) +# +# If a new field has NO behavior-preserving default, it is NOT additive: bump +# SCHEMA_VERSION and add a migration step instead. +# ----------------------------------------------------------------------------- + + +def resolve_schema_version(serializer: "JSONSerializer", migrate=None) -> int: + """Read the envelope ``schema_version`` and enforce the RFC 0005 reader rules. + + Returns the loaded version (``0`` for a legacy / absent stamp). Behavior vs the + current ``SCHEMA_VERSION``: + + - ``loaded > current`` -> hard error (model written by a newer stochtree). + - ``loaded == current`` -> parse directly (caller proceeds). + - ``0 < loaded < current`` -> caller runs the migration ladder (no rungs exist yet + at ``SCHEMA_VERSION == 1``; this is the hook for future ``migrate_vN_to_vN+1``). + - ``loaded == 0`` -> legacy model; handled by field-presence default-filling on parse. + """ + loaded = serializer.get_integer_or_default("schema_version", 0) + if loaded > SCHEMA_VERSION: + raise ValueError( + f"This model was serialized with schema_version={loaded}, but this " + f"installation of stochtree supports up to schema_version={SCHEMA_VERSION}. " + "Please upgrade stochtree to load it." + ) + if loaded < SCHEMA_VERSION and migrate is not None: + # Dispatch owns the migration: upgrade the JSON in place to the current + # schema before the (v1-only) parser runs. + migrate(serializer, loaded) + return loaded + + +def infer_platform_v0(serializer: "JSONSerializer", default: str) -> str: + """Infer the writer platform of a legacy (v0) envelope from structural + fingerprints. R wrote ``preprocessor_metadata`` and a top-level + ``rfx_unique_group_ids``; Python wrote ``covariate_preprocessor``. Called + before the v0->v1 renames, so the legacy keys are still present. Falls back + to ``default`` (the loading platform) when no fingerprint is decisive.""" + if serializer.contains("covariate_preprocessor"): + return "python" + if serializer.contains("preprocessor_metadata") or serializer.contains( + "rfx_unique_group_ids" + ): + return "R" + return default + + +def enforce_cross_platform_gate( + serializer: "JSONSerializer", loader_platform: str +) -> bool: + """Return ``True`` iff this is a cross-platform load (writer ``platform`` != + loader). For a cross-platform load, verify the model is portable (all-numeric + covariates) and rfx-compatible (integer group ids); otherwise raise a clear, + actionable error. Same-platform loads return ``False`` and ignore the flags. + + The portability flag is peeked from the ``covariate_preprocessor`` JSON via a + plain parse -- the foreign native preprocessor is never reconstructed. An + absent portability flag is treated as non-portable (e.g. legacy v0 models), + so a cross-platform load of such a model is refused rather than mis-handled. + """ + writer_platform = serializer.get_string_or_default("platform", loader_platform) + if writer_platform == loader_platform: + return False + portable = False + non_portable_columns = "" + if serializer.contains("covariate_preprocessor"): + prep = json.loads(serializer.get_string("covariate_preprocessor")) + portable = bool(prep.get("cross_platform_portable", False)) + non_portable_columns = prep.get("non_portable_columns", "") + compatible = serializer.get_boolean_or_default( + "cross_platform_compatible", True, subfolder_name="random_effects" + ) + if not portable: + raise ValueError( + f"Cannot load a model written on platform '{writer_platform}' from " + f"'{loader_platform}': it was fit with non-numeric covariates " + f"(columns: {non_portable_columns or 'unknown'}), which are not " + f"cross-platform portable. Load it on '{writer_platform}', or " + "refit / re-save with all-numeric covariates." + ) + if not compatible: + raise ValueError( + f"Cannot load a model written on platform '{writer_platform}' from " + f"'{loader_platform}': it has string-labeled random effects, which are " + "not cross-platform compatible (only integer-valued group ids are). " + f"Load it on '{writer_platform}'." + ) + return True + class JSONSerializer: """ @@ -41,17 +162,25 @@ def load_from_json_string(self, json_string: str) -> None: """ self.json_cpp.LoadFromString(json_string) - def add_forest(self, forest_samples: ForestContainer) -> None: + def add_forest( + self, forest_samples: ForestContainer, forest_label: str = "" + ) -> str: """Adds a container of forest samples to a json object Parameters ---------- forest_samples : ForestContainer Samples of a tree ensemble + forest_label : str, optional + Key under the ``forests`` subfolder to store this container. Defaults + to an auto-numbered ``forest_`` label when empty. """ - forest_label = self.json_cpp.AddForest(forest_samples.forest_container_cpp) + forest_label = self.json_cpp.AddForest( + forest_samples.forest_container_cpp, forest_label + ) self.num_forests += 1 self.forest_labels.append(forest_label) + return forest_label def add_random_effects(self, rfx_container: RandomEffectsContainer) -> None: """Adds a container of random effect samples to a json object @@ -300,6 +429,52 @@ def get_string(self, field_name: str, subfolder_name: str = None) -> str: else: return self.json_cpp.ExtractStringSubfolder(subfolder_name, field_name) + def contains(self, field_name: str, subfolder_name: str = None) -> bool: + """Whether the json object contains ``field_name`` (optionally under ``subfolder_name``).""" + if subfolder_name is None: + return self.json_cpp.ContainsField(field_name) + return self.json_cpp.ContainsFieldSubfolder(subfolder_name, field_name) + + def get_scalar_or_default(self, field_name, default, subfolder_name=None): + """Read a numeric field, returning ``default`` if absent (augmentation; see RFC 0005).""" + if self.contains(field_name, subfolder_name): + return self.get_scalar(field_name, subfolder_name) + return default + + def get_integer_or_default(self, field_name, default, subfolder_name=None): + """Read an integer field, returning ``default`` if absent.""" + if self.contains(field_name, subfolder_name): + return self.get_integer(field_name, subfolder_name) + return default + + def get_boolean_or_default(self, field_name, default, subfolder_name=None): + """Read a boolean field, returning ``default`` if absent.""" + if self.contains(field_name, subfolder_name): + return self.get_boolean(field_name, subfolder_name) + return default + + def get_string_or_default(self, field_name, default, subfolder_name=None): + """Read a string field, returning ``default`` if absent.""" + if self.contains(field_name, subfolder_name): + return self.get_string(field_name, subfolder_name) + return default + + def rename_field(self, old_name, new_name, subfolder_name=None): + """Rename ``old_name`` to ``new_name`` (optionally under ``subfolder_name``). + + No-op if ``old_name`` is absent. Used by JSON schema migrations (RFC 0005).""" + if subfolder_name is None: + self.json_cpp.RenameField(old_name, new_name) + else: + self.json_cpp.RenameFieldSubfolder(subfolder_name, old_name, new_name) + + def erase_field(self, field_name, subfolder_name=None): + """Erase ``field_name`` (optionally under ``subfolder_name``). No-op if absent.""" + if subfolder_name is None: + self.json_cpp.EraseField(field_name) + else: + self.json_cpp.EraseFieldSubfolder(subfolder_name, field_name) + def get_numeric_vector( self, field_name: str, subfolder_name: str = None ) -> np.array: diff --git a/test/R/testthat/fixtures/bart_categorical_rfx_v1.json b/test/R/testthat/fixtures/bart_categorical_rfx_v1.json new file mode 100644 index 00000000..6cd55066 --- /dev/null +++ b/test/R/testthat/fixtures/bart_categorical_rfx_v1.json @@ -0,0 +1 @@ +{"covariate_preprocessor":"{\"cross_platform_portable\":false,\"feature_types\":[0.0,0.0,0.0,0.0,1.0,1.0,1.0,1.0],\"forests\":{},\"non_portable_columns\":\"cat\",\"num_forests\":0,\"num_numeric_vars\":4,\"num_ordered_cat_vars\":0,\"num_random_effects\":0,\"num_unordered_cat_vars\":1,\"numeric_vars\":[\"X1\",\"X2\",\"X3\",\"X4\"],\"original_var_indices\":[1.0,2.0,3.0,4.0,5.0,5.0,5.0,5.0],\"random_effects\":{},\"unordered_cat_vars\":[\"cat\"],\"unordered_unique_level_keys\":{\"key_1\":\"cat\"},\"unordered_unique_levels\":{\"cat\":[\"a\",\"b\",\"c\"]}}","forests":{"mean_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.04885873017931268,0.16521956681326327],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[5,-1,-1],"threshold":[0.27422900901325475,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.03945338751414626,0.12351140116166949],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.8709056605344936,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.10151557492097157,0.3167937750161585],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.6876814769913645,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.3701310664754703,-0.08204989580054663],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.4872431174865849,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.03537089897802479,-0.06025009351828198],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.44556619948076925,0.0,0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.26654929772933583,0.17172519008932544],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[5,-1,-1],"threshold":[0.27422900901325475,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0,0,0],"category_list_end":[0,0,0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0,1],"is_log_scale":false,"leaf_parents":[1],"leaf_value":[0.0,-0.03945338751414626,-0.05641427152400197,-0.15939539039531248,-0.05018754060977407],"leaf_vector":[],"leaf_vector_begin":[0,0,0,0,0],"leaf_vector_end":[0,0,0,0,0],"leaves":[2,3,4],"left":[1,3,-1,-1,-1],"node_deleted":[false,false,false,false,false],"node_type":[1,1,0,0,0],"num_deleted_nodes":0,"num_nodes":5,"output_dimension":1,"parent":[-1,0,0,1,1],"right":[2,4,-1,-1,-1],"split_index":[0,5,-1,-1,-1],"threshold":[0.8709056605344936,0.17422101564004022,0.0,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0,0,0],"category_list_end":[0,0,0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0,2],"is_log_scale":false,"leaf_parents":[2],"leaf_value":[0.0,0.30920223914846195,0.3167937750161585,0.022987854048784527,0.05991049965964413],"leaf_vector":[],"leaf_vector_begin":[0,0,0,0,0],"leaf_vector_end":[0,0,0,0,0],"leaves":[1,3,4],"left":[1,-1,3,-1,-1],"node_deleted":[false,false,false,false,false],"node_type":[1,0,1,0,0],"num_deleted_nodes":0,"num_nodes":5,"output_dimension":1,"parent":[-1,0,0,2,2],"right":[2,-1,4,-1,-1],"split_index":[2,-1,3,-1,-1],"threshold":[0.6876814769913645,0.0,0.622359896390653,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.37264067308884924,0.1362034348747787],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.4872431174865849,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.3376828311301402,-0.0732615175250862],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.44556619948076925,0.0,0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":true,"num_samples":2,"num_trees":5,"output_dimension":1}},"has_rfx":true,"has_rfx_basis":true,"include_mean_forest":true,"include_variance_forest":false,"keep_every":1.0,"num_basis":0.0,"num_burnin":0.0,"num_chains":1.0,"num_covariates":5.0,"num_forests":1,"num_gfr":0.0,"num_mcmc":2.0,"num_random_effects":1,"num_rfx_basis":1.0,"num_samples":2.0,"outcome_mean":1.6024014438950713,"outcome_model":{"link":"identity","outcome":"continuous"},"outcome_scale":0.9636847924729185,"parameters":{"sigma2_global_samples":[0.9745642523804949,0.7549123776587986],"sigma2_leaf_samples":[0.024972323629083196,0.033870876351265254]},"platform":"R","probit_outcome_model":false,"random_effects":{"cross_platform_compatible":true,"random_effect_container_0":{"alpha":[0.09899066592578562,0.3494759091955891],"alpha_size":2,"beta":[-0.008172338863651646,-0.1277405670683955,0.019722328408684944,-0.9686028380594023,-0.15760825239060994,0.6867878312661562],"beta_size":6,"num_components":1,"num_groups":3,"num_samples":2,"sigma_xi":[0.41510710134476175,3.364648878043308],"xi":[-0.0825566611480171,-1.2904304246641194,0.199234222986953,-2.771586860704925,-0.4509845979180278,1.9651936319358305]},"random_effect_groupids_0":[3,2,2,3,1,1,2,3,3,2,3,2,1,2,1,3,3,1,2,1,2,2,2,1,2,3,3,3,2,2,1,1,1,1,3,3,2,3,3,2,3,3,2,1,1,2,1,2,1,2,1,1,3,2,2,3,3,3,2,2,1,3,1,2,3,3,2,1,3,2,1,3,3,3,3,2,3,2,2,2,3,3,2,3,3,2,1,1,3,1,1,3,1,3,3,3,1,2,2,3,2,1,3,2,1,3,1,2,1,3,1,2,2,3,1,1,1,1,3,3,3,1,1,2,3,3,1,1,3,3,3,1,1,2,2,3,1,1,2,1,1,2,3,1,3,3,3,3,2,3,3,3,3,1,1,3,1,2,1,1,3,1,3,2,3,3,2,1,2,1,2,1,3,1,3,3,1,2,3,1,2,2,3,3,3,3,1,1,1,3,2,3,3,3,2,2,1,3,1,1],"random_effect_label_mapper_0":{"keys":[1,2,3],"values":[0,1,2]},"rfx_unique_group_ids":["0","1","2"]},"requires_basis":false,"rfx_model_spec":"custom","sample_sigma2_global":true,"sample_sigma2_leaf":true,"schema_version":1,"sigma2_init":1.0,"standardize":true,"stochtree_version":"0.4.5.9000"} diff --git a/test/R/testthat/fixtures/bart_categorical_v1.json b/test/R/testthat/fixtures/bart_categorical_v1.json new file mode 100644 index 00000000..d45880fe --- /dev/null +++ b/test/R/testthat/fixtures/bart_categorical_v1.json @@ -0,0 +1 @@ +{"covariate_preprocessor":"{\"cross_platform_portable\":false,\"feature_types\":[0.0,0.0,0.0,0.0,1.0,1.0,1.0,1.0],\"forests\":{},\"non_portable_columns\":\"cat\",\"num_forests\":0,\"num_numeric_vars\":4,\"num_ordered_cat_vars\":0,\"num_random_effects\":0,\"num_unordered_cat_vars\":1,\"numeric_vars\":[\"X1\",\"X2\",\"X3\",\"X4\"],\"original_var_indices\":[1.0,2.0,3.0,4.0,5.0,5.0,5.0,5.0],\"random_effects\":{},\"unordered_cat_vars\":[\"cat\"],\"unordered_unique_level_keys\":{\"key_1\":\"cat\"},\"unordered_unique_levels\":{\"cat\":[\"a\",\"b\",\"c\"]}}","forests":{"mean_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.017061578535734295],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.08530894408963365,-0.04529504965015144],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[4,-1,-1],"threshold":[0.07098034051274182,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.07763961072881047,0.13790499237755058],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.3667947828764807,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.03920587132242699,-0.19710728829018462],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.5123362130865561,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.3149406257363808,-0.014805486318148464],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.07219663472846179,0.0,0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[-0.017061578535734295,0.1305161516465159,-0.21949066221956012],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[5,-1,-1],"threshold":[0.8878203633235752,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.1550408340374099,-0.05249768256596441],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[4,-1,-1],"threshold":[0.07098034051274182,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.10122203151934948,0.22364185021384073],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.3667947828764807,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.09869413033222156,-0.12087997026320184],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.5123362130865561,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.285756784752209,-0.17508187698397512],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.07219663472846179,0.0,0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":true,"num_samples":2,"num_trees":5,"output_dimension":1}},"has_rfx":false,"has_rfx_basis":false,"include_mean_forest":true,"include_variance_forest":false,"keep_every":1.0,"num_basis":0.0,"num_burnin":0.0,"num_chains":1.0,"num_covariates":5.0,"num_forests":1,"num_gfr":0.0,"num_mcmc":2.0,"num_random_effects":0,"num_rfx_basis":0.0,"num_samples":2.0,"outcome_mean":0.544417932509007,"outcome_model":{"link":"identity","outcome":"continuous"},"outcome_scale":0.5978139073725202,"parameters":{"sigma2_global_samples":[0.36130025767646373,0.33823993227566773],"sigma2_leaf_samples":[0.020744694747098117,0.02275522115811404]},"platform":"R","probit_outcome_model":false,"random_effects":{},"requires_basis":false,"rfx_model_spec":"custom","sample_sigma2_global":true,"sample_sigma2_leaf":true,"schema_version":1,"sigma2_init":1.0,"standardize":true,"stochtree_version":"0.4.5.9000"} diff --git a/test/R/testthat/fixtures/bart_numeric_rfx_v1.json b/test/R/testthat/fixtures/bart_numeric_rfx_v1.json new file mode 100644 index 00000000..66503276 --- /dev/null +++ b/test/R/testthat/fixtures/bart_numeric_rfx_v1.json @@ -0,0 +1 @@ +{"covariate_preprocessor":"{\"cross_platform_portable\":true,\"feature_types\":[0.0,0.0,0.0,0.0],\"forests\":{},\"num_forests\":0,\"num_numeric_vars\":4,\"num_ordered_cat_vars\":0,\"num_random_effects\":0,\"num_unordered_cat_vars\":0,\"numeric_vars\":[\"x1\",\"x2\",\"x3\",\"x4\"],\"original_var_indices\":[1.0,2.0,3.0,4.0],\"random_effects\":{}}","forests":{"mean_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.10398776894645842,0.3415516697449695],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.8755700733589689,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.393933000256073,0.14593194815452287],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.12959924560830197,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.1669651626670278,-0.159082050890466],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.6495255142223529,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.05553342078640208],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.0466306364628901,-0.4419197925089475],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.8413677847082969,0.0,0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.06619830762916759,0.020132599743703128],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.8755700733589689,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.21662164665949526,0.1906812964426531],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.12959924560830197,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.18216391128515777,0.12993467801851288],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.6495255142223529,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.05553342078640208,-0.17590414540382715,-0.12373351995578129],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.6758393757946902,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.09277726622386032,-0.07184424491648865],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.8413677847082969,0.0,0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":true,"num_samples":2,"num_trees":5,"output_dimension":1}},"has_rfx":true,"has_rfx_basis":true,"include_mean_forest":true,"include_variance_forest":false,"keep_every":1.0,"num_basis":0.0,"num_burnin":0.0,"num_chains":1.0,"num_covariates":4.0,"num_forests":1,"num_gfr":0.0,"num_mcmc":2.0,"num_random_effects":1,"num_rfx_basis":1.0,"num_samples":2.0,"outcome_mean":1.4314553647577828,"outcome_model":{"link":"identity","outcome":"continuous"},"outcome_scale":0.9953923262928258,"parameters":{"sigma2_global_samples":[1.0520956665390389,0.9549955776088448],"sigma2_leaf_samples":[0.0760200192138279,0.029617163545580794]},"platform":"R","probit_outcome_model":false,"random_effects":{"cross_platform_compatible":true,"random_effect_container_0":{"alpha":[-0.03488326820065428,-0.22303164725589206],"alpha_size":2,"beta":[0.02233899021643641,0.06737993283505501,0.04350676700140397,-0.7613226436987385,-0.0872348023098768,1.1735956110057288],"beta_size":6,"num_components":1,"num_groups":3,"num_samples":2,"sigma_xi":[1.7341768697134379,84.86639792467857],"xi":[-0.6403926973796972,-1.9315831431698025,-1.247210173976415,3.413518453841872,0.3911319464443055,-5.262013823801523]},"random_effect_groupids_0":[3,1,3,3,3,1,3,3,2,1,3,3,1,3,3,2,2,1,2,2,3,1,1,3,1,3,3,1,3,1,3,2,2,3,3,1,1,2,2,3,2,1,1,3,3,2,3,2,2,1,1,3,2,2,1,3,2,2,3,1,1,1,2,1,3,3,1,1,2,1,3,3,2,2,1,1,1,1,1,1,3,1,1,1,2,2,2,2,1,1,3,3,3,1,1,1,1,1,3,3,2,2,3,3,1,2,1,3,1,2,2,1,3,3,2,1,2,1,3,1,1,2,2,3,2,1,1,3,2,1,1,1,3,3,1,3,3,1,2,3,2,2,2,1,3,2,2,3,2,2,1,1,1,1,2,2,3,2,1,3,2,2,2,2,2,2,3,2,2,2,2,1,2,1,1,1,3,2,3,2,1,3,1,2,3,1,3,2,2,2,3,3,3,3,1,1,2,2,2,3],"random_effect_label_mapper_0":{"keys":[1,2,3],"values":[0,1,2]},"rfx_unique_group_ids":["0","1","2"]},"requires_basis":false,"rfx_model_spec":"custom","sample_sigma2_global":true,"sample_sigma2_leaf":true,"schema_version":1,"sigma2_init":1.0,"standardize":true,"stochtree_version":"0.4.5.9000"} diff --git a/test/R/testthat/fixtures/bart_numeric_v1.json b/test/R/testthat/fixtures/bart_numeric_v1.json new file mode 100644 index 00000000..f4d95467 --- /dev/null +++ b/test/R/testthat/fixtures/bart_numeric_v1.json @@ -0,0 +1 @@ +{"covariate_preprocessor":"{\"cross_platform_portable\":true,\"feature_types\":[0.0,0.0,0.0,0.0],\"forests\":{},\"num_forests\":0,\"num_numeric_vars\":4,\"num_ordered_cat_vars\":0,\"num_random_effects\":0,\"num_unordered_cat_vars\":0,\"numeric_vars\":[\"x1\",\"x2\",\"x3\",\"x4\"],\"original_var_indices\":[1.0,2.0,3.0,4.0],\"random_effects\":{}}","forests":{"mean_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.32901356315743796,-0.014177636313433323],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.2925176053251436,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.2851356971980384,0.23605771644206075],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.11522111727888279,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.14413194101152957,0.11131571895114768],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.5776559593749006,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.11585280618689299,-0.05766133256763008],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.3820197371814679,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.022812999553732442,0.10797234765718894],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.7579355009907407,0.0,0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[1,2],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.15644027824076226,-0.32901356315743796,-0.014177636313433323],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[0],"left":[-1,-1,-1],"node_deleted":[false,true,true],"node_type":[0,0,0],"num_deleted_nodes":2,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[-1,-1,-1],"split_index":[1,-1,-1],"threshold":[0.2925176053251436,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.01704987390916035,0.10211403131121345],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.11522111727888279,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.016713396965898637,-0.005654774396837854],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.5776559593749006,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.18835940659522954,0.014876447457371192],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.3820197371814679,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.13632315914754894,0.02809620256863898],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.7579355009907407,0.0,0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":true,"num_samples":2,"num_trees":5,"output_dimension":1}},"has_rfx":false,"has_rfx_basis":false,"include_mean_forest":true,"include_variance_forest":false,"keep_every":1.0,"num_basis":0.0,"num_burnin":0.0,"num_chains":1.0,"num_covariates":4.0,"num_forests":1,"num_gfr":0.0,"num_mcmc":2.0,"num_random_effects":0,"num_rfx_basis":0.0,"num_samples":2.0,"outcome_mean":0.461344881073025,"outcome_model":{"link":"identity","outcome":"continuous"},"outcome_scale":0.5577278661513684,"parameters":{"sigma2_global_samples":[0.3175247925040276,0.31866777190978646],"sigma2_leaf_samples":[0.024398345415268436,0.010637685556095814]},"platform":"R","probit_outcome_model":false,"random_effects":{},"requires_basis":false,"rfx_model_spec":"custom","sample_sigma2_global":true,"sample_sigma2_leaf":true,"schema_version":1,"sigma2_init":1.0,"standardize":true,"stochtree_version":"0.4.5.9000"} diff --git a/test/R/testthat/fixtures/bcf_categorical_rfx_v1.json b/test/R/testthat/fixtures/bcf_categorical_rfx_v1.json new file mode 100644 index 00000000..748c363f --- /dev/null +++ b/test/R/testthat/fixtures/bcf_categorical_rfx_v1.json @@ -0,0 +1 @@ +{"adaptive_coding":false,"binary_treatment":true,"covariate_preprocessor":"{\"cross_platform_portable\":false,\"feature_types\":[0.0,0.0,0.0,0.0,1.0,1.0,1.0,1.0],\"forests\":{},\"non_portable_columns\":\"cat\",\"num_forests\":0,\"num_numeric_vars\":4,\"num_ordered_cat_vars\":0,\"num_random_effects\":0,\"num_unordered_cat_vars\":1,\"numeric_vars\":[\"X1\",\"X2\",\"X3\",\"X4\"],\"original_var_indices\":[1.0,2.0,3.0,4.0,5.0,5.0,5.0,5.0],\"random_effects\":{},\"unordered_cat_vars\":[\"cat\"],\"unordered_unique_level_keys\":{\"key_1\":\"cat\"},\"unordered_unique_levels\":{\"cat\":[\"a\",\"b\",\"c\"]}}","forests":{"prognostic_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.18042191796498597,0.14047774670432467],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.4616803744859633,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.10436330193087084],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.16597401833681572,0.07696406509506483],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.6899133834808194,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.06480934720758251,-0.12957216359536286],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.7949844600943714,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.06046252219775554],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.09302395553402311,0.25031565693857366],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.4616803744859633,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.16069409702480608],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.0704356781702906,-0.005616050820683134],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.6899133834808194,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.03747929674622692,-0.03051655704455757],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.7949844600943714,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[-0.06046252219775554,-0.022905151342982658,-0.029874482427494763],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.46062767151312395,0.0,0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":true,"num_samples":2,"num_trees":5,"output_dimension":1},"treatment_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":false,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.04305927855061697],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_1":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.03923006687748638],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.45181258791982515,-0.22697836071850705],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.47464878288334866,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-1.152771790552996,0.4534692296619075],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.4584331002318041,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.04138323786084819,0.08322740872509223],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.05013707480230672,0.0,0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":false,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.0790692704986196],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.03923006687748638,0.007039266104045444,0.0021325998606110223],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.36804937725120246,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.21576056910586558,-0.1177746726198096],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.47464878288334866,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.8093524060863997,0.3414264266995864],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.4584331002318041,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.26560782304190944,0.029350410740014748],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.05013707480230672,0.0,0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":false,"num_samples":2,"num_trees":5,"output_dimension":1}},"has_rfx":true,"has_rfx_basis":true,"include_variance_forest":false,"internal_propensity_model":false,"keep_every":1.0,"multivariate_treatment":false,"num_burnin":0.0,"num_chains":1.0,"num_covariates":5.0,"num_forests":2,"num_gfr":0.0,"num_mcmc":2.0,"num_random_effects":1,"num_rfx_basis":1.0,"num_samples":2.0,"outcome_mean":3.921308031798132,"outcome_model":{"link":"identity","outcome":"continuous"},"outcome_scale":1.8319846506714386,"parameters":{"sigma2_global_samples":[2.0467014547409557,1.8233185481189764],"sigma2_leaf_mu_samples":[0.01003183440180122,0.015045473555574139],"tau_0_samples":[1.0546996529117894,1.0853532686488738]},"platform":"R","probit_outcome_model":false,"propensity_covariate":"prognostic","random_effects":{"cross_platform_compatible":true,"random_effect_container_0":{"alpha":[0.2621025363500752,0.3341739893329059],"alpha_size":2,"beta":[-0.04665788724779571,-0.3997942918501097,-0.08460238729291594,-0.7937697797518433,-0.44966844950731394,0.031044459183398408],"beta_size":6,"num_components":1,"num_groups":3,"num_samples":2,"sigma_xi":[1.1025406325688607,5.228908106049238],"xi":[-0.1780138715845064,-1.525335456182414,-0.32278355055640295,-2.3753188611010825,-1.3456117587277323,0.09289908902057531]},"random_effect_groupids_0":[1,1,3,3,3,1,1,2,2,2,2,3,3,3,1,1,1,3,2,1,3,3,1,2,1,3,2,1,2,3,3,2,1,1,1,3,3,3,2,2,3,2,3,3,2,3,2,2,2,1,3,3,1,2,1,2,3,3,2,3,2,1,1,2,1,1,2,2,1,3,3,1,3,2,3,2,2,1,1,1,2,1,2,1,3,2,1,3,1,3,2,1,1,2,2,1,1,2,2,2,1,1,3,3,2,2,2,2,2,3,3,3,3,2,2,3,2,2,1,2,1,1,1,3,2,3,2,2,2,3,2,2,2,2,1,3,1,1,2,1,2,2,2,2,3,2,3,1,3,1,2,3,1,2,2,2,3,3,1,3,1,3,2,2,2,2,2,2,2,1,3,2,2,1,2,2,2,1,2,3,3,3,3,2,3,3,3,2,1,3,2,2,1,2,1,2,1,1,3,1],"random_effect_label_mapper_0":{"keys":[1,2,3],"values":[0,1,2]},"rfx_unique_group_ids":["0","1","2"]},"rfx_model_spec":"custom","sample_sigma2_global":true,"sample_sigma2_leaf_mu":true,"sample_sigma2_leaf_tau":false,"sample_tau_0":true,"schema_version":1,"sigma2_init":1.0,"standardize":true,"stochtree_version":"0.4.5.9000","tau_0_dim":1.0,"treatment_dim":1.0} diff --git a/test/R/testthat/fixtures/bcf_categorical_v1.json b/test/R/testthat/fixtures/bcf_categorical_v1.json new file mode 100644 index 00000000..c19e0f92 --- /dev/null +++ b/test/R/testthat/fixtures/bcf_categorical_v1.json @@ -0,0 +1 @@ +{"adaptive_coding":false,"binary_treatment":true,"covariate_preprocessor":"{\"cross_platform_portable\":false,\"feature_types\":[0.0,0.0,0.0,0.0,1.0,1.0,1.0,1.0],\"forests\":{},\"non_portable_columns\":\"cat\",\"num_forests\":0,\"num_numeric_vars\":4,\"num_ordered_cat_vars\":0,\"num_random_effects\":0,\"num_unordered_cat_vars\":1,\"numeric_vars\":[\"X1\",\"X2\",\"X3\",\"X4\"],\"original_var_indices\":[1.0,2.0,3.0,4.0,5.0,5.0,5.0,5.0],\"random_effects\":{},\"unordered_cat_vars\":[\"cat\"],\"unordered_unique_level_keys\":{\"key_1\":\"cat\"},\"unordered_unique_levels\":{\"cat\":[\"a\",\"b\",\"c\"]}}","forests":{"prognostic_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.08185884324192944,-0.09032844732472756],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.1435486555427689,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.06513786304395654,0.10093467614509707],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[6,-1,-1],"threshold":[0.7575362642591073,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.033266427001844144,-0.01051291792405894],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.3448163274707924,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.0708482292745689,0.3026612379343009],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[5,-1,-1],"threshold":[0.9290980944262642,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.1827111605322551,-0.06964636226737422],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.1739479007690341,0.0,0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.03384388411041029,-0.26978034512630317],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.1435486555427689,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.05435789742607878,0.2485274532024364],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[6,-1,-1],"threshold":[0.7575362642591073,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.25639644708075604,-0.016842395630418467],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.3448163274707924,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.0708323950437144,0.11391163738176342],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[5,-1,-1],"threshold":[0.9290980944262642,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[1,2],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.11113864608970511,-0.1827111605322551,-0.06964636226737422],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[0],"left":[-1,-1,-1],"node_deleted":[false,true,true],"node_type":[0,0,0],"num_deleted_nodes":2,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[-1,-1,-1],"split_index":[3,-1,-1],"threshold":[0.1739479007690341,0.0,0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":true,"num_samples":2,"num_trees":5,"output_dimension":1},"treatment_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":false,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.014375090427459443],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.23310619332050408,0.3288498438243571],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.9304524291293423,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.8825036896648653,0.4488325693894204],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.7085016606520307,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.2583025985683647],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_4":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.1595787245014281],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":false,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.14473212756623458],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[1,2],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.3015155105762455,0.23310619332050408,0.3288498438243571],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[0],"left":[-1,-1,-1],"node_deleted":[false,true,true],"node_type":[0,0,0],"num_deleted_nodes":2,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[-1,-1,-1],"split_index":[0,-1,-1],"threshold":[0.9304524291293423,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-1.0215242398532187,0.36033339746836357],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.7085016606520307,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.12459536722728352],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_4":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.07906792652994463],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":false,"num_samples":2,"num_trees":5,"output_dimension":1}},"has_rfx":false,"has_rfx_basis":false,"include_variance_forest":false,"internal_propensity_model":false,"keep_every":1.0,"multivariate_treatment":false,"num_burnin":0.0,"num_chains":1.0,"num_covariates":5.0,"num_forests":2,"num_gfr":0.0,"num_mcmc":2.0,"num_random_effects":0,"num_rfx_basis":0.0,"num_samples":2.0,"outcome_mean":3.09448822211111,"outcome_model":{"link":"identity","outcome":"continuous"},"outcome_scale":1.5734824706294417,"parameters":{"sigma2_global_samples":[1.6264105674577767,1.2955876018382568],"sigma2_leaf_mu_samples":[0.013509470826724414,0.04768520217272847],"tau_0_samples":[0.9766582387380113,1.063746794930983]},"platform":"R","probit_outcome_model":false,"propensity_covariate":"prognostic","random_effects":{},"rfx_model_spec":"custom","sample_sigma2_global":true,"sample_sigma2_leaf_mu":true,"sample_sigma2_leaf_tau":false,"sample_tau_0":true,"schema_version":1,"sigma2_init":1.0,"standardize":true,"stochtree_version":"0.4.5.9000","tau_0_dim":1.0,"treatment_dim":1.0} diff --git a/test/R/testthat/fixtures/bcf_numeric_rfx_v1.json b/test/R/testthat/fixtures/bcf_numeric_rfx_v1.json new file mode 100644 index 00000000..67decb60 --- /dev/null +++ b/test/R/testthat/fixtures/bcf_numeric_rfx_v1.json @@ -0,0 +1 @@ +{"adaptive_coding":false,"binary_treatment":true,"covariate_preprocessor":"{\"cross_platform_portable\":true,\"feature_types\":[0.0,0.0,0.0,0.0],\"forests\":{},\"num_forests\":0,\"num_numeric_vars\":4,\"num_ordered_cat_vars\":0,\"num_random_effects\":0,\"num_unordered_cat_vars\":0,\"numeric_vars\":[\"x1\",\"x2\",\"x3\",\"x4\"],\"original_var_indices\":[1.0,2.0,3.0,4.0],\"random_effects\":{}}","forests":{"prognostic_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.04999730647890461,-0.2559939734381046],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.8083109755109188,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.05016763647951462,-0.5834583391278416],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.9616577341451928,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.04023690879835688,-0.11194477723028749],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.564331380304315,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.24228867316220834,0.1315983683090614],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.690642008868686,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.5481340310419853,0.3636745930405258],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.5602792062257769,0.0,0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.09793183999326693,-0.1942066798624267],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.8083109755109188,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.1122392940263486,0.211852073327657],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.9616577341451928,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.15738089293399693,0.04155210579114442],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.564331380304315,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.22490220901411725,0.21052785443290276],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.690642008868686,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.609847012639031,0.4825622697663648],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.5602792062257769,0.0,0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":true,"num_samples":2,"num_trees":5,"output_dimension":1},"treatment_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":false,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.4613304621747731,0.025993938931986545],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.35946943253124536,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.04091454062665775],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_2":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.10745477237989225],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_3":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.12294179771791022],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.25684601015586944,0.07001498591711426],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.2951668903783159,0.0,0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":false,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[1,2],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.05248780731116516,-0.4613304621747731,0.025993938931986545],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[0],"left":[-1,-1,-1],"node_deleted":[false,true,true],"node_type":[0,0,0],"num_deleted_nodes":2,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[-1,-1,-1],"split_index":[0,-1,-1],"threshold":[0.35946943253124536,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.08312118064713782],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_2":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.2095705607903301],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_3":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.3159601748378549],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.2482335894083703,0.10034080199309886],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.2951668903783159,0.0,0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":false,"num_samples":2,"num_trees":5,"output_dimension":1}},"has_rfx":true,"has_rfx_basis":true,"include_variance_forest":false,"internal_propensity_model":false,"keep_every":1.0,"multivariate_treatment":false,"num_burnin":0.0,"num_chains":1.0,"num_covariates":4.0,"num_forests":2,"num_gfr":0.0,"num_mcmc":2.0,"num_random_effects":1,"num_rfx_basis":1.0,"num_samples":2.0,"outcome_mean":4.118612918169948,"outcome_model":{"link":"identity","outcome":"continuous"},"outcome_scale":1.7652956665631405,"parameters":{"sigma2_global_samples":[2.0100038253004957,1.8076094801631222],"sigma2_leaf_mu_samples":[0.05405366627035948,0.035066256009818726],"tau_0_samples":[0.48819230875699604,0.5710709607064383]},"platform":"R","probit_outcome_model":false,"propensity_covariate":"prognostic","random_effects":{"cross_platform_compatible":true,"random_effect_container_0":{"alpha":[0.04583696792688299,0.33363462410000166],"alpha_size":2,"beta":[0.006042635992374171,0.06845241081558724,0.03898226145560701,-0.28439045964316767,-0.3738888528207721,0.3411314482639395],"beta_size":6,"num_components":1,"num_groups":3,"num_samples":2,"sigma_xi":[0.4747173717678096,1.0214815534633148],"xi":[0.13182887668340332,1.4933887190090618,0.8504546268808553,-0.8524009173518098,-1.1206536306876378,1.0224701623345027]},"random_effect_groupids_0":[1,2,2,3,2,3,2,1,2,1,2,3,1,1,2,1,1,3,3,3,3,2,3,1,1,3,2,1,3,3,1,2,1,3,3,1,2,1,2,2,2,1,3,3,3,3,2,1,3,1,3,2,1,3,1,3,1,2,1,1,2,2,2,1,1,2,1,3,2,1,2,1,2,3,2,1,2,2,3,1,2,3,3,2,2,3,3,2,2,3,1,1,3,2,3,2,3,2,3,2,1,1,1,2,3,3,2,1,1,3,2,2,1,3,1,3,3,1,2,3,1,2,2,3,2,2,3,2,3,3,1,2,1,2,2,3,2,3,3,1,1,1,1,3,1,1,1,3,1,1,3,3,3,3,1,3,3,1,1,2,3,3,2,3,1,1,2,2,2,2,3,2,2,1,1,2,1,3,1,2,1,2,2,3,2,3,2,1,2,1,1,1,3,2,3,3,2,1,2,2],"random_effect_label_mapper_0":{"keys":[1,2,3],"values":[0,1,2]},"rfx_unique_group_ids":["0","1","2"]},"rfx_model_spec":"custom","sample_sigma2_global":true,"sample_sigma2_leaf_mu":true,"sample_sigma2_leaf_tau":false,"sample_tau_0":true,"schema_version":1,"sigma2_init":1.0,"standardize":true,"stochtree_version":"0.4.5.9000","tau_0_dim":1.0,"treatment_dim":1.0} diff --git a/test/R/testthat/fixtures/bcf_numeric_v1.json b/test/R/testthat/fixtures/bcf_numeric_v1.json new file mode 100644 index 00000000..ec1dccab --- /dev/null +++ b/test/R/testthat/fixtures/bcf_numeric_v1.json @@ -0,0 +1 @@ +{"adaptive_coding":false,"binary_treatment":true,"covariate_preprocessor":"{\"cross_platform_portable\":true,\"feature_types\":[0.0,0.0,0.0,0.0],\"forests\":{},\"num_forests\":0,\"num_numeric_vars\":4,\"num_ordered_cat_vars\":0,\"num_random_effects\":0,\"num_unordered_cat_vars\":0,\"numeric_vars\":[\"x1\",\"x2\",\"x3\",\"x4\"],\"original_var_indices\":[1.0,2.0,3.0,4.0],\"random_effects\":{}}","forests":{"prognostic_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.5618062379345891,0.47702890518298136],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.5689465328005905,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.9189056364164913,0.1111349737242411],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.020017584242135084,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.08925678668109703,-0.1138856825816371],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.23320559444666064,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.053504061914274995,0.6445509714918987],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.9639442340854351,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.22755369336741674,0.051709278471769655],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.2345037608406191,0.0,0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.7182763363756342,0.3483698277123991],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.5689465328005905,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.25587065370682294,0.11704265523381838],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.020017584242135084,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.06179992839316321,-0.10503778410634956],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.23320559444666064,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.14581791807992797,-0.12569129970855972],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.9639442340854351,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0,0,0],"category_list_end":[0,0,0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0,2],"is_log_scale":false,"leaf_parents":[2],"leaf_value":[0.0,-0.08638914809982157,0.051709278471769655,-0.11058979604056965,-0.03422302603004217],"leaf_vector":[],"leaf_vector_begin":[0,0,0,0,0],"leaf_vector_end":[0,0,0,0,0],"leaves":[1,3,4],"left":[1,-1,3,-1,-1],"node_deleted":[false,false,false,false,false],"node_type":[1,0,1,0,0],"num_deleted_nodes":0,"num_nodes":5,"output_dimension":1,"parent":[-1,0,0,2,2],"right":[2,-1,4,-1,-1],"split_index":[1,-1,3,-1,-1],"threshold":[0.2345037608406191,0.0,0.1622819171885667,0.0,0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":true,"num_samples":2,"num_trees":5,"output_dimension":1},"treatment_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":false,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.20414561784411026],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_1":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.03340624496368301],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_2":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.21714067410182286],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_3":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.015063814554423101],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_4":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.11637575476193755],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":false,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.2854486782185377],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[-0.03340624496368301,-0.23655771356069463,0.24254377942764702],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.7282710648421039,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.37469320075240115],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_3":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.05785626081608691],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_4":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.03769590839794458],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":false,"num_samples":2,"num_trees":5,"output_dimension":1}},"has_rfx":false,"has_rfx_basis":false,"include_variance_forest":false,"internal_propensity_model":false,"keep_every":1.0,"multivariate_treatment":false,"num_burnin":0.0,"num_chains":1.0,"num_covariates":4.0,"num_forests":2,"num_gfr":0.0,"num_mcmc":2.0,"num_random_effects":0,"num_rfx_basis":0.0,"num_samples":2.0,"outcome_mean":3.092255997584547,"outcome_model":{"link":"identity","outcome":"continuous"},"outcome_scale":1.7116093779622612,"parameters":{"sigma2_global_samples":[1.6657406096210687,1.5427612478466444],"sigma2_leaf_mu_samples":[0.08208314375530636,0.06921795188683333],"tau_0_samples":[0.6904678515495952,0.973707336073473]},"platform":"R","probit_outcome_model":false,"propensity_covariate":"prognostic","random_effects":{},"rfx_model_spec":"custom","sample_sigma2_global":true,"sample_sigma2_leaf_mu":true,"sample_sigma2_leaf_tau":false,"sample_tau_0":true,"schema_version":1,"sigma2_init":1.0,"standardize":true,"stochtree_version":"0.4.5.9000","tau_0_dim":1.0,"treatment_dim":1.0} diff --git a/test/R/testthat/fixtures/generate_v1_fixtures.R b/test/R/testthat/fixtures/generate_v1_fixtures.R new file mode 100644 index 00000000..fd3e7db1 --- /dev/null +++ b/test/R/testthat/fixtures/generate_v1_fixtures.R @@ -0,0 +1,109 @@ +# Generate the v1 (schema_version=1) golden-fixture matrix for the unified envelope. +# +# Run once to (re)mint the checked-in v1 fixtures: +# NOT_CRAN=true Rscript -e "devtools::load_all('.'); \ +# source('test/R/testthat/fixtures/generate_v1_fixtures.R')" +# +# The matrix is {bart, bcf} x {numeric, categorical} x {no-rfx, rfx} = 8 models, +# covering the format-relevant axes: forest naming (BART vs BCF), the covariate +# preprocessor (identity vs categorical encoding -> the cross-platform-portable +# axis), and random effects (which add the random_effects subfolder, including +# the relocated rfx_unique_group_ids). +# +# Fixtures only need to be structurally complete and loadable, not statistically +# meaningful, so the ensembles are tiny (5 trees, 2 MCMC) to keep the checked-in +# files small. The legacy v0 fixtures (bart_mcmc.json / bcf_mcmc.json) are kept +# separately and exercise the v0 -> v1 migration path. + +out_dir <- "test/R/testthat/fixtures" + +NUM_MCMC <- 2 +SMALL <- list(num_trees = 5) +N <- 200 +P <- 4 # numeric covariate columns (categorical models add one "cat" column) + +# Superseded ad-hoc fixtures from the first cut (replaced by the matrix below). +SUPERSEDED <- c("bart_mcmc_v1.json", "bart_rfx_v1.json", "bcf_mcmc_v1.json") + +make_covariates <- function(categorical) { + X_num <- matrix(runif(N * P), ncol = P) + if (!categorical) { + return(list(X = X_num, x0 = X_num[, 1])) + } + X <- data.frame(X_num) + X$cat <- factor(sample(c("a", "b", "c"), N, replace = TRUE)) + list(X = X, x0 = X_num[, 1]) +} + +bart_fixture <- function(seed, rfx, categorical) { + set.seed(seed) + cv <- make_covariates(categorical) + y <- cv$x0 + rnorm(N, 0, 0.5) + args <- list( + X_train = cv$X, + num_gfr = 0, + num_burnin = 0, + num_mcmc = NUM_MCMC, + mean_forest_params = SMALL + ) + if (rfx) { + g <- sample(0:2, N, replace = TRUE) + y <- y + g + args$rfx_group_ids_train <- g + args$rfx_basis_train <- matrix(1, nrow = N, ncol = 1) + } + args$y_train <- y + saveBARTModelToJsonString(do.call(bart, args)) +} + +bcf_fixture <- function(seed, rfx, categorical) { + set.seed(seed) + cv <- make_covariates(categorical) + pi_x <- 0.25 + 0.5 * cv$x0 + Z <- rbinom(N, 1, pi_x) + y <- pi_x * 5 + Z * cv$x0 * 2 + rnorm(N) + args <- list( + X_train = cv$X, + Z_train = Z, + propensity_train = pi_x, + num_gfr = 0, + num_burnin = 0, + num_mcmc = NUM_MCMC, + prognostic_forest_params = SMALL, + treatment_effect_forest_params = SMALL + ) + if (rfx) { + g <- sample(0:2, N, replace = TRUE) + y <- y + g + args$rfx_group_ids_train <- g + args$rfx_basis_train <- matrix(1, nrow = N, ncol = 1) + } + args$y_train <- y + saveBCFModelToJsonString(do.call(bcf, args)) +} + +seed <- 110 +for (model in c("bart", "bcf")) { + fn <- if (model == "bart") bart_fixture else bcf_fixture + for (categorical in c(FALSE, TRUE)) { + for (rfx in c(FALSE, TRUE)) { + kind <- if (categorical) "categorical" else "numeric" + suffix <- if (rfx) "_rfx" else "" + name <- paste0(model, "_", kind, suffix, "_v1.json") + js <- fn(seed, rfx, categorical) + seed <- seed + 1 + ver <- jsonlite::fromJSON(js, simplifyVector = FALSE)$schema_version + stopifnot(ver == 1) + writeLines(js, file.path(out_dir, name)) + cat("wrote", name, "(schema_version =", ver, ")\n") + } + } +} + +for (old in SUPERSEDED) { + p <- file.path(out_dir, old) + if (file.exists(p)) { + file.remove(p) + cat("removed superseded", old, "\n") + } +} diff --git a/test/R/testthat/test-bart.R b/test/R/testthat/test-bart.R index 595ddf70..84ba168f 100644 --- a/test/R/testthat/test-bart.R +++ b/test/R/testthat/test-bart.R @@ -464,9 +464,9 @@ test_that("BART Predictions", { # Check that cached predictions agree with results of predict() function train_preds <- predict(bart_model, X = X_train) - train_preds_mean_cached <- bart_model$y_hat_train + train_preds_mean_cached <- bart_model$samples$y_hat_train() train_preds_mean_recomputed <- train_preds$mean_forest_predictions - train_preds_variance_cached <- bart_model$sigma2_x_hat_train + train_preds_variance_cached <- bart_model$samples$variance_forest_predictions_train() train_preds_variance_recomputed <- train_preds$variance_forest_predictions # Assertion @@ -734,13 +734,13 @@ test_that("Cloglog Binary BART", { ) # Check model outputs - expect_true(!is.null(bart_model$y_hat_train)) - expect_true(!is.null(bart_model$y_hat_test)) - expect_true(is.null(bart_model$cloglog_cutpoint_samples)) - expect_equal(nrow(bart_model$y_hat_train), n_train) - expect_equal(ncol(bart_model$y_hat_train), 10) - expect_equal(nrow(bart_model$y_hat_test), n_test) - expect_equal(ncol(bart_model$y_hat_test), 10) + expect_true(!is.null(bart_model$samples$y_hat_train())) + expect_true(!is.null(bart_model$samples$y_hat_test())) + expect_true(is.null(bart_model$samples$cloglog_cutpoint_samples())) + expect_equal(nrow(bart_model$samples$y_hat_train()), n_train) + expect_equal(ncol(bart_model$samples$y_hat_train()), 10) + expect_equal(nrow(bart_model$samples$y_hat_test()), n_test) + expect_equal(ncol(bart_model$samples$y_hat_test()), 10) # Predict from model on linear scale (terms = "y_hat" for single matrix return) expect_no_error({ @@ -844,11 +844,11 @@ test_that("Cloglog Binary BART with GFR", { ) # Check model outputs - expect_true(!is.null(bart_model$y_hat_train)) - expect_true(!is.null(bart_model$y_hat_test)) - expect_true(is.null(bart_model$cloglog_cutpoint_samples)) - expect_equal(nrow(bart_model$y_hat_train), n_train) - expect_equal(ncol(bart_model$y_hat_train), 10) + expect_true(!is.null(bart_model$samples$y_hat_train())) + expect_true(!is.null(bart_model$samples$y_hat_test())) + expect_true(is.null(bart_model$samples$cloglog_cutpoint_samples())) + expect_equal(nrow(bart_model$samples$y_hat_train()), n_train) + expect_equal(ncol(bart_model$samples$y_hat_train()), 10) }) test_that("Cloglog Ordinal BART", { @@ -903,16 +903,16 @@ test_that("Cloglog Ordinal BART", { ) # Check model outputs - expect_true(!is.null(bart_model$y_hat_train)) - expect_true(!is.null(bart_model$y_hat_test)) - expect_true(!is.null(bart_model$cloglog_cutpoint_samples)) - expect_equal(nrow(bart_model$y_hat_train), n_train) - expect_equal(ncol(bart_model$y_hat_train), 10) - expect_equal(nrow(bart_model$y_hat_test), n_test) - expect_equal(ncol(bart_model$y_hat_test), 10) + expect_true(!is.null(bart_model$samples$y_hat_train())) + expect_true(!is.null(bart_model$samples$y_hat_test())) + expect_true(!is.null(bart_model$samples$cloglog_cutpoint_samples())) + expect_equal(nrow(bart_model$samples$y_hat_train()), n_train) + expect_equal(ncol(bart_model$samples$y_hat_train()), 10) + expect_equal(nrow(bart_model$samples$y_hat_test()), n_test) + expect_equal(ncol(bart_model$samples$y_hat_test()), 10) # 3 categories means 2 cutpoint rows - expect_equal(nrow(bart_model$cloglog_cutpoint_samples), 2) - expect_equal(ncol(bart_model$cloglog_cutpoint_samples), 10) + expect_equal(nrow(bart_model$samples$cloglog_cutpoint_samples()), 2) + expect_equal(ncol(bart_model$samples$cloglog_cutpoint_samples()), 10) expect_equal(bart_model$model_params$cloglog_num_categories, 3) # Predict from model on linear scale @@ -1023,14 +1023,14 @@ test_that("Cloglog Ordinal BART with GFR", { ) # Check model outputs - expect_true(!is.null(bart_model$y_hat_train)) - expect_true(!is.null(bart_model$y_hat_test)) - expect_true(!is.null(bart_model$cloglog_cutpoint_samples)) - expect_equal(nrow(bart_model$y_hat_train), n_train) - expect_equal(ncol(bart_model$y_hat_train), 10) + expect_true(!is.null(bart_model$samples$y_hat_train())) + expect_true(!is.null(bart_model$samples$y_hat_test())) + expect_true(!is.null(bart_model$samples$cloglog_cutpoint_samples())) + expect_equal(nrow(bart_model$samples$y_hat_train()), n_train) + expect_equal(ncol(bart_model$samples$y_hat_train()), 10) # 3 categories means 2 cutpoint rows - expect_equal(nrow(bart_model$cloglog_cutpoint_samples), 2) - expect_equal(ncol(bart_model$cloglog_cutpoint_samples), 10) + expect_equal(nrow(bart_model$samples$cloglog_cutpoint_samples()), 2) + expect_equal(ncol(bart_model$samples$cloglog_cutpoint_samples()), 10) expect_equal(bart_model$model_params$cloglog_num_categories, 3) }) @@ -1073,10 +1073,10 @@ test_that("Cloglog BART multi-chain", { ) # Check model outputs: 2 chains x 10 MCMC = 20 total samples - expect_true(!is.null(bart_model$y_hat_train)) - expect_true(is.null(bart_model$cloglog_cutpoint_samples)) - expect_equal(nrow(bart_model$y_hat_train), n_train) - expect_equal(ncol(bart_model$y_hat_train), 20) + expect_true(!is.null(bart_model$samples$y_hat_train())) + expect_true(is.null(bart_model$samples$cloglog_cutpoint_samples())) + expect_equal(nrow(bart_model$samples$y_hat_train()), n_train) + expect_equal(ncol(bart_model$samples$y_hat_train()), 20) }) test_that("Cloglog BART JSON round-trip", { @@ -1141,8 +1141,8 @@ test_that("Cloglog BART JSON round-trip", { bart_model$model_params$cloglog_num_categories ) expect_equal( - bart_model_loaded$cloglog_cutpoint_samples, - bart_model$cloglog_cutpoint_samples + bart_model_loaded$samples$cloglog_cutpoint_samples(), + bart_model$samples$cloglog_cutpoint_samples() ) # Check that forest predictions survive round-trip diff --git a/test/R/testthat/test-bcf.R b/test/R/testthat/test-bcf.R index 36502d70..a40e3325 100644 --- a/test/R/testthat/test-bcf.R +++ b/test/R/testthat/test-bcf.R @@ -656,9 +656,12 @@ test_that("BCF Predictions", { Z = Z_train, propensity = pi_train ) - train_preds_mean_cached <- bcf_model$y_hat_train + train_preds_mean_cached <- extractParameter(bcf_model, "y_hat_train") train_preds_mean_recomputed <- train_preds$y_hat - train_preds_variance_cached <- bcf_model$sigma2_x_hat_train + train_preds_variance_cached <- extractParameter( + bcf_model, + "sigma2_x_train" + ) train_preds_variance_recomputed <- train_preds$variance_forest_predictions # Assertion @@ -844,8 +847,10 @@ test_that("BCF internal propensity model works with data frame covariates", { num_mcmc = 5 ) ) - expect_true(!is.null(bcf_model$tau_hat_train)) - expect_true(!is.null(bcf_model$tau_hat_test)) + tau_hat_train <- extractParameter(bcf_model, "tau_hat_train") + tau_hat_test <- extractParameter(bcf_model, "tau_hat_test") + expect_true(!is.null(tau_hat_train)) + expect_true(!is.null(tau_hat_test)) }) test_that("BCF JSON serialization roundtrip covers all deserialization paths", { @@ -946,8 +951,13 @@ test_that("BCF factor-valued treatment handling", { # Factor treatment should run without error and emit an informative message expect_message( suppressWarnings(bcf( - X_train = X, y_train = y, Z_train = Z_factor_binary, - propensity_train = pi_X, num_gfr = 0, num_burnin = 5, num_mcmc = 5 + X_train = X, + y_train = y, + Z_train = Z_factor_binary, + propensity_train = pi_X, + num_gfr = 0, + num_burnin = 5, + num_mcmc = 5 )), regexp = "Z_train is a factor" ) @@ -961,8 +971,13 @@ test_that("BCF factor-valued treatment handling", { expect_message( suppressWarnings(bcf( - X_train = X, y_train = y, Z_train = Z_factor_logical, - propensity_train = pi_X, num_gfr = 0, num_burnin = 5, num_mcmc = 5 + X_train = X, + y_train = y, + Z_train = Z_factor_logical, + propensity_train = pi_X, + num_gfr = 0, + num_burnin = 5, + num_mcmc = 5 )), regexp = "Z_train is a factor" ) @@ -971,8 +986,13 @@ test_that("BCF factor-valued treatment handling", { Z_factor_categorical <- factor(sample(c("A", "B", "C"), n, replace = TRUE)) expect_error( bcf( - X_train = X, y_train = y, Z_train = Z_factor_categorical, - propensity_train = pi_X, num_gfr = 0, num_burnin = 5, num_mcmc = 5 + X_train = X, + y_train = y, + Z_train = Z_factor_categorical, + propensity_train = pi_X, + num_gfr = 0, + num_burnin = 5, + num_mcmc = 5 ), regexp = "exactly 2 levels" ) @@ -980,8 +1000,13 @@ test_that("BCF factor-valued treatment handling", { # predict.bcfmodel should also handle factor Z, raising a warning suppressMessages( bcf_model <- bcf( - X_train = X, y_train = y, Z_train = Z_numeric, - propensity_train = pi_X, num_gfr = 0, num_burnin = 5, num_mcmc = 5 + X_train = X, + y_train = y, + Z_train = Z_numeric, + propensity_train = pi_X, + num_gfr = 0, + num_burnin = 5, + num_mcmc = 5 ) ) expect_warning( @@ -1014,40 +1039,62 @@ test_that("Warmstart BCF reuses internal propensity model", { test_inds <- (n_train + 1):n X_train <- X[train_inds, ] - X_test <- X[test_inds, ] + X_test <- X[test_inds, ] Z_train <- Z[train_inds] - Z_test <- Z[test_inds] + Z_test <- Z[test_inds] y_train <- y[train_inds] # Fit first model without propensity — triggers internal propensity BART m1 <- bcf( - X_train = X_train, Z_train = Z_train, y_train = y_train, - X_test = X_test, Z_test = Z_test, - num_gfr = 5, num_burnin = 0, num_mcmc = 10 + X_train = X_train, + Z_train = Z_train, + y_train = y_train, + X_test = X_test, + Z_test = Z_test, + num_gfr = 5, + num_burnin = 0, + num_mcmc = 10 ) expect_true(m1$model_params$internal_propensity_model) # Propensity predictions from the first model's propensity BART - pi_train_m1 <- predict(m1$bart_propensity_model, X = X_train, terms = "y_hat", type = "mean") + pi_train_m1 <- predict( + m1$bart_propensity_model, + X = X_train, + terms = "y_hat", + type = "mean" + ) # Warm-start second model from first — propensity model should be reused m1_json <- saveBCFModelToJsonString(m1) m2 <- bcf( - X_train = X_train, Z_train = Z_train, y_train = y_train, - X_test = X_test, Z_test = Z_test, - num_gfr = 0, num_burnin = 0, num_mcmc = 10, + X_train = X_train, + Z_train = Z_train, + y_train = y_train, + X_test = X_test, + Z_test = Z_test, + num_gfr = 0, + num_burnin = 0, + num_mcmc = 10, previous_model_json = m1_json, previous_model_warmstart_sample_num = 10L ) expect_true(m2$model_params$internal_propensity_model) # Propensity model reused: predictions on train set should be identical - pi_train_m2 <- predict(m2$bart_propensity_model, X = X_train, terms = "y_hat", type = "mean") + pi_train_m2 <- predict( + m2$bart_propensity_model, + X = X_train, + terms = "y_hat", + type = "mean" + ) expect_equal(pi_train_m1, pi_train_m2) # Output shapes should be correct - expect_equal(dim(m2$y_hat_train), c(n_train, 10)) - expect_equal(dim(m2$y_hat_test), c(n_test, 10)) + y_hat_train <- extractParameter(m2, "y_hat_train") + y_hat_test <- extractParameter(m2, "y_hat_test") + expect_equal(dim(y_hat_train), c(n_train, 10)) + expect_equal(dim(y_hat_test), c(n_test, 10)) }) test_that("predict.bcfmodel works with data frame X when internal propensity model is used", { @@ -1068,9 +1115,9 @@ test_that("predict.bcfmodel works with data frame X when internal propensity mod test_inds <- 1:40 train_inds <- 41:n X_train <- as.data.frame(X[train_inds, ]) - X_test <- as.data.frame(X[test_inds, ]) + X_test <- as.data.frame(X[test_inds, ]) Z_train <- Z[train_inds] - Z_test <- Z[test_inds] + Z_test <- Z[test_inds] y_train <- y[train_inds] # No propensity_train provided — internal propensity model is fitted @@ -1105,48 +1152,92 @@ test_that("predict(terms='tau') == tau_hat_test, tau==cate without treatment RFX n_train <- 160 train_inds <- seq_len(n_train) - test_inds <- seq(n_train + 1, n) - X_train <- X[train_inds, ]; X_test <- X[test_inds, ] - Z_train <- Z[train_inds]; Z_test <- Z[test_inds] + test_inds <- seq(n_train + 1, n) + X_train <- X[train_inds, ] + X_test <- X[test_inds, ] + Z_train <- Z[train_inds] + Z_test <- Z[test_inds] y_train <- y[train_inds] - pi_train <- pi_x[train_inds]; pi_test <- pi_x[test_inds] + pi_train <- pi_x[train_inds] + pi_test <- pi_x[test_inds] # Fit BCF with sample_intercept = TRUE (default) bcf_model <- bcf( - X_train = X_train, Z_train = Z_train, y_train = y_train, - propensity_train = pi_train, X_test = X_test, Z_test = Z_test, - propensity_test = pi_test, num_gfr = 5, num_burnin = 0, num_mcmc = 10 + X_train = X_train, + Z_train = Z_train, + y_train = y_train, + propensity_train = pi_train, + X_test = X_test, + Z_test = Z_test, + propensity_test = pi_test, + num_gfr = 5, + num_burnin = 0, + num_mcmc = 10 ) # predict(terms = "tau") must match tau_hat_test exactly - tau_from_predict <- predict(bcf_model, X = X_test, Z = Z_test, - propensity = pi_test, terms = "tau") - expect_equal(tau_from_predict, bcf_model$tau_hat_test) + tau_from_predict <- predict( + bcf_model, + X = X_test, + Z = Z_test, + propensity = pi_test, + terms = "tau" + ) + tau_hat_test <- extractParameter(bcf_model, "tau_hat_test") + expect_equal(tau_from_predict, tau_hat_test) # predict(terms = "tau") == predict(terms = "cate") when no treatment RFX - cate_from_predict <- predict(bcf_model, X = X_test, Z = Z_test, - propensity = pi_test, terms = "cate") + cate_from_predict <- predict( + bcf_model, + X = X_test, + Z = Z_test, + propensity = pi_test, + terms = "cate" + ) expect_equal(tau_from_predict, cate_from_predict) # y_hat_test = mu_hat_test + Z_test * tau_hat_test (stored attributes decompose) - expected_y <- bcf_model$mu_hat_test + sweep(bcf_model$tau_hat_test, 1, as.numeric(Z_test), "*") - expect_equal(bcf_model$y_hat_test, expected_y) + tau_hat_test <- extractParameter(bcf_model, "tau_hat_test") + y_hat_test <- extractParameter(bcf_model, "y_hat_test") + mu_hat_test <- extractParameter(bcf_model, "mu_hat_test") + expected_y <- mu_hat_test + + sweep(tau_hat_test, 1, as.numeric(Z_test), "*") + expect_equal(y_hat_test, expected_y) # y_hat_train = mu_hat_train + Z_train * tau_hat_train - expected_y_train <- bcf_model$mu_hat_train + sweep(bcf_model$tau_hat_train, 1, as.numeric(Z_train), "*") - expect_equal(bcf_model$y_hat_train, expected_y_train) + tau_hat_train <- extractParameter(bcf_model, "tau_hat_train") + y_hat_train <- extractParameter(bcf_model, "y_hat_train") + mu_hat_train <- extractParameter(bcf_model, "mu_hat_train") + expected_y_train <- mu_hat_train + + sweep(tau_hat_train, 1, as.numeric(Z_train), "*") + expect_equal(y_hat_train, expected_y_train) # With sample_intercept = FALSE, tau includes only the forest; decomposition still holds bcf_no_int <- bcf( - X_train = X_train, Z_train = Z_train, y_train = y_train, - propensity_train = pi_train, X_test = X_test, Z_test = Z_test, - propensity_test = pi_test, num_gfr = 5, num_burnin = 0, num_mcmc = 10, + X_train = X_train, + Z_train = Z_train, + y_train = y_train, + propensity_train = pi_train, + X_test = X_test, + Z_test = Z_test, + propensity_test = pi_test, + num_gfr = 5, + num_burnin = 0, + num_mcmc = 10, treatment_effect_forest_params = list(sample_intercept = FALSE) ) - tau_no_int <- predict(bcf_no_int, X = X_test, Z = Z_test, - propensity = pi_test, terms = "tau") - expect_equal(tau_no_int, bcf_no_int$tau_hat_test) - expected_y_no_int <- bcf_no_int$mu_hat_test + - sweep(bcf_no_int$tau_hat_test, 1, as.numeric(Z_test), "*") - expect_equal(bcf_no_int$y_hat_test, expected_y_no_int) + tau_no_int <- predict( + bcf_no_int, + X = X_test, + Z = Z_test, + propensity = pi_test, + terms = "tau" + ) + tau_hat_test <- extractParameter(bcf_no_int, "tau_hat_test") + mu_hat_test <- extractParameter(bcf_no_int, "mu_hat_test") + y_hat_test <- extractParameter(bcf_no_int, "y_hat_test") + expect_equal(tau_no_int, tau_hat_test) + expected_y_no_int <- mu_hat_test + + sweep(tau_hat_test, 1, as.numeric(Z_test), "*") + expect_equal(y_hat_test, expected_y_no_int) }) diff --git a/test/R/testthat/test-forest-container.R b/test/R/testthat/test-forest-container.R index b78c8c8b..87c96574 100644 --- a/test/R/testthat/test-forest-container.R +++ b/test/R/testthat/test-forest-container.R @@ -162,8 +162,8 @@ test_that("Collapse forests", { general_params = general_param_list ) - # Extract the mean forest container - mean_forest_container <- bart_model$mean_forests + # Extract the mean forest container (a standalone deep copy we can mutate) + mean_forest_container <- bart_model$samples$materialize_mean_forest() # Predict from the original container pred_orig <- mean_forest_container$predict(forest_dataset_test) @@ -208,8 +208,8 @@ test_that("Collapse forests", { general_params = general_param_list ) - # Extract the mean forest container - mean_forest_container <- bart_model$mean_forests + # Extract the mean forest container (a standalone deep copy we can mutate) + mean_forest_container <- bart_model$samples$materialize_mean_forest() # Predict from the original container pred_orig <- mean_forest_container$predict(forest_dataset_test) @@ -254,8 +254,8 @@ test_that("Collapse forests", { general_params = general_param_list ) - # Extract the mean forest container - mean_forest_container <- bart_model$mean_forests + # Extract the mean forest container (a standalone deep copy we can mutate) + mean_forest_container <- bart_model$samples$materialize_mean_forest() # Predict from the original container pred_orig <- mean_forest_container$predict(forest_dataset_test) @@ -286,8 +286,8 @@ test_that("Collapse forests", { general_params = general_param_list ) - # Extract the mean forest container - mean_forest_container <- bart_model$mean_forests + # Extract the mean forest container (a standalone deep copy we can mutate) + mean_forest_container <- bart_model$samples$materialize_mean_forest() # Predict from the original container pred_orig <- mean_forest_container$predict(forest_dataset_test) diff --git a/test/R/testthat/test-init-params.R b/test/R/testthat/test-init-params.R new file mode 100644 index 00000000..89e72a46 --- /dev/null +++ b/test/R/testthat/test-init-params.R @@ -0,0 +1,184 @@ +# Regression tests that user-supplied initialization / calibration parameters +# are threaded into the C++ sampler rather than silently dropped. Each test would +# fail if the corresponding parameter were ignored on the C++ path. + +test_that("BART honors user-supplied sigma2_global_init", { + skip_on_cran() + set.seed(1) + n <- 200 + p <- 3 + X <- matrix(runif(n * p), ncol = p) + y <- X[, 1] + rnorm(n) + fit <- function(s) { + general_params <- list(standardize = FALSE, random_seed = 1) + if (!is.null(s)) { + general_params$sigma2_global_init <- s + } + bart( + X_train = X, y_train = y, + num_gfr = 0, num_burnin = 0, num_mcmc = 5, + general_params = general_params + ) + } + m_set <- fit(9.0) + m_default <- fit(NULL) + # User-supplied global variance init must be honored, not hardcoded to 1.0. + expect_equal(m_set$model_params$sigma2_init, 9.0, tolerance = 1e-8) + expect_false(isTRUE(all.equal(m_default$model_params$sigma2_init, 9.0))) +}) + +test_that("BART honors user-supplied var_forest_leaf_init", { + skip_on_cran() + set.seed(2) + n <- 200 + p <- 3 + X <- matrix(runif(n * p), ncol = p) + # Heteroskedastic outcome so the variance forest is meaningfully fit. + y <- X[, 1] + rnorm(n) * exp(0.5 * X[, 2]) + fit <- function(v) { + variance_forest_params <- list(num_trees = 50) + if (!is.null(v)) { + variance_forest_params$var_forest_leaf_init <- v + } + m <- bart( + X_train = X, y_train = y, X_test = X, + num_gfr = 0, num_burnin = 0, num_mcmc = 3, + general_params = list(standardize = FALSE, random_seed = 1), + mean_forest_params = list(num_trees = 50), + variance_forest_params = variance_forest_params + ) + m$samples$variance_forest_predictions_test() + } + # Same seed; differing only by the variance-forest leaf init must change output. + expect_false(isTRUE(all.equal(fit(0.05), fit(2.0)))) +}) + +test_that("BCF honors delta_max for probit treatment leaf calibration", { + skip_on_cran() + set.seed(3) + n <- 300 + p <- 5 + X <- matrix(runif(n * p), ncol = p) + Z <- rbinom(n, 1, 0.5) + propensity <- rep(0.5, n) + prob <- pnorm(X[, 1] - 0.5 + 0.5 * Z) + y <- rbinom(n, 1, prob) + num_trees_tau <- 50 + fit <- function(dm) { + m <- bcf( + X_train = X, Z_train = Z, y_train = y, propensity_train = propensity, + num_gfr = 0, num_burnin = 0, num_mcmc = 5, + general_params = list( + random_seed = 1, + outcome_model = OutcomeModel(outcome = "binary", link = "probit") + ), + treatment_effect_forest_params = list( + num_trees = num_trees_tau, + delta_max = dm + ) + ) + m$model_params$sigma2_leaf_tau + } + expected <- function(dm) { + p_coverage <- 0.6827 + q_quantile <- qnorm((p_coverage + 1) / 2) + ((dm / (q_quantile * dnorm(0)))^2) / num_trees_tau + } + v_05 <- fit(0.5) + v_09 <- fit(0.9) + # delta_max must drive the treatment-effect leaf scale calibration. + expect_false(isTRUE(all.equal(v_05, v_09))) + expect_equal(v_05, expected(0.5), tolerance = 1e-8) + expect_equal(v_09, expected(0.9), tolerance = 1e-8) +}) + +test_that("BART honors observation_weights", { + skip_on_cran() + set.seed(4) + n <- 200 + p <- 3 + X <- matrix(runif(n * p), ncol = p) + y <- X[, 1] + rnorm(n) + w <- runif(n, 0.1, 2.0) + fit <- function(weights) { + bart( + X_train = X, y_train = y, X_test = X, + num_gfr = 0, num_burnin = 0, num_mcmc = 5, + general_params = list(standardize = FALSE, random_seed = 1), + observation_weights_train = weights + )$samples$y_hat_test() + } + # Non-uniform observation weights must change the fit (same seed). + expect_false(isTRUE(all.equal(fit(NULL), fit(w)))) +}) + +test_that("BCF honors observation_weights", { + skip_on_cran() + set.seed(5) + n <- 200 + p <- 5 + X <- matrix(runif(n * p), ncol = p) + Z <- rbinom(n, 1, 0.5) + propensity <- rep(0.5, n) + y <- X[, 1] + Z * X[, 2] + rnorm(n) + w <- runif(n, 0.1, 2.0) + fit <- function(weights) { + bcf( + X_train = X, Z_train = Z, y_train = y, propensity_train = propensity, + X_test = X, Z_test = Z, propensity_test = propensity, + num_gfr = 0, num_burnin = 0, num_mcmc = 5, + general_params = list(standardize = FALSE, random_seed = 1), + observation_weights_train = weights + )$samples$y_hat_test() + } + expect_false(isTRUE(all.equal(fit(NULL), fit(w)))) +}) + +test_that("BCF internal propensity model is reproducible with random_seed", { + skip_on_cran() + set.seed(7) + n <- 200 + p <- 5 + X <- matrix(runif(n * p), ncol = p) + Z <- rbinom(n, 1, 0.5) + y <- X[, 1] + Z * X[, 2] + rnorm(n) + fit <- function() { + # No propensity_train -> BCF estimates it with an internal BART model. + bcf( + X_train = X, Z_train = Z, y_train = y, + num_gfr = 0, num_burnin = 0, num_mcmc = 10, + general_params = list(random_seed = 99) + )$samples$y_hat_train() + } + # With a fixed random_seed, the internally-estimated propensity (and hence the + # full fit) must be reproducible across runs. + expect_equal(fit(), fit()) +}) + +test_that("verbose prints sampler progress", { + skip_on_cran() + set.seed(8) + n <- 100 + p <- 3 + X <- matrix(runif(n * p), ncol = p) + y <- X[, 1] + rnorm(n) + # verbose = TRUE prints GFR/MCMC progress to the console. + expect_output( + bart( + X_train = X, y_train = y, + num_gfr = 5, num_burnin = 0, num_mcmc = 20, + general_params = list(random_seed = 1, verbose = TRUE) + ), + "Running GFR sampler" + ) + # verbose = FALSE is silent. Assign the result inside capture.output so the + # returned model is not auto-printed (its print method itself mentions GFR/MCMC). + silent_output <- capture.output( + fit_silent <- bart( + X_train = X, y_train = y, + num_gfr = 5, num_burnin = 0, num_mcmc = 20, + general_params = list(random_seed = 1, verbose = FALSE) + ) + ) + expect_false(any(grepl("Running GFR sampler|GFR:|MCMC:", silent_output))) +}) diff --git a/test/R/testthat/test-json-primitives.R b/test/R/testthat/test-json-primitives.R new file mode 100644 index 00000000..82557fbb --- /dev/null +++ b/test/R/testthat/test-json-primitives.R @@ -0,0 +1,37 @@ +# Unit tests for the JSON key primitives (rename/erase) used by schema migrations. + +test_that("CppJson rename_field / erase_field (top-level)", { + skip_on_cran() + j <- createCppJson() + j$add_scalar("old", 3.5) + + j$rename_field("old", "new") + expect_false(j$contains("old")) + expect_true(j$contains("new")) + expect_equal(j$get_scalar("new"), 3.5) + + j$erase_field("new") + expect_false(j$contains("new")) +}) + +test_that("CppJson rename_field / erase_field (subfolder)", { + skip_on_cran() + j <- createCppJson() + j$add_scalar("old", 7.0, subfolder_name = "sub") + + j$rename_field("old", "new", subfolder_name = "sub") + expect_false(j$contains("old", subfolder_name = "sub")) + expect_true(j$contains("new", subfolder_name = "sub")) + expect_equal(j$get_scalar("new", subfolder_name = "sub"), 7.0) + + j$erase_field("new", subfolder_name = "sub") + expect_false(j$contains("new", subfolder_name = "sub")) +}) + +test_that("CppJson rename/erase of a missing field is a no-op", { + skip_on_cran() + j <- createCppJson() + expect_silent(j$rename_field("nope", "whatever")) + expect_silent(j$erase_field("nope")) + expect_false(j$contains("whatever")) +}) diff --git a/test/R/testthat/test-multi-chain.R b/test/R/testthat/test-multi-chain.R index 98530bad..d4c74474 100644 --- a/test/R/testthat/test-multi-chain.R +++ b/test/R/testthat/test-multi-chain.R @@ -60,9 +60,9 @@ test_that("BART multi-chain: sample counts with no GFR", { general_params = list(num_chains = n_chains, num_threads = 1) ) expected <- n_chains * n_mcmc - expect_length(m$sigma2_global_samples, expected) - expect_equal(dim(m$y_hat_train), c(d$n_train, expected)) - expect_equal(dim(m$y_hat_test), c(d$n_test, expected)) + expect_length(m$samples$global_var_samples(), expected) + expect_equal(dim(m$samples$y_hat_train()), c(d$n_train, expected)) + expect_equal(dim(m$samples$y_hat_test()), c(d$n_test, expected)) }) test_that("BART multi-chain: sample counts with GFR warm-start", { @@ -75,8 +75,8 @@ test_that("BART multi-chain: sample counts with GFR warm-start", { general_params = list(num_chains = n_chains, num_threads = 1) ) expected <- n_chains * n_mcmc - expect_length(m$sigma2_global_samples, expected) - expect_equal(dim(m$y_hat_train), c(d$n_train, expected)) + expect_length(m$samples$global_var_samples(), expected) + expect_equal(dim(m$samples$y_hat_train()), c(d$n_train, expected)) }) test_that("BART multi-chain: leaf-scale sample count", { @@ -92,7 +92,7 @@ test_that("BART multi-chain: leaf-scale sample count", { ), mean_forest_params = list(sample_sigma2_leaf = TRUE) ) - expect_length(m$sigma2_leaf_samples, n_chains * n_mcmc) + expect_length(m$samples$leaf_scale_samples(), n_chains * n_mcmc) }) test_that("BART multi-chain: chain independence (no GFR)", { @@ -104,8 +104,8 @@ test_that("BART multi-chain: chain independence (no GFR)", { num_gfr = 0, num_burnin = 0, num_mcmc = n_mcmc, general_params = list(num_chains = 2, num_threads = 1) ) - chain1 <- m$sigma2_global_samples[seq_len(n_mcmc)] - chain2 <- m$sigma2_global_samples[seq(n_mcmc + 1, 2 * n_mcmc)] + chain1 <- m$samples$global_var_samples()[seq_len(n_mcmc)] + chain2 <- m$samples$global_var_samples()[seq(n_mcmc + 1, 2 * n_mcmc)] expect_false(isTRUE(all.equal(chain1, chain2)), label = "Chains should produce distinct sigma2 samples") }) @@ -119,8 +119,8 @@ test_that("BART multi-chain: chain independence (with GFR)", { num_gfr = n_gfr, num_burnin = 5, num_mcmc = n_mcmc, general_params = list(num_chains = 2, num_threads = 1) ) - chain1 <- m$sigma2_global_samples[seq_len(n_mcmc)] - chain2 <- m$sigma2_global_samples[seq(n_mcmc + 1, 2 * n_mcmc)] + chain1 <- m$samples$global_var_samples()[seq_len(n_mcmc)] + chain2 <- m$samples$global_var_samples()[seq(n_mcmc + 1, 2 * n_mcmc)] expect_false(isTRUE(all.equal(chain1, chain2))) }) @@ -189,8 +189,8 @@ test_that("BART multi-chain: sigma2 samples are finite and positive with GFR", { num_gfr = 6, num_burnin = 10, num_mcmc = 10, general_params = list(num_chains = 3, num_threads = 1) ) - expect_true(all(is.finite(m$sigma2_global_samples))) - expect_true(all(m$sigma2_global_samples > 0)) + expect_true(all(is.finite(m$samples$global_var_samples()))) + expect_true(all(m$samples$global_var_samples() > 0)) }) # --------------------------------------------------------------------------- @@ -209,10 +209,10 @@ test_that("BCF multi-chain: sample counts with no GFR", { general_params = list(num_chains = n_chains, num_threads = 1) ) expected <- n_chains * n_mcmc - expect_length(m$sigma2_global_samples, expected) - expect_equal(dim(m$tau_hat_train), c(d$n_train, expected)) - expect_equal(dim(m$mu_hat_train), c(d$n_train, expected)) - expect_equal(dim(m$tau_hat_test), c(d$n_test, expected)) + expect_length(m$samples$global_var_samples(), expected) + expect_equal(dim(m$samples$tau_forest_predictions_train()), c(d$n_train, expected)) + expect_equal(dim(m$samples$mu_forest_predictions_train()), c(d$n_train, expected)) + expect_equal(dim(m$samples$tau_forest_predictions_test()), c(d$n_test, expected)) }) test_that("BCF multi-chain: sample counts with GFR warm-start", { @@ -227,13 +227,13 @@ test_that("BCF multi-chain: sample counts with GFR warm-start", { general_params = list(num_chains = n_chains, num_threads = 1, adaptive_coding = TRUE) ) expected <- n_chains * n_mcmc - expect_length(m$sigma2_global_samples, expected) - expect_equal(dim(m$tau_hat_train), c(d$n_train, expected)) - expect_equal(dim(m$mu_hat_train), c(d$n_train, expected)) + expect_length(m$samples$global_var_samples(), expected) + expect_equal(dim(m$samples$tau_forest_predictions_train()), c(d$n_train, expected)) + expect_equal(dim(m$samples$mu_forest_predictions_train()), c(d$n_train, expected)) # BCF-specific scalar parameter arrays - expect_length(m$b_0_samples, expected) - expect_length(m$b_1_samples, expected) - expect_length(m$sigma2_leaf_mu_samples, expected) + expect_length(m$samples$b0_samples(), expected) + expect_length(m$samples$b1_samples(), expected) + expect_length(m$samples$leaf_scale_mu_samples(), expected) }) test_that("BCF multi-chain: chain independence (no GFR)", { @@ -247,8 +247,8 @@ test_that("BCF multi-chain: chain independence (no GFR)", { num_gfr = 0, num_burnin = 10, num_mcmc = n_mcmc, general_params = list(num_chains = 2, num_threads = 1) ) - chain1 <- m$sigma2_global_samples[seq_len(n_mcmc)] - chain2 <- m$sigma2_global_samples[seq(n_mcmc + 1, 2 * n_mcmc)] + chain1 <- m$samples$global_var_samples()[seq_len(n_mcmc)] + chain2 <- m$samples$global_var_samples()[seq(n_mcmc + 1, 2 * n_mcmc)] expect_false(isTRUE(all.equal(chain1, chain2)), label = "BCF chains should produce distinct sigma2 samples") }) @@ -264,8 +264,8 @@ test_that("BCF multi-chain: chain independence (with GFR)", { num_gfr = 4, num_burnin = 5, num_mcmc = n_mcmc, general_params = list(num_chains = 2, num_threads = 1) ) - chain1 <- m$sigma2_global_samples[seq_len(n_mcmc)] - chain2 <- m$sigma2_global_samples[seq(n_mcmc + 1, 2 * n_mcmc)] + chain1 <- m$samples$global_var_samples()[seq_len(n_mcmc)] + chain2 <- m$samples$global_var_samples()[seq(n_mcmc + 1, 2 * n_mcmc)] expect_false(isTRUE(all.equal(chain1, chain2))) }) @@ -281,11 +281,11 @@ test_that("BCF multi-chain: all samples finite with GFR + multiple chains", { num_gfr = 6, num_burnin = 20, num_mcmc = 10, general_params = list(num_chains = 3, num_threads = 1, adaptive_coding = TRUE) ) - expect_true(all(is.finite(m$sigma2_global_samples)), + expect_true(all(is.finite(m$samples$global_var_samples())), label = "sigma2 samples must be finite (no chain-transition blowup)") - expect_true(all(m$sigma2_global_samples > 0)) - expect_true(all(is.finite(m$b_0_samples))) - expect_true(all(is.finite(m$b_1_samples))) + expect_true(all(m$samples$global_var_samples() > 0)) + expect_true(all(is.finite(m$samples$b0_samples()))) + expect_true(all(is.finite(m$samples$b1_samples()))) }) test_that("BCF multi-chain: extractParameter dimensions", { diff --git a/test/R/testthat/test-observation-weights.R b/test/R/testthat/test-observation-weights.R index 2e2dfda4..eef4d14d 100644 --- a/test/R/testthat/test-observation-weights.R +++ b/test/R/testthat/test-observation-weights.R @@ -39,13 +39,13 @@ test_that("BART: uniform weights produce identical predictions to no weights", { set.seed(1) m2 <- bart( X_train = d$X_train, y_train = d$y_train, X_test = d$X_test, - observation_weights = rep(1.0, d$n_train), + observation_weights_train = rep(1.0, d$n_train), num_gfr = 0, num_burnin = 0, num_mcmc = num_mcmc, general_params = list(random_seed = 1L) ) - expect_equal(m1$y_hat_train, m2$y_hat_train) - expect_equal(m1$y_hat_test, m2$y_hat_test) + expect_equal(m1$samples$y_hat_train(), m2$samples$y_hat_train()) + expect_equal(m1$samples$y_hat_test(), m2$samples$y_hat_test()) }) test_that("BART: non-uniform weights run and produce correct output shape", { @@ -57,12 +57,12 @@ test_that("BART: non-uniform weights run and produce correct output shape", { expect_no_error( m <- bart( X_train = d$X_train, y_train = d$y_train, X_test = d$X_test, - observation_weights = weights, + observation_weights_train = weights, num_gfr = 0, num_burnin = 0, num_mcmc = num_mcmc ) ) - expect_equal(dim(m$y_hat_train), c(d$n_train, num_mcmc)) - expect_equal(dim(m$y_hat_test), c(d$n_test, num_mcmc)) + expect_equal(dim(m$samples$y_hat_train()), c(d$n_train, num_mcmc)) + expect_equal(dim(m$samples$y_hat_test()), c(d$n_test, num_mcmc)) }) test_that("BART: all-zero weights (prior mode) run with num_gfr = 0", { @@ -73,40 +73,40 @@ test_that("BART: all-zero weights (prior mode) run with num_gfr = 0", { expect_no_error( m <- bart( X_train = d$X_train, y_train = d$y_train, - observation_weights = rep(0.0, d$n_train), + observation_weights_train = rep(0.0, d$n_train), num_gfr = 0, num_burnin = 0, num_mcmc = num_mcmc ) ) - expect_equal(dim(m$y_hat_train), c(d$n_train, num_mcmc)) + expect_equal(dim(m$samples$y_hat_train()), c(d$n_train, num_mcmc)) }) -test_that("BART: non-numeric observation_weights raises error", { +test_that("BART: non-numeric observation_weights_train raises error", { skip_on_cran() d <- make_bart_data() expect_error( bart( X_train = d$X_train, y_train = d$y_train, - observation_weights = as.character(rep(1, d$n_train)), + observation_weights_train = as.character(rep(1, d$n_train)), num_gfr = 0, num_burnin = 0, num_mcmc = 5 ), "numeric" ) }) -test_that("BART: wrong-length observation_weights raises error", { +test_that("BART: wrong-length observation_weights_train raises error", { skip_on_cran() d <- make_bart_data() expect_error( bart( X_train = d$X_train, y_train = d$y_train, - observation_weights = rep(1.0, d$n_train + 1), + observation_weights_train = rep(1.0, d$n_train + 1), num_gfr = 0, num_burnin = 0, num_mcmc = 5 ), "nrow" ) }) -test_that("BART: negative observation_weights raises error", { +test_that("BART: negative observation_weights_train raises error", { skip_on_cran() d <- make_bart_data() weights <- rep(1.0, d$n_train) @@ -114,7 +114,7 @@ test_that("BART: negative observation_weights raises error", { expect_error( bart( X_train = d$X_train, y_train = d$y_train, - observation_weights = weights, + observation_weights_train = weights, num_gfr = 0, num_burnin = 0, num_mcmc = 5 ), "negative" @@ -127,21 +127,21 @@ test_that("BART: all-zero weights with num_gfr > 0 raises error", { expect_error( bart( X_train = d$X_train, y_train = d$y_train, - observation_weights = rep(0.0, d$n_train), + observation_weights_train = rep(0.0, d$n_train), num_gfr = 5, num_burnin = 0, num_mcmc = 10 ), "num_gfr" ) }) -test_that("BART: observation_weights with cloglog outcome raises error", { +test_that("BART: observation_weights_train with cloglog outcome raises error", { skip_on_cran() d <- make_bart_data() y_ord <- sample(1:3, d$n_train, replace = TRUE) expect_error( bart( X_train = d$X_train, y_train = y_ord, - observation_weights = rep(1.0, d$n_train), + observation_weights_train = rep(1.0, d$n_train), num_gfr = 0, num_burnin = 0, num_mcmc = 5, general_params = list(outcome_model = OutcomeModel(outcome = "ordinal", link = "cloglog")) ), @@ -149,17 +149,17 @@ test_that("BART: observation_weights with cloglog outcome raises error", { ) }) -test_that("BART: observation_weights with variance forest raises warning", { +test_that("BART: observation_weights_train with variance forest raises error", { skip_on_cran() d <- make_bart_data() - expect_warning( + expect_error( bart( X_train = d$X_train, y_train = d$y_train, - observation_weights = rep(1.0, d$n_train), + observation_weights_train = rep(1.0, d$n_train), num_gfr = 0, num_burnin = 0, num_mcmc = 5, variance_forest_params = list(num_trees = 5) ), - "variance forest" + "not compatible with a variance forest" ) }) @@ -182,13 +182,13 @@ test_that("BCF: uniform weights produce identical predictions to no weights", { X_train = d$X_train, Z_train = d$Z_train, y_train = d$y_train, propensity_train = d$pi_train, X_test = d$X_test, Z_test = d$Z_test, propensity_test = d$pi_test, - observation_weights = rep(1.0, d$n_train), + observation_weights_train = rep(1.0, d$n_train), num_gfr = 0, num_burnin = 0, num_mcmc = num_mcmc, general_params = list(random_seed = 1L) ) - expect_equal(m1$y_hat_train, m2$y_hat_train) - expect_equal(m1$tau_hat_train, m2$tau_hat_train) + expect_equal(m1$samples$y_hat_train(), m2$samples$y_hat_train()) + expect_equal(m1$samples$tau_forest_predictions_train(), m2$samples$tau_forest_predictions_train()) }) test_that("BCF: non-uniform weights run and produce correct output shape", { @@ -202,17 +202,17 @@ test_that("BCF: non-uniform weights run and produce correct output shape", { X_train = d$X_train, Z_train = d$Z_train, y_train = d$y_train, propensity_train = d$pi_train, X_test = d$X_test, Z_test = d$Z_test, propensity_test = d$pi_test, - observation_weights = weights, + observation_weights_train = weights, num_gfr = 0, num_burnin = 0, num_mcmc = num_mcmc ) ) - expect_equal(dim(m$y_hat_train), c(d$n_train, num_mcmc)) - expect_equal(dim(m$tau_hat_train), c(d$n_train, num_mcmc)) - expect_equal(dim(m$y_hat_test), c(d$n_test, num_mcmc)) - expect_equal(dim(m$tau_hat_test), c(d$n_test, num_mcmc)) + expect_equal(dim(m$samples$y_hat_train()), c(d$n_train, num_mcmc)) + expect_equal(dim(m$samples$tau_forest_predictions_train()), c(d$n_train, num_mcmc)) + expect_equal(dim(m$samples$y_hat_test()), c(d$n_test, num_mcmc)) + expect_equal(dim(m$samples$tau_forest_predictions_test()), c(d$n_test, num_mcmc)) }) -test_that("BCF: negative observation_weights raises error", { +test_that("BCF: negative observation_weights_train raises error", { skip_on_cran() d <- make_bcf_data() weights <- rep(1.0, d$n_train) @@ -221,7 +221,7 @@ test_that("BCF: negative observation_weights raises error", { bcf( X_train = d$X_train, Z_train = d$Z_train, y_train = d$y_train, propensity_train = d$pi_train, - observation_weights = weights, + observation_weights_train = weights, num_gfr = 0, num_burnin = 0, num_mcmc = 5 ), "negative" @@ -235,14 +235,14 @@ test_that("BCF: all-zero weights with num_gfr > 0 raises error", { bcf( X_train = d$X_train, Z_train = d$Z_train, y_train = d$y_train, propensity_train = d$pi_train, - observation_weights = rep(0.0, d$n_train), + observation_weights_train = rep(0.0, d$n_train), num_gfr = 5, num_burnin = 0, num_mcmc = 10 ), "num_gfr" ) }) -test_that("BCF: observation_weights with cloglog outcome raises error", { +test_that("BCF: observation_weights_train with cloglog outcome raises error", { skip_on_cran() d <- make_bcf_data() y_bin <- rbinom(d$n_train, 1, 0.5) @@ -250,10 +250,35 @@ test_that("BCF: observation_weights with cloglog outcome raises error", { bcf( X_train = d$X_train, Z_train = d$Z_train, y_train = y_bin, propensity_train = d$pi_train, - observation_weights = rep(1.0, d$n_train), + observation_weights_train = rep(1.0, d$n_train), num_gfr = 0, num_burnin = 0, num_mcmc = 5, general_params = list(outcome_model = OutcomeModel(outcome = "binary", link = "cloglog")) ), "cloglog" ) }) + +test_that("BART: deprecated observation_weights alias still works and warns", { + skip_on_cran() + d <- make_bart_data() + weights <- runif(d$n_train, 0.5, 2.0) + set.seed(1) + m_new <- bart( + X_train = d$X_train, y_train = d$y_train, X_test = d$X_test, + observation_weights_train = weights, + num_gfr = 0, num_burnin = 0, num_mcmc = 10, + general_params = list(random_seed = 1L) + ) + set.seed(1) + expect_warning( + m_old <- bart( + X_train = d$X_train, y_train = d$y_train, X_test = d$X_test, + observation_weights = weights, + num_gfr = 0, num_burnin = 0, num_mcmc = 10, + general_params = list(random_seed = 1L) + ), + "deprecated" + ) + # Deprecated alias must produce the same result as the new parameter. + expect_equal(m_new$samples$y_hat_test(), m_old$samples$y_hat_test()) +}) diff --git a/test/R/testthat/test-predict.R b/test/R/testthat/test-predict.R index bda5c980..8d029176 100644 --- a/test/R/testthat/test-predict.R +++ b/test/R/testthat/test-predict.R @@ -1048,16 +1048,22 @@ test_that("BART cloglog ordinal: sample posterior predictive", { }) test_that("BART cloglog ordinal: probability transform correctness (K=4)", { + # K=4 is the minimal case that exposes cumulative-product bugs; + # K=3 is accidentally correct via the residual-last-class formula. skip_on_cran() set.seed(123) - n <- 500; p <- 3; n_categories <- 4L + n <- 500 + p <- 3 + n_categories <- 4L X <- matrix(runif(n * p), ncol = p) beta <- rep(1 / sqrt(p), p) true_lambda <- X %*% beta # Balanced cutpoints: ~25% per category so all 4 are reliably observed. - # Unbalanced cutpoints give ~0.3% K=4 probability, often no K=4 training obs. + # Unbalanced cutpoints (e.g., c(-2, -0.5, 0.5)) give ~0.3% K=4 probability + # and can produce no K=4 training observations, which would degrade to K=3. gamma_true <- c(-2.1, -1.7, -1.2) + # Sequential DGP: P(Y=k) = prod_{j= 0)) - row_sums_r <- apply(p_model_r, c(1,3), sum) - expect_equal(row_sums_r, matrix(1, nrow=n_test, ncol=num_mcmc), tolerance=1e-10) + gamma_samples <- bart_model$samples$cloglog_cutpoint_samples() # (K-1) x num_mcmc + expect_equal(dim(gamma_samples), c(n_categories - 1L, num_mcmc)) + + f_hat <- predict( + bart_model, + X = X_test, + scale = "linear", + terms = "mean_forest" + ) + expect_equal(dim(f_hat), c(n_test, num_mcmc)) + + p_manual <- assemble_probs(f_hat, gamma_samples, n_categories) + + p_model <- predict( + bart_model, + X = X_test, + scale = "probability", + terms = "y_hat" + ) + expect_equal(dim(p_model), c(n_test, n_categories, num_mcmc)) + + expect_equal(p_manual, p_model, tolerance = 1e-10) + expect_true(all(p_model >= 0)) + row_sums <- apply(p_model, c(1, 3), sum) + expect_equal( + row_sums, + matrix(1, nrow = n_test, ncol = num_mcmc), + tolerance = 1e-10 + ) }) test_that("BART gaussian: posterior interval and contrast", { diff --git a/test/R/testthat/test-random-effects.R b/test/R/testthat/test-random-effects.R index bafa75ab..c1c93d00 100644 --- a/test/R/testthat/test-random-effects.R +++ b/test/R/testthat/test-random-effects.R @@ -374,20 +374,20 @@ test_that("Random Effects BART with Offset Numbering", { mean_forest_params = mean_forest_param_list, random_effects_params = rfx_param_list ) - rfx_output <- capture.output(print(bart_model$rfx_samples)) + rfx_output <- capture.output(print(bart_model$samples$materialize_rfx())) expect_true(any(grepl("Random effects container", rfx_output, fixed = TRUE))) expect_true(any(grepl( - bart_model$rfx_samples$num_samples(), + bart_model$samples$materialize_rfx()$num_samples(), rfx_output, fixed = TRUE ))) expect_true(any(grepl( - bart_model$rfx_samples$num_components(), + bart_model$samples$materialize_rfx()$num_components(), rfx_output, fixed = TRUE ))) expect_true(any(grepl( - bart_model$rfx_samples$num_groups(), + bart_model$samples$materialize_rfx()$num_groups(), rfx_output, fixed = TRUE ))) @@ -583,20 +583,20 @@ test_that("Random Effects BCF with Default Numbering", { num_mcmc = 10, random_effects_params = rfx_param_list ) - rfx_output <- capture.output(print(bcf_model$rfx_samples)) + rfx_output <- capture.output(print(bcf_model$samples$materialize_rfx())) expect_true(any(grepl("Random effects container", rfx_output, fixed = TRUE))) expect_true(any(grepl( - bcf_model$rfx_samples$num_samples(), + bcf_model$samples$materialize_rfx()$num_samples(), rfx_output, fixed = TRUE ))) expect_true(any(grepl( - bcf_model$rfx_samples$num_components(), + bcf_model$samples$materialize_rfx()$num_components(), rfx_output, fixed = TRUE ))) expect_true(any(grepl( - bcf_model$rfx_samples$num_groups(), + bcf_model$samples$materialize_rfx()$num_groups(), rfx_output, fixed = TRUE ))) @@ -792,20 +792,20 @@ test_that("Random Effects BCF with Offset Numbering", { num_mcmc = 10, random_effects_params = rfx_param_list ) - rfx_output <- capture.output(print(bcf_model$rfx_samples)) + rfx_output <- capture.output(print(bcf_model$samples$materialize_rfx())) expect_true(any(grepl("Random effects container", rfx_output, fixed = TRUE))) expect_true(any(grepl( - bcf_model$rfx_samples$num_samples(), + bcf_model$samples$materialize_rfx()$num_samples(), rfx_output, fixed = TRUE ))) expect_true(any(grepl( - bcf_model$rfx_samples$num_components(), + bcf_model$samples$materialize_rfx()$num_components(), rfx_output, fixed = TRUE ))) expect_true(any(grepl( - bcf_model$rfx_samples$num_groups(), + bcf_model$samples$materialize_rfx()$num_groups(), rfx_output, fixed = TRUE ))) diff --git a/test/R/testthat/test-samples-container.R b/test/R/testthat/test-samples-container.R new file mode 100644 index 00000000..6c47fef6 --- /dev/null +++ b/test/R/testthat/test-samples-container.R @@ -0,0 +1,144 @@ +test_that("BARTSamples round-trips a fitted model's forests and parameters", { + skip_on_cran() + set.seed(1) + n <- 100 + p <- 4 + X <- matrix(runif(n * p), n, p) + y <- X[, 1] * 2 + rnorm(n, 0, 0.5) + m <- bart(X_train = X, y_train = y, num_gfr = 0, num_burnin = 0, num_mcmc = 10) + + # Rebuild an independent BARTSamples from the model's serialized form. + sc <- createBARTSamplesFromJson( + createCppJsonString(saveBARTModelToJsonString(m)) + ) + + # Scalars / counts match the model + expect_equal(sc$num_samples(), m$model_params$num_samples) + expect_true(sc$has_mean_forest()) + expect_false(sc$has_variance_forest()) + expect_equal(sc$y_bar(), m$model_params$outcome_mean) + + # Materialized mean forest predicts identically to the model's forest (faithful deep copy) + model_mean_forest <- m$samples$materialize_mean_forest() + fc <- sc$materialize_mean_forest() + expect_false(is.null(fc)) + expect_equal(fc$num_samples(), model_mean_forest$num_samples()) + ds <- createForestDataset(X) + expect_equal(model_mean_forest$predict(ds), fc$predict(ds)) + + # Parameter traces round-trip + if (length(m$samples$global_var_samples()) > 0) { + expect_equal( + as.numeric(sc$global_var_samples()), + as.numeric(m$samples$global_var_samples()) + ) + } +}) + +test_that("BARTSamples merge concatenates draws", { + skip_on_cran() + set.seed(2) + n <- 80 + p <- 3 + X <- matrix(runif(n * p), n, p) + y <- X[, 1] + rnorm(n, 0, 0.5) + m <- bart(X_train = X, y_train = y, num_gfr = 0, num_burnin = 0, num_mcmc = 8) + + build <- function() { + createBARTSamplesFromJson(createCppJsonString(saveBARTModelToJsonString(m))) + } + # Two containers from the same model share standardization (merge guards against mismatch). + a <- build() + b <- build() + na <- a$num_samples() + nb <- b$num_samples() + + a$merge(b) + + expect_equal(a$num_samples(), na + nb) + expect_equal(a$materialize_mean_forest()$num_samples(), na + nb) +}) + +test_that("BCFSamples wraps a fitted model's forests and parameters", { + skip_on_cran() + set.seed(3) + n <- 120 + p <- 4 + X <- matrix(runif(n * p), n, p) + pi_x <- 0.3 + 0.4 * X[, 2] + Z <- rbinom(n, 1, pi_x) + y <- 1 + 2 * X[, 1] + 1.5 * X[, 3] * Z + 0.5 * rnorm(n) + m <- bcf( + X_train = X, Z_train = Z, y_train = y, propensity_train = pi_x, + num_gfr = 0, num_burnin = 0, num_mcmc = 10 + ) + + model_mu_forest <- m$samples$materialize_mu_forest() + model_tau_forest <- m$samples$materialize_tau_forest() + sc <- createBCFSamplesFromJson( + createCppJsonString(saveBCFModelToJsonString(m)) + ) + + expect_equal(sc$num_samples(), m$model_params$num_samples) + expect_equal(sc$treatment_dim(), 1L) + expect_true(sc$has_mu_forest()) + expect_true(sc$has_tau_forest()) + expect_false(sc$has_variance_forest()) + + # Both forests materialize as deep copies. (Prediction parity is awkward to check directly here: + # BCF trains its forests on propensity-augmented covariates, and the treatment forest needs a + # basis -- so verify the copies via sample count.) + fc_mu <- sc$materialize_mu_forest() + fc_tau <- sc$materialize_tau_forest() + expect_false(is.null(fc_mu)) + expect_false(is.null(fc_tau)) + expect_equal(fc_mu$num_samples(), model_mu_forest$num_samples()) + expect_equal(fc_tau$num_samples(), model_tau_forest$num_samples()) + + # Global error variance round-trips + if (length(m$samples$global_var_samples()) > 0) { + expect_equal( + as.numeric(sc$global_var_samples()), + as.numeric(m$samples$global_var_samples()) + ) + } +}) + +test_that("direct forest access is a hard error pointing at the extraction path", { + skip_on_cran() + set.seed(7) + n <- 60 + p <- 3 + X <- matrix(runif(n * p), n, p) + y <- X[, 1] + rnorm(n, 0, 0.5) + m <- bart(X_train = X, y_train = y, num_gfr = 0, num_burnin = 0, num_mcmc = 5) + + # Removed fields raise, and the message names the supported extraction call. + expect_error(m$mean_forests, "extractForest") + expect_error(m$variance_forests, "extractForest") + + # Non-forest fields still read through normally. + expect_false(is.null(m$model_params)) + expect_false(is.null(m$samples)) + + # The documented extraction path works. + expect_false(is.null(m$samples$materialize_mean_forest())) +}) + +test_that("BCF direct forest access is a hard error", { + skip_on_cran() + set.seed(8) + n <- 100 + p <- 3 + X <- matrix(runif(n * p), n, p) + pi_x <- 0.3 + 0.4 * X[, 2] + Z <- rbinom(n, 1, pi_x) + y <- 1 + 2 * X[, 1] + 1.5 * X[, 3] * Z + 0.5 * rnorm(n) + m <- bcf( + X_train = X, Z_train = Z, y_train = y, propensity_train = pi_x, + num_gfr = 0, num_burnin = 0, num_mcmc = 5 + ) + expect_error(m$forests_mu, "extractForest") + expect_error(m$forests_tau, "extractForest") + expect_false(is.null(m$samples$materialize_tau_forest())) +}) diff --git a/test/R/testthat/test-serialization-compat.R b/test/R/testthat/test-serialization-compat.R index ff7c7e5b..0b6ff0d0 100644 --- a/test/R/testthat/test-serialization-compat.R +++ b/test/R/testthat/test-serialization-compat.R @@ -334,3 +334,185 @@ test_that("BCF loads with multiple missing optional fields simultaneously", { # At least the preprocessor_metadata warning should fire expect_true(any(grepl("preprocessor|preprocess", warns, ignore.case = TRUE))) }) + +test_that("a future schema_version is refused with a clear error", { + skip_on_cran() + set.seed(1) + n <- 80 + p <- 3 + X <- matrix(runif(n * p), ncol = p) + y <- X[, 1] + rnorm(n) + m <- bart( + X_train = X, y_train = y, + num_gfr = 0, num_burnin = 0, num_mcmc = 5, + general_params = list(random_seed = 1) + ) + js <- saveBARTModelToJsonString(m) + # Bump the stamped schema_version to a value this install cannot read. + js_future <- sub('("schema_version"\\s*:\\s*)[0-9]+', "\\199", js, perl = TRUE) + expect_false(identical(js, js_future)) # confirm the substitution matched + expect_error(createBARTModelFromJsonString(js_future), "schema_version") +}) + +# =========================================================================== +# v1 (unified-envelope) golden-fixture matrix snapshot tests +# =========================================================================== +# +# Matrix: {bart, bcf} x {numeric, categorical} x {no-rfx, rfx}. These lock the +# on-disk schema_version=1 format directly (named forest keys, +# covariate_preprocessor, and rfx_unique_group_ids relocated into +# random_effects), whereas the legacy v0 fixtures above guard the v0 -> v1 +# migration path. Regenerate with test/R/testthat/fixtures/generate_v1_fixtures.R. + +make_v1_covariates <- function(categorical, k, seed = 7) { + set.seed(seed) + X_num <- matrix(runif(k * 4), ncol = 4) + if (!categorical) { + return(X_num) + } + X <- data.frame(X_num) + X$cat <- factor( + sample(c("a", "b", "c"), k, replace = TRUE), + levels = c("a", "b", "c") + ) + X +} + +for (.categorical in c(FALSE, TRUE)) { + for (.rfx in c(FALSE, TRUE)) { + local({ + categorical <- .categorical + rfx <- .rfx + kind <- if (categorical) "categorical" else "numeric" + sfx <- if (rfx) "_rfx" else "" + + test_that(sprintf("BART v1 fixture loads and predicts (%s%s)", kind, sfx), { + skip_on_cran() + skip_if_not_installed("jsonlite") + obj <- read_fixture_json(sprintf("bart_%s%s_v1.json", kind, sfx)) + expect_equal(obj$schema_version, 1) + expect_equal(names(obj$forests), "mean_forest") + if (rfx) { + # rfx unique group ids live in the random_effects subfolder, not top-level + expect_true("rfx_unique_group_ids" %in% names(obj$random_effects)) + expect_false("rfx_unique_group_ids" %in% names(obj)) + } + m <- createBARTModelFromJsonString(write_json_string(obj)) + expect_equal(m$model_params$has_rfx, rfx) + k <- 12 + X <- make_v1_covariates(categorical, k) + args <- list(object = m, X = X) + if (rfx) { + args$rfx_group_ids <- (0:(k - 1)) %% 3 + args$rfx_basis <- matrix(1, nrow = k, ncol = 1) + } + preds <- do.call(predict, args) + expect_equal(nrow(preds$y_hat), k) + }) + + test_that(sprintf("BCF v1 fixture loads and predicts (%s%s)", kind, sfx), { + skip_on_cran() + skip_if_not_installed("jsonlite") + obj <- read_fixture_json(sprintf("bcf_%s%s_v1.json", kind, sfx)) + expect_equal(obj$schema_version, 1) + expect_true(all( + c("prognostic_forest", "treatment_forest") %in% names(obj$forests) + )) + if (rfx) { + expect_true("rfx_unique_group_ids" %in% names(obj$random_effects)) + } + m <- createBCFModelFromJsonString(write_json_string(obj)) + expect_equal(m$model_params$has_rfx, rfx) + k <- 12 + X <- make_v1_covariates(categorical, k) + set.seed(3) + Z <- rbinom(k, 1, 0.5) + pi <- rep(0.5, k) + args <- list(object = m, X = X, Z = Z, propensity = pi) + if (rfx) { + args$rfx_group_ids <- (0:(k - 1)) %% 3 + args$rfx_basis <- matrix(1, nrow = k, ncol = 1) + } + preds <- do.call(predict, args) + expect_equal(nrow(preds$y_hat), k) + }) + }) + } +} + +# =========================================================================== +# WS-E: cross-platform load gate +# =========================================================================== +# +# A cross-platform load succeeds for portable (all-numeric, integer-rfx) models +# and is refused with a clear error otherwise; same-platform loads ignore the +# flags. We relabel a fixture's `platform` as the other platform to drive the +# gate from within R's own suite (the gate peeks generic flags; the foreign +# preprocessor body is ignored). + +read_fixture_raw <- function(fixture_name) { + paste( + readLines(testthat::test_path("fixtures", fixture_name), warn = FALSE), + collapse = "" + ) +} + +as_foreign <- function(fixture_name, compat = NULL) { + cj <- createCppJsonString(read_fixture_raw(fixture_name)) + cj$erase_field("platform") + cj$add_string("platform", "python") + if (!is.null(compat)) { + cj$add_boolean( + "cross_platform_compatible", + compat, + subfolder_name = "random_effects" + ) + } + cj$return_json_string() +} + +test_that("numeric BART loads cross-platform and predicts", { + skip_on_cran() + skip_if_not_installed("jsonlite") + m <- createBARTModelFromJsonString(as_foreign("bart_numeric_v1.json")) + set.seed(0) + X <- matrix(runif(8 * 4), ncol = 4) + expect_equal(nrow(predict(m, X = X)$y_hat), 8) +}) + +test_that("numeric BCF loads cross-platform and predicts", { + skip_on_cran() + skip_if_not_installed("jsonlite") + m <- createBCFModelFromJsonString(as_foreign("bcf_numeric_v1.json")) + set.seed(0) + X <- matrix(runif(8 * 4), ncol = 4) + Z <- rbinom(8, 1, 0.5) + preds <- predict(m, X = X, Z = Z, propensity = rep(0.5, 8)) + expect_equal(nrow(preds$y_hat), 8) +}) + +test_that("non-portable models are refused cross-platform", { + skip_on_cran() + skip_if_not_installed("jsonlite") + expect_error( + createBARTModelFromJsonString(as_foreign("bart_categorical_v1.json")), + "non-numeric covariates" + ) + expect_error( + createBCFModelFromJsonString(as_foreign("bcf_categorical_v1.json")), + "non-numeric covariates" + ) + expect_error( + createBARTModelFromJsonString( + as_foreign("bart_numeric_rfx_v1.json", compat = FALSE) + ), + "random effects" + ) +}) + +test_that("same-platform categorical load is not refused", { + skip_on_cran() + skip_if_not_installed("jsonlite") + m <- createBARTModelFromJsonString(read_fixture_raw("bart_categorical_v1.json")) + expect_s3_class(m, "bartmodel") +}) diff --git a/test/R/testthat/test-serialization.R b/test/R/testthat/test-serialization.R index 1b31b918..6a6b6836 100644 --- a/test/R/testthat/test-serialization.R +++ b/test/R/testthat/test-serialization.R @@ -209,7 +209,7 @@ test_that("BCF JSON uses canonical field names (sigma2_init, b1_samples, b0_samp expect_false(grepl('"b_0_samples"', json_string)) }) -test_that("BCF JSON deserialization handles legacy field names with warnings", { +test_that("BCF pre-v1 (v0) JSON is migrated to canonical field names before parsing", { skip_on_cran() set.seed(2) @@ -230,28 +230,27 @@ test_that("BCF JSON deserialization handles legacy field names with warnings", { ) preds_orig <- predict(bcf_model, X_test, Z_test, pi_test) - # Simulate a legacy JSON by replacing canonical names with old names - json_new <- saveBCFModelToJsonString(bcf_model) - json_legacy <- gsub('"sigma2_init"', '"initial_sigma2"', json_new, fixed = TRUE) - json_legacy <- gsub('"b1_samples"', '"b_1_samples"', json_legacy, fixed = TRUE) - json_legacy <- gsub('"b0_samples"', '"b_0_samples"', json_legacy, fixed = TRUE) - - # Loading a legacy JSON should emit deprecation warnings for all renamed fields - all_warnings <- character(0) - withCallingHandlers( - bcf_legacy <- createBCFModelFromJsonString(json_legacy), - warning = function(w) { - all_warnings <<- c(all_warnings, conditionMessage(w)) - invokeRestart("muffleWarning") - } - ) - expect_true(any(grepl("initial_sigma2.*deprecated|deprecated.*initial_sigma2", all_warnings))) - expect_true(any(grepl("b_1_samples.*deprecated|b_0_samples.*deprecated|deprecated.*b_[01]_samples", all_warnings))) - - # Predictions must still match - preds_legacy <- predict(bcf_legacy, X_test, Z_test, pi_test) - expect_equal(rowMeans(preds_legacy[["y_hat"]]), rowMeans(preds_orig[["y_hat"]])) - expect_equal(rowMeans(preds_legacy[["tau_hat"]]), rowMeans(preds_orig[["tau_hat"]])) + # Synthesize a pre-v1 (schema_version 0) JSON: legacy parameter-trace field names + # (initial_sigma2, b_0_samples, b_1_samples) plus a v0 schema stamp. The v0 -> v1 + # migration owns renaming these back to the canonical names *before* the C++ layer + # parses them -- there is no read-time fallback, so the reloaded model must simply + # predict identically to the original. + json_v1 <- saveBCFModelToJsonString(bcf_model) + json_v0 <- gsub('"sigma2_init"', '"initial_sigma2"', json_v1, fixed = TRUE) + json_v0 <- gsub('"b1_samples"', '"b_1_samples"', json_v0, fixed = TRUE) + json_v0 <- gsub('"b0_samples"', '"b_0_samples"', json_v0, fixed = TRUE) + json_v0 <- gsub('"schema_version":1', '"schema_version":0', json_v0, fixed = TRUE) + + bcf_migrated <- createBCFModelFromJsonString(json_v0) + + # Migration must have restored the canonical adaptive-coding traces (else the C++ + # layer would see empty b0/b1 arrays and predict would read out of bounds). + expect_length(bcf_migrated$samples$b0_samples(), bcf_model$model_params$num_samples) + expect_length(bcf_migrated$samples$b1_samples(), bcf_model$model_params$num_samples) + + preds_migrated <- predict(bcf_migrated, X_test, Z_test, pi_test) + expect_equal(rowMeans(preds_migrated[["y_hat"]]), rowMeans(preds_orig[["y_hat"]])) + expect_equal(rowMeans(preds_migrated[["tau_hat"]]), rowMeans(preds_orig[["tau_hat"]])) }) # Fields that the BCF JSON contract must preserve across a round-trip. These are diff --git a/test/cpp/test_category_tracker.cpp b/test/cpp/test_category_tracker.cpp index 05f28b2a..4ef249be 100644 --- a/test/cpp/test_category_tracker.cpp +++ b/test/cpp/test_category_tracker.cpp @@ -5,22 +5,19 @@ #include #include #include -#include -#include #include TEST(CategorySampleTracker, BasicOperations) { // Create a vector of categorical data - std::vector category_data { - 3, 4, 3, 2, 2, 4, 3, 3, 3, 4, 3, 4 - }; + std::vector category_data{ + 3, 4, 3, 2, 2, 4, 3, 3, 3, 4, 3, 4}; // Create a CategorySamplerTracker StochTree::CategorySampleTracker category_tracker = StochTree::CategorySampleTracker(category_data); // Extract the label map std::map label_map = category_tracker.GetLabelMap(); - std::map expected_label_map {{2, 0}, {3, 1}, {4, 2}}; + std::map expected_label_map{{2, 0}, {3, 1}, {4, 2}}; // Check that the map was constructed as expected ASSERT_EQ(label_map[2], 0); diff --git a/test/cpp/test_cutpoints.cpp b/test/cpp/test_cutpoints.cpp index 5e71a498..42fbe388 100644 --- a/test/cpp/test_cutpoints.cpp +++ b/test/cpp/test_cutpoints.cpp @@ -3,7 +3,6 @@ #include #include #include -#include #include #include diff --git a/test/cpp/test_data.cpp b/test/cpp/test_data.cpp index 62215cff..18a23a08 100644 --- a/test/cpp/test_data.cpp +++ b/test/cpp/test_data.cpp @@ -8,8 +8,6 @@ #include #include #include -#include -#include TEST(Data, ReadFromSmallDatasetRowMajor) { // Load test data @@ -25,14 +23,14 @@ TEST(Data, ReadFromSmallDatasetRowMajor) { dataset.AddCovariates(test_dataset.covariates.data(), n, test_dataset.x_cols, test_dataset.row_major); dataset.AddBasis(test_dataset.omega.data(), test_dataset.n, test_dataset.omega_cols, test_dataset.row_major); StochTree::ColumnVector residual = StochTree::ColumnVector(test_dataset.outcome.data(), n); - + // Compute average value for each feature, compared to their known values std::vector total; std::vector average; total.resize(p, 0.); average.resize(p, 0.); for (int j = 0; j < p; j++) { - for (data_size_t i = 0; i < n; i++) { + for (data_size_t i = 0; i < n; i++) { total[j] += dataset.CovariateValue(i, j); } } @@ -60,14 +58,14 @@ TEST(Data, ReadFromMediumDatasetRowMajor) { dataset.AddCovariates(test_dataset.covariates.data(), n, test_dataset.x_cols, test_dataset.row_major); dataset.AddBasis(test_dataset.omega.data(), test_dataset.n, test_dataset.omega_cols, test_dataset.row_major); StochTree::ColumnVector residual = StochTree::ColumnVector(test_dataset.outcome.data(), n); - + // Compute average value for each feature, compared to their known values std::vector total; std::vector average; total.resize(p, 0.); average.resize(p, 0.); for (int j = 0; j < p; j++) { - for (data_size_t i = 0; i < n; i++) { + for (data_size_t i = 0; i < n; i++) { total[j] += dataset.CovariateValue(i, j); } } @@ -80,4 +78,3 @@ TEST(Data, ReadFromMediumDatasetRowMajor) { EXPECT_NEAR(0.4863596, average[3], 0.0001); EXPECT_NEAR(0.4413101, average[4], 0.0001); } - diff --git a/test/cpp/test_forest.cpp b/test/cpp/test_forest.cpp index d870997c..cabc3289 100644 --- a/test/cpp/test_forest.cpp +++ b/test/cpp/test_forest.cpp @@ -8,8 +8,6 @@ #include #include #include -#include -#include TEST(Forest, UnivariateForestConstruction) { int num_trees = 2; @@ -45,13 +43,13 @@ TEST(Forest, UnivariateForestMerge) { dataset.AddCovariates(test_dataset.covariates.data(), n, test_dataset.x_cols, test_dataset.row_major); dataset.AddBasis(test_dataset.omega.data(), test_dataset.n, test_dataset.omega_cols, test_dataset.row_major); StochTree::ColumnVector residual = StochTree::ColumnVector(test_dataset.outcome.data(), n); - + // Create a small ensemble int output_dim = 1; int num_trees = 2; bool is_leaf_constant = true; StochTree::TreeEnsemble ensemble1(num_trees, output_dim, is_leaf_constant); - + // Create another small ensemble StochTree::TreeEnsemble ensemble2(num_trees, output_dim, is_leaf_constant); @@ -63,8 +61,8 @@ TEST(Forest, UnivariateForestMerge) { tree->ExpandNode(0, 1, tree_split, -2.5, 2.5); // Run predict on the supplied covariates and check the result for the first forest - std::vector result(n*output_dim); - std::vector expected_pred = {7.5,2.5,-7.5,7.5,7.5,7.5,2.5,7.5,-7.5,-2.5}; + std::vector result(n * output_dim); + std::vector expected_pred = {7.5, 2.5, -7.5, 7.5, 7.5, 7.5, 2.5, 7.5, -7.5, -2.5}; ensemble1.PredictInplace(dataset.GetCovariates(), result, 0); for (int i = 0; i < n; i++) { ASSERT_NEAR(expected_pred[i], result[i], 0.01); @@ -79,8 +77,8 @@ TEST(Forest, UnivariateForestMerge) { tree->ExpandNode(0, 2, tree_split, -0.5, 0.5); // Run predict on the supplied covariates and check the result for the second forest - result = std::vector(n*output_dim); - expected_pred = std::vector{1.5,-1.5,-0.5,-0.5,1.5,1.5,-1.5,1.5,-0.5,-1.5}; + result = std::vector(n * output_dim); + expected_pred = std::vector{1.5, -1.5, -0.5, -0.5, 1.5, 1.5, -1.5, 1.5, -0.5, -1.5}; ensemble2.PredictInplace(dataset.GetCovariates(), result, 0); for (int i = 0; i < n; i++) { ASSERT_NEAR(expected_pred[i], result[i], 0.01); @@ -88,8 +86,8 @@ TEST(Forest, UnivariateForestMerge) { // Merge the second forest into the first ensemble1.MergeForest(ensemble2); - result = std::vector(n*output_dim); - expected_pred = std::vector{9,1,-8,7,9,9,1,9,-8,-4}; + result = std::vector(n * output_dim); + expected_pred = std::vector{9, 1, -8, 7, 9, 9, 1, 9, -8, -4}; ensemble1.PredictInplace(dataset.GetCovariates(), result, 0); for (int i = 0; i < n; i++) { ASSERT_NEAR(expected_pred[i], result[i], 0.01); @@ -108,13 +106,13 @@ TEST(Forest, UnivariateForestAdd) { dataset.AddCovariates(test_dataset.covariates.data(), n, test_dataset.x_cols, test_dataset.row_major); dataset.AddBasis(test_dataset.omega.data(), test_dataset.n, test_dataset.omega_cols, test_dataset.row_major); StochTree::ColumnVector residual = StochTree::ColumnVector(test_dataset.outcome.data(), n); - + // Create a small ensemble int output_dim = 1; int num_trees = 2; bool is_leaf_constant = true; StochTree::TreeEnsemble ensemble1(num_trees, output_dim, is_leaf_constant); - + // Create another small ensemble StochTree::TreeEnsemble ensemble2(num_trees, output_dim, is_leaf_constant); @@ -126,8 +124,8 @@ TEST(Forest, UnivariateForestAdd) { tree->ExpandNode(0, 1, tree_split, -2.5, 2.5); // Run predict on the supplied covariates and check the result for the first forest - std::vector result(n*output_dim); - std::vector expected_pred = {7.5,2.5,-7.5,7.5,7.5,7.5,2.5,7.5,-7.5,-2.5}; + std::vector result(n * output_dim); + std::vector expected_pred = {7.5, 2.5, -7.5, 7.5, 7.5, 7.5, 2.5, 7.5, -7.5, -2.5}; ensemble1.PredictInplace(dataset.GetCovariates(), result, 0); for (int i = 0; i < n; i++) { ASSERT_NEAR(expected_pred[i], result[i], 0.01); @@ -137,8 +135,8 @@ TEST(Forest, UnivariateForestAdd) { ensemble1.AddValueToLeaves(1.0); // Run predict on the supplied covariates and check the result for the first forest - result = std::vector(n*output_dim); - expected_pred = std::vector{9.5,4.5,-5.5,9.5,9.5,9.5,4.5,9.5,-5.5,-0.5}; + result = std::vector(n * output_dim); + expected_pred = std::vector{9.5, 4.5, -5.5, 9.5, 9.5, 9.5, 4.5, 9.5, -5.5, -0.5}; ensemble1.PredictInplace(dataset.GetCovariates(), result, 0); for (int i = 0; i < n; i++) { ASSERT_NEAR(expected_pred[i], result[i], 0.01); @@ -153,8 +151,8 @@ TEST(Forest, UnivariateForestAdd) { tree->ExpandNode(0, 2, tree_split, -0.5, 0.5); // Run predict on the supplied covariates and check the result for the second forest - result = std::vector(n*output_dim); - expected_pred = std::vector{1.5,-1.5,-0.5,-0.5,1.5,1.5,-1.5,1.5,-0.5,-1.5}; + result = std::vector(n * output_dim); + expected_pred = std::vector{1.5, -1.5, -0.5, -0.5, 1.5, 1.5, -1.5, 1.5, -0.5, -1.5}; ensemble2.PredictInplace(dataset.GetCovariates(), result, 0); for (int i = 0; i < n; i++) { ASSERT_NEAR(expected_pred[i], result[i], 0.01); @@ -164,8 +162,8 @@ TEST(Forest, UnivariateForestAdd) { ensemble2.AddValueToLeaves(-1.0); // Run predict on the supplied covariates and check the result for the first forest - result = std::vector(n*output_dim); - expected_pred = std::vector{-0.5,-3.5,-2.5,-2.5,-0.5,-0.5,-3.5,-0.5,-2.5,-3.5}; + result = std::vector(n * output_dim); + expected_pred = std::vector{-0.5, -3.5, -2.5, -2.5, -0.5, -0.5, -3.5, -0.5, -2.5, -3.5}; ensemble2.PredictInplace(dataset.GetCovariates(), result, 0); for (int i = 0; i < n; i++) { ASSERT_NEAR(expected_pred[i], result[i], 0.01); @@ -173,8 +171,8 @@ TEST(Forest, UnivariateForestAdd) { // Merge the second forest into the first ensemble1.MergeForest(ensemble2); - result = std::vector(n*output_dim); - expected_pred = std::vector{9,1,-8,7,9,9,1,9,-8,-4}; + result = std::vector(n * output_dim); + expected_pred = std::vector{9, 1, -8, 7, 9, 9, 1, 9, -8, -4}; ensemble1.PredictInplace(dataset.GetCovariates(), result, 0); for (int i = 0; i < n; i++) { ASSERT_NEAR(expected_pred[i], result[i], 0.01); @@ -182,11 +180,10 @@ TEST(Forest, UnivariateForestAdd) { // Merge the second forest into the first ensemble1.MultiplyLeavesByValue(2.0); - result = std::vector(n*output_dim); - expected_pred = std::vector{18,2,-16,14,18,18,2,18,-16,-8}; + result = std::vector(n * output_dim); + expected_pred = std::vector{18, 2, -16, 14, 18, 18, 2, 18, -16, -8}; ensemble1.PredictInplace(dataset.GetCovariates(), result, 0); for (int i = 0; i < n; i++) { ASSERT_NEAR(expected_pred[i], result[i], 0.01); } } - diff --git a/test/cpp/test_json.cpp b/test/cpp/test_json.cpp index 77a28faf..8da2c8e6 100644 --- a/test/cpp/test_json.cpp +++ b/test/cpp/test_json.cpp @@ -4,10 +4,12 @@ */ #include #include +#include +#include +#include #include #include #include -#include #include TEST(Json, TreeUnivariateLeaf) { @@ -15,7 +17,7 @@ TEST(Json, TreeUnivariateLeaf) { StochTree::Tree tree; StochTree::TreeSplit split; tree.Init(1); - + // Perform three splits split = StochTree::TreeSplit(0.5); tree.ExpandNode(0, 0, split, 0., 0.); @@ -23,10 +25,10 @@ TEST(Json, TreeUnivariateLeaf) { tree.ExpandNode(1, 1, split, 0., 0.); split = StochTree::TreeSplit(0.6); tree.ExpandNode(3, 2, split, 0., 0.); - + // Prune node 3 to a leaf tree.CollapseToLeaf(3, 0.); - + // Write to json nlohmann::json tree_json = tree.to_json(); @@ -43,20 +45,20 @@ TEST(Json, TreeUnivariateLeafCategoricalSplit) { StochTree::Tree tree; StochTree::TreeSplit split; tree.Init(1); - + // Perform three splits - std::vector split_categories_1{1,3,5,7}; + std::vector split_categories_1{1, 3, 5, 7}; split = StochTree::TreeSplit(split_categories_1); tree.ExpandNode(0, 0, split, 0., 0.); - std::vector split_categories_2{2,3,5}; + std::vector split_categories_2{2, 3, 5}; split = StochTree::TreeSplit(split_categories_2); tree.ExpandNode(1, 1, split, 0., 0.); split = StochTree::TreeSplit(0.6); tree.ExpandNode(3, 2, split, 0., 0.); - + // Prune node 3 to a leaf tree.CollapseToLeaf(3, 0.); - + // Write to json nlohmann::json tree_json = tree.to_json(); @@ -80,7 +82,7 @@ TEST(Json, TreeMultivariateLeaf) { std::vector leaf_values5(tree_dim, 345235636.4); std::vector leaf_values6(tree_dim, 10023.1); tree.Init(tree_dim); - + // Perform three splits split = StochTree::TreeSplit(0.5); tree.ExpandNode(0, 0, split, leaf_values1, leaf_values2); @@ -88,10 +90,10 @@ TEST(Json, TreeMultivariateLeaf) { tree.ExpandNode(1, 1, split, leaf_values3, leaf_values4); split = StochTree::TreeSplit(0.6); tree.ExpandNode(1, 1, split, leaf_values5, leaf_values6); - + // Prune node 3 to a leaf tree.CollapseToLeaf(3, leaf_values3); - + // Write to json nlohmann::json tree_json = tree.to_json(); @@ -104,6 +106,191 @@ TEST(Json, TreeMultivariateLeaf) { ASSERT_EQ(tree, tree_parsed); } +TEST(Json, BARTSamplesRoundTrip) { + // Build a minimal BARTSamples: a 2-sample mean forest container + parameter traces + scalars. + StochTree::BARTSamples samples; + samples.mean_forests = std::make_unique( + /*num_trees=*/1, /*output_dimension=*/1, /*is_leaf_constant=*/true, /*is_exponentiated=*/false); + // Build two retained samples via AddSample (InitializeRoot does NOT count as a retained sample -- + // it sets num_samples_ back to 0 -- so two AddSample calls give a consistent 2-sample container). + StochTree::TreeEnsemble ens0(/*num_trees=*/1, /*output_dimension=*/1, /*is_leaf_constant=*/true); + ens0.SetLeafValue(0.5); + samples.mean_forests->AddSample(ens0); // sample 0 + StochTree::TreeEnsemble ens1(/*num_trees=*/1, /*output_dimension=*/1, /*is_leaf_constant=*/true); + ens1.SetLeafValue(-0.25); + samples.mean_forests->AddSample(ens1); // sample 1 + samples.global_error_variance_samples = {1.1, 2.2}; + samples.leaf_scale_samples = {0.3, 0.4}; + samples.y_bar = 3.14; + samples.y_std = 2.71; + samples.num_samples = 2; + + // Round-trip through the samples-owned subtree + nlohmann::json obj; + samples.AppendToJson(obj); + StochTree::BARTSamples restored; + restored.FromJson(obj); + + // Round-trip must reproduce the original samples exactly (compare against the source object, + // not re-typed literals -- ToJson is const, so `samples` is untouched and is the source of truth). + EXPECT_EQ(restored.num_samples, samples.num_samples); + EXPECT_DOUBLE_EQ(restored.y_bar, samples.y_bar); + EXPECT_DOUBLE_EQ(restored.y_std, samples.y_std); + // Every per-draw parameter trace should carry exactly one entry per retained sample. + ASSERT_EQ(restored.global_error_variance_samples.size(), static_cast(samples.num_samples)); + ASSERT_EQ(restored.leaf_scale_samples.size(), static_cast(samples.num_samples)); + EXPECT_EQ(restored.global_error_variance_samples, samples.global_error_variance_samples); + EXPECT_EQ(restored.leaf_scale_samples, samples.leaf_scale_samples); + + // Mean forest survives (byte-level: re-serialized JSON identical); variance forest stays absent + ASSERT_NE(restored.mean_forests, nullptr); + EXPECT_EQ(restored.mean_forests->NumSamples(), samples.mean_forests->NumSamples()); + EXPECT_EQ(samples.mean_forests->to_json(), restored.mean_forests->to_json()); + EXPECT_EQ(restored.variance_forests, nullptr); + + // num_forests counter matches the number of forests actually written (mean only -> 1) + EXPECT_EQ(obj.at("num_forests").get(), 1); +} + +TEST(Json, BCFSamplesRoundTrip) { + // Build a minimal BCFSamples: 2-sample prognostic (mu) + treatment (tau) forests, the full set of + // BCF parameter traces (univariate treatment), and scalars including treatment_dim. + StochTree::BCFSamples samples; + samples.mu_forests = std::make_unique(1, 1, true, false); + samples.tau_forests = std::make_unique(1, 1, false, false); + StochTree::TreeEnsemble mu0(1, 1, true), mu1(1, 1, true); + mu0.SetLeafValue(0.5); + mu1.SetLeafValue(-0.25); + samples.mu_forests->AddSample(mu0); + samples.mu_forests->AddSample(mu1); + StochTree::TreeEnsemble tau0(1, 1, false), tau1(1, 1, false); + tau0.SetLeafValue(1.0); + tau1.SetLeafValue(0.8); + samples.tau_forests->AddSample(tau0); + samples.tau_forests->AddSample(tau1); + + samples.global_error_variance_samples = {1.1, 2.2}; + samples.leaf_scale_mu_samples = {0.3, 0.4}; + samples.leaf_scale_tau_samples = {0.05, 0.06}; + samples.b0_samples = {-0.5, -0.4}; + samples.b1_samples = {0.5, 0.6}; + samples.tau_0_samples = {0.1, 0.2}; // univariate: treatment_dim x num_samples = 1 x 2 + samples.y_bar = 3.14; + samples.y_std = 2.71; + samples.num_samples = 2; + samples.treatment_dim = 1; + + // Round-trip through the samples-owned subtree + nlohmann::json obj; + samples.AppendToJson(obj); + StochTree::BCFSamples restored; + restored.FromJson(obj); + + // Scalars survive exactly (compared against the source object, not re-typed literals) + EXPECT_EQ(restored.num_samples, samples.num_samples); + EXPECT_EQ(restored.treatment_dim, samples.treatment_dim); + EXPECT_DOUBLE_EQ(restored.y_bar, samples.y_bar); + EXPECT_DOUBLE_EQ(restored.y_std, samples.y_std); + + // Every per-draw parameter trace carries one entry per retained sample, and round-trips exactly + ASSERT_EQ(restored.global_error_variance_samples.size(), static_cast(samples.num_samples)); + ASSERT_EQ(restored.leaf_scale_mu_samples.size(), static_cast(samples.num_samples)); + ASSERT_EQ(restored.leaf_scale_tau_samples.size(), static_cast(samples.num_samples)); + ASSERT_EQ(restored.b0_samples.size(), static_cast(samples.num_samples)); + ASSERT_EQ(restored.b1_samples.size(), static_cast(samples.num_samples)); + ASSERT_EQ(restored.tau_0_samples.size(), + static_cast(samples.num_samples * samples.treatment_dim)); + EXPECT_EQ(restored.global_error_variance_samples, samples.global_error_variance_samples); + EXPECT_EQ(restored.leaf_scale_mu_samples, samples.leaf_scale_mu_samples); + EXPECT_EQ(restored.leaf_scale_tau_samples, samples.leaf_scale_tau_samples); + EXPECT_EQ(restored.b0_samples, samples.b0_samples); + EXPECT_EQ(restored.b1_samples, samples.b1_samples); + EXPECT_EQ(restored.tau_0_samples, samples.tau_0_samples); + + // Both forests survive (byte-level: re-serialized JSON identical); no variance forest present + ASSERT_NE(restored.mu_forests, nullptr); + ASSERT_NE(restored.tau_forests, nullptr); + EXPECT_EQ(restored.mu_forests->NumSamples(), samples.mu_forests->NumSamples()); + EXPECT_EQ(restored.tau_forests->NumSamples(), samples.tau_forests->NumSamples()); + EXPECT_EQ(samples.mu_forests->to_json(), restored.mu_forests->to_json()); + EXPECT_EQ(samples.tau_forests->to_json(), restored.tau_forests->to_json()); + EXPECT_EQ(restored.variance_forests, nullptr); + + // num_forests counter matches the number written (prognostic + treatment -> 2) + EXPECT_EQ(obj.at("num_forests").get(), 2); +} + +TEST(Json, BARTSamplesMerge) { + // Two single-draw chains with matching structure/standardization. + auto make_chain = [](double mu_leaf, double global_var, double leaf_scale) { + StochTree::BARTSamples s; + s.mean_forests = std::make_unique(1, 1, true, false); + StochTree::TreeEnsemble e(1, 1, true); + e.SetLeafValue(mu_leaf); + s.mean_forests->AddSample(e); + s.global_error_variance_samples = {global_var}; + s.leaf_scale_samples = {leaf_scale}; + s.num_samples = 1; + s.y_bar = 1.0; + s.y_std = 2.0; + return s; + }; + StochTree::BARTSamples a = make_chain(0.5, 1.1, 0.3); + StochTree::BARTSamples b = make_chain(-0.25, 2.2, 0.4); + + a.Merge(b); + + // Draw counts add; draw order is preserved (a's draw, then b's) + EXPECT_EQ(a.num_samples, 2); + ASSERT_NE(a.mean_forests, nullptr); + EXPECT_EQ(a.mean_forests->NumSamples(), 2); + // Merged forest's second sample is a deep copy of b's first sample + EXPECT_EQ(a.mean_forests->GetEnsemble(1)->to_json(), b.mean_forests->GetEnsemble(0)->to_json()); + ASSERT_EQ(a.global_error_variance_samples.size(), 2u); + EXPECT_DOUBLE_EQ(a.global_error_variance_samples[0], 1.1); + EXPECT_DOUBLE_EQ(a.global_error_variance_samples[1], 2.2); + ASSERT_EQ(a.leaf_scale_samples.size(), 2u); + EXPECT_DOUBLE_EQ(a.leaf_scale_samples[0], 0.3); + EXPECT_DOUBLE_EQ(a.leaf_scale_samples[1], 0.4); +} + +TEST(Json, BCFSamplesMerge) { + auto make_chain = [](double mu_leaf, double tau_leaf, double gv, double tau0) { + StochTree::BCFSamples s; + s.mu_forests = std::make_unique(1, 1, true, false); + s.tau_forests = std::make_unique(1, 1, false, false); + StochTree::TreeEnsemble mu(1, 1, true); + mu.SetLeafValue(mu_leaf); + s.mu_forests->AddSample(mu); + StochTree::TreeEnsemble tau(1, 1, false); + tau.SetLeafValue(tau_leaf); + s.tau_forests->AddSample(tau); + s.global_error_variance_samples = {gv}; + s.tau_0_samples = {tau0}; + s.num_samples = 1; + s.treatment_dim = 1; + s.y_bar = 1.0; + s.y_std = 2.0; + return s; + }; + StochTree::BCFSamples a = make_chain(0.5, 1.0, 1.1, 0.1); + StochTree::BCFSamples b = make_chain(-0.25, 0.8, 2.2, 0.2); + + a.Merge(b); + + EXPECT_EQ(a.num_samples, 2); + ASSERT_NE(a.mu_forests, nullptr); + ASSERT_NE(a.tau_forests, nullptr); + EXPECT_EQ(a.mu_forests->NumSamples(), 2); + EXPECT_EQ(a.tau_forests->NumSamples(), 2); + EXPECT_EQ(a.tau_forests->GetEnsemble(1)->to_json(), b.tau_forests->GetEnsemble(0)->to_json()); + ASSERT_EQ(a.global_error_variance_samples.size(), 2u); + EXPECT_DOUBLE_EQ(a.global_error_variance_samples[1], 2.2); + ASSERT_EQ(a.tau_0_samples.size(), 2u); + EXPECT_DOUBLE_EQ(a.tau_0_samples[0], 0.1); + EXPECT_DOUBLE_EQ(a.tau_0_samples[1], 0.2); +} + TEST(Json, TreeMultivariateLeafCategoricalSplit) { // Initialize tree StochTree::Tree tree; @@ -116,20 +303,20 @@ TEST(Json, TreeMultivariateLeafCategoricalSplit) { std::vector leaf_values5(tree_dim, 345235636.4); std::vector leaf_values6(tree_dim, 10023.1); tree.Init(tree_dim); - + // Perform three splits - std::vector split_categories_1{1,3,5,7}; + std::vector split_categories_1{1, 3, 5, 7}; split = StochTree::TreeSplit(split_categories_1); tree.ExpandNode(0, 0, split, leaf_values1, leaf_values2); - std::vector split_categories_2{2,3,5}; + std::vector split_categories_2{2, 3, 5}; split = StochTree::TreeSplit(split_categories_2); tree.ExpandNode(1, 1, split, leaf_values3, leaf_values4); split = StochTree::TreeSplit(0.6); tree.ExpandNode(1, 1, split, leaf_values5, leaf_values6); - + // Prune node 3 to a leaf tree.CollapseToLeaf(3, leaf_values3); - + // Write to json nlohmann::json tree_json = tree.to_json(); diff --git a/test/cpp/test_linear_regression.cpp b/test/cpp/test_linear_regression.cpp new file mode 100644 index 00000000..92ee8610 --- /dev/null +++ b/test/cpp/test_linear_regression.cpp @@ -0,0 +1,233 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +TEST(LinearRegression, UnivariateDegeneratePosteriorMeanCorrectness) { + // Test that the posterior mean of the regression coefficient is correct in a degenerate case where the outcome has no variance + // and the prior variance is nearly infinite (i.e. the posterior mean should equal the OLS estimate). + + // Generate data + std::mt19937 gen(1234); + int n = 100; + std::vector x(n, 0.0); + std::vector y(n, 0.0); + for (int i = 0; i < n; i++) { + x[i] = StochTree::standard_uniform_draw_53bit(gen); + y[i] = 2.0 * x[i]; + } + double sigma2 = 1e-6; // near-zero outcome variance + double tau2 = 1e6; // near-infinite prior variance + + // Sample from the regression model + int num_samples = 1000; + std::vector beta_samples(num_samples); + for (int i = 0; i < num_samples; i++) { + beta_samples[i] = StochTree::sample_univariate_gaussian_regression_coefficient(y.data(), x.data(), sigma2, tau2, n, gen); + } + + // Check posterior mean is close to true value (which should also be the OLS estimate without noise) + double beta_mean = std::accumulate(beta_samples.begin(), beta_samples.end(), 0.0) / num_samples; + double ols_estimate = 2.0; + EXPECT_NEAR(beta_mean, ols_estimate, 1e-2); +} + +TEST(LinearRegression, UnivariatePosteriorMeanCorrectness) { + // Test that the sampled regression coefficients average out close to the expected posterior mean with enough samples + + // Generate data + std::mt19937 gen(1234); + int n = 100; + std::vector x(n, 0.0); + std::vector y(n, 0.0); + for (int i = 0; i < n; i++) { + x[i] = StochTree::standard_uniform_draw_53bit(gen); + y[i] = 2.0 * x[i] + StochTree::sample_standard_normal(0.0, 0.1, gen); + } + double sigma2 = 1; + double tau2 = 1; + + // Compute the "true" posterior mean analytically for comparison + double sum_xx = 0.0; + double sum_yx = 0.0; + for (int i = 0; i < n; i++) { + sum_xx += x[i] * x[i]; + sum_yx += y[i] * x[i]; + } + double post_mean = (sum_yx * sigma2) / (sum_xx * tau2 + sigma2); + + // Draw many samples from the posterior and compute their average + int num_samples = 10000; + std::vector beta_samples(num_samples); + for (int i = 0; i < num_samples; i++) { + beta_samples[i] = StochTree::sample_univariate_gaussian_regression_coefficient(y.data(), x.data(), sigma2, tau2, n, gen); + } + double beta_mean = std::accumulate(beta_samples.begin(), beta_samples.end(), 0.0) / num_samples; + EXPECT_NEAR(beta_mean, post_mean, 1e-2); +} + +TEST(LinearRegression, BivariatePosteriorMeanCorrectness) { + // Test that the sampled regression coefficients average out close to the expected posterior mean with enough samples + + // Generate data + std::mt19937 gen(1234); + int n = 100; + std::vector x1(n, 0.0); + std::vector x2(n, 0.0); + std::vector y(n, 0.0); + for (int i = 0; i < n; i++) { + x1[i] = StochTree::standard_uniform_draw_53bit(gen); + x2[i] = StochTree::standard_uniform_draw_53bit(gen); + y[i] = 2.0 * x1[i] + 3.0 * x2[i] + StochTree::sample_standard_normal(0.0, 0.1, gen); + } + double sigma2 = 1; + double prior_variance_11 = 1; + double prior_variance_12 = 0.5; + double prior_variance_22 = 1; + + // Compute the "true" posterior mean analytically for comparison + double det_prior_var = prior_variance_11 * prior_variance_22 - prior_variance_12 * prior_variance_12; + double inv_prior_var_11 = prior_variance_22 / det_prior_var; + double inv_prior_var_12 = -prior_variance_12 / det_prior_var; + double inv_prior_var_22 = prior_variance_11 / det_prior_var; + double sum_x1x1 = 0.0; + double sum_x2x2 = 0.0; + double sum_x1x2 = 0.0; + double sum_yx1 = 0.0; + double sum_yx2 = 0.0; + for (int i = 0; i < n; i++) { + sum_x1x1 += x1[i] * x1[i]; + sum_x2x2 += x2[i] * x2[i]; + sum_x1x2 += x1[i] * x2[i]; + sum_yx1 += y[i] * x1[i]; + sum_yx2 += y[i] * x2[i]; + } + double post_var_pre_inv_11 = inv_prior_var_11 + sum_x1x1 / sigma2; + double post_var_pre_inv_12 = inv_prior_var_12 + sum_x1x2 / sigma2; + double post_var_pre_inv_22 = inv_prior_var_22 + sum_x2x2 / sigma2; + double det_post_var_pre_inv = post_var_pre_inv_11 * post_var_pre_inv_22 - post_var_pre_inv_12 * post_var_pre_inv_12; + double post_var_11 = post_var_pre_inv_22 / det_post_var_pre_inv; + double post_var_12 = -post_var_pre_inv_12 / det_post_var_pre_inv; + double post_var_22 = post_var_pre_inv_11 / det_post_var_pre_inv; + double post_mean_1 = post_var_11 * (sum_yx1 / sigma2) + post_var_12 * (sum_yx2 / sigma2); + double post_mean_2 = post_var_12 * (sum_yx1 / sigma2) + post_var_22 * (sum_yx2 / sigma2); + + // Draw many samples from the posterior and compute their average + int num_samples = 10000; + double beta_mean_1_sum = 0.0; + double beta_mean_2_sum = 0.0; + std::vector beta_samples(num_samples * 2); + for (int i = 0; i < num_samples; i++) { + StochTree::sample_general_bivariate_gaussian_regression_coefficients(beta_samples.data() + 2 * i, y.data(), x1.data(), x2.data(), sigma2, prior_variance_11, prior_variance_12, prior_variance_22, n, gen); + beta_mean_1_sum += beta_samples[2 * i]; + beta_mean_2_sum += beta_samples[2 * i + 1]; + } + double beta_mean_1 = beta_mean_1_sum / num_samples; + double beta_mean_2 = beta_mean_2_sum / num_samples; + EXPECT_NEAR(beta_mean_1, post_mean_1, 1e-2); + EXPECT_NEAR(beta_mean_2, post_mean_2, 1e-2); +} + +TEST(LinearRegression, BivariateMatchWhenDiagonalPrior) { + // Test that the sampled regression coefficients for the general bivariate and specialized diagonal bivariate samplers are close to each other with enough samples when the covariance is diagonal + + // Generate data + std::mt19937 gen(1234); + int n = 100; + std::vector x1(n, 0.0); + std::vector x2(n, 0.0); + std::vector y(n, 0.0); + for (int i = 0; i < n; i++) { + x1[i] = StochTree::standard_uniform_draw_53bit(gen); + x2[i] = StochTree::standard_uniform_draw_53bit(gen); + y[i] = 2.0 * x1[i] + 3.0 * x2[i] + StochTree::sample_standard_normal(0.0, 0.1, gen); + } + double sigma2 = 1; + double prior_variance_11 = 1; + double prior_variance_12 = 0; + double prior_variance_22 = 1; + + // Draw many samples from the posterior and compute their average + int num_samples = 10000; + double beta_mean_1_sum_general = 0.0; + double beta_mean_2_sum_general = 0.0; + double beta_mean_1_sum_diagonal = 0.0; + double beta_mean_2_sum_diagonal = 0.0; + std::vector beta_samples_general(num_samples * 2); + std::vector beta_samples_diagonal(num_samples * 2); + for (int i = 0; i < num_samples; i++) { + StochTree::sample_general_bivariate_gaussian_regression_coefficients(beta_samples_general.data() + 2 * i, y.data(), x1.data(), x2.data(), sigma2, prior_variance_11, prior_variance_12, prior_variance_22, n, gen); + StochTree::sample_diagonal_bivariate_gaussian_regression_coefficients(beta_samples_diagonal.data() + 2 * i, y.data(), x1.data(), x2.data(), sigma2, prior_variance_11, prior_variance_22, n, gen); + beta_mean_1_sum_general += beta_samples_general[2 * i]; + beta_mean_2_sum_general += beta_samples_general[2 * i + 1]; + beta_mean_1_sum_diagonal += beta_samples_diagonal[2 * i]; + beta_mean_2_sum_diagonal += beta_samples_diagonal[2 * i + 1]; + } + double beta_mean_1_general = beta_mean_1_sum_general / num_samples; + double beta_mean_2_general = beta_mean_2_sum_general / num_samples; + double beta_mean_1_diagonal = beta_mean_1_sum_diagonal / num_samples; + double beta_mean_2_diagonal = beta_mean_2_sum_diagonal / num_samples; + EXPECT_NEAR(beta_mean_1_general, beta_mean_1_diagonal, 1e-2); + EXPECT_NEAR(beta_mean_2_general, beta_mean_2_diagonal, 1e-2); +} + +TEST(LinearRegression, MultivariateBivariateMatch) { + // Test that the sampled regression coefficients for the general bivariate and multivariate samplers are close to each other with enough samples when covariates are bivariate + + // Generate data + std::mt19937 gen(1234); + int n = 100; + std::vector x1(n, 0.0); + std::vector x2(n, 0.0); + std::vector y(n, 0.0); + Eigen::VectorXd y_eigen(n); + Eigen::MatrixXd X_eigen(n, 2); + double x1_elem, x2_elem, y_elem; + for (int i = 0; i < n; i++) { + x1_elem = StochTree::standard_uniform_draw_53bit(gen); + x2_elem = StochTree::standard_uniform_draw_53bit(gen); + y_elem = 2.0 * x1_elem + 3.0 * x2_elem + StochTree::sample_standard_normal(0.0, 0.1, gen); + x1[i] = x1_elem; + x2[i] = x2_elem; + y[i] = y_elem; + y_eigen(i) = y_elem; + X_eigen(i, 0) = x1_elem; + X_eigen(i, 1) = x2_elem; + } + double sigma2 = 1; + double prior_variance_11 = 1; + double prior_variance_12 = 0; + double prior_variance_22 = 1; + Eigen::MatrixXd prior_variance(2, 2); + prior_variance(0, 0) = prior_variance_11; + prior_variance(0, 1) = prior_variance_12; + prior_variance(1, 0) = prior_variance_12; + prior_variance(1, 1) = prior_variance_22; + + // Draw many samples from the posterior and compute their average + int num_samples = 10000; + double beta_mean_1_sum_bivariate = 0.0; + double beta_mean_2_sum_bivariate = 0.0; + double beta_mean_1_sum_multivariate = 0.0; + double beta_mean_2_sum_multivariate = 0.0; + std::vector beta_samples_bivariate(num_samples * 2); + Eigen::VectorXd beta(2); + for (int i = 0; i < num_samples; i++) { + StochTree::sample_general_bivariate_gaussian_regression_coefficients(beta_samples_bivariate.data() + 2 * i, y.data(), x1.data(), x2.data(), sigma2, prior_variance_11, prior_variance_12, prior_variance_22, n, gen); + beta = StochTree::sample_general_gaussian_regression_coefficients(y_eigen, X_eigen, sigma2, prior_variance, n, gen); + beta_mean_1_sum_bivariate += beta_samples_bivariate[2 * i]; + beta_mean_2_sum_bivariate += beta_samples_bivariate[2 * i + 1]; + beta_mean_1_sum_multivariate += beta(0); + beta_mean_2_sum_multivariate += beta(1); + } + double beta_mean_1_bivariate = beta_mean_1_sum_bivariate / num_samples; + double beta_mean_2_bivariate = beta_mean_2_sum_bivariate / num_samples; + double beta_mean_1_multivariate = beta_mean_1_sum_multivariate / num_samples; + double beta_mean_2_multivariate = beta_mean_2_sum_multivariate / num_samples; + EXPECT_NEAR(beta_mean_1_bivariate, beta_mean_1_multivariate, 1e-2); + EXPECT_NEAR(beta_mean_2_bivariate, beta_mean_2_multivariate, 1e-2); +} diff --git a/test/cpp/test_predict.cpp b/test/cpp/test_predict.cpp index 8f35adfa..ed72ae17 100644 --- a/test/cpp/test_predict.cpp +++ b/test/cpp/test_predict.cpp @@ -1,14 +1,10 @@ -/*! - * Test of the ensemble prediction method - */ +/*! Test of the ensemble prediction method */ #include #include #include #include #include #include -#include -#include /*! \brief Test forest prediction procedures for trees with constants in leaf nodes */ TEST(Ensemble, PredictConstant) { @@ -23,7 +19,7 @@ TEST(Ensemble, PredictConstant) { dataset.AddCovariates(test_dataset.covariates.data(), n, test_dataset.x_cols, test_dataset.row_major); dataset.AddBasis(test_dataset.omega.data(), test_dataset.n, test_dataset.omega_cols, test_dataset.row_major); StochTree::ColumnVector residual = StochTree::ColumnVector(test_dataset.outcome.data(), n); - + // Create a small ensemble int output_dim = 1; int num_trees = 2; @@ -38,8 +34,8 @@ TEST(Ensemble, PredictConstant) { tree->ExpandNode(0, 1, tree_split, -2.5, 2.5); // Run predict on the supplied covariates and check the result - std::vector result(n*output_dim); - std::vector expected_pred = {7.5,2.5,-7.5,7.5,7.5,7.5,2.5,7.5,-7.5,-2.5}; + std::vector result(n * output_dim); + std::vector expected_pred = {7.5, 2.5, -7.5, 7.5, 7.5, 7.5, 2.5, 7.5, -7.5, -2.5}; ensemble.PredictInplace(dataset.GetCovariates(), result, 0); for (int i = 0; i < n; i++) { ASSERT_NEAR(expected_pred[i], result[i], 0.01); @@ -59,7 +55,7 @@ TEST(Ensemble, PredictUnivariateRegression) { dataset.AddCovariates(test_dataset.covariates.data(), n, test_dataset.x_cols, test_dataset.row_major); dataset.AddBasis(test_dataset.omega.data(), test_dataset.n, test_dataset.omega_cols, test_dataset.row_major); StochTree::ColumnVector residual = StochTree::ColumnVector(test_dataset.outcome.data(), n); - + // Create a small ensemble int output_dim = 1; int num_trees = 2; @@ -71,10 +67,11 @@ TEST(Ensemble, PredictUnivariateRegression) { StochTree::TreeSplit tree_split = StochTree::TreeSplit(0.5); tree->ExpandNode(0, 0, tree_split, -5., 5.); tree = ensemble.GetTree(1); - tree->ExpandNode(0, 1, tree_split, -2.5, 2.5);; + tree->ExpandNode(0, 1, tree_split, -2.5, 2.5); + ; // Run predict on the supplied covariates and check the result - std::vector result(n*output_dim); + std::vector result(n * output_dim); std::vector expected_pred = {7.3351256, 0.8511415, -1.5396290, 5.7172741, 4.7433491, 4.5919388, 1.0123031, 2.4834167, -6.5187785, -1.4611208}; ensemble.PredictInplace(dataset.GetCovariates(), dataset.GetBasis(), result, 0); for (int i = 0; i < n; i++) { @@ -94,7 +91,7 @@ TEST(Ensemble, PredictMultivariateRegression) { dataset.AddCovariates(test_dataset.covariates.data(), n, test_dataset.x_cols, test_dataset.row_major); dataset.AddBasis(test_dataset.omega.data(), test_dataset.n, test_dataset.omega_cols, test_dataset.row_major); StochTree::ColumnVector residual = StochTree::ColumnVector(test_dataset.outcome.data(), n); - + // Create a small ensemble int output_dim = 2; int num_trees = 2; diff --git a/test/cpp/test_random_effects.cpp b/test/cpp/test_random_effects.cpp index 2908ced5..f703d9e7 100644 --- a/test/cpp/test_random_effects.cpp +++ b/test/cpp/test_random_effects.cpp @@ -8,8 +8,6 @@ #include #include #include -#include -#include TEST(RandomEffects, Setup) { // Load test data @@ -22,7 +20,7 @@ TEST(RandomEffects, Setup) { StochTree::RandomEffectsDataset dataset = StochTree::RandomEffectsDataset(); dataset.AddBasis(test_dataset.rfx_basis.data(), test_dataset.n, test_dataset.rfx_basis_cols, test_dataset.row_major); dataset.AddGroupLabels(test_dataset.rfx_groups); - + // Construct tracker, model state, and container StochTree::RandomEffectsTracker tracker = StochTree::RandomEffectsTracker(test_dataset.rfx_groups); StochTree::MultivariateRegressionRandomEffectsModel model = StochTree::MultivariateRegressionRandomEffectsModel(test_dataset.rfx_basis_cols, test_dataset.rfx_num_groups); @@ -30,7 +28,7 @@ TEST(RandomEffects, Setup) { // Check the internal label map of the RandomEffectsTracker std::map label_map = tracker.GetLabelMap(); - std::map expected_label_map {{1, 0}, {2, 1}}; + std::map expected_label_map{{1, 0}, {2, 1}}; ASSERT_EQ(label_map, expected_label_map); } @@ -45,7 +43,7 @@ TEST(RandomEffects, Construction) { StochTree::RandomEffectsDataset dataset = StochTree::RandomEffectsDataset(); dataset.AddBasis(test_dataset.rfx_basis.data(), test_dataset.n, test_dataset.rfx_basis_cols, test_dataset.row_major); dataset.AddGroupLabels(test_dataset.rfx_groups); - + // Construct tracker, model state, and container StochTree::RandomEffectsTracker tracker = StochTree::RandomEffectsTracker(test_dataset.rfx_groups); StochTree::MultivariateRegressionRandomEffectsModel model = StochTree::MultivariateRegressionRandomEffectsModel(test_dataset.rfx_basis_cols, test_dataset.rfx_num_groups); @@ -65,7 +63,7 @@ TEST(RandomEffects, Construction) { model.SetGroupParameter(xi0, 0); model.SetGroupParameter(xi1, 1); model.SetGroupParameterCovariance(sigma); - + // Push to the container container.AddSample(model); @@ -83,17 +81,17 @@ TEST(RandomEffects, Construction) { // Check data in the container std::vector alpha_retrieved = container.GetAlpha(); - std::vector alpha_expected {1.5, 2.0}; + std::vector alpha_expected{1.5, 2.0}; for (int i = 0; i < alpha_expected.size(); i++) { ASSERT_EQ(alpha_retrieved[i], alpha_expected[i]); } std::vector xi_retrieved = container.GetXi(); - std::vector xi_expected {2, 4, 1, 3}; + std::vector xi_expected{2, 4, 1, 3}; for (int i = 0; i < xi_expected.size(); i++) { ASSERT_EQ(xi_retrieved[i], xi_expected[i]); } std::vector beta_retrieved = container.GetBeta(); - std::vector beta_expected {3, 6, 2, 6}; + std::vector beta_expected{3, 6, 2, 6}; for (int i = 0; i < beta_expected.size(); i++) { ASSERT_EQ(beta_retrieved[i], beta_expected[i]); } @@ -111,7 +109,7 @@ TEST(RandomEffects, Computation) { StochTree::RandomEffectsDataset dataset = StochTree::RandomEffectsDataset(); dataset.AddBasis(test_dataset.rfx_basis.data(), test_dataset.n, test_dataset.rfx_basis_cols, test_dataset.row_major); dataset.AddGroupLabels(test_dataset.rfx_groups); - + // Construct tracker, model state, and container StochTree::RandomEffectsTracker tracker = StochTree::RandomEffectsTracker(test_dataset.rfx_groups); StochTree::MultivariateRegressionRandomEffectsModel model = StochTree::MultivariateRegressionRandomEffectsModel(test_dataset.rfx_basis_cols, test_dataset.rfx_num_groups); @@ -134,7 +132,7 @@ TEST(RandomEffects, Computation) { model.SetGroupParameter(xi2, 2); model.SetGroupParameterCovariance(sigma); double sigma2 = 1.; - + // Compute the posterior mean for the group parameters Eigen::VectorXd xi0_mean = model.GroupParameterMean(dataset, residual, tracker, sigma2, 0); Eigen::VectorXd xi1_mean = model.GroupParameterMean(dataset, residual, tracker, sigma2, 1); @@ -167,7 +165,7 @@ TEST(RandomEffects, Predict) { StochTree::RandomEffectsDataset dataset = StochTree::RandomEffectsDataset(); dataset.AddBasis(test_dataset.rfx_basis.data(), test_dataset.n, test_dataset.rfx_basis_cols, test_dataset.row_major); dataset.AddGroupLabels(test_dataset.rfx_groups); - + // Construct tracker, model state, and container StochTree::RandomEffectsTracker tracker = StochTree::RandomEffectsTracker(test_dataset.rfx_groups); StochTree::MultivariateRegressionRandomEffectsModel model = StochTree::MultivariateRegressionRandomEffectsModel(test_dataset.rfx_basis_cols, test_dataset.rfx_num_groups); @@ -187,7 +185,7 @@ TEST(RandomEffects, Predict) { model.SetGroupParameter(xi0, 0); model.SetGroupParameter(xi1, 1); model.SetGroupParameterCovariance(sigma); - + // Push to the container container.AddSample(model); @@ -205,13 +203,31 @@ TEST(RandomEffects, Predict) { // Predict from the container int num_samples = 2; - std::vector output(n*num_samples); + std::vector output(n * num_samples); container.Predict(dataset, label_mapper, output); // Check predictions - std::vector output_expected { - 3, 6, 3, 6, 3, 6, 3, 6, 3, 6, - 2, 6, 2, 6, 2, 6, 2, 6, 2, 6, + std::vector output_expected{ + 3, + 6, + 3, + 6, + 3, + 6, + 3, + 6, + 3, + 6, + 2, + 6, + 2, + 6, + 2, + 6, + 2, + 6, + 2, + 6, }; for (int i = 0; i < output.size(); i++) { ASSERT_EQ(output[i], output_expected[i]); @@ -229,7 +245,7 @@ TEST(RandomEffects, Serialization) { StochTree::RandomEffectsDataset dataset = StochTree::RandomEffectsDataset(); dataset.AddBasis(test_dataset.rfx_basis.data(), test_dataset.n, test_dataset.rfx_basis_cols, test_dataset.row_major); dataset.AddGroupLabels(test_dataset.rfx_groups); - + // Construct tracker, model state, and container StochTree::RandomEffectsTracker tracker = StochTree::RandomEffectsTracker(test_dataset.rfx_groups); StochTree::MultivariateRegressionRandomEffectsModel model = StochTree::MultivariateRegressionRandomEffectsModel(test_dataset.rfx_basis_cols, test_dataset.rfx_num_groups); @@ -249,7 +265,7 @@ TEST(RandomEffects, Serialization) { model.SetGroupParameter(xi0, 0); model.SetGroupParameter(xi1, 1); model.SetGroupParameterCovariance(sigma); - + // Push to the container container.AddSample(model); diff --git a/test/cpp/test_sorted_partition_tracker.cpp b/test/cpp/test_sorted_partition_tracker.cpp index 99dbbe5d..752cf33d 100644 --- a/test/cpp/test_sorted_partition_tracker.cpp +++ b/test/cpp/test_sorted_partition_tracker.cpp @@ -5,7 +5,6 @@ #include #include #include -#include #include TEST(SortedNodeSampleTracker, BasicOperations) { @@ -52,23 +51,23 @@ TEST(SortedNodeSampleTracker, BasicOperations) { ASSERT_EQ(sample_node_mapper.GetNodeId(7, 0), 2); ASSERT_EQ(sample_node_mapper.GetNodeId(8, 0), 1); ASSERT_EQ(sample_node_mapper.GetNodeId(9, 0), 1); - + // Check that node begin and node end haven't changed for root node, but that the indices have been sifted ASSERT_EQ(sorted_node_sampler_tracker.NodeBegin(0, 0), 0); ASSERT_EQ(sorted_node_sampler_tracker.NodeEnd(0, 0), n); - std::vector expected_result{2,8,9,4,1,6,3,7,0,5}; + std::vector expected_result{2, 8, 9, 4, 1, 6, 3, 7, 0, 5}; ASSERT_EQ(sorted_node_sampler_tracker.NodeIndices(0, 0), expected_result); - + // Check node begin and node end for left node ASSERT_EQ(sorted_node_sampler_tracker.NodeBegin(1, 0), 0); ASSERT_EQ(sorted_node_sampler_tracker.NodeEnd(1, 0), 3); - expected_result = {2,8,9}; + expected_result = {2, 8, 9}; ASSERT_EQ(sorted_node_sampler_tracker.NodeIndices(1, 0), expected_result); - + // Check node begin and node end for right node ASSERT_EQ(sorted_node_sampler_tracker.NodeBegin(2, 0), 3); ASSERT_EQ(sorted_node_sampler_tracker.NodeEnd(2, 0), n); - expected_result = {4,1,6,3,7,0,5}; + expected_result = {4, 1, 6, 3, 7, 0, 5}; ASSERT_EQ(sorted_node_sampler_tracker.NodeIndices(2, 0), expected_result); // Partition right node based on X[,1] <= 0.5 @@ -77,24 +76,24 @@ TEST(SortedNodeSampleTracker, BasicOperations) { // Check that node begin and node end haven't changed for old right node, but that the indices have been sifted ASSERT_EQ(sorted_node_sampler_tracker.NodeBegin(2, 0), 3); ASSERT_EQ(sorted_node_sampler_tracker.NodeEnd(2, 0), n); - expected_result = {1,6,4,3,7,0,5}; + expected_result = {1, 6, 4, 3, 7, 0, 5}; ASSERT_EQ(sorted_node_sampler_tracker.NodeIndices(2, 0), expected_result); // Check same indices for feature 1 ASSERT_EQ(sorted_node_sampler_tracker.NodeBegin(2, 1), 3); ASSERT_EQ(sorted_node_sampler_tracker.NodeEnd(2, 1), n); - expected_result = {6,1,3,0,7,4,5}; + expected_result = {6, 1, 3, 0, 7, 4, 5}; ASSERT_EQ(sorted_node_sampler_tracker.NodeIndices(2, 1), expected_result); - + // Check node begin and node end for new left node ASSERT_EQ(sorted_node_sampler_tracker.NodeBegin(3, 1), 3); ASSERT_EQ(sorted_node_sampler_tracker.NodeEnd(3, 1), 5); - expected_result = {6,1}; + expected_result = {6, 1}; ASSERT_EQ(sorted_node_sampler_tracker.NodeIndices(3, 1), expected_result); - + // Check node begin and node end for new right node ASSERT_EQ(sorted_node_sampler_tracker.NodeBegin(4, 1), 5); ASSERT_EQ(sorted_node_sampler_tracker.NodeEnd(4, 1), n); - expected_result = {3,0,7,4,5}; + expected_result = {3, 0, 7, 4, 5}; ASSERT_EQ(sorted_node_sampler_tracker.NodeIndices(4, 1), expected_result); } diff --git a/test/cpp/test_tree.cpp b/test/cpp/test_tree.cpp index c30302c2..7deecf48 100644 --- a/test/cpp/test_tree.cpp +++ b/test/cpp/test_tree.cpp @@ -7,8 +7,6 @@ #include #include #include -#include -#include TEST(Tree, UnivariateTreeConstruction) { StochTree::Tree tree; @@ -33,10 +31,10 @@ TEST(Tree, UnivariateTreeCopyConstruction) { StochTree::Tree tree_2; StochTree::TreeSplit split; tree_1.Init(1); - + // Check max depth ASSERT_EQ(tree_1.MaxLeafDepth(), 0); - + // Perform two splits split = StochTree::TreeSplit(0.5); tree_1.ExpandNode(0, 0, split, 0., 0.); @@ -46,7 +44,7 @@ TEST(Tree, UnivariateTreeCopyConstruction) { ASSERT_EQ(tree_1.MaxLeafDepth(), 2); ASSERT_EQ(tree_1.NumValidNodes(), 5); ASSERT_EQ(tree_1.NumLeafParents(), 1); - + // Check leaves std::vector leaves = tree_1.GetLeaves(); for (int i = 0; i < leaves.size(); i++) { @@ -57,7 +55,7 @@ TEST(Tree, UnivariateTreeCopyConstruction) { for (int i = 0; i < leaf_parents.size(); i++) { ASSERT_TRUE(tree_1.IsLeafParent(leaf_parents[i])); } - + // Perform another split split = StochTree::TreeSplit(0.6); tree_1.ExpandNode(3, 2, split, 0., 0.); @@ -65,7 +63,7 @@ TEST(Tree, UnivariateTreeCopyConstruction) { ASSERT_EQ(tree_1.NumValidNodes(), 7); ASSERT_EQ(tree_1.NumLeaves(), 4); ASSERT_EQ(tree_1.NumLeafParents(), 1); - + // Check leaves leaves = tree_1.GetLeaves(); for (int i = 0; i < leaves.size(); i++) { @@ -83,7 +81,7 @@ TEST(Tree, UnivariateTreeCopyConstruction) { ASSERT_EQ(tree_1.NumValidNodes(), 5); ASSERT_EQ(tree_1.NumLeaves(), 3); ASSERT_EQ(tree_1.NumLeafParents(), 1); - + // Check leaves leaves = tree_1.GetLeaves(); for (int i = 0; i < leaves.size(); i++) { @@ -106,12 +104,12 @@ TEST(Tree, UnivariateTreeCategoricalSplitConstruction) { StochTree::Tree tree; tree.Init(1); ASSERT_EQ(tree.LeafValue(0), 0.); - tree.ExpandNode(0, 0, std::vector{1,4,6}, 0., 0.); + tree.ExpandNode(0, 0, std::vector{1, 4, 6}, 0., 0.); ASSERT_EQ(tree.NumNodes(), 3); ASSERT_EQ(tree.NodeType(0), StochTree::TreeNodeType::kCategoricalSplitNode); tree.CollapseToLeaf(0, 0.); ASSERT_EQ(tree.NumValidNodes(), 1); - tree.ExpandNode(0, 0, std::vector{2,3,5}, 0., 0.); + tree.ExpandNode(0, 0, std::vector{2, 3, 5}, 0., 0.); ASSERT_EQ(tree.NodeType(0), StochTree::TreeNodeType::kCategoricalSplitNode); ASSERT_EQ(tree.NumValidNodes(), 3); ASSERT_EQ(tree.NumLeaves(), 2); @@ -148,12 +146,12 @@ TEST(Tree, MultivariateTreeCategoricalSplitConstruction) { std::vector leaf_values(tree_dim, 0.); tree.Init(tree_dim); ASSERT_EQ(tree.LeafVector(0), leaf_values); - tree.ExpandNode(0, 0, std::vector{1,4,6}, leaf_values, leaf_values); + tree.ExpandNode(0, 0, std::vector{1, 4, 6}, leaf_values, leaf_values); ASSERT_EQ(tree.NumNodes(), 3); ASSERT_EQ(tree.NodeType(0), StochTree::TreeNodeType::kCategoricalSplitNode); tree.CollapseToLeaf(0, leaf_values); ASSERT_EQ(tree.NumValidNodes(), 1); - tree.ExpandNode(0, 0, std::vector{2,3,5}, leaf_values, leaf_values); + tree.ExpandNode(0, 0, std::vector{2, 3, 5}, leaf_values, leaf_values); ASSERT_EQ(tree.NodeType(0), StochTree::TreeNodeType::kCategoricalSplitNode); ASSERT_EQ(tree.NumValidNodes(), 3); ASSERT_EQ(tree.NumLeaves(), 2); @@ -167,7 +165,7 @@ TEST(Tree, SparseLeafRepresentation) { StochTree::Tree tree; tree.Init(1); tree.ExpandNode(0, 0, 0.5, 0., 0.); - + // Load test data StochTree::TestUtils::TestDataset test_dataset; test_dataset = StochTree::TestUtils::LoadSmallDatasetUnivariateBasis(); @@ -180,6 +178,6 @@ TEST(Tree, SparseLeafRepresentation) { // Predict leaf indices of each observation in `dataset` std::vector leaf_index_preds(n); tree.PredictLeafIndexInplace(&dataset, leaf_index_preds, 0, 0); - std::vector leaf_index_expected{1,1,0,1,1,1,1,1,0,0}; + std::vector leaf_index_expected{1, 1, 0, 1, 1, 1, 1, 1, 0, 0}; ASSERT_EQ(leaf_index_expected, leaf_index_preds); } diff --git a/test/cpp/test_unsorted_partition_tracker.cpp b/test/cpp/test_unsorted_partition_tracker.cpp index 067e870f..02e5fb93 100644 --- a/test/cpp/test_unsorted_partition_tracker.cpp +++ b/test/cpp/test_unsorted_partition_tracker.cpp @@ -2,11 +2,12 @@ #include #include #include +#include #include #include #include -#include -#include +#include +#include #include TEST(UnsortedNodeSampleTracker, BasicOperations) { @@ -40,13 +41,13 @@ TEST(UnsortedNodeSampleTracker, BasicOperations) { StochTree::TreeSplit tree_split = StochTree::TreeSplit(0.5); node_sample_tracker.PartitionTreeNode(dataset.GetCovariates(), 0, 0, 1, 2, 0, tree_split); sample_node_mapper.AddSplit(dataset.GetCovariates(), tree_split, 0, 0, 0, 1, 2); - + // Check that node begin and node end haven't changed for root node, but that the indices have been sifted ASSERT_EQ(node_sample_tracker.NodeBegin(0, 0), 0); ASSERT_EQ(node_sample_tracker.NodeEnd(0, 0), n); std::vector expected_result{2, 8, 9, 0, 1, 3, 4, 5, 6, 7}; ASSERT_EQ(node_sample_tracker.TreeNodeIndices(0, 0), expected_result); - + // Check that terminal nodes are updated for for every observation ASSERT_EQ(sample_node_mapper.GetNodeId(0, 0), 2); ASSERT_EQ(sample_node_mapper.GetNodeId(1, 0), 2); @@ -58,13 +59,13 @@ TEST(UnsortedNodeSampleTracker, BasicOperations) { ASSERT_EQ(sample_node_mapper.GetNodeId(7, 0), 2); ASSERT_EQ(sample_node_mapper.GetNodeId(8, 0), 1); ASSERT_EQ(sample_node_mapper.GetNodeId(9, 0), 1); - + // Check node begin and node end for left node ASSERT_EQ(node_sample_tracker.NodeBegin(0, 1), 0); ASSERT_EQ(node_sample_tracker.NodeEnd(0, 1), 3); expected_result = {2, 8, 9}; ASSERT_EQ(node_sample_tracker.TreeNodeIndices(0, 1), expected_result); - + // Check node begin and node end for right node ASSERT_EQ(node_sample_tracker.NodeBegin(0, 2), 3); ASSERT_EQ(node_sample_tracker.NodeEnd(0, 2), n); @@ -81,13 +82,13 @@ TEST(UnsortedNodeSampleTracker, BasicOperations) { ASSERT_EQ(node_sample_tracker.NodeEnd(0, 2), n); expected_result = {1, 6, 0, 3, 4, 5, 7}; ASSERT_EQ(node_sample_tracker.TreeNodeIndices(0, 2), expected_result); - + // Check node begin and node end for new left node ASSERT_EQ(node_sample_tracker.NodeBegin(0, 3), 3); ASSERT_EQ(node_sample_tracker.NodeEnd(0, 3), 5); expected_result = {1, 6}; ASSERT_EQ(node_sample_tracker.TreeNodeIndices(0, 3), expected_result); - + // Check node begin and node end for new right node ASSERT_EQ(node_sample_tracker.NodeBegin(0, 4), 5); ASSERT_EQ(node_sample_tracker.NodeEnd(0, 4), n); @@ -100,3 +101,81 @@ TEST(UnsortedNodeSampleTracker, BasicOperations) { ASSERT_FALSE(node_sample_tracker.IsValidNode(0, 3)); ASSERT_FALSE(node_sample_tracker.IsValidNode(0, 4)); } + +// Regression test for a bug in FeatureUnsortedPartition::ReconstituteFromTree. +// +// The reconstitution used to iterate nodes by numeric id, assuming every parent has a +// SMALLER id than its children (so the parent always set a node's index bounds before the +// node itself was processed). That holds for grow-only trees, but the Tree free-list +// RECYCLES deleted node ids after prune+regrow, so a split node can end up with a parent +// of *higher* id. Iterating by id then partitioned that node from stale bounds and +// mis-assigned observations. The fix traverses top-down (BFS from root). This test builds +// a tree whose ids are recycled into a non-topological layout and asserts the +// reconstituted unsorted partition matches a direct evaluation of the tree. +TEST(UnsortedNodeSampleTracker, ReconstituteWithRecycledNodeIds) { + int n = 12; + int p = 2; + // Row-major covariates (n x p), chosen so the splits below route observations to + // several different leaves. + std::vector covariates_raw = { + 0.10, 0.20, 0.30, 0.80, 0.40, 0.40, 0.20, 0.90, 0.60, 0.10, 0.70, 0.60, + 0.80, 0.30, 0.90, 0.70, 0.55, 0.95, 0.65, 0.05, 0.85, 0.50, 0.95, 0.25}; + StochTree::ForestDataset dataset; + dataset.AddCovariates(covariates_raw.data(), n, p, /*row_major=*/true); + std::vector residual_raw(n, 0.0); + StochTree::ColumnVector residual(residual_raw.data(), n); + std::vector feature_types(p, StochTree::FeatureType::kNumeric); + + // Build a single-tree forest, then grow/prune/regrow so node ids get recycled into a + // non-topological layout (a split node whose parent has a larger id). + StochTree::TreeEnsemble forest(/*num_trees=*/1, /*output_dim=*/1, /*is_leaf_constant=*/true); + StochTree::Tree* tree = forest.GetTree(0); + tree->ExpandNode(0, 0, 0.5, 0.0, 0.0); // 0 -> {1, 2} + tree->ExpandNode(1, 0, 0.25, 0.0, 0.0); // 1 -> {3, 4} + tree->ExpandNode(2, 0, 0.75, 0.0, 0.0); // 2 -> {5, 6} + tree->CollapseToLeaf(1, 0.0); // prune node 1: frees ids 3, 4 + tree->ExpandNode(5, 1, 0.5, 0.0, 0.0); // 5 -> recycled {3, 4} + tree->ExpandNode(3, 0, 0.6, 0.0, 0.0); // 3 -> new {7, 8}; node 3 is a split whose parent (5) has a LARGER id + + // Guard: confirm the setup actually produced the non-topological condition under test + // (some internal/split node has a parent with a larger id). If this ever fails, the + // Tree id-recycling policy changed and the test setup needs to be revisited. + bool has_parent_with_larger_id = false; + for (int nid = 0; nid < tree->NumNodes(); nid++) { + if (tree->IsDeleted(nid) || tree->IsLeaf(nid)) continue; + int parent = tree->Parent(nid); + if (parent != StochTree::Tree::kInvalidNodeId && parent > nid) { + has_parent_with_larger_id = true; + } + } + ASSERT_TRUE(has_parent_with_larger_id) + << "Test setup did not produce a node whose parent has a larger id"; + + // Reconstitute the tracker from the forest (this exercises ReconstituteFromTree). + StochTree::ForestTracker tracker(dataset.GetCovariates(), feature_types, /*num_trees=*/1, n); + tracker.ReconstituteFromForest(forest, dataset, residual, /*is_mean_model=*/true); + + // Ground truth membership: evaluate the tree directly for every observation. + Eigen::MatrixXd& covariates = dataset.GetCovariates(); + std::map> expected_membership; + for (int i = 0; i < n; i++) { + int leaf = StochTree::EvaluateTree(*tree, covariates, i); + expected_membership[leaf].insert(i); + } + + // Every leaf's reconstituted membership must match the ground truth exactly. + int total_assigned = 0; + for (int leaf : tree->GetLeaves()) { + std::set reconstituted; + auto it_begin = tracker.UnsortedNodeBeginIterator(0, leaf); + auto it_end = tracker.UnsortedNodeEndIterator(0, leaf); + for (auto it = it_begin; it != it_end; ++it) { + reconstituted.insert(static_cast(*it)); + } + total_assigned += static_cast(reconstituted.size()); + EXPECT_EQ(reconstituted, expected_membership[leaf]) + << "Reconstituted membership mismatch for leaf node " << leaf; + } + // Every observation must be assigned to exactly one leaf. + EXPECT_EQ(total_assigned, n); +} diff --git a/test/cpp/testutils.cpp b/test/cpp/testutils.cpp index e6d967a5..4ae1266a 100644 --- a/test/cpp/testutils.cpp +++ b/test/cpp/testutils.cpp @@ -2,18 +2,14 @@ #include #include #include -#include -#include -#include -#include namespace StochTree { -namespace TestUtils{ +namespace TestUtils { TestDataset LoadSmallDatasetUnivariateBasis() { TestDataset output; - + // Data dimensions output.n = 10; output.x_cols = 5; @@ -24,29 +20,29 @@ TestDataset LoadSmallDatasetUnivariateBasis() { output.rfx_basis.resize(output.n, output.rfx_basis_cols); output.rfx_groups.resize(output.n); output.outcome.resize(output.n); - + // Covariates output.covariates << 0.766969853, 0.83894646, 0.63649772, 0.6747788934, 0.27398269, - 0.634970996, 0.15237997, 0.3800786, 0.6457891271, 0.21604451, - 0.229598754, 0.12461481, 0.81407372, 0.364336529, 0.45160373, - 0.741084778, 0.53356288, 0.58940162, 0.9995219493, 0.19142269, - 0.618177813, 0.88876378, 0.51174404, 0.8827708189, 0.12730742, - 0.858657839, 0.9271676, 0.5115294, 0.67865624, 0.28658962, - 0.719224842, 0.0546961, 0.42850897, 0.260336376, 0.1371501, - 0.747422328, 0.87172033, 0.98791964, 0.4018020707, 0.29145664, - 0.3158837, 0.39253551, 0.83610831, 0.0101785748, 0.1955386, - 0.419554105, 0.5586495, 0.19908607, 0.4873921743, 0.35568569; - + 0.634970996, 0.15237997, 0.3800786, 0.6457891271, 0.21604451, + 0.229598754, 0.12461481, 0.81407372, 0.364336529, 0.45160373, + 0.741084778, 0.53356288, 0.58940162, 0.9995219493, 0.19142269, + 0.618177813, 0.88876378, 0.51174404, 0.8827708189, 0.12730742, + 0.858657839, 0.9271676, 0.5115294, 0.67865624, 0.28658962, + 0.719224842, 0.0546961, 0.42850897, 0.260336376, 0.1371501, + 0.747422328, 0.87172033, 0.98791964, 0.4018020707, 0.29145664, + 0.3158837, 0.39253551, 0.83610831, 0.0101785748, 0.1955386, + 0.419554105, 0.5586495, 0.19908607, 0.4873921743, 0.35568569; + // Leaf regression basis output.omega << 0.97801674, 0.34045661, 0.20528387, 0.76230322, 0.63244655, 0.61225851, 0.40492125, 0.33112223, 0.86917047, 0.58444831; - + // Outcome - output.outcome << 2.158854445, 1.175387297, 0.40481061, 1.751578365, 0.299641379, - 0.347249942, 0.546179903, 1.164750138, 3.389946886, -0.605464414; - + output.outcome << 2.158854445, 1.175387297, 0.40481061, 1.751578365, 0.299641379, + 0.347249942, 0.546179903, 1.164750138, 3.389946886, -0.605464414; + // Random effects regression basis (i.e. constant, intercept-only RFX model) output.rfx_basis << 1, 1, 1, 1, 1, 1, 1, 1, 1, 1; - + // Random effects group labels for (int i = 0; i < output.n; i++) { if (i % 2 == 0) { @@ -62,7 +58,7 @@ TestDataset LoadSmallDatasetUnivariateBasis() { TestDataset LoadSmallDatasetMultivariateBasis() { TestDataset output; - + // Data dimensions output.n = 10; output.x_cols = 5; @@ -73,38 +69,38 @@ TestDataset LoadSmallDatasetMultivariateBasis() { output.rfx_basis.resize(output.n, output.rfx_basis_cols); output.rfx_groups.resize(output.n); output.outcome.resize(output.n); - + // Covariates output.covariates << 0.766969853, 0.83894646, 0.63649772, 0.6747788934, 0.27398269, - 0.634970996, 0.15237997, 0.3800786, 0.6457891271, 0.21604451, - 0.229598754, 0.12461481, 0.81407372, 0.364336529, 0.45160373, - 0.741084778, 0.53356288, 0.58940162, 0.9995219493, 0.19142269, - 0.618177813, 0.88876378, 0.51174404, 0.8827708189, 0.12730742, - 0.858657839, 0.9271676, 0.5115294, 0.67865624, 0.28658962, - 0.719224842, 0.0546961, 0.42850897, 0.260336376, 0.1371501, - 0.747422328, 0.87172033, 0.98791964, 0.4018020707, 0.29145664, - 0.3158837, 0.39253551, 0.83610831, 0.0101785748, 0.1955386, - 0.419554105, 0.5586495, 0.19908607, 0.4873921743, 0.35568569; - + 0.634970996, 0.15237997, 0.3800786, 0.6457891271, 0.21604451, + 0.229598754, 0.12461481, 0.81407372, 0.364336529, 0.45160373, + 0.741084778, 0.53356288, 0.58940162, 0.9995219493, 0.19142269, + 0.618177813, 0.88876378, 0.51174404, 0.8827708189, 0.12730742, + 0.858657839, 0.9271676, 0.5115294, 0.67865624, 0.28658962, + 0.719224842, 0.0546961, 0.42850897, 0.260336376, 0.1371501, + 0.747422328, 0.87172033, 0.98791964, 0.4018020707, 0.29145664, + 0.3158837, 0.39253551, 0.83610831, 0.0101785748, 0.1955386, + 0.419554105, 0.5586495, 0.19908607, 0.4873921743, 0.35568569; + // Leaf regression basis - output.omega << 0.97801674, 0.3707159, - 0.34045661, 0.1312134, - 0.20528387, 0.5614470, - 0.76230322, 0.2276504, - 0.63244655, 0.9029984, - 0.61225851, 0.7448547, - 0.40492125, 0.2549813, - 0.33112223, 0.5295535, - 0.86917047, 0.5584614, - 0.58444831, 0.2365117; - + output.omega << 0.97801674, 0.3707159, + 0.34045661, 0.1312134, + 0.20528387, 0.5614470, + 0.76230322, 0.2276504, + 0.63244655, 0.9029984, + 0.61225851, 0.7448547, + 0.40492125, 0.2549813, + 0.33112223, 0.5295535, + 0.86917047, 0.5584614, + 0.58444831, 0.2365117; + // Outcome - output.outcome << 2.158854445, 1.175387297, 0.40481061, 1.751578365, 0.299641379, - 0.347249942, 0.546179903, 1.164750138, 3.389946886, -0.605464414; - + output.outcome << 2.158854445, 1.175387297, 0.40481061, 1.751578365, 0.299641379, + 0.347249942, 0.546179903, 1.164750138, 3.389946886, -0.605464414; + // Random effects regression basis (i.e. constant, intercept-only RFX model) output.rfx_basis << 1, 1, 1, 1, 1, 1, 1, 1, 1, 1; - + // Random effects group labels for (int i = 0; i < output.n; i++) { if (i % 2 == 0) { @@ -120,7 +116,7 @@ TestDataset LoadSmallDatasetMultivariateBasis() { TestDataset LoadSmallRFXDatasetMultivariateBasis() { TestDataset output; - + // Data dimensions output.n = 10; output.x_cols = 5; @@ -131,49 +127,49 @@ TestDataset LoadSmallRFXDatasetMultivariateBasis() { output.rfx_basis.resize(output.n, output.rfx_basis_cols); output.rfx_groups.resize(output.n); output.outcome.resize(output.n); - + // Covariates output.covariates << 0.766969853, 0.83894646, 0.63649772, 0.6747788934, 0.27398269, - 0.634970996, 0.15237997, 0.3800786, 0.6457891271, 0.21604451, - 0.229598754, 0.12461481, 0.81407372, 0.364336529, 0.45160373, - 0.741084778, 0.53356288, 0.58940162, 0.9995219493, 0.19142269, - 0.618177813, 0.88876378, 0.51174404, 0.8827708189, 0.12730742, - 0.858657839, 0.9271676, 0.5115294, 0.67865624, 0.28658962, - 0.719224842, 0.0546961, 0.42850897, 0.260336376, 0.1371501, - 0.747422328, 0.87172033, 0.98791964, 0.4018020707, 0.29145664, - 0.3158837, 0.39253551, 0.83610831, 0.0101785748, 0.1955386, - 0.419554105, 0.5586495, 0.19908607, 0.4873921743, 0.35568569; - + 0.634970996, 0.15237997, 0.3800786, 0.6457891271, 0.21604451, + 0.229598754, 0.12461481, 0.81407372, 0.364336529, 0.45160373, + 0.741084778, 0.53356288, 0.58940162, 0.9995219493, 0.19142269, + 0.618177813, 0.88876378, 0.51174404, 0.8827708189, 0.12730742, + 0.858657839, 0.9271676, 0.5115294, 0.67865624, 0.28658962, + 0.719224842, 0.0546961, 0.42850897, 0.260336376, 0.1371501, + 0.747422328, 0.87172033, 0.98791964, 0.4018020707, 0.29145664, + 0.3158837, 0.39253551, 0.83610831, 0.0101785748, 0.1955386, + 0.419554105, 0.5586495, 0.19908607, 0.4873921743, 0.35568569; + // Leaf regression basis - output.omega << 0.97801674, 0.3707159, - 0.34045661, 0.1312134, - 0.20528387, 0.5614470, - 0.76230322, 0.2276504, - 0.63244655, 0.9029984, - 0.61225851, 0.7448547, - 0.40492125, 0.2549813, - 0.33112223, 0.5295535, - 0.86917047, 0.5584614, - 0.58444831, 0.2365117; - + output.omega << 0.97801674, 0.3707159, + 0.34045661, 0.1312134, + 0.20528387, 0.5614470, + 0.76230322, 0.2276504, + 0.63244655, 0.9029984, + 0.61225851, 0.7448547, + 0.40492125, 0.2549813, + 0.33112223, 0.5295535, + 0.86917047, 0.5584614, + 0.58444831, 0.2365117; + // Outcome - output.outcome << 2.158854445, 1.175387297, 0.40481061, 1.751578365, 0.299641379, - 0.347249942, 0.546179903, 1.164750138, 3.389946886, -0.605464414; - + output.outcome << 2.158854445, 1.175387297, 0.40481061, 1.751578365, 0.299641379, + 0.347249942, 0.546179903, 1.164750138, 3.389946886, -0.605464414; + // Random effects regression basis (i.e. constant, intercept-only RFX model) - output.rfx_basis << 1, 0.3707159, - 1, 0.1312134, - 1, 0.5614470, - 1, 0.2276504, - 1, 0.9029984, - 1, 0.7448547, - 1, 0.2549813, - 1, 0.5295535, - 1, 0.5584614, - 1, 0.2365117; - + output.rfx_basis << 1, 0.3707159, + 1, 0.1312134, + 1, 0.5614470, + 1, 0.2276504, + 1, 0.9029984, + 1, 0.7448547, + 1, 0.2549813, + 1, 0.5295535, + 1, 0.5584614, + 1, 0.2365117; + // Random effects group labels - output.rfx_groups = {1,2,3,1,2,3,1,2,3,1}; + output.rfx_groups = {1, 2, 3, 1, 2, 3, 1, 2, 3, 1}; // for (int i = 0; i < output.n; i++) { // if (i % 2 == 0) { // output.rfx_groups[i] = 1; @@ -188,7 +184,7 @@ TestDataset LoadSmallRFXDatasetMultivariateBasis() { TestDataset LoadMediumDatasetUnivariateBasis() { TestDataset output; - + // Data dimensions output.n = 100; output.x_cols = 5; @@ -199,148 +195,148 @@ TestDataset LoadMediumDatasetUnivariateBasis() { output.rfx_basis.resize(output.n, output.rfx_basis_cols); output.rfx_groups.resize(output.n); output.outcome.resize(output.n); - + // Covariates output.covariates << 0.766969853, 0.83894646, 0.63649772, 0.6747788934, 0.27398269, - 0.634970996, 0.15237997, 0.3800786, 0.6457891271, 0.21604451, - 0.229598754, 0.12461481, 0.81407372, 0.364336529, 0.45160373, - 0.741084778, 0.53356288, 0.58940162, 0.9995219493, 0.19142269, - 0.618177813, 0.88876378, 0.51174404, 0.8827708189, 0.12730742, - 0.858657839, 0.9271676, 0.5115294, 0.67865624, 0.28658962, - 0.719224842, 0.0546961, 0.42850897, 0.260336376, 0.1371501, - 0.747422328, 0.87172033, 0.98791964, 0.4018020707, 0.29145664, - 0.3158837, 0.39253551, 0.83610831, 0.0101785748, 0.1955386, - 0.419554105, 0.5586495, 0.19908607, 0.4873921743, 0.35568569, - 0.012786428, 0.46925501, 0.25363201, 0.3429851863, 0.2071495, - 0.887479904, 0.66166194, 0.31100105, 0.2895678403, 0.00117005, - 0.147758652, 0.14108789, 0.0361254, 0.4790630946, 0.47336526, - 0.899947367, 0.03730855, 0.33408769, 0.368503517, 0.30600202, - 0.527616998, 0.22344076, 0.20325828, 0.9296060419, 0.34518043, - 0.947085596, 0.85906392, 0.35535464, 0.529360628, 0.8781696, - 0.716097994, 0.9149628, 0.11689428, 0.1157865208, 0.31602707, - 0.433331308, 0.53848417, 0.34146036, 0.4967994317, 0.12822296, - 0.420861259, 0.28802486, 0.62324752, 0.2045601751, 0.06909585, - 0.275279159, 0.69079999, 0.29498051, 0.0082852058, 0.45247107, - 0.909681016, 0.35067747, 0.66813255, 0.3866910117, 0.65315347, - 0.828031845, 0.74096924, 0.33982958, 0.0009472317, 0.65103292, - 0.261653444, 0.43179244, 0.89632155, 0.8636559783, 0.93461464, - 0.209384357, 0.12561389, 0.69809409, 0.4752417156, 0.34963379, - 0.737655852, 0.42078584, 0.09970929, 0.5218528947, 0.36737846, - 0.975034732, 0.69977514, 0.33918481, 0.5443784453, 0.35411297, - 0.053533786, 0.98021485, 0.71035393, 0.189234901, 0.73372176, - 0.364139644, 0.47595789, 0.24620073, 0.4284725219, 0.46145259, - 0.696115067, 0.18095114, 0.66919045, 0.9517078404, 0.31686943, - 0.920878008, 0.89758374, 0.21445324, 0.5666448742, 0.29554824, - 0.397853079, 0.12019741, 0.10775046, 0.0799620333, 0.20065807, - 0.322087545, 0.68342919, 0.29873607, 0.0044371644, 0.66733723, - 0.661407114, 0.0558764, 0.10688295, 0.067841246, 0.52254161, - 0.593253554, 0.40498486, 0.97342655, 0.1917967587, 0.2078643, - 0.392762915, 0.91608107, 0.98894976, 0.3599016496, 0.70576753, - 0.758995247, 0.19899099, 0.95978035, 0.8000916124, 0.8356055, - 0.105617762, 0.12135206, 0.47523114, 0.3594282658, 0.71053726, - 0.754330984, 0.803395, 0.11297253, 0.5072350584, 0.05109695, - 0.410083859, 0.13842349, 0.3671543, 0.262290115, 0.76582706, - 0.498883172, 0.52094766, 0.23674406, 0.8919167451, 0.26313017, - 0.315790046, 0.57934811, 0.96794023, 0.7292640421, 0.63874656, - 0.969918807, 0.86839672, 0.17867962, 0.797609952, 0.3123159, - 0.291589217, 0.37982099, 0.92081884, 0.3760313739, 0.30599535, - 0.874146047, 0.64472863, 0.74944373, 0.0179410274, 0.06637048, - 0.006168369, 0.36819005, 0.48640614, 0.5182905369, 0.37514676, - 0.018794786, 0.50404546, 0.30706335, 0.239409535, 0.78368968, - 0.218041312, 0.08232156, 0.910968, 0.236348928, 0.08734924, - 0.240712896, 0.81851635, 0.75910757, 0.7666831033, 0.51030368, - 0.32422135, 0.37234399, 0.4268269, 0.0688136201, 0.52522145, - 0.737050103, 0.55333162, 0.35681609, 0.5527229193, 0.45528166, - 0.666105454, 0.44928217, 0.93068357, 0.2682658806, 0.47992145, - 0.072705164, 0.24379538, 0.36250275, 0.2693803106, 0.88583253, - 0.393483048, 0.7180344, 0.88936403, 0.9690254654, 0.41720031, - 0.726532397, 0.15675097, 0.14675637, 0.973136256, 0.86701643, - 0.206543021, 0.70612692, 0.9923119, 0.1270776591, 0.43317344, - 0.392393596, 0.6581254, 0.51121301, 0.8005079071, 0.16056554, - 0.326374607, 0.48817642, 0.68630408, 0.9265561129, 0.48683193, - 0.761818521, 0.71751337, 0.83854992, 0.134206275, 0.25700676, - 0.930924999, 0.37469277, 0.42861545, 0.7379696709, 0.9670993, - 0.601101112, 0.56631699, 0.85690728, 0.0792362478, 0.23640603, - 0.294070227, 0.02818223, 0.83060893, 0.8203584203, 0.17647972, - 0.393978659, 0.88639966, 0.80788018, 0.4202279691, 0.75344798, - 0.381183787, 0.98751161, 0.13933232, 0.5427466533, 0.15809025, - 0.203872876, 0.31032719, 0.53000948, 0.6001499062, 0.43581315, - 0.355075927, 0.10865708, 0.21823445, 0.5707600345, 0.84459087, - 0.415892882, 0.09056941, 0.85957968, 0.9296874236, 0.39317951, - 0.885163931, 0.60617414, 0.22888755, 0.9225545505, 0.41601782, - 0.803631177, 0.63855664, 0.4968153, 0.4970232591, 0.28230652, - 0.755692566, 0.36382158, 0.31492054, 0.9853899847, 0.45864754, - 0.761099141, 0.88094342, 0.82542666, 0.977985516, 0.5416208, - 0.536037115, 0.19298885, 0.67674639, 0.213044832, 0.29409245, - 0.050087478, 0.56597845, 0.22309031, 0.7668617836, 0.02385271, - 0.847882026, 0.86580035, 0.8381724, 0.618777399, 0.4707389, - 0.280194086, 0.95490103, 0.27399251, 0.5894525715, 0.17181438, - 0.261382768, 0.96124295, 0.33737123, 0.3545607659, 0.36367031, - 0.465759262, 0.17167592, 0.87114988, 0.4175856721, 0.16020522, - 0.982323635, 0.30892377, 0.96513595, 0.376671114, 0.9411435, - 0.851789546, 0.42260807, 0.37396782, 0.0759502219, 0.41219659, - 0.23932738, 0.70124641, 0.08544481, 0.8599137105, 0.35298377, - 0.985171556, 0.48493665, 0.92919919, 0.3128095574, 0.84388465, - 0.936608667, 0.70159722, 0.23570122, 0.5124408882, 0.99478731, - 0.328337863, 0.83252833, 0.29078719, 0.7531193637, 0.49378383, - 0.504403078, 0.72845174, 0.12801659, 0.5383322216, 0.12559066, - 0.906952623, 0.36801267, 0.13168735, 0.9791060984, 0.14008791, - 0.454210506, 0.67248289, 0.4041049, 0.234963659, 0.92138674, - 0.499037576, 0.7534805, 0.4168877, 0.6275620307, 0.24189188, - 0.707788941, 0.91990553, 0.56701198, 0.1408275496, 0.80566006, - 0.694437274, 0.69339343, 0.42296251, 0.8271595608, 0.53699966, - 0.447118821, 0.97512181, 0.16431204, 0.3697280197, 0.38753206, - 0.885936489, 0.94468978, 0.48918779, 0.3676202064, 0.06938232, - 0.593980148, 0.28140352, 0.27760537, 0.2819242389, 0.8730862, - 0.04248501, 0.45279893, 0.69760642, 0.0949480394, 0.42568701, - 0.35842742, 0.68098838, 0.82745029, 0.5315801166, 0.31104918, - 0.724621041, 0.28763999, 0.48743089, 0.8648093319, 0.93792148, - 0.961828358, 0.5548953, 0.7250596, 0.249875583, 0.90661302, - 0.251438316, 0.86021024, 0.65037498, 0.209739062, 0.07886205, - 0.699615913, 0.12223695, 0.20393331, 0.6357937951, 0.81502268, - 0.391076967, 0.25143855, 0.16091307, 0.6037441837, 0.50651534, - 0.343597198, 0.82570727, 0.62455707, 0.6284155636, 0.17288776, - 0.451352309, 0.29346835, 0.12641623, 0.1194773833, 0.88849468; - + 0.634970996, 0.15237997, 0.3800786, 0.6457891271, 0.21604451, + 0.229598754, 0.12461481, 0.81407372, 0.364336529, 0.45160373, + 0.741084778, 0.53356288, 0.58940162, 0.9995219493, 0.19142269, + 0.618177813, 0.88876378, 0.51174404, 0.8827708189, 0.12730742, + 0.858657839, 0.9271676, 0.5115294, 0.67865624, 0.28658962, + 0.719224842, 0.0546961, 0.42850897, 0.260336376, 0.1371501, + 0.747422328, 0.87172033, 0.98791964, 0.4018020707, 0.29145664, + 0.3158837, 0.39253551, 0.83610831, 0.0101785748, 0.1955386, + 0.419554105, 0.5586495, 0.19908607, 0.4873921743, 0.35568569, + 0.012786428, 0.46925501, 0.25363201, 0.3429851863, 0.2071495, + 0.887479904, 0.66166194, 0.31100105, 0.2895678403, 0.00117005, + 0.147758652, 0.14108789, 0.0361254, 0.4790630946, 0.47336526, + 0.899947367, 0.03730855, 0.33408769, 0.368503517, 0.30600202, + 0.527616998, 0.22344076, 0.20325828, 0.9296060419, 0.34518043, + 0.947085596, 0.85906392, 0.35535464, 0.529360628, 0.8781696, + 0.716097994, 0.9149628, 0.11689428, 0.1157865208, 0.31602707, + 0.433331308, 0.53848417, 0.34146036, 0.4967994317, 0.12822296, + 0.420861259, 0.28802486, 0.62324752, 0.2045601751, 0.06909585, + 0.275279159, 0.69079999, 0.29498051, 0.0082852058, 0.45247107, + 0.909681016, 0.35067747, 0.66813255, 0.3866910117, 0.65315347, + 0.828031845, 0.74096924, 0.33982958, 0.0009472317, 0.65103292, + 0.261653444, 0.43179244, 0.89632155, 0.8636559783, 0.93461464, + 0.209384357, 0.12561389, 0.69809409, 0.4752417156, 0.34963379, + 0.737655852, 0.42078584, 0.09970929, 0.5218528947, 0.36737846, + 0.975034732, 0.69977514, 0.33918481, 0.5443784453, 0.35411297, + 0.053533786, 0.98021485, 0.71035393, 0.189234901, 0.73372176, + 0.364139644, 0.47595789, 0.24620073, 0.4284725219, 0.46145259, + 0.696115067, 0.18095114, 0.66919045, 0.9517078404, 0.31686943, + 0.920878008, 0.89758374, 0.21445324, 0.5666448742, 0.29554824, + 0.397853079, 0.12019741, 0.10775046, 0.0799620333, 0.20065807, + 0.322087545, 0.68342919, 0.29873607, 0.0044371644, 0.66733723, + 0.661407114, 0.0558764, 0.10688295, 0.067841246, 0.52254161, + 0.593253554, 0.40498486, 0.97342655, 0.1917967587, 0.2078643, + 0.392762915, 0.91608107, 0.98894976, 0.3599016496, 0.70576753, + 0.758995247, 0.19899099, 0.95978035, 0.8000916124, 0.8356055, + 0.105617762, 0.12135206, 0.47523114, 0.3594282658, 0.71053726, + 0.754330984, 0.803395, 0.11297253, 0.5072350584, 0.05109695, + 0.410083859, 0.13842349, 0.3671543, 0.262290115, 0.76582706, + 0.498883172, 0.52094766, 0.23674406, 0.8919167451, 0.26313017, + 0.315790046, 0.57934811, 0.96794023, 0.7292640421, 0.63874656, + 0.969918807, 0.86839672, 0.17867962, 0.797609952, 0.3123159, + 0.291589217, 0.37982099, 0.92081884, 0.3760313739, 0.30599535, + 0.874146047, 0.64472863, 0.74944373, 0.0179410274, 0.06637048, + 0.006168369, 0.36819005, 0.48640614, 0.5182905369, 0.37514676, + 0.018794786, 0.50404546, 0.30706335, 0.239409535, 0.78368968, + 0.218041312, 0.08232156, 0.910968, 0.236348928, 0.08734924, + 0.240712896, 0.81851635, 0.75910757, 0.7666831033, 0.51030368, + 0.32422135, 0.37234399, 0.4268269, 0.0688136201, 0.52522145, + 0.737050103, 0.55333162, 0.35681609, 0.5527229193, 0.45528166, + 0.666105454, 0.44928217, 0.93068357, 0.2682658806, 0.47992145, + 0.072705164, 0.24379538, 0.36250275, 0.2693803106, 0.88583253, + 0.393483048, 0.7180344, 0.88936403, 0.9690254654, 0.41720031, + 0.726532397, 0.15675097, 0.14675637, 0.973136256, 0.86701643, + 0.206543021, 0.70612692, 0.9923119, 0.1270776591, 0.43317344, + 0.392393596, 0.6581254, 0.51121301, 0.8005079071, 0.16056554, + 0.326374607, 0.48817642, 0.68630408, 0.9265561129, 0.48683193, + 0.761818521, 0.71751337, 0.83854992, 0.134206275, 0.25700676, + 0.930924999, 0.37469277, 0.42861545, 0.7379696709, 0.9670993, + 0.601101112, 0.56631699, 0.85690728, 0.0792362478, 0.23640603, + 0.294070227, 0.02818223, 0.83060893, 0.8203584203, 0.17647972, + 0.393978659, 0.88639966, 0.80788018, 0.4202279691, 0.75344798, + 0.381183787, 0.98751161, 0.13933232, 0.5427466533, 0.15809025, + 0.203872876, 0.31032719, 0.53000948, 0.6001499062, 0.43581315, + 0.355075927, 0.10865708, 0.21823445, 0.5707600345, 0.84459087, + 0.415892882, 0.09056941, 0.85957968, 0.9296874236, 0.39317951, + 0.885163931, 0.60617414, 0.22888755, 0.9225545505, 0.41601782, + 0.803631177, 0.63855664, 0.4968153, 0.4970232591, 0.28230652, + 0.755692566, 0.36382158, 0.31492054, 0.9853899847, 0.45864754, + 0.761099141, 0.88094342, 0.82542666, 0.977985516, 0.5416208, + 0.536037115, 0.19298885, 0.67674639, 0.213044832, 0.29409245, + 0.050087478, 0.56597845, 0.22309031, 0.7668617836, 0.02385271, + 0.847882026, 0.86580035, 0.8381724, 0.618777399, 0.4707389, + 0.280194086, 0.95490103, 0.27399251, 0.5894525715, 0.17181438, + 0.261382768, 0.96124295, 0.33737123, 0.3545607659, 0.36367031, + 0.465759262, 0.17167592, 0.87114988, 0.4175856721, 0.16020522, + 0.982323635, 0.30892377, 0.96513595, 0.376671114, 0.9411435, + 0.851789546, 0.42260807, 0.37396782, 0.0759502219, 0.41219659, + 0.23932738, 0.70124641, 0.08544481, 0.8599137105, 0.35298377, + 0.985171556, 0.48493665, 0.92919919, 0.3128095574, 0.84388465, + 0.936608667, 0.70159722, 0.23570122, 0.5124408882, 0.99478731, + 0.328337863, 0.83252833, 0.29078719, 0.7531193637, 0.49378383, + 0.504403078, 0.72845174, 0.12801659, 0.5383322216, 0.12559066, + 0.906952623, 0.36801267, 0.13168735, 0.9791060984, 0.14008791, + 0.454210506, 0.67248289, 0.4041049, 0.234963659, 0.92138674, + 0.499037576, 0.7534805, 0.4168877, 0.6275620307, 0.24189188, + 0.707788941, 0.91990553, 0.56701198, 0.1408275496, 0.80566006, + 0.694437274, 0.69339343, 0.42296251, 0.8271595608, 0.53699966, + 0.447118821, 0.97512181, 0.16431204, 0.3697280197, 0.38753206, + 0.885936489, 0.94468978, 0.48918779, 0.3676202064, 0.06938232, + 0.593980148, 0.28140352, 0.27760537, 0.2819242389, 0.8730862, + 0.04248501, 0.45279893, 0.69760642, 0.0949480394, 0.42568701, + 0.35842742, 0.68098838, 0.82745029, 0.5315801166, 0.31104918, + 0.724621041, 0.28763999, 0.48743089, 0.8648093319, 0.93792148, + 0.961828358, 0.5548953, 0.7250596, 0.249875583, 0.90661302, + 0.251438316, 0.86021024, 0.65037498, 0.209739062, 0.07886205, + 0.699615913, 0.12223695, 0.20393331, 0.6357937951, 0.81502268, + 0.391076967, 0.25143855, 0.16091307, 0.6037441837, 0.50651534, + 0.343597198, 0.82570727, 0.62455707, 0.6284155636, 0.17288776, + 0.451352309, 0.29346835, 0.12641623, 0.1194773833, 0.88849468; + // Leaf regression basis output.omega << 0.97801674, 0.34045661, 0.20528387, 0.76230322, 0.63244655, 0.61225851, 0.40492125, 0.33112223, - 0.86917047, 0.58444831, 0.33316433, 0.62217709, 0.96820668, 0.20778425, 0.23764591, 0.94193115, - 0.03869153, 0.60847765, 0.51535811, 0.81554404, 0.78515289, 0.23337815, 0.16730957, 0.02168331, - 0.08699654, 0.34067049, 0.93141264, 0.03679176, 0.4364772, 0.2644173, 0.23717182, 0.59084776, - 0.63438143, 0.57132227, 0.17568721, 0.15552373, 0.8625478, 0.02466334, 0.47269628, 0.97782225, - 0.90593388, 0.82272111, 0.67374992, 0.47619752, 0.5276532, 0.75182919, 0.09559243, 0.5126907, - 0.45892102, 0.11357212, 0.77861167, 0.78424907, 0.84693988, 0.38814934, 0.01010333, 0.10064384, - 0.68664865, 0.1264298, 0.14314708, 0.62679815, 0.71101772, 0.43504811, 0.8868721, 0.95098048, - 0.38291537, 0.71337451, 0.12109764, 0.68943347, 0.89878588, 0.67524475, 0.95549402, 0.58758459, - 0.68558459, 0.16794963, 0.23680754, 0.40289479, 0.98291039, 0.87276966, 0.76995475, 0.55282963, - 0.12448394, 0.5479543, 0.8718802, 0.14515363, 0.71311006, 0.39196408, 0.94504373, 0.44020353, - 0.24090674, 0.52675625, 0.86674581, 0.90576332, 0.09167602, 0.74795585, 0.26901811, 0.544173, - 0.03336554, 0.8314331, 0.27185696, 0.83434459; - + 0.86917047, 0.58444831, 0.33316433, 0.62217709, 0.96820668, 0.20778425, 0.23764591, 0.94193115, + 0.03869153, 0.60847765, 0.51535811, 0.81554404, 0.78515289, 0.23337815, 0.16730957, 0.02168331, + 0.08699654, 0.34067049, 0.93141264, 0.03679176, 0.4364772, 0.2644173, 0.23717182, 0.59084776, + 0.63438143, 0.57132227, 0.17568721, 0.15552373, 0.8625478, 0.02466334, 0.47269628, 0.97782225, + 0.90593388, 0.82272111, 0.67374992, 0.47619752, 0.5276532, 0.75182919, 0.09559243, 0.5126907, + 0.45892102, 0.11357212, 0.77861167, 0.78424907, 0.84693988, 0.38814934, 0.01010333, 0.10064384, + 0.68664865, 0.1264298, 0.14314708, 0.62679815, 0.71101772, 0.43504811, 0.8868721, 0.95098048, + 0.38291537, 0.71337451, 0.12109764, 0.68943347, 0.89878588, 0.67524475, 0.95549402, 0.58758459, + 0.68558459, 0.16794963, 0.23680754, 0.40289479, 0.98291039, 0.87276966, 0.76995475, 0.55282963, + 0.12448394, 0.5479543, 0.8718802, 0.14515363, 0.71311006, 0.39196408, 0.94504373, 0.44020353, + 0.24090674, 0.52675625, 0.86674581, 0.90576332, 0.09167602, 0.74795585, 0.26901811, 0.544173, + 0.03336554, 0.8314331, 0.27185696, 0.83434459; + // Outcome output.outcome << 2.158854445, 1.175387297, 0.40481061, 1.751578365, 0.299641379, 0.347249942, 0.546179903, - 1.164750138, 3.389946886, -0.605464414, 1.271432631, 2.203609096, 2.192327323, 0.746140817, - 3.009233058, -0.292800298, 1.752730639, 1.824961588, 2.055603702, -0.153889672, 0.248010541, - 1.099472562, 0.822333874, 1.291797503, 0.877720106, 2.365239601, 0.685716301, 1.445624363, - 1.342180906, 0.148136818, -1.157010472, 2.186988614, 1.523371203, 1.740153725, 0.73351857, - 0.449967161, 1.25200968, 1.155083428, 1.580760814, 3.025557265, 1.488059405, -0.069025021, - 1.100181892, 1.014150762, 0.418207324, 3.210834777, 1.658875834, 2.215173806, 1.351802193, - 1.33331705, 2.357354695, -1.449598055, 1.042660314, 0.404779346, 1.35048031, -0.58922199, - -0.281044393, 0.128478258, 0.006620112, 1.237840372, 1.0999817, 2.245489523, 2.114281687, - 1.337789336, 0.668884629, 2.275744698, 1.483665856, 0.577564239, -0.557180209, 3.810578895, - 0.946494502, 1.464014296, 0.793749131, 2.735140925, 2.037714409, 1.530792369, 1.857142205, - 1.015348805, -0.91839562, 1.924546112, -0.218826033, 1.761318971, 0.928338732, 1.109589807, - 2.165307398, 2.258640565, 1.147428989, 0.332872857, 0.373646084, 0.520770108, 1.857996323, - -1.971537882, 0.962010578, 1.552073631, 0.459464684, -0.149159276, 0.203079262, -0.453721958, 2.152977755, 0.948865461; - + 1.164750138, 3.389946886, -0.605464414, 1.271432631, 2.203609096, 2.192327323, 0.746140817, + 3.009233058, -0.292800298, 1.752730639, 1.824961588, 2.055603702, -0.153889672, 0.248010541, + 1.099472562, 0.822333874, 1.291797503, 0.877720106, 2.365239601, 0.685716301, 1.445624363, + 1.342180906, 0.148136818, -1.157010472, 2.186988614, 1.523371203, 1.740153725, 0.73351857, + 0.449967161, 1.25200968, 1.155083428, 1.580760814, 3.025557265, 1.488059405, -0.069025021, + 1.100181892, 1.014150762, 0.418207324, 3.210834777, 1.658875834, 2.215173806, 1.351802193, + 1.33331705, 2.357354695, -1.449598055, 1.042660314, 0.404779346, 1.35048031, -0.58922199, + -0.281044393, 0.128478258, 0.006620112, 1.237840372, 1.0999817, 2.245489523, 2.114281687, + 1.337789336, 0.668884629, 2.275744698, 1.483665856, 0.577564239, -0.557180209, 3.810578895, + 0.946494502, 1.464014296, 0.793749131, 2.735140925, 2.037714409, 1.530792369, 1.857142205, + 1.015348805, -0.91839562, 1.924546112, -0.218826033, 1.761318971, 0.928338732, 1.109589807, + 2.165307398, 2.258640565, 1.147428989, 0.332872857, 0.373646084, 0.520770108, 1.857996323, + -1.971537882, 0.962010578, 1.552073631, 0.459464684, -0.149159276, 0.203079262, -0.453721958, 2.152977755, 0.948865461; + // Random effects regression basis (i.e. constant, intercept-only RFX model) - output.rfx_basis << 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1; - + output.rfx_basis << 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1; + // Random effects group labels for (int i = 0; i < output.n; i++) { if (i % 2 == 0) { @@ -354,6 +350,6 @@ TestDataset LoadMediumDatasetUnivariateBasis() { return output; } -} +} // namespace TestUtils } // namespace StochTree diff --git a/test/cpp/testutils.h b/test/cpp/testutils.h index 09b43133..1b10b441 100644 --- a/test/cpp/testutils.h +++ b/test/cpp/testutils.h @@ -39,8 +39,8 @@ TestDataset LoadSmallRFXDatasetMultivariateBasis(); /*! Creates a modest dataset (100 observations) */ TestDataset LoadMediumDatasetUnivariateBasis(); -} // namespace TestUtils +} // namespace TestUtils -} // namespace StochTree +} // namespace StochTree #endif // STOCHTREE_TESTUTILS_H_ diff --git a/test/cross_language/generate_predictions.py b/test/cross_language/generate_predictions.py new file mode 100644 index 00000000..f6a81adc --- /dev/null +++ b/test/cross_language/generate_predictions.py @@ -0,0 +1,449 @@ +#!/usr/bin/env python3 +""" +Generate cross-language parity fixtures using Python stochtree. + +Fixture layout for each scenario +----------------------------------------- + .json metadata: seed, n_train, n_test, p, num_gfr, num_burnin, + num_mcmc, and any model-specific params + _.csv (n_obs x num_mcmc) float64 matrix, no header, comma-separated + rows = observations, columns = posterior draws + +Scenarios +--------- + bart_basic continuous BART, no RFX + terms: yhat_train, yhat_test + bart_rfx continuous BART with group random effects [TODO] + terms: yhat_train, yhat_test, rfx_train, rfx_test + bcf_basic BCF with sample_tau_0=True, no variance forest [TODO] + terms: yhat_train, yhat_test, tau_train, tau_test, + mu_train, mu_test + bcf_varforest BCF with variance forest [TODO] + terms: same as bcf_basic + sigma2x_train, sigma2x_test + +Usage +----- + python test/cross_language/generate_predictions.py [--output-dir DIR] +""" +import argparse +import json +import os + +import numpy as np + +# Fixed seed used for both data generation and model sampling. +# R side must use the same seed value. +GLOBAL_SEED = 42 + + +def write_matrix(path: str, mat: np.ndarray) -> None: + """Write (n, S) float64 matrix as CSV with no header.""" + np.savetxt(path, mat, delimiter=",", fmt="%.17g") + + +def write_metadata(path: str, meta: dict) -> None: + with open(path, "w") as f: + json.dump(meta, f, indent=2) + + +def write_model_json(path: str, model) -> None: + """Serialize a fitted model so the other language can load it (WS-E cross-load).""" + with open(path, "w") as f: + f.write(model.to_json()) + + +# --------------------------------------------------------------------------- +# Scenario: bart_basic +# --------------------------------------------------------------------------- + +def scenario_bart_basic(output_dir: str) -> None: + """Continuous BART, no RFX. Writes X_train, y_train, X_test, and prediction matrices.""" + from stochtree import BARTModel + + n = 250 + p = 5 + n_test = 50 + n_train = n - n_test + num_gfr = 10 + num_burnin = 10 + num_mcmc = 50 + + # Generate data — written to fixture so R reads identical inputs + rng = np.random.default_rng(GLOBAL_SEED) + X = rng.uniform(size=(n, p)) + f_x = np.sin(np.pi * X[:, 0]) * X[:, 1] + y = f_x + rng.normal(scale=0.5, size=n) + X_train, X_test = X[:n_train], X[n_train:] + y_train = y[:n_train] + + write_matrix(os.path.join(output_dir, "bart_basic_X_train.csv"), X_train) + write_matrix(os.path.join(output_dir, "bart_basic_y_train.csv"), y_train.reshape(-1, 1)) + write_matrix(os.path.join(output_dir, "bart_basic_X_test.csv"), X_test) + + # Fit model and write predictions + model = BARTModel() + model.sample( + X_train, y_train, + num_gfr=num_gfr, + num_burnin=num_burnin, + num_mcmc=num_mcmc, + general_params={"random_seed": GLOBAL_SEED}, + ) + write_matrix( + os.path.join(output_dir, "bart_basic_yhat_train.csv"), + model.predict(X_train, terms="y_hat"), + ) + write_matrix( + os.path.join(output_dir, "bart_basic_yhat_test.csv"), + model.predict(X_test, terms="y_hat"), + ) + + meta = { + "scenario": "bart_basic", + "seed": GLOBAL_SEED, + "n": n, + "p": p, + "n_train": n_train, + "n_test": n_test, + "num_gfr": num_gfr, + "num_burnin": num_burnin, + "num_mcmc": num_mcmc, + "params": {}, + } + write_model_json(os.path.join(output_dir, "bart_basic_py_model.json"), model) + write_metadata(os.path.join(output_dir, "bart_basic.json"), meta) + + print(" bart_basic — OK") + + +# --------------------------------------------------------------------------- +# Scenario: bcf_basic +# --------------------------------------------------------------------------- + +def scenario_bcf_basic(output_dir: str) -> None: + """BCF with sample_tau_0=True, no variance forest, no RFX. + + DGP: mu(X) and tau(X) are step functions of X[:,0] and X[:,1] respectively, + propensity pi(X) is a step function of X[:,0]. Matches the example in the + predict.bcfmodel Rd docs so the scenario is easy to reason about. + + Terms written: yhat, tau (= tau_0 + tau(X)), mu (prognostic forest). + """ + from stochtree import BCFModel + + n = 500 + p = 5 + n_test = 100 + n_train = n - n_test + num_gfr = 10 + num_burnin = 10 + num_mcmc = 50 + + rng = np.random.default_rng(GLOBAL_SEED) + X = rng.uniform(size=(n, p)) + + # Step-function propensity and DGP (canonical BCF example) + pi_x = ( + ((X[:, 0] >= 0.00) & (X[:, 0] < 0.25)) * 0.2 + + ((X[:, 0] >= 0.25) & (X[:, 0] < 0.50)) * 0.4 + + ((X[:, 0] >= 0.50) & (X[:, 0] < 0.75)) * 0.6 + + ((X[:, 0] >= 0.75) & (X[:, 0] <= 1.00)) * 0.8 + ) + Z = rng.binomial(1, pi_x).astype(float) + mu_x = ( + ((X[:, 0] >= 0.00) & (X[:, 0] < 0.25)) * (-7.5) + + ((X[:, 0] >= 0.25) & (X[:, 0] < 0.50)) * (-2.5) + + ((X[:, 0] >= 0.50) & (X[:, 0] < 0.75)) * 2.5 + + ((X[:, 0] >= 0.75) & (X[:, 0] <= 1.00)) * 7.5 + ) + tau_x = ( + ((X[:, 1] >= 0.00) & (X[:, 1] < 0.25)) * 0.5 + + ((X[:, 1] >= 0.25) & (X[:, 1] < 0.50)) * 1.0 + + ((X[:, 1] >= 0.50) & (X[:, 1] < 0.75)) * 1.5 + + ((X[:, 1] >= 0.75) & (X[:, 1] <= 1.00)) * 2.0 + ) + y = mu_x + tau_x * Z + rng.normal(scale=1.0, size=n) + + X_train, X_test = X[:n_train], X[n_train:] + Z_train, Z_test = Z[:n_train], Z[n_train:] + pi_train, pi_test = pi_x[:n_train], pi_x[n_train:] + y_train = y[:n_train] + + # Write data fixtures (R reads identical inputs) + write_matrix(os.path.join(output_dir, "bcf_basic_X_train.csv"), X_train) + write_matrix(os.path.join(output_dir, "bcf_basic_Z_train.csv"), Z_train.reshape(-1, 1)) + write_matrix(os.path.join(output_dir, "bcf_basic_y_train.csv"), y_train.reshape(-1, 1)) + write_matrix(os.path.join(output_dir, "bcf_basic_pi_train.csv"), pi_train.reshape(-1, 1)) + write_matrix(os.path.join(output_dir, "bcf_basic_X_test.csv"), X_test) + write_matrix(os.path.join(output_dir, "bcf_basic_Z_test.csv"), Z_test.reshape(-1, 1)) + write_matrix(os.path.join(output_dir, "bcf_basic_pi_test.csv"), pi_test.reshape(-1, 1)) + + model = BCFModel() + model.sample( + X_train, Z_train, y_train, pi_train, + num_gfr=num_gfr, + num_burnin=num_burnin, + num_mcmc=num_mcmc, + general_params={"random_seed": GLOBAL_SEED}, + ) + + for split, X_s, Z_s, pi_s in [ + ("train", X_train, Z_train, pi_train), + ("test", X_test, Z_test, pi_test), + ]: + # terms="all" → dict; keys are y_hat, tau_hat, mu_hat, ... + import warnings + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + preds = model.predict(X_s, Z_s, propensity=pi_s) + write_matrix( + os.path.join(output_dir, f"bcf_basic_yhat_{split}.csv"), preds["y_hat"] + ) + write_matrix( + os.path.join(output_dir, f"bcf_basic_tau_{split}.csv"), preds["tau_hat"] + ) + write_matrix( + os.path.join(output_dir, f"bcf_basic_mu_{split}.csv"), preds["mu_hat"] + ) + + meta = { + "scenario": "bcf_basic", + "seed": GLOBAL_SEED, + "n": n, + "p": p, + "n_train": n_train, + "n_test": n_test, + "num_gfr": num_gfr, + "num_burnin": num_burnin, + "num_mcmc": num_mcmc, + "params": {"sample_intercept": True, "num_trees_variance": 0}, + } + write_model_json(os.path.join(output_dir, "bcf_basic_py_model.json"), model) + write_metadata(os.path.join(output_dir, "bcf_basic.json"), meta) + + print(" bcf_basic — OK") + + +# --------------------------------------------------------------------------- +# Scenario: bart_rfx +# --------------------------------------------------------------------------- + +def scenario_bart_rfx(output_dir: str) -> None: + """Continuous BART with intercept-only group random effects. + + Group IDs are 1-indexed integers (1..num_groups) so R's factor() conversion + produces the same integers Python uses. + Terms written: yhat_train, yhat_test, rfx_train, rfx_test. + """ + from stochtree import BARTModel + + n = 300 + p = 5 + n_test = 60 + n_train = n - n_test + num_groups = 8 + num_gfr = 10 + num_burnin = 10 + num_mcmc = 50 + + rng = np.random.default_rng(GLOBAL_SEED) + X = rng.uniform(size=(n, p)) + group_ids = rng.integers(1, num_groups + 1, size=n) # 1..num_groups + group_effects = rng.normal(scale=2.0, size=num_groups + 1) # index 0 unused + f_x = np.sin(np.pi * X[:, 0]) * X[:, 1] + y = f_x + group_effects[group_ids] + rng.normal(scale=0.5, size=n) + + X_train, X_test = X[:n_train], X[n_train:] + g_train, g_test = group_ids[:n_train], group_ids[n_train:] + y_train = y[:n_train] + + write_matrix(os.path.join(output_dir, "bart_rfx_X_train.csv"), X_train) + write_matrix(os.path.join(output_dir, "bart_rfx_y_train.csv"), y_train.reshape(-1, 1)) + write_matrix(os.path.join(output_dir, "bart_rfx_group_train.csv"), g_train.reshape(-1, 1)) + write_matrix(os.path.join(output_dir, "bart_rfx_X_test.csv"), X_test) + write_matrix(os.path.join(output_dir, "bart_rfx_group_test.csv"), g_test.reshape(-1, 1)) + + model = BARTModel() + model.sample( + X_train, y_train, + rfx_group_ids_train=g_train, + num_gfr=num_gfr, + num_burnin=num_burnin, + num_mcmc=num_mcmc, + general_params={"random_seed": GLOBAL_SEED}, + random_effects_params={"model_spec": "intercept_only"}, + ) + + for split, X_s, g_s in [("train", X_train, g_train), ("test", X_test, g_test)]: + import warnings + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + preds = model.predict(X_s, rfx_group_ids=g_s.astype(np.int64)) + write_matrix( + os.path.join(output_dir, f"bart_rfx_yhat_{split}.csv"), preds["y_hat"] + ) + write_matrix( + os.path.join(output_dir, f"bart_rfx_rfx_{split}.csv"), preds["rfx_predictions"] + ) + + meta = { + "scenario": "bart_rfx", + "seed": GLOBAL_SEED, + "n": n, + "p": p, + "n_train": n_train, + "n_test": n_test, + "num_groups": num_groups, + "num_gfr": num_gfr, + "num_burnin": num_burnin, + "num_mcmc": num_mcmc, + "params": {"model_spec": "intercept_only"}, + } + write_model_json(os.path.join(output_dir, "bart_rfx_py_model.json"), model) + write_metadata(os.path.join(output_dir, "bart_rfx.json"), meta) + + print(" bart_rfx — OK") + + +# --------------------------------------------------------------------------- +# Scenario: bcf_varforest +# --------------------------------------------------------------------------- + +def scenario_bcf_varforest(output_dir: str) -> None: + """BCF with sample_tau_0=True AND a variance forest (num_trees=50). + + Same DGP and data layout as bcf_basic. Uses a different RNG seed offset + so the actual data differs from bcf_basic. + Terms written: yhat, tau, mu, sigma2x (variance forest). + """ + from stochtree import BCFModel + + SEED = GLOBAL_SEED + 1 # distinct from bcf_basic + + n = 500 + p = 5 + n_test = 100 + n_train = n - n_test + num_gfr = 10 + num_burnin = 10 + num_mcmc = 50 + num_trees_variance = 50 + + rng = np.random.default_rng(SEED) + X = rng.uniform(size=(n, p)) + + pi_x = ( + ((X[:, 0] >= 0.00) & (X[:, 0] < 0.25)) * 0.2 + + ((X[:, 0] >= 0.25) & (X[:, 0] < 0.50)) * 0.4 + + ((X[:, 0] >= 0.50) & (X[:, 0] < 0.75)) * 0.6 + + ((X[:, 0] >= 0.75) & (X[:, 0] <= 1.00)) * 0.8 + ) + Z = rng.binomial(1, pi_x).astype(float) + mu_x = ( + ((X[:, 0] >= 0.00) & (X[:, 0] < 0.25)) * (-7.5) + + ((X[:, 0] >= 0.25) & (X[:, 0] < 0.50)) * (-2.5) + + ((X[:, 0] >= 0.50) & (X[:, 0] < 0.75)) * 2.5 + + ((X[:, 0] >= 0.75) & (X[:, 0] <= 1.00)) * 7.5 + ) + tau_x = ( + ((X[:, 1] >= 0.00) & (X[:, 1] < 0.25)) * 0.5 + + ((X[:, 1] >= 0.25) & (X[:, 1] < 0.50)) * 1.0 + + ((X[:, 1] >= 0.50) & (X[:, 1] < 0.75)) * 1.5 + + ((X[:, 1] >= 0.75) & (X[:, 1] <= 1.00)) * 2.0 + ) + y = mu_x + tau_x * Z + rng.normal(scale=1.0, size=n) + + X_train, X_test = X[:n_train], X[n_train:] + Z_train, Z_test = Z[:n_train], Z[n_train:] + pi_train, pi_test = pi_x[:n_train], pi_x[n_train:] + y_train = y[:n_train] + + write_matrix(os.path.join(output_dir, "bcf_varforest_X_train.csv"), X_train) + write_matrix(os.path.join(output_dir, "bcf_varforest_Z_train.csv"), Z_train.reshape(-1, 1)) + write_matrix(os.path.join(output_dir, "bcf_varforest_y_train.csv"), y_train.reshape(-1, 1)) + write_matrix(os.path.join(output_dir, "bcf_varforest_pi_train.csv"), pi_train.reshape(-1, 1)) + write_matrix(os.path.join(output_dir, "bcf_varforest_X_test.csv"), X_test) + write_matrix(os.path.join(output_dir, "bcf_varforest_Z_test.csv"), Z_test.reshape(-1, 1)) + write_matrix(os.path.join(output_dir, "bcf_varforest_pi_test.csv"), pi_test.reshape(-1, 1)) + + model = BCFModel() + model.sample( + X_train, Z_train, y_train, pi_train, + num_gfr=num_gfr, + num_burnin=num_burnin, + num_mcmc=num_mcmc, + general_params={"random_seed": SEED}, + variance_forest_params={"num_trees": num_trees_variance}, + ) + + for split, X_s, Z_s, pi_s in [ + ("train", X_train, Z_train, pi_train), + ("test", X_test, Z_test, pi_test), + ]: + import warnings + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + preds = model.predict(X_s, Z_s, propensity=pi_s) + write_matrix( + os.path.join(output_dir, f"bcf_varforest_yhat_{split}.csv"), preds["y_hat"] + ) + write_matrix( + os.path.join(output_dir, f"bcf_varforest_tau_{split}.csv"), preds["tau_hat"] + ) + write_matrix( + os.path.join(output_dir, f"bcf_varforest_mu_{split}.csv"), preds["mu_hat"] + ) + write_matrix( + os.path.join(output_dir, f"bcf_varforest_sigma2x_{split}.csv"), + preds["variance_forest_predictions"], + ) + + meta = { + "scenario": "bcf_varforest", + "seed": SEED, + "n": n, + "p": p, + "n_train": n_train, + "n_test": n_test, + "num_gfr": num_gfr, + "num_burnin": num_burnin, + "num_mcmc": num_mcmc, + "params": { + "sample_intercept": True, + "num_trees_variance": num_trees_variance, + }, + } + write_model_json(os.path.join(output_dir, "bcf_varforest_py_model.json"), model) + write_metadata(os.path.join(output_dir, "bcf_varforest.json"), meta) + + print(" bcf_varforest — OK") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main() -> None: + parser = argparse.ArgumentParser( + description="Generate cross-language parity fixtures" + ) + parser.add_argument( + "--output-dir", + default="test/cross_language/fixtures", + help="Directory to write fixture files (default: test/cross_language/fixtures)", + ) + args = parser.parse_args() + + os.makedirs(args.output_dir, exist_ok=True) + print(f"Writing fixtures to: {args.output_dir}") + + scenario_bart_basic(args.output_dir) + scenario_bcf_basic(args.output_dir) + scenario_bart_rfx(args.output_dir) + scenario_bcf_varforest(args.output_dir) + + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/test/cross_language/run_local.sh b/test/cross_language/run_local.sh new file mode 100644 index 00000000..80d4ab7f --- /dev/null +++ b/test/cross_language/run_local.sh @@ -0,0 +1,36 @@ +#!/usr/bin/env bash +# run_local.sh — Run cross-language parity tests locally. +# +# Usage: +# bash test/cross_language/run_local.sh [fixture-dir] +# +# Prerequisites: +# Python pip install -e . (or activate a venv with stochtree installed) +# R R CMD INSTALL . (builds and installs stochtree from source) +# +# For rapid local dev where stochtree is not installed in R, use devtools: +# STOCHTREE_R_DEV=1 bash test/cross_language/run_local.sh +# This sets STOCHTREE_R_DEV and STOCHTREE_REPO_ROOT so that verify_predictions.R +# calls devtools::load_all() instead of library(stochtree). + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +FIXTURE_DIR="${1:-$SCRIPT_DIR/fixtures}" + +mkdir -p "$FIXTURE_DIR" + +echo "==> Generating Python predictions..." +python "$SCRIPT_DIR/generate_predictions.py" --output-dir "$FIXTURE_DIR" + +echo "==> Verifying R predictions match Python..." +# cd to repo root so that devtools::load_all(".") resolves correctly in dev mode +cd "$REPO_ROOT" +export STOCHTREE_REPO_ROOT="$REPO_ROOT" +Rscript "$SCRIPT_DIR/verify_predictions.R" "$FIXTURE_DIR" + +echo "==> Verifying Python loads R-written models (Python<-R cross-load)..." +python "$SCRIPT_DIR/verify_r_models.py" "$FIXTURE_DIR" + +echo "==> Done" diff --git a/test/cross_language/verify_predictions.R b/test/cross_language/verify_predictions.R new file mode 100644 index 00000000..457f9934 --- /dev/null +++ b/test/cross_language/verify_predictions.R @@ -0,0 +1,385 @@ +#!/usr/bin/env Rscript +# verify_predictions.R +# +# Load fixtures generated by generate_predictions.py, reproduce the same +# predictions in R with the same seed and parameters, then compare. +# Exits non-zero if any parity check fails. +# +# Usage: +# Rscript test/cross_language/verify_predictions.R [fixture-dir] +# +# For local dev with devtools::load_all instead of an installed package: +# STOCHTREE_R_DEV=1 Rscript test/cross_language/verify_predictions.R [fixture-dir] +# +# Fixture schema: see generate_predictions.py + +if (nzchar(Sys.getenv("STOCHTREE_R_DEV"))) { + repo_root <- Sys.getenv("STOCHTREE_REPO_ROOT", ".") + devtools::load_all(repo_root, quiet = TRUE) +} else { + library(stochtree) +} +library(jsonlite) + +args <- commandArgs(trailingOnly = TRUE) +fixture_dir <- if (length(args) > 0) args[[1]] else "test/cross_language/fixtures" + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +read_matrix <- function(path) { + as.matrix(read.table(path, sep = ",", header = FALSE)) +} + +read_text <- function(path) { + paste(readLines(path, warn = FALSE), collapse = "") +} + +check_parity <- function(label, python_mat, r_mat, tol = 1e-10) { + stopifnot(identical(dim(python_mat), dim(r_mat))) + max_diff <- max(abs(python_mat - r_mat)) + if (max_diff <= tol) { + cat(sprintf(" PASS %-40s max|diff| = %.2e\n", label, max_diff)) + } else { + cat(sprintf(" FAIL %-40s max|diff| = %.2e (tol %.2e)\n", + label, max_diff, tol)) + stop(paste("Parity check failed:", label)) + } +} + +scenario_exists <- function(fixture_dir, name) { + file.exists(file.path(fixture_dir, paste0(name, ".json"))) +} + +# WS-E cross-load correctness. Cross-language prediction *equality* is confounded +# by ~1-ULP differences between R and numpy CSV float parsing (which flip +# split-boundary observations), so we assert the real serialization guarantee +# directly: the foreign-written forest round-trips into R bit-exactly. A loaded +# forest container is re-serialized and compared to the original named forest at +# full precision. +check_forest_bitexact <- function(label, forest_samples, envelope, forest_name) { + reser <- createCppJson() + reser$add_forest(forest_samples) + got <- jsonlite::fromJSON(reser$return_json_string(), simplifyVector = FALSE) + got <- got$forests$forest_0 + want <- envelope$forests[[forest_name]] + if (identical(got, want)) { + cat(sprintf(" PASS %-44s (forest bit-exact across platforms)\n", label)) + } else { + cat(sprintf(" FAIL %-44s (forest did NOT round-trip bit-exactly)\n", label)) + stop(paste("Cross-load forest round-trip failed:", label)) + } +} + +# Soft prediction sanity check: cross-language predictions agree except for rare +# split-boundary flips from float-parsing noise, so we only report the magnitude +# and guard against gross errors (e.g. a transpose), not exact equality. +soft_pred_info <- function(label, a, b) { + md <- max(abs(a - b)) + # Rows differing beyond float-exchange noise are split-boundary flips from + # ~1-ULP R/numpy CSV parsing differences -- inherently rare and scale-free to + # count. A large fraction would indicate a real cross-load bug (the magnitude + # alone is scale-dependent and not a reliable signal). + frac <- mean(apply(abs(a - b), 1, max) > 1e-6) + cat(sprintf( + " INFO %-40s max|diff|=%.2e, boundary-flipped rows=%.1f%%\n", + label, md, 100 * frac + )) + if (frac > 0.15) { + stop(sprintf( + "Cross-load: %.1f%% of rows differ for %s -- likely a real bug, not float noise", + 100 * frac, label + )) + } +} + +# --------------------------------------------------------------------------- +# Scenario: bart_basic +# --------------------------------------------------------------------------- + +scenario_bart_basic <- function(fixture_dir) { + if (!scenario_exists(fixture_dir, "bart_basic")) { + message("Skipping bart_basic — no fixture found") + return(invisible(NULL)) + } + + cat("--- bart_basic ---\n") + meta <- fromJSON(file.path(fixture_dir, "bart_basic.json")) + + # Load the exact same data Python used (written to fixture by generator) + X_train <- read_matrix(file.path(fixture_dir, "bart_basic_X_train.csv")) + y_train <- as.numeric(read_matrix(file.path(fixture_dir, "bart_basic_y_train.csv"))) + X_test <- read_matrix(file.path(fixture_dir, "bart_basic_X_test.csv")) + + # Fit with the same hyperparameters and seed + model <- bart( + X_train = X_train, y_train = y_train, + num_gfr = meta$num_gfr, + num_burnin = meta$num_burnin, + num_mcmc = meta$num_mcmc, + general_params = list(random_seed = meta$seed) + ) + writeLines( + saveBARTModelToJsonString(model), + file.path(fixture_dir, "bart_basic_r_model.json") + ) + + py_yhat_train <- read_matrix(file.path(fixture_dir, "bart_basic_yhat_train.csv")) + r_yhat_train <- predict(model, X_train, terms = "y_hat") + check_parity("bart_basic / yhat_train", py_yhat_train, r_yhat_train) + + py_yhat_test <- read_matrix(file.path(fixture_dir, "bart_basic_yhat_test.csv")) + r_yhat_test <- predict(model, X_test, terms = "y_hat") + check_parity("bart_basic / yhat_test", py_yhat_test, r_yhat_test) + + # WS-E cross-load: R loads the Python-written model; assert the forest round- + # trips bit-exactly, then sanity-check predictions. + py_path <- file.path(fixture_dir, "bart_basic_py_model.json") + envelope <- fromJSON(read_text(py_path), simplifyVector = FALSE) + py_model <- createBARTModelFromJsonString(read_text(py_path)) + check_forest_bitexact( + "bart_basic / R<-py mean_forest", py_model$mean_forests, envelope, "mean_forest" + ) + soft_pred_info("bart_basic / R<-py yhat_train", py_yhat_train, + predict(py_model, X_train, terms = "y_hat")) +} + + +# --------------------------------------------------------------------------- +# Scenario: bcf_basic +# --------------------------------------------------------------------------- + +scenario_bcf_basic <- function(fixture_dir) { + if (!scenario_exists(fixture_dir, "bcf_basic")) { + message("Skipping bcf_basic — no fixture found") + return(invisible(NULL)) + } + + cat("--- bcf_basic ---\n") + meta <- fromJSON(file.path(fixture_dir, "bcf_basic.json")) + + X_train <- read_matrix(file.path(fixture_dir, "bcf_basic_X_train.csv")) + Z_train <- as.numeric(read_matrix(file.path(fixture_dir, "bcf_basic_Z_train.csv"))) + y_train <- as.numeric(read_matrix(file.path(fixture_dir, "bcf_basic_y_train.csv"))) + pi_train <- as.numeric(read_matrix(file.path(fixture_dir, "bcf_basic_pi_train.csv"))) + X_test <- read_matrix(file.path(fixture_dir, "bcf_basic_X_test.csv")) + Z_test <- as.numeric(read_matrix(file.path(fixture_dir, "bcf_basic_Z_test.csv"))) + pi_test <- as.numeric(read_matrix(file.path(fixture_dir, "bcf_basic_pi_test.csv"))) + + model <- bcf( + X_train = X_train, Z_train = Z_train, + y_train = y_train, propensity_train = pi_train, + num_gfr = meta$num_gfr, + num_burnin = meta$num_burnin, + num_mcmc = meta$num_mcmc, + general_params = list(random_seed = meta$seed) + ) + writeLines( + saveBCFModelToJsonString(model), + file.path(fixture_dir, "bcf_basic_r_model.json") + ) + + for (split in c("train", "test")) { + X_s <- if (split == "train") X_train else X_test + Z_s <- if (split == "train") Z_train else Z_test + pi_s <- if (split == "train") pi_train else pi_test + + check_parity( + sprintf("bcf_basic / yhat_%s", split), + read_matrix(file.path(fixture_dir, sprintf("bcf_basic_yhat_%s.csv", split))), + predict(model, X_s, Z_s, propensity = pi_s, terms = "y_hat") + ) + check_parity( + sprintf("bcf_basic / tau_%s", split), + read_matrix(file.path(fixture_dir, sprintf("bcf_basic_tau_%s.csv", split))), + predict(model, X_s, Z_s, propensity = pi_s, terms = "tau") + ) + check_parity( + sprintf("bcf_basic / mu_%s", split), + read_matrix(file.path(fixture_dir, sprintf("bcf_basic_mu_%s.csv", split))), + predict(model, X_s, Z_s, propensity = pi_s, terms = "mu") + ) + } + + # WS-E cross-load: R loads the Python-written model; assert forests round-trip + # bit-exactly, then sanity-check predictions. + py_path <- file.path(fixture_dir, "bcf_basic_py_model.json") + envelope <- fromJSON(read_text(py_path), simplifyVector = FALSE) + py_model <- createBCFModelFromJsonString(read_text(py_path)) + check_forest_bitexact( + "bcf_basic / R<-py prognostic", py_model$forests_mu, envelope, "prognostic_forest" + ) + check_forest_bitexact( + "bcf_basic / R<-py treatment", py_model$forests_tau, envelope, "treatment_forest" + ) + soft_pred_info( + "bcf_basic / R<-py yhat_train", + read_matrix(file.path(fixture_dir, "bcf_basic_yhat_train.csv")), + predict(py_model, X_train, Z_train, propensity = pi_train, terms = "y_hat") + ) +} + + +# --------------------------------------------------------------------------- +# Scenario: bart_rfx +# --------------------------------------------------------------------------- + +scenario_bart_rfx <- function(fixture_dir) { + if (!scenario_exists(fixture_dir, "bart_rfx")) { + message("Skipping bart_rfx — no fixture found") + return(invisible(NULL)) + } + + cat("--- bart_rfx ---\n") + meta <- fromJSON(file.path(fixture_dir, "bart_rfx.json")) + + X_train <- read_matrix(file.path(fixture_dir, "bart_rfx_X_train.csv")) + y_train <- as.numeric(read_matrix(file.path(fixture_dir, "bart_rfx_y_train.csv"))) + g_train <- as.integer(as.numeric(read_matrix(file.path(fixture_dir, "bart_rfx_group_train.csv")))) + X_test <- read_matrix(file.path(fixture_dir, "bart_rfx_X_test.csv")) + g_test <- as.integer(as.numeric(read_matrix(file.path(fixture_dir, "bart_rfx_group_test.csv")))) + + model <- bart( + X_train = X_train, y_train = y_train, + rfx_group_ids_train = g_train, + num_gfr = meta$num_gfr, + num_burnin = meta$num_burnin, + num_mcmc = meta$num_mcmc, + general_params = list(random_seed = meta$seed), + random_effects_params = list(model_spec = "intercept_only") + ) + writeLines( + saveBARTModelToJsonString(model), + file.path(fixture_dir, "bart_rfx_r_model.json") + ) + + for (split in c("train", "test")) { + X_s <- if (split == "train") X_train else X_test + g_s <- if (split == "train") g_train else g_test + + check_parity( + sprintf("bart_rfx / yhat_%s", split), + read_matrix(file.path(fixture_dir, sprintf("bart_rfx_yhat_%s.csv", split))), + predict(model, X_s, rfx_group_ids = g_s, terms = "y_hat") + ) + check_parity( + sprintf("bart_rfx / rfx_%s", split), + read_matrix(file.path(fixture_dir, sprintf("bart_rfx_rfx_%s.csv", split))), + predict(model, X_s, rfx_group_ids = g_s, terms = "rfx") + ) + } + + # WS-E cross-load: R loads the Python-written rfx model. Python never writes + # rfx_unique_group_ids, so R reconstructs the levels from the rfx label mapper. + py_path <- file.path(fixture_dir, "bart_rfx_py_model.json") + envelope <- fromJSON(read_text(py_path), simplifyVector = FALSE) + py_model <- createBARTModelFromJsonString(read_text(py_path)) + check_forest_bitexact( + "bart_rfx / R<-py mean_forest", py_model$mean_forests, envelope, "mean_forest" + ) + soft_pred_info( + "bart_rfx / R<-py yhat_train", + read_matrix(file.path(fixture_dir, "bart_rfx_yhat_train.csv")), + predict(py_model, X_train, + rfx_group_ids = g_train, rfx_basis = matrix(1, nrow(X_train), 1), + terms = "y_hat" + ) + ) +} + + +# --------------------------------------------------------------------------- +# Scenario: bcf_varforest +# --------------------------------------------------------------------------- + +scenario_bcf_varforest <- function(fixture_dir) { + if (!scenario_exists(fixture_dir, "bcf_varforest")) { + message("Skipping bcf_varforest — no fixture found") + return(invisible(NULL)) + } + + cat("--- bcf_varforest ---\n") + meta <- fromJSON(file.path(fixture_dir, "bcf_varforest.json")) + + X_train <- read_matrix(file.path(fixture_dir, "bcf_varforest_X_train.csv")) + Z_train <- as.numeric(read_matrix(file.path(fixture_dir, "bcf_varforest_Z_train.csv"))) + y_train <- as.numeric(read_matrix(file.path(fixture_dir, "bcf_varforest_y_train.csv"))) + pi_train <- as.numeric(read_matrix(file.path(fixture_dir, "bcf_varforest_pi_train.csv"))) + X_test <- read_matrix(file.path(fixture_dir, "bcf_varforest_X_test.csv")) + Z_test <- as.numeric(read_matrix(file.path(fixture_dir, "bcf_varforest_Z_test.csv"))) + pi_test <- as.numeric(read_matrix(file.path(fixture_dir, "bcf_varforest_pi_test.csv"))) + + model <- bcf( + X_train = X_train, Z_train = Z_train, + y_train = y_train, propensity_train = pi_train, + num_gfr = meta$num_gfr, + num_burnin = meta$num_burnin, + num_mcmc = meta$num_mcmc, + general_params = list(random_seed = meta$seed), + variance_forest_params = list(num_trees = meta$params$num_trees_variance) + ) + writeLines( + saveBCFModelToJsonString(model), + file.path(fixture_dir, "bcf_varforest_r_model.json") + ) + + for (split in c("train", "test")) { + X_s <- if (split == "train") X_train else X_test + Z_s <- if (split == "train") Z_train else Z_test + pi_s <- if (split == "train") pi_train else pi_test + + check_parity( + sprintf("bcf_varforest / yhat_%s", split), + read_matrix(file.path(fixture_dir, sprintf("bcf_varforest_yhat_%s.csv", split))), + predict(model, X_s, Z_s, propensity = pi_s, terms = "y_hat") + ) + check_parity( + sprintf("bcf_varforest / tau_%s", split), + read_matrix(file.path(fixture_dir, sprintf("bcf_varforest_tau_%s.csv", split))), + predict(model, X_s, Z_s, propensity = pi_s, terms = "tau") + ) + check_parity( + sprintf("bcf_varforest / mu_%s", split), + read_matrix(file.path(fixture_dir, sprintf("bcf_varforest_mu_%s.csv", split))), + predict(model, X_s, Z_s, propensity = pi_s, terms = "mu") + ) + check_parity( + sprintf("bcf_varforest / sigma2x_%s", split), + read_matrix(file.path(fixture_dir, sprintf("bcf_varforest_sigma2x_%s.csv", split))), + predict(model, X_s, Z_s, propensity = pi_s, terms = "variance_forest") + ) + } + + # WS-E cross-load: R loads the Python-written model; assert forests round-trip + # bit-exactly (incl. the variance forest), then sanity-check predictions. + py_path <- file.path(fixture_dir, "bcf_varforest_py_model.json") + envelope <- fromJSON(read_text(py_path), simplifyVector = FALSE) + py_model <- createBCFModelFromJsonString(read_text(py_path)) + check_forest_bitexact( + "bcf_varforest / R<-py prognostic", py_model$forests_mu, envelope, "prognostic_forest" + ) + check_forest_bitexact( + "bcf_varforest / R<-py treatment", py_model$forests_tau, envelope, "treatment_forest" + ) + check_forest_bitexact( + "bcf_varforest / R<-py variance", py_model$forests_variance, envelope, "variance_forest" + ) + soft_pred_info( + "bcf_varforest / R<-py yhat_train", + read_matrix(file.path(fixture_dir, "bcf_varforest_yhat_train.csv")), + predict(py_model, X_train, Z_train, propensity = pi_train, terms = "y_hat") + ) +} + + +# --------------------------------------------------------------------------- +# Dispatch +# --------------------------------------------------------------------------- + +scenario_bart_basic(fixture_dir) +scenario_bcf_basic(fixture_dir) +scenario_bart_rfx(fixture_dir) +scenario_bcf_varforest(fixture_dir) + +cat("Cross-language parity check complete\n") diff --git a/test/cross_language/verify_r_models.py b/test/cross_language/verify_r_models.py new file mode 100644 index 00000000..165cbe81 --- /dev/null +++ b/test/cross_language/verify_r_models.py @@ -0,0 +1,166 @@ +#!/usr/bin/env python3 +"""Verify the Python<-R cross-load direction. + +Companion to verify_predictions.R, which writes ``_r_model.json`` for +each scenario. Here we load each R-written model in Python and assert the real +cross-platform serialization guarantee: the foreign forest reconstructs +**bit-exactly** (re-serialize the loaded forest, compare to the original named +forest at full precision). Predictions are only sanity-checked -- cross-language +prediction *equality* is confounded by ~1-ULP differences between R and numpy +CSV float parsing, which flip rare split-boundary observations -- so we bound the +fraction of flipped rows rather than asserting equality. + +Usage: + python test/cross_language/verify_r_models.py [fixture-dir] +""" + +import json +import os +import sys + +import numpy as np + +from stochtree import BARTModel, BCFModel +from stochtree.serialization import JSONSerializer + + +def read_matrix(path): + return np.loadtxt(path, delimiter=",", ndmin=2) + + +def read_text(path): + with open(path) as f: + return f.read() + + +def check_forest_bitexact(label, forest_container, envelope, forest_name): + s = JSONSerializer() + s.add_forest(forest_container) + got = json.loads(s.return_json_string())["forests"]["forest_0"] + want = envelope["forests"][forest_name] + if got == want: + print(f" PASS {label:<42} (forest bit-exact across platforms)") + else: + print(f" FAIL {label:<42} (forest did NOT round-trip bit-exactly)") + raise SystemExit(f"Cross-load forest round-trip failed: {label}") + + +def soft_pred_info(label, a, b): + diff = np.abs(a - b) + md = float(diff.max()) + frac = float(np.mean(diff.max(axis=1) > 1e-6)) + print( + f" INFO {label:<38} max|diff|={md:.2e}, boundary-flipped rows={100 * frac:.1f}%" + ) + if frac > 0.15: + raise SystemExit( + f"Cross-load: {100 * frac:.1f}% of rows differ for {label} -- likely a real bug" + ) + + +def has_r_model(fixture_dir, name): + return os.path.exists(os.path.join(fixture_dir, name + "_r_model.json")) + + +def scenario_bart_basic(fixture_dir): + name = "bart_basic" + if not has_r_model(fixture_dir, name): + print(f"Skipping {name} -- no R model fixture") + return + print(f"--- {name} ---") + r_path = os.path.join(fixture_dir, name + "_r_model.json") + envelope = json.loads(read_text(r_path)) + m = BARTModel() + m.from_json(read_text(r_path)) + check_forest_bitexact( + f"{name} / py<-R mean_forest", m.forest_container_mean, envelope, "mean_forest" + ) + X = read_matrix(os.path.join(fixture_dir, f"{name}_X_train.csv")) + py = read_matrix(os.path.join(fixture_dir, f"{name}_yhat_train.csv")) + soft_pred_info(f"{name} / py<-R yhat_train", m.predict(X, terms="y_hat"), py) + + +def scenario_bcf_basic(fixture_dir): + name = "bcf_basic" + if not has_r_model(fixture_dir, name): + print(f"Skipping {name} -- no R model fixture") + return + print(f"--- {name} ---") + r_path = os.path.join(fixture_dir, name + "_r_model.json") + envelope = json.loads(read_text(r_path)) + m = BCFModel() + m.from_json(read_text(r_path)) + check_forest_bitexact( + f"{name} / py<-R prognostic", m.forest_container_mu, envelope, "prognostic_forest" + ) + check_forest_bitexact( + f"{name} / py<-R treatment", m.forest_container_tau, envelope, "treatment_forest" + ) + X = read_matrix(os.path.join(fixture_dir, f"{name}_X_train.csv")) + Z = read_matrix(os.path.join(fixture_dir, f"{name}_Z_train.csv")).ravel() + pi = read_matrix(os.path.join(fixture_dir, f"{name}_pi_train.csv")).ravel() + py = read_matrix(os.path.join(fixture_dir, f"{name}_yhat_train.csv")) + soft_pred_info(f"{name} / py<-R yhat_train", m.predict(X, Z, pi, terms="y_hat"), py) + + +def scenario_bart_rfx(fixture_dir): + name = "bart_rfx" + if not has_r_model(fixture_dir, name): + print(f"Skipping {name} -- no R model fixture") + return + print(f"--- {name} ---") + r_path = os.path.join(fixture_dir, name + "_r_model.json") + envelope = json.loads(read_text(r_path)) + m = BARTModel() + m.from_json(read_text(r_path)) + check_forest_bitexact( + f"{name} / py<-R mean_forest", m.forest_container_mean, envelope, "mean_forest" + ) + X = read_matrix(os.path.join(fixture_dir, f"{name}_X_train.csv")) + g = read_matrix(os.path.join(fixture_dir, f"{name}_group_train.csv")).ravel().astype(np.int32) + basis = np.ones((X.shape[0], 1)) + py = read_matrix(os.path.join(fixture_dir, f"{name}_yhat_train.csv")) + soft_pred_info( + f"{name} / py<-R yhat_train", + m.predict(X, rfx_group_ids=g, rfx_basis=basis, terms="y_hat"), + py, + ) + + +def scenario_bcf_varforest(fixture_dir): + name = "bcf_varforest" + if not has_r_model(fixture_dir, name): + print(f"Skipping {name} -- no R model fixture") + return + print(f"--- {name} ---") + r_path = os.path.join(fixture_dir, name + "_r_model.json") + envelope = json.loads(read_text(r_path)) + m = BCFModel() + m.from_json(read_text(r_path)) + check_forest_bitexact( + f"{name} / py<-R prognostic", m.forest_container_mu, envelope, "prognostic_forest" + ) + check_forest_bitexact( + f"{name} / py<-R treatment", m.forest_container_tau, envelope, "treatment_forest" + ) + check_forest_bitexact( + f"{name} / py<-R variance", m.forest_container_variance, envelope, "variance_forest" + ) + X = read_matrix(os.path.join(fixture_dir, f"{name}_X_train.csv")) + Z = read_matrix(os.path.join(fixture_dir, f"{name}_Z_train.csv")).ravel() + pi = read_matrix(os.path.join(fixture_dir, f"{name}_pi_train.csv")).ravel() + py = read_matrix(os.path.join(fixture_dir, f"{name}_yhat_train.csv")) + soft_pred_info(f"{name} / py<-R yhat_train", m.predict(X, Z, pi, terms="y_hat"), py) + + +def main(): + fixture_dir = sys.argv[1] if len(sys.argv) > 1 else "test/cross_language/fixtures" + scenario_bart_basic(fixture_dir) + scenario_bcf_basic(fixture_dir) + scenario_bart_rfx(fixture_dir) + scenario_bcf_varforest(fixture_dir) + print("Cross-language (Python<-R) check complete") + + +if __name__ == "__main__": + main() diff --git a/test/python/fixtures/bart_categorical_rfx_v1.json b/test/python/fixtures/bart_categorical_rfx_v1.json new file mode 100644 index 00000000..a709ea64 --- /dev/null +++ b/test/python/fixtures/bart_categorical_rfx_v1.json @@ -0,0 +1 @@ +{"covariate_preprocessor":"{\"cross_platform_portable\":false,\"forests\":{},\"is_fitted\":true,\"non_portable_columns\":\"4\",\"num_forests\":0,\"num_onehot_features\":1,\"num_ordinal_features\":0,\"num_original_features\":5,\"num_random_effects\":0,\"onehot_categories_list\":{\"cats_0\":[\"a\",\"b\",\"c\"]},\"onehot_dtype_list\":{\"dtype_0\":\"str\"},\"onehot_feature_index\":[-1,-1,-1,-1,0],\"ordinal_feature_index\":[-1,-1,-1,-1,-1],\"original_feature_indices\":[0,1,2,3,4,4,4],\"original_feature_types\":[\"float\",\"float\",\"float\",\"float\",\"category\"],\"processed_feature_types\":[0,0,0,0,1,1,1],\"random_effects\":{}}","forests":{"mean_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.5236632364273408,0.05066670237417858],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.3034134423993235,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.18586932454411767,0.23613016513770765],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.2044247659188199,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.01695052115399959,0.4972825661177521],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.966054982767047,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.19629841299356368,-0.03304791231944931],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.17128146157459764,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.3531608850598189,-0.17052356285511602],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.31546743080873424,0.0,0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.6967626943975539,-0.15549632083222184],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.3034134423993235,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.24775613076837677,0.12915057719343231],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.2044247659188199,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.050907143001618985,-0.036836883144944765],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.966054982767047,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.09743693408295125,-0.07023382687589891],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.17128146157459764,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.20130360407134362,-0.1367153008782409],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.31546743080873424,0.0,0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":true,"num_samples":2,"num_trees":5,"output_dimension":1}},"has_rfx":true,"has_rfx_basis":true,"include_mean_forest":true,"include_variance_forest":false,"keep_every":1,"num_basis":0,"num_burnin":0,"num_chains":1,"num_forests":1,"num_gfr":0,"num_mcmc":2,"num_random_effects":1,"num_rfx_basis":1.0,"num_samples":2,"outcome_mean":1.5073242708159946,"outcome_model":{"link":"identity","outcome":"continuous"},"outcome_scale":1.01099404839377,"parameters":{"sigma2_global_samples":[0.9753108405816692,0.45786116282820194],"sigma2_leaf_samples":[0.06004940980963514,0.11600181269218321]},"platform":"python","probit_outcome_model":false,"random_effects":{"cross_platform_compatible":true,"random_effect_container_0":{"alpha":[-0.5068091214498942,-0.5918214818198173],"alpha_size":2,"beta":[-0.33261075738654383,0.5363340784095936,0.8811658557516957,-0.6551318133231622,0.2859689662887389,1.472379330318233],"beta_size":6,"num_components":1,"num_groups":3,"num_samples":2,"sigma_xi":[2.143303169938896,2.2735120090136895],"xi":[0.6562840787770381,-1.0582565619088178,-1.7386542949953838,1.1069753860719438,-0.48320139615311136,-2.4878774690481844]},"random_effect_groupids_0":[0,1,2],"random_effect_label_mapper_0":{"keys":[0,1,2],"values":[0,1,2]}},"requires_basis":false,"rfx_model_spec":"custom","sample_sigma2_global":true,"sample_sigma2_leaf":true,"schema_version":1,"sigma2_init":1.0,"standardize":true,"stochtree_version":"0.4.5.dev0"} \ No newline at end of file diff --git a/test/python/fixtures/bart_categorical_v1.json b/test/python/fixtures/bart_categorical_v1.json new file mode 100644 index 00000000..86dfd259 --- /dev/null +++ b/test/python/fixtures/bart_categorical_v1.json @@ -0,0 +1 @@ +{"covariate_preprocessor":"{\"cross_platform_portable\":false,\"forests\":{},\"is_fitted\":true,\"non_portable_columns\":\"4\",\"num_forests\":0,\"num_onehot_features\":1,\"num_ordinal_features\":0,\"num_original_features\":5,\"num_random_effects\":0,\"onehot_categories_list\":{\"cats_0\":[\"a\",\"b\",\"c\"]},\"onehot_dtype_list\":{\"dtype_0\":\"str\"},\"onehot_feature_index\":[-1,-1,-1,-1,0],\"ordinal_feature_index\":[-1,-1,-1,-1,-1],\"original_feature_indices\":[0,1,2,3,4,4,4],\"original_feature_types\":[\"float\",\"float\",\"float\",\"float\",\"category\"],\"processed_feature_types\":[0,0,0,0,1,1,1],\"random_effects\":{}}","forests":{"mean_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.6971117541939098,0.1485112085002089],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.3222570769557214,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.2330086619337179,0.20665351652683173],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.20644166551501808,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.002712252352572681,0.17529524881429925],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.9676494332327997,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.26794055711412756,-0.002620057011379868],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.19377879466894937,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.22329666759608918,-0.10931033340909625],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.3162852092413055,0.0,0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.4771251923240624,0.07186884742394903],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.3222570769557214,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.06704097270692166,0.13417061645089282],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.20644166551501808,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.056950152988485526,0.48702482405695957],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.9676494332327997,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.39604703348728415,0.03279105393200044],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.19377879466894937,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.03279147924368678,0.004550508243007983],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.3162852092413055,0.0,0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":true,"num_samples":2,"num_trees":5,"output_dimension":1}},"has_rfx":false,"has_rfx_basis":false,"include_mean_forest":true,"include_variance_forest":false,"keep_every":1,"num_basis":0,"num_burnin":0,"num_chains":1,"num_forests":1,"num_gfr":0,"num_mcmc":2,"num_random_effects":0,"num_rfx_basis":0.0,"num_samples":2,"outcome_mean":0.4667257063943834,"outcome_model":{"link":"identity","outcome":"continuous"},"outcome_scale":0.5584801778547619,"parameters":{"sigma2_global_samples":[0.08116406662051638,0.07542315150857114],"sigma2_leaf_samples":[0.05280929164207198,0.06286707931310706]},"platform":"python","probit_outcome_model":false,"random_effects":{},"requires_basis":false,"rfx_model_spec":"custom","sample_sigma2_global":true,"sample_sigma2_leaf":true,"schema_version":1,"sigma2_init":1.0,"standardize":true,"stochtree_version":"0.4.5.dev0"} \ No newline at end of file diff --git a/test/python/fixtures/bart_numeric_rfx_v1.json b/test/python/fixtures/bart_numeric_rfx_v1.json new file mode 100644 index 00000000..27c6510d --- /dev/null +++ b/test/python/fixtures/bart_numeric_rfx_v1.json @@ -0,0 +1 @@ +{"covariate_preprocessor":"{\"cross_platform_portable\":true,\"forests\":{},\"is_fitted\":true,\"num_forests\":0,\"num_onehot_features\":0,\"num_ordinal_features\":0,\"num_original_features\":4,\"num_random_effects\":0,\"onehot_feature_index\":[-1,-1,-1,-1],\"ordinal_feature_index\":[-1,-1,-1,-1],\"original_feature_indices\":[0,1,2,3],\"original_feature_types\":[\"float\",\"float\",\"float\",\"float\"],\"processed_feature_types\":[0,0,0,0],\"random_effects\":{}}","forests":{"mean_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.5008973537723335,0.06568246632817014],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.30135352239689017,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.1715161281492864,0.2307279448228717],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.22542524665970431,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.0018273458653479703,0.29781182134116324],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.9616875040728227,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.33359997983568873,-0.003500378032897418],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.1697221241628432,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.3964843736503252,-0.16269010421516955],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.33348295584540016,0.0,0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[1,2],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.47941685561498,-0.5008973537723335,0.06568246632817014],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[0],"left":[-1,-1,-1],"node_deleted":[false,true,true],"node_type":[0,0,0],"num_deleted_nodes":2,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[-1,-1,-1],"split_index":[0,-1,-1],"threshold":[0.30135352239689017,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.0063995190150948275,0.2948315810130058],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.22542524665970431,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.061412935969462436,0.31336000970212124],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.9616875040728227,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.393033301226236,0.0753495256706962],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.1697221241628432,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.08274455207030383,-0.029725309567544527],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.33348295584540016,0.0,0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":true,"num_samples":2,"num_trees":5,"output_dimension":1}},"has_rfx":true,"has_rfx_basis":true,"include_mean_forest":true,"include_variance_forest":false,"keep_every":1,"num_basis":0,"num_burnin":0,"num_chains":1,"num_forests":1,"num_gfr":0,"num_mcmc":2,"num_random_effects":1,"num_rfx_basis":1.0,"num_samples":2,"outcome_mean":1.4613251529417188,"outcome_model":{"link":"identity","outcome":"continuous"},"outcome_scale":0.9897050045474932,"parameters":{"sigma2_global_samples":[0.8921235974634758,0.40359007385112583],"sigma2_leaf_samples":[0.05405785658558741,0.11603385528268308]},"platform":"python","probit_outcome_model":false,"random_effects":{"cross_platform_compatible":true,"random_effect_container_0":{"alpha":[-0.47975708077148904,-0.5572337729517062],"alpha_size":2,"beta":[-0.31485693379087776,0.5077060788486469,0.8341317090377967,-0.8051578975061499,0.28630877230413593,1.3240694102843011],"beta_size":6,"num_components":1,"num_groups":3,"num_samples":2,"sigma_xi":[2.143303169938896,2.365375856286227],"xi":[0.6562840787770381,-1.0582565619088178,-1.7386542949953838,1.4449194154926617,-0.5138036964047215,-2.3761470940115728]},"random_effect_groupids_0":[0,1,2],"random_effect_label_mapper_0":{"keys":[0,1,2],"values":[0,1,2]}},"requires_basis":false,"rfx_model_spec":"custom","sample_sigma2_global":true,"sample_sigma2_leaf":true,"schema_version":1,"sigma2_init":1.0,"standardize":true,"stochtree_version":"0.4.5.dev0"} \ No newline at end of file diff --git a/test/python/fixtures/bart_numeric_v1.json b/test/python/fixtures/bart_numeric_v1.json new file mode 100644 index 00000000..ffeb98f0 --- /dev/null +++ b/test/python/fixtures/bart_numeric_v1.json @@ -0,0 +1 @@ +{"covariate_preprocessor":"{\"cross_platform_portable\":true,\"forests\":{},\"is_fitted\":true,\"num_forests\":0,\"num_onehot_features\":0,\"num_ordinal_features\":0,\"num_original_features\":4,\"num_random_effects\":0,\"onehot_feature_index\":[-1,-1,-1,-1],\"ordinal_feature_index\":[-1,-1,-1,-1],\"original_feature_indices\":[0,1,2,3],\"original_feature_types\":[\"float\",\"float\",\"float\",\"float\"],\"processed_feature_types\":[0,0,0,0],\"random_effects\":{}}","forests":{"mean_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.6832572750522327,0.1980969982846796],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.30197913266705073,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.3723160377701503,0.26905792877180085],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.20368987279319398,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.005774008731624321,0.21758763841600556],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.9594887140651809,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.2802917297957185,0.009722260647103848],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.16993231200237638,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.23379446205793675,-0.10126767939050017],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.3142445605979907,0.0,0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.49902396179527675,0.11356951117639683],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.30197913266705073,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.22769271433086855,0.2137733566349887],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.20368987279319398,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.04532973924859999,0.48682620384986813],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.9594887140651809,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.39816645673242573,0.03294040272001218],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.16993231200237638,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.03862634838803235,0.008234656890208771],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[2,-1,-1],"threshold":[0.3142445605979907,0.0,0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":true,"num_samples":2,"num_trees":5,"output_dimension":1}},"has_rfx":false,"has_rfx_basis":false,"include_mean_forest":true,"include_variance_forest":false,"keep_every":1,"num_basis":0,"num_burnin":0,"num_chains":1,"num_forests":1,"num_gfr":0,"num_mcmc":2,"num_random_effects":0,"num_rfx_basis":0.0,"num_samples":2,"outcome_mean":0.43845557755270703,"outcome_model":{"link":"identity","outcome":"continuous"},"outcome_scale":0.5743156003426314,"parameters":{"sigma2_global_samples":[0.08482665033320452,0.07928485372848336],"sigma2_leaf_samples":[0.06207584360036479,0.07264822656080219]},"platform":"python","probit_outcome_model":false,"random_effects":{},"requires_basis":false,"rfx_model_spec":"custom","sample_sigma2_global":true,"sample_sigma2_leaf":true,"schema_version":1,"sigma2_init":1.0,"standardize":true,"stochtree_version":"0.4.5.dev0"} \ No newline at end of file diff --git a/test/python/fixtures/bcf_categorical_rfx_v1.json b/test/python/fixtures/bcf_categorical_rfx_v1.json new file mode 100644 index 00000000..54b2235d --- /dev/null +++ b/test/python/fixtures/bcf_categorical_rfx_v1.json @@ -0,0 +1 @@ +{"adaptive_coding":false,"binary_treatment":true,"covariate_preprocessor":"{\"cross_platform_portable\":false,\"forests\":{},\"is_fitted\":true,\"non_portable_columns\":\"4\",\"num_forests\":0,\"num_onehot_features\":1,\"num_ordinal_features\":0,\"num_original_features\":5,\"num_random_effects\":0,\"onehot_categories_list\":{\"cats_0\":[\"a\",\"b\",\"c\"]},\"onehot_dtype_list\":{\"dtype_0\":\"str\"},\"onehot_feature_index\":[-1,-1,-1,-1,0],\"ordinal_feature_index\":[-1,-1,-1,-1,-1],\"original_feature_indices\":[0,1,2,3,4,4,4],\"original_feature_types\":[\"float\",\"float\",\"float\",\"float\",\"category\"],\"processed_feature_types\":[0,0,0,0,1,1,1],\"random_effects\":{}}","forests":{"prognostic_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.8762260279918795,0.26922390975141164],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.30246647484937217,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.23045979874096123,0.24184512731269148],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.21118767831452812,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.03638313371749197,0.015504116821562641],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[6,-1,-1],"threshold":[0.9682615757193975,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.1880534949148276,-0.06319815700947139],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.1708921179415913,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.3486892071832219,-0.14293431258601577],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.32118815577385745,0.0,0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.8723462340040467,-0.021630719513607743],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.30246647484937217,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.03477658766254907,0.18009376107363206],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.21118767831452812,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.14242284642868772,0.08066114714109177],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[6,-1,-1],"threshold":[0.9682615757193975,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.2828745135428721,-0.08834399605258317],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.1708921179415913,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[1,2],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.06306355714376406,0.3486892071832219,-0.14293431258601577],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[0],"left":[-1,-1,-1],"node_deleted":[false,true,true],"node_type":[0,0,0],"num_deleted_nodes":2,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[-1,-1,-1],"split_index":[3,-1,-1],"threshold":[0.32118815577385745,0.0,0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":true,"num_samples":2,"num_trees":5,"output_dimension":1},"treatment_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":false,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.09185433353906003],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_1":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.01077244222638806],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_2":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.005171987178458667],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_3":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.06464553222868631],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_4":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.053019043897354436],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":false,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.13523254891054842],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_1":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.10065600486078201],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.005171987178458667,-0.006536970164820224,0.900892099710261],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.7972774171116785,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.06464553222868631,-0.010390302238034126,-0.09890009656979093],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[4,-1,-1],"threshold":[0.13645522566068502,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.07366124713731273],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":false,"num_samples":2,"num_trees":5,"output_dimension":1}},"has_rfx":true,"has_rfx_basis":true,"include_variance_forest":false,"internal_propensity_model":false,"keep_every":1.0,"multivariate_treatment":false,"num_burnin":0.0,"num_chains":1.0,"num_forests":2,"num_gfr":0.0,"num_mcmc":2.0,"num_random_effects":1,"num_rfx_basis":1.0,"num_samples":2.0,"outcome_mean":4.017200556271111,"outcome_model":{"link":"identity","outcome":"continuous"},"outcome_scale":1.7810629310983845,"parameters":{"sigma2_global_samples":[2.0795979493246746,1.8537860128415273],"sigma2_leaf_mu_samples":[0.07721739558039187,0.06960609966247516],"tau_0_samples":[0.38373158055850004,0.5346683892557097]},"platform":"python","probit_outcome_model":false,"propensity_covariate":"prognostic","random_effects":{"cross_platform_compatible":true,"random_effect_container_0":{"alpha":[0.3911236610586947,0.47105940708392946],"alpha_size":2,"beta":[-0.11659116100074,0.20061898922807053,-0.029557885069272347,-0.6210265297498089,0.0037494938099618164,0.45614160376084395],"beta_size":6,"num_components":1,"num_groups":3,"num_samples":2,"sigma_xi":[0.18310378031923932,1.5486022800538237],"xi":[-0.29809283510271584,0.512929820418009,-0.07557171302105573,-1.3183613795003981,0.007959704770939354,0.9683313758333933]},"random_effect_groupids_0":[0,1,2],"random_effect_label_mapper_0":{"keys":[0,1,2],"values":[0,1,2]}},"rfx_model_spec":"custom","sample_sigma2_global":true,"sample_sigma2_leaf_mu":true,"sample_sigma2_leaf_tau":false,"sample_tau_0":true,"schema_version":1,"sigma2_init":1.0,"standardize":true,"stochtree_version":"0.4.5.dev0","tau_0_dim":1.0,"treatment_dim":1.0} \ No newline at end of file diff --git a/test/python/fixtures/bcf_categorical_v1.json b/test/python/fixtures/bcf_categorical_v1.json new file mode 100644 index 00000000..4a8c9320 --- /dev/null +++ b/test/python/fixtures/bcf_categorical_v1.json @@ -0,0 +1 @@ +{"adaptive_coding":false,"binary_treatment":true,"covariate_preprocessor":"{\"cross_platform_portable\":false,\"forests\":{},\"is_fitted\":true,\"non_portable_columns\":\"4\",\"num_forests\":0,\"num_onehot_features\":1,\"num_ordinal_features\":0,\"num_original_features\":5,\"num_random_effects\":0,\"onehot_categories_list\":{\"cats_0\":[\"a\",\"b\",\"c\"]},\"onehot_dtype_list\":{\"dtype_0\":\"str\"},\"onehot_feature_index\":[-1,-1,-1,-1,0],\"ordinal_feature_index\":[-1,-1,-1,-1,-1],\"original_feature_indices\":[0,1,2,3,4,4,4],\"original_feature_types\":[\"float\",\"float\",\"float\",\"float\",\"category\"],\"processed_feature_types\":[0,0,0,0,1,1,1],\"random_effects\":{}}","forests":{"prognostic_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.9565634585008379,0.33935050347748597],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.3040795568062851,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.08998514452963734,0.20469331635530935],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.2098111859420376,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.057017520982358134,-0.005867169914194989],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[6,-1,-1],"threshold":[0.9682615757193975,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.20145467033722309,-0.055314876526326856],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.17270313633619544,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.24069716414324235,-0.09751834749269524],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.31997132352299656,0.0,0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.7541392713319439,0.19453661236715042],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.3040795568062851,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.08977748396204221,0.06039886756567847],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.2098111859420376,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.06429419793231894,-0.02729518714911613],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[6,-1,-1],"threshold":[0.9682615757193975,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.13613450605148203,0.028308615415648764],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.17270313633619544,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.059459853284013,-0.02598450900609771],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.31997132352299656,0.0,0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":true,"num_samples":2,"num_trees":5,"output_dimension":1},"treatment_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":false,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.0935524343600449],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_1":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.010845414446811755],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_2":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.005303586156079363],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_3":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.06612290145593955],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_4":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.05435497151895199],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":false,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.1338656624842498],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_1":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.10525809346171741],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.005303586156079363,-0.5267253239940609,0.10356377483727927],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.2811273351907882,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.02274567689192029],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_4":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.024050070408488358],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":false,"num_samples":2,"num_trees":5,"output_dimension":1}},"has_rfx":false,"has_rfx_basis":false,"include_variance_forest":false,"internal_propensity_model":false,"keep_every":1.0,"multivariate_treatment":false,"num_burnin":0.0,"num_chains":1.0,"num_forests":2,"num_gfr":0.0,"num_mcmc":2.0,"num_random_effects":0,"num_rfx_basis":0.0,"num_samples":2.0,"outcome_mean":2.903950515727082,"outcome_model":{"link":"identity","outcome":"continuous"},"outcome_scale":1.6197243675523427,"parameters":{"sigma2_global_samples":[1.5132003284229594,1.2035803811622228],"sigma2_leaf_mu_samples":[0.08095398840681357,0.05346669433775765],"tau_0_samples":[0.4070121634806551,0.691378622690623]},"platform":"python","probit_outcome_model":false,"propensity_covariate":"prognostic","random_effects":{},"rfx_model_spec":"custom","sample_sigma2_global":true,"sample_sigma2_leaf_mu":true,"sample_sigma2_leaf_tau":false,"sample_tau_0":true,"schema_version":1,"sigma2_init":1.0,"standardize":true,"stochtree_version":"0.4.5.dev0","tau_0_dim":1.0,"treatment_dim":1.0} \ No newline at end of file diff --git a/test/python/fixtures/bcf_numeric_rfx_v1.json b/test/python/fixtures/bcf_numeric_rfx_v1.json new file mode 100644 index 00000000..8937b8bb --- /dev/null +++ b/test/python/fixtures/bcf_numeric_rfx_v1.json @@ -0,0 +1 @@ +{"adaptive_coding":false,"binary_treatment":true,"covariate_preprocessor":"{\"cross_platform_portable\":true,\"forests\":{},\"is_fitted\":true,\"num_forests\":0,\"num_onehot_features\":0,\"num_ordinal_features\":0,\"num_original_features\":4,\"num_random_effects\":0,\"onehot_feature_index\":[-1,-1,-1,-1],\"ordinal_feature_index\":[-1,-1,-1,-1],\"original_feature_indices\":[0,1,2,3],\"original_feature_types\":[\"float\",\"float\",\"float\",\"float\"],\"processed_feature_types\":[0,0,0,0],\"random_effects\":{}}","forests":{"prognostic_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-1.043404515993365,0.20990324416009074],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.3006881086247295,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.06676759524633163,0.18252123189731895],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.20217194019257026,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.018109013604513686,0.6220350158546815],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[4,-1,-1],"threshold":[0.7306474989029743,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.34254048621898464,-0.0034513868282639412],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.16924491461794333,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.4246715204116045,-0.1899375892145438],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.31191955214441724,0.0,0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.9412308662280215,-0.09559171796714556],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.3006881086247295,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[1,2],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.23668400833515235,-0.06676759524633163,0.18252123189731895],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[0],"left":[-1,-1,-1],"node_deleted":[false,true,true],"node_type":[0,0,0],"num_deleted_nodes":2,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[-1,-1,-1],"split_index":[3,-1,-1],"threshold":[0.20217194019257026,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.15168663734441903,0.5087795936477899],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[4,-1,-1],"threshold":[0.7306474989029743,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.25350364854170676,0.07524671377967958],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.16924491461794333,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.15627520065524705,-0.01730916927881411],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.31191955214441724,0.0,0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":true,"num_samples":2,"num_trees":5,"output_dimension":1},"treatment_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":false,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.09224496404724696],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_1":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.010814630599505062],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_2":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.005224081909475082],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_3":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.06522453047325263],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.05794482691226333,0.09836832077464629],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.39310739324787647,0.0,0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":false,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.03548678232894472],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_1":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.06402288571951621],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.005224081909475082,-0.09674569644037001,0.49296326165709664],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.7950057873047063,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.06522453047325263,0.07100854854510232,-0.037164282447073274],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[1,-1,-1],"threshold":[0.1371907717282265,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[1,2],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.07299148230557802,-0.05794482691226333,0.09836832077464629],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[0],"left":[-1,-1,-1],"node_deleted":[false,true,true],"node_type":[0,0,0],"num_deleted_nodes":2,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[-1,-1,-1],"split_index":[3,-1,-1],"threshold":[0.39310739324787647,0.0,0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":false,"num_samples":2,"num_trees":5,"output_dimension":1}},"has_rfx":true,"has_rfx_basis":true,"include_variance_forest":false,"internal_propensity_model":false,"keep_every":1.0,"multivariate_treatment":false,"num_burnin":0.0,"num_chains":1.0,"num_forests":2,"num_gfr":0.0,"num_mcmc":2.0,"num_random_effects":1,"num_rfx_basis":1.0,"num_samples":2.0,"outcome_mean":4.0113342061929105,"outcome_model":{"link":"identity","outcome":"continuous"},"outcome_scale":1.7135934445263885,"parameters":{"sigma2_global_samples":[1.7222465199677461,1.7080313197868549],"sigma2_leaf_mu_samples":[0.12454102960469965,0.1221270136982205],"tau_0_samples":[0.446759110307799,0.7304653258128424]},"platform":"python","probit_outcome_model":false,"propensity_covariate":"prognostic","random_effects":{"cross_platform_compatible":true,"random_effect_container_0":{"alpha":[0.19048519121434435,0.3855818094724671],"alpha_size":2,"beta":[-0.056782270694166845,0.09770553492186376,-0.014395292205211357,-0.6517055555733899,-0.13775356883705012,0.40380891101593847],"beta_size":6,"num_components":1,"num_groups":3,"num_samples":2,"sigma_xi":[0.18310378031923932,2.1406672018693493],"xi":[-0.29809283510271584,0.512929820418009,-0.07557171302105573,-1.6901875025303177,-0.35726158613529346,1.0472716842332597]},"random_effect_groupids_0":[0,1,2],"random_effect_label_mapper_0":{"keys":[0,1,2],"values":[0,1,2]}},"rfx_model_spec":"custom","sample_sigma2_global":true,"sample_sigma2_leaf_mu":true,"sample_sigma2_leaf_tau":false,"sample_tau_0":true,"schema_version":1,"sigma2_init":1.0,"standardize":true,"stochtree_version":"0.4.5.dev0","tau_0_dim":1.0,"treatment_dim":1.0} \ No newline at end of file diff --git a/test/python/fixtures/bcf_numeric_v1.json b/test/python/fixtures/bcf_numeric_v1.json new file mode 100644 index 00000000..034170e7 --- /dev/null +++ b/test/python/fixtures/bcf_numeric_v1.json @@ -0,0 +1 @@ +{"adaptive_coding":false,"binary_treatment":true,"covariate_preprocessor":"{\"cross_platform_portable\":true,\"forests\":{},\"is_fitted\":true,\"num_forests\":0,\"num_onehot_features\":0,\"num_ordinal_features\":0,\"num_original_features\":4,\"num_random_effects\":0,\"onehot_feature_index\":[-1,-1,-1,-1],\"ordinal_feature_index\":[-1,-1,-1,-1],\"original_feature_indices\":[0,1,2,3],\"original_feature_types\":[\"float\",\"float\",\"float\",\"float\"],\"processed_feature_types\":[0,0,0,0],\"random_effects\":{}}","forests":{"prognostic_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.897780108721336,0.20421147875283965],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.30413278204822336,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.07408307052501317,0.192299277205924],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.2118028910967174,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.0325425461248192,0.9239642076871497],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[4,-1,-1],"threshold":[0.7329067050688588,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.30531919867027324,0.0007457885051878342],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.17247594256794246,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.30000310270637565,-0.11681592725175163],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.3213683786227764,0.0,0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":true,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.7063086753741229,-0.1482376587510231],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.30413278204822336,0.0,0.0]},"tree_1":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.010075734873151923,0.0756690354611048],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.2118028910967174,0.0,0.0]},"tree_2":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.09935979406648472,0.45894366009414983],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[4,-1,-1],"threshold":[0.7329067050688588,0.0,0.0]},"tree_3":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,-0.393483585944056,0.03510503339182133],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[0,-1,-1],"threshold":[0.17247594256794246,0.0,0.0]},"tree_4":{"category_list":[],"category_list_begin":[0,0,0],"category_list_end":[0,0,0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[0],"is_log_scale":false,"leaf_parents":[0],"leaf_value":[0.0,0.08108804869491043,-0.0969172777123195],"leaf_vector":[],"leaf_vector_begin":[0,0,0],"leaf_vector_end":[0,0,0],"leaves":[1,2],"left":[1,-1,-1],"node_deleted":[false,false,false],"node_type":[1,0,0],"num_deleted_nodes":0,"num_nodes":3,"output_dimension":1,"parent":[-1,0,0],"right":[2,-1,-1],"split_index":[3,-1,-1],"threshold":[0.3213683786227764,0.0,0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":true,"num_samples":2,"num_trees":5,"output_dimension":1},"treatment_forest":{"forest_0":{"is_exponentiated":false,"is_leaf_constant":false,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.09471389567192955],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_1":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.010942852167361065],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_2":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.005446363569340318],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_3":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.06770651680574899],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_4":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.05579614726795498],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]}},"forest_1":{"is_exponentiated":false,"is_leaf_constant":false,"num_trees":5,"output_dimension":1,"tree_0":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.13619091425618626],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_1":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.009271298361936308],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_2":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.08249637094828131],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_3":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[-0.02381204588624635],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]},"tree_4":{"category_list":[],"category_list_begin":[0],"category_list_end":[0],"deleted_nodes":[],"has_categorical_split":false,"internal_nodes":[],"is_log_scale":false,"leaf_parents":[],"leaf_value":[0.15337595934963608],"leaf_vector":[],"leaf_vector_begin":[0],"leaf_vector_end":[0],"leaves":[0],"left":[-1],"node_deleted":[false],"node_type":[0],"num_deleted_nodes":0,"num_nodes":1,"output_dimension":1,"parent":[-1],"right":[-1],"split_index":[-1],"threshold":[0.0]}},"initialized":true,"is_exponentiated":false,"is_leaf_constant":false,"num_samples":2,"num_trees":5,"output_dimension":1}},"has_rfx":false,"has_rfx_basis":false,"include_variance_forest":false,"internal_propensity_model":false,"keep_every":1.0,"multivariate_treatment":false,"num_burnin":0.0,"num_chains":1.0,"num_forests":2,"num_gfr":0.0,"num_mcmc":2.0,"num_random_effects":0,"num_rfx_basis":0.0,"num_samples":2.0,"outcome_mean":3.0009607735656543,"outcome_model":{"link":"identity","outcome":"continuous"},"outcome_scale":1.6084942607224353,"parameters":{"sigma2_global_samples":[1.6120674724735369,1.73203723999291],"sigma2_leaf_mu_samples":[0.12768452733418903,0.06869728695335575],"tau_0_samples":[0.5699599873636958,0.9458850926099168]},"platform":"python","probit_outcome_model":false,"propensity_covariate":"prognostic","random_effects":{},"rfx_model_spec":"custom","sample_sigma2_global":true,"sample_sigma2_leaf_mu":true,"sample_sigma2_leaf_tau":false,"sample_tau_0":true,"schema_version":1,"sigma2_init":1.0,"standardize":true,"stochtree_version":"0.4.5.dev0","tau_0_dim":1.0,"treatment_dim":1.0} \ No newline at end of file diff --git a/test/python/fixtures/generate_v1_fixtures.py b/test/python/fixtures/generate_v1_fixtures.py new file mode 100644 index 00000000..d8972bd4 --- /dev/null +++ b/test/python/fixtures/generate_v1_fixtures.py @@ -0,0 +1,108 @@ +"""Generate the v1 (schema_version=1) golden-fixture matrix for the unified envelope. + +Run once to (re)mint the checked-in v1 fixtures: + + python test/python/fixtures/generate_v1_fixtures.py + +The matrix is {bart, bcf} x {numeric, categorical} x {no-rfx, rfx} = 8 models, +covering the format-relevant axes: forest naming (BART vs BCF), the covariate +preprocessor (identity vs categorical encoding -> the cross-platform-portable +axis), and random effects (which add the random_effects subfolder). + +Fixtures only need to be structurally complete and loadable, not statistically +meaningful, so the ensembles are tiny (5 trees, 2 MCMC) to keep the checked-in +files small (a few KB each). The legacy v0 fixtures (bart_mcmc.json / +bcf_mcmc.json) are kept separately and exercise the v0 -> v1 migration path. +""" + +import json +from pathlib import Path + +import numpy as np +import pandas as pd + +from stochtree import BARTModel, BCFModel + +OUT = Path(__file__).parent + +_NUM_MCMC = 2 +_SMALL = {"num_trees": 5} +N = 200 +P = 4 # numeric covariate columns (categorical models add one "cat" column) + +# Superseded ad-hoc fixtures from the first cut (replaced by the matrix below). +_SUPERSEDED = ("bart_mcmc_v1.json", "bart_rfx_v1.json", "bcf_mcmc_v1.json") + + +def _make_covariates(rng, n, categorical): + """Return (X, x0) where x0 is a numeric column used to build the outcome.""" + X_num = rng.uniform(size=(n, P)) + if not categorical: + return X_num, X_num[:, 0] + X = pd.DataFrame(X_num, columns=[f"x{i}" for i in range(P)]) + X["cat"] = pd.Categorical(rng.choice(["a", "b", "c"], size=n)) + return X, X_num[:, 0] + + +def _bart(seed, *, rfx, categorical): + rng = np.random.default_rng(seed) + X, x0 = _make_covariates(rng, N, categorical) + y = x0 + rng.normal(scale=0.5, size=N) + kw = dict(num_gfr=0, num_burnin=0, num_mcmc=_NUM_MCMC, mean_forest_params=_SMALL) + if rfx: + g = rng.integers(0, 3, size=N) + y = y + g + kw["rfx_group_ids_train"] = g + kw["rfx_basis_train"] = np.ones((N, 1)) + m = BARTModel() + m.sample(X_train=X, y_train=y, **kw) + return m.to_json() + + +def _bcf(seed, *, rfx, categorical): + rng = np.random.default_rng(seed) + X, x0 = _make_covariates(rng, N, categorical) + pi = 0.25 + 0.5 * x0 + Z = rng.binomial(1, pi).astype(float) + y = pi * 5 + Z * x0 * 2 + rng.normal(size=N) + kw = dict( + num_gfr=0, + num_burnin=0, + num_mcmc=_NUM_MCMC, + prognostic_forest_params=_SMALL, + treatment_effect_forest_params=_SMALL, + ) + if rfx: + g = rng.integers(0, 3, size=N) + y = y + g + kw["rfx_group_ids_train"] = g + kw["rfx_basis_train"] = np.ones((N, 1)) + m = BCFModel() + m.sample(X_train=X, Z_train=Z, y_train=y, propensity_train=pi, **kw) + return m.to_json() + + +def main(): + seed = 110 + for model, fn in (("bart", _bart), ("bcf", _bcf)): + for categorical in (False, True): + for rfx in (False, True): + kind = "categorical" if categorical else "numeric" + suffix = "_rfx" if rfx else "" + name = f"{model}_{kind}{suffix}_v1.json" + js = fn(seed, rfx=rfx, categorical=categorical) + seed += 1 + payload = json.loads(js) + assert payload["schema_version"] == 1, f"{name}: expected schema_version=1" + (OUT / name).write_text(js) + print(f"wrote {name} (schema_version={payload['schema_version']})") + + for old in _SUPERSEDED: + p = OUT / old + if p.exists(): + p.unlink() + print(f"removed superseded {old}") + + +if __name__ == "__main__": + main() diff --git a/test/python/test_bart.py b/test/python/test_bart.py index 50864f51..81a3d0ed 100644 --- a/test/python/test_bart.py +++ b/test/python/test_bart.py @@ -1233,8 +1233,8 @@ def outcome_mean(X): num_gfr=num_gfr, num_burnin=num_burnin, num_mcmc=num_mcmc, - general_params={"outcome_model": OutcomeModel(outcome="binary", link="probit"), - "sample_sigma2_global": False} + general_params={"outcome_model": OutcomeModel(outcome="binary", link="probit"), + "sample_sigma2_global": False}, ) # Assertions @@ -1250,8 +1250,8 @@ def outcome_mean(X): num_gfr=num_gfr, num_burnin=num_burnin, num_mcmc=num_mcmc, - general_params={"outcome_model": OutcomeModel(outcome="binary", link="probit"), - "sample_sigma2_global": False} + general_params={"outcome_model": OutcomeModel(outcome="binary", link="probit"), + "sample_sigma2_global": False}, ) # Assertions @@ -1412,6 +1412,16 @@ def test_cloglog_binary_bart_with_gfr(self): assert bart_model.y_hat_train.shape == (n_train, num_mcmc) assert bart_model.y_hat_test.shape == (n_test, num_mcmc) + # Structural validity: GFR cloglog warm-starts with unconverged latent variables, + # making correlation-based checks unreliable across platforms and seeds. Instead + # verify that predictions are finite and in a valid range — corrupted residuals + # from a reconstitute_from_forest bug would produce NaN or extreme clipping. + p_hat_mean = bart_model.predict( + X=X_test, type="mean", scale="probability", terms="y_hat" + ) + assert np.all(np.isfinite(p_hat_mean)) + assert np.all((p_hat_mean >= 0.0) & (p_hat_mean <= 1.0)) + def test_cloglog_ordinal_bart(self): # RNG random_seed = 101 @@ -1579,6 +1589,17 @@ def test_cloglog_ordinal_bart_with_gfr(self): assert bart_model.y_hat_test.shape == (n_test, num_mcmc) assert bart_model.cloglog_cutpoint_samples.shape == (2, num_mcmc) + # Structural validity: GFR cloglog warm-starts with unconverged latent variables, + # making correlation-based checks unreliable across platforms and seeds. Instead + # verify that predictions are finite and valid — corrupted residuals from a + # reconstitute_from_forest bug would produce NaN or extreme clipping. + preds_mean_prob = bart_model.predict( + X=X_test, type="mean", scale="probability", terms="y_hat" + ) + assert np.all(np.isfinite(preds_mean_prob)) + assert np.all((preds_mean_prob >= 0.0) & (preds_mean_prob <= 1.0)) + assert np.allclose(preds_mean_prob.sum(axis=1), 1.0, atol=1e-6) + def test_categorical_covariates_mean_only(self): """A mean-only BART model with categorical (one-hot expanded) covariates must sample and predict without error. diff --git a/test/python/test_bcf.py b/test/python/test_bcf.py index 022ceb3a..b7b8abab 100644 --- a/test/python/test_bcf.py +++ b/test/python/test_bcf.py @@ -605,10 +605,10 @@ def test_multivariate_bcf(self): # Assertions assert bcf_model.y_hat_train.shape == (n_train, num_mcmc) assert bcf_model.mu_hat_train.shape == (n_train, num_mcmc) - assert bcf_model.tau_hat_train.shape == (n_train, num_mcmc, treatment_dim) + assert bcf_model.tau_hat_train.shape == (n_train, treatment_dim, num_mcmc) assert bcf_model.y_hat_test.shape == (n_test, num_mcmc) assert bcf_model.mu_hat_test.shape == (n_test, num_mcmc) - assert bcf_model.tau_hat_test.shape == (n_test, num_mcmc, treatment_dim) + assert bcf_model.tau_hat_test.shape == (n_test, treatment_dim, num_mcmc) # Check overall prediction method bcf_preds = bcf_model.predict(X_test, Z_test, pi_test) @@ -617,7 +617,7 @@ def test_multivariate_bcf(self): bcf_preds["mu_hat"], bcf_preds["y_hat"], ) - assert tau_hat.shape == (n_test, num_mcmc, treatment_dim) + assert tau_hat.shape == (n_test, treatment_dim, num_mcmc) assert mu_hat.shape == (n_test, num_mcmc) assert y_hat.shape == (n_test, num_mcmc) @@ -625,7 +625,7 @@ def test_multivariate_bcf(self): tau_hat = bcf_model.predict( X=X_test, Z=Z_test, propensity=pi_test, terms="cate" ) - assert tau_hat.shape == (n_test, num_mcmc, treatment_dim) + assert tau_hat.shape == (n_test, treatment_dim, num_mcmc) # Run BCF without test set and with propensity score with pytest.warns(UserWarning): @@ -646,7 +646,7 @@ def test_multivariate_bcf(self): # Assertions assert bcf_model.y_hat_train.shape == (n_train, num_mcmc) assert bcf_model.mu_hat_train.shape == (n_train, num_mcmc) - assert bcf_model.tau_hat_train.shape == (n_train, num_mcmc, treatment_dim) + assert bcf_model.tau_hat_train.shape == (n_train, treatment_dim, num_mcmc) # Check overall prediction method bcf_preds = bcf_model.predict(X_test, Z_test, pi_test) @@ -655,7 +655,7 @@ def test_multivariate_bcf(self): bcf_preds["mu_hat"], bcf_preds["y_hat"], ) - assert tau_hat.shape == (n_test, num_mcmc, treatment_dim) + assert tau_hat.shape == (n_test, treatment_dim, num_mcmc) assert mu_hat.shape == (n_test, num_mcmc) assert y_hat.shape == (n_test, num_mcmc) @@ -663,7 +663,7 @@ def test_multivariate_bcf(self): tau_hat = bcf_model.predict( X=X_test, Z=Z_test, propensity=pi_test, terms="cate" ) - assert tau_hat.shape == (n_test, num_mcmc, treatment_dim) + assert tau_hat.shape == (n_test, treatment_dim, num_mcmc) # Run BCF with test set and without propensity score with pytest.warns(UserWarning): diff --git a/test/python/test_continuation.py b/test/python/test_continuation.py new file mode 100644 index 00000000..7a9662a5 --- /dev/null +++ b/test/python/test_continuation.py @@ -0,0 +1,205 @@ +import numpy as np +import pytest + +from stochtree import BARTModel, BCFModel + + +def _make_bcf_data(seed=202, n=200, p=5): + rng = np.random.default_rng(seed) + X = rng.uniform(0, 1, (n, p)) + propensity = 0.3 + 0.4 * X[:, 1] + Z = rng.binomial(1, propensity).astype(np.float64).reshape(-1, 1) + pi = propensity.reshape(-1, 1) + mu = 1.0 + 2.0 * X[:, 0] + tau = 1.5 * X[:, 2] + y = mu + tau * Z[:, 0] + 0.5 * rng.standard_normal(n) + return X, Z, y, pi + + +def _make_data(seed=101, n=200, p=5): + rng = np.random.default_rng(seed) + X = rng.uniform(0, 1, (n, p)) + y = 2.0 * X[:, 0] + 0.5 * rng.standard_normal(n) + return X, y + + +class TestBARTContinuation: + def test_continue_matches_one_shot(self): + # A continued chain (N then +M) must match a single run of N+M to within + # floating-point noise. The residual is reconstructed at continue-time, so the + # two chains differ only by ~1e-15 rounding, not by RNG/leaf-cache state. + X, y = _make_data() + N, M, seed = 10, 8, 1234 + + cont = BARTModel() + cont.sample( + X_train=X, y_train=y, num_gfr=0, num_burnin=0, num_mcmc=N, + general_params={"random_seed": seed}, + ) + cont.continue_sampling(X_train=X, y_train=y, num_mcmc=M, num_burnin=0) + + one_shot = BARTModel() + one_shot.sample( + X_train=X, y_train=y, num_gfr=0, num_burnin=0, num_mcmc=N + M, + general_params={"random_seed": seed}, + ) + + assert cont.num_samples == one_shot.num_samples == N + M + pred_cont = cont.predict(X, terms="y_hat") + pred_one_shot = one_shot.predict(X, terms="y_hat") + np.testing.assert_allclose(pred_cont, pred_one_shot, atol=1e-10, rtol=0) + + # Sampled scalars should match too (global variance is sampled by default) + np.testing.assert_allclose( + cont.global_var_samples, one_shot.global_var_samples, atol=1e-10, rtol=0 + ) + + def test_continuation_is_deterministic(self): + # Two continuations from the same fitted model produce identical results. + X, y = _make_data() + + def run(): + m = BARTModel() + m.sample( + X_train=X, y_train=y, num_gfr=0, num_burnin=0, num_mcmc=10, + general_params={"random_seed": 1234}, + ) + m.continue_sampling(X_train=X, y_train=y, num_mcmc=8, num_burnin=0) + return m.predict(X, terms="y_hat") + + np.testing.assert_array_equal(run(), run()) + + def test_override_seed_changes_draws(self): + # Supplying a new random_seed re-seeds the continued draws (no resume). + X, y = _make_data() + m = BARTModel() + m.sample( + X_train=X, y_train=y, num_gfr=0, num_burnin=0, num_mcmc=10, + general_params={"random_seed": 1234}, + ) + resumed = BARTModel() + resumed.sample( + X_train=X, y_train=y, num_gfr=0, num_burnin=0, num_mcmc=10, + general_params={"random_seed": 1234}, + ) + m.continue_sampling(X_train=X, y_train=y, num_mcmc=8, num_burnin=0) + resumed.continue_sampling( + X_train=X, y_train=y, num_mcmc=8, num_burnin=0, random_seed=999 + ) + # The newly drawn samples should differ between resume and override. + assert not np.allclose( + m.predict(X, terms="y_hat")[:, 10:], + resumed.predict(X, terms="y_hat")[:, 10:], + ) + + def test_continuation_guards(self): + X, y = _make_data() + # Variance forest is not yet supported for continuation. + var_model = BARTModel() + var_model.sample( + X_train=X, y_train=y, num_gfr=0, num_burnin=0, num_mcmc=5, + variance_forest_params={"num_trees": 20}, + ) + with pytest.raises(NotImplementedError): + var_model.continue_sampling(X_train=X, y_train=y, num_mcmc=5) + + +class TestBCFContinuation: + def test_continue_matches_one_shot(self): + # A continued BCF chain (N then +M) must match a single run of N+M to within + # floating-point noise (bit-identical RNG resume + warm-start residual reconstruction). + X, Z, y, pi = _make_bcf_data() + N, M, seed = 10, 8, 1234 + + cont = BCFModel() + cont.sample( + X_train=X, Z_train=Z, y_train=y, propensity_train=pi, + num_gfr=0, num_burnin=0, num_mcmc=N, + general_params={"random_seed": seed}, + ) + cont.continue_sampling( + X_train=X, Z_train=Z, y_train=y, propensity_train=pi, num_mcmc=M, num_burnin=0 + ) + + one_shot = BCFModel() + one_shot.sample( + X_train=X, Z_train=Z, y_train=y, propensity_train=pi, + num_gfr=0, num_burnin=0, num_mcmc=N + M, + general_params={"random_seed": seed}, + ) + + assert cont.num_samples == one_shot.num_samples == N + M + pred_cont = cont.predict(X, Z, propensity=pi, terms="y_hat") + pred_one_shot = one_shot.predict(X, Z, propensity=pi, terms="y_hat") + np.testing.assert_allclose(pred_cont, pred_one_shot, atol=1e-9, rtol=0) + + # Sampled scalars should match too (global variance + prognostic leaf scale + tau_0). + np.testing.assert_allclose( + cont.global_var_samples, one_shot.global_var_samples, atol=1e-9, rtol=0 + ) + if getattr(cont, "sample_sigma2_leaf_mu", False): + np.testing.assert_allclose( + cont.leaf_scale_mu_samples, one_shot.leaf_scale_mu_samples, atol=1e-9, rtol=0 + ) + if getattr(cont, "sample_tau_0", False): + np.testing.assert_allclose( + cont.tau_0_samples, one_shot.tau_0_samples, atol=1e-9, rtol=0 + ) + + def test_continuation_is_deterministic(self): + X, Z, y, pi = _make_bcf_data() + + def run(): + m = BCFModel() + m.sample( + X_train=X, Z_train=Z, y_train=y, propensity_train=pi, + num_gfr=0, num_burnin=0, num_mcmc=10, + general_params={"random_seed": 1234}, + ) + m.continue_sampling( + X_train=X, Z_train=Z, y_train=y, propensity_train=pi, num_mcmc=8, num_burnin=0 + ) + return m.predict(X, Z, propensity=pi, terms="y_hat") + + np.testing.assert_array_equal(run(), run()) + + def test_override_seed_changes_draws(self): + X, Z, y, pi = _make_bcf_data() + m = BCFModel() + m.sample( + X_train=X, Z_train=Z, y_train=y, propensity_train=pi, + num_gfr=0, num_burnin=0, num_mcmc=10, + general_params={"random_seed": 1234}, + ) + resumed = BCFModel() + resumed.sample( + X_train=X, Z_train=Z, y_train=y, propensity_train=pi, + num_gfr=0, num_burnin=0, num_mcmc=10, + general_params={"random_seed": 1234}, + ) + m.continue_sampling( + X_train=X, Z_train=Z, y_train=y, propensity_train=pi, num_mcmc=8, num_burnin=0 + ) + resumed.continue_sampling( + X_train=X, Z_train=Z, y_train=y, propensity_train=pi, num_mcmc=8, num_burnin=0, + random_seed=999, + ) + # The newly drawn samples should differ between resume and override. + assert not np.allclose( + m.predict(X, Z, propensity=pi, terms="y_hat")[:, 10:], + resumed.predict(X, Z, propensity=pi, terms="y_hat")[:, 10:], + ) + + def test_continuation_guards(self): + X, Z, y, pi = _make_bcf_data() + # Variance forest is not yet supported for continuation. + var_model = BCFModel() + var_model.sample( + X_train=X, Z_train=Z, y_train=y, propensity_train=pi, + num_gfr=0, num_burnin=0, num_mcmc=5, + variance_forest_params={"num_trees": 20}, + ) + with pytest.raises(NotImplementedError): + var_model.continue_sampling( + X_train=X, Z_train=Z, y_train=y, propensity_train=pi, num_mcmc=5 + ) diff --git a/test/python/test_forest_container.py b/test/python/test_forest_container.py index ab2e8aea..ec5d6ac2 100644 --- a/test/python/test_forest_container.py +++ b/test/python/test_forest_container.py @@ -129,7 +129,7 @@ def outcome_mean(X): ) # Extract the mean forest container - mean_forest_container = bart_model.forest_container_mean + mean_forest_container = bart_model.extract_forest('mean') # Predict from the original container pred_orig = mean_forest_container.predict(forest_dataset_test) @@ -175,7 +175,7 @@ def outcome_mean(X): ) # Extract the mean forest container - mean_forest_container = bart_model.forest_container_mean + mean_forest_container = bart_model.extract_forest('mean') # Predict from the original container pred_orig = mean_forest_container.predict(forest_dataset_test) @@ -221,7 +221,7 @@ def outcome_mean(X): ) # Extract the mean forest container - mean_forest_container = bart_model.forest_container_mean + mean_forest_container = bart_model.extract_forest('mean') # Predict from the original container pred_orig = mean_forest_container.predict(forest_dataset_test) @@ -254,7 +254,7 @@ def outcome_mean(X): ) # Extract the mean forest container - mean_forest_container = bart_model.forest_container_mean + mean_forest_container = bart_model.extract_forest('mean') # Predict from the original container pred_orig = mean_forest_container.predict(forest_dataset_test) diff --git a/test/python/test_init_params.py b/test/python/test_init_params.py new file mode 100644 index 00000000..7f516172 --- /dev/null +++ b/test/python/test_init_params.py @@ -0,0 +1,205 @@ +import numpy as np +from scipy.stats import norm + +from stochtree import BARTModel, BCFModel, OutcomeModel + + +class TestInitParamsHonored: + """Regression tests that user-supplied initialization / calibration parameters + are threaded into the C++ sampler rather than silently dropped. + + Each test would fail if the corresponding parameter were ignored on the + C++ path (as it was before these fixes).""" + + def test_bart_sigma2_init_honored(self): + rng = np.random.default_rng(1) + n, p = 200, 3 + X = rng.uniform(size=(n, p)) + y = X[:, 0] + rng.normal(size=n) + + def fit(s): + m = BARTModel() + general_params = {"standardize": False, "random_seed": 1} + if s is not None: + general_params["sigma2_init"] = s + m.sample( + X_train=X, + y_train=y, + num_gfr=0, + num_burnin=0, + num_mcmc=5, + general_params=general_params, + ) + return m + + m_set = fit(9.0) + m_default = fit(None) + # The user-supplied global variance init must be honored, not hardcoded to 1.0. + assert np.isclose(m_set.sigma2_init, 9.0) + assert not np.isclose(m_default.sigma2_init, 9.0) + + def test_bart_var_forest_leaf_init_honored(self): + rng = np.random.default_rng(2) + n, p = 200, 3 + X = rng.uniform(size=(n, p)) + # Heteroskedastic outcome so the variance forest is meaningfully fit. + y = X[:, 0] + rng.normal(size=n) * np.exp(0.5 * X[:, 1]) + + def fit(v): + m = BARTModel() + variance_forest_params = {"num_trees": 50} + if v is not None: + variance_forest_params["var_forest_leaf_init"] = v + m.sample( + X_train=X, + y_train=y, + X_test=X, + num_gfr=0, + num_burnin=0, + num_mcmc=3, + general_params={"standardize": False, "random_seed": 1}, + mean_forest_params={"num_trees": 50}, + variance_forest_params=variance_forest_params, + ) + return np.asarray(m.sigma2_x_train) + + # Same seed; differing only by the variance-forest leaf init must change output. + assert not np.allclose(fit(0.05), fit(2.0)) + + def test_bcf_delta_max_honored(self): + rng = np.random.default_rng(3) + n, p = 300, 5 + X = rng.uniform(size=(n, p)) + Z = rng.binomial(1, 0.5, size=n).astype(float) + propensity = np.full(n, 0.5) + prob = norm.cdf(X[:, 0] - 0.5 + 0.5 * Z) + y = rng.binomial(1, prob).astype(float) + num_trees_tau = 50 + + def fit(delta_max): + m = BCFModel() + m.sample( + X_train=X, + Z_train=Z, + y_train=y, + propensity_train=propensity, + num_gfr=0, + num_burnin=0, + num_mcmc=5, + general_params={ + "random_seed": 1, + "outcome_model": OutcomeModel(outcome="binary", link="probit"), + }, + treatment_effect_forest_params={ + "num_trees": num_trees_tau, + "delta_max": delta_max, + }, + ) + return m.sigma2_leaf_tau_init + + def expected(delta_max): + p_coverage = 0.6827 + q_quantile = norm.ppf((p_coverage + 1) / 2) + return ((delta_max / (q_quantile * norm.pdf(0))) ** 2) / num_trees_tau + + v_05 = fit(0.5) + v_09 = fit(0.9) + # delta_max must drive the treatment-effect leaf scale calibration. + assert not np.isclose(v_05, v_09) + assert np.isclose(v_05, expected(0.5)) + assert np.isclose(v_09, expected(0.9)) + + def test_bart_observation_weights_honored(self): + rng = np.random.default_rng(4) + n, p = 200, 3 + X = rng.uniform(size=(n, p)) + y = X[:, 0] + rng.normal(size=n) + w = rng.uniform(0.1, 2.0, size=n) + + def fit(weights): + m = BARTModel() + m.sample( + X_train=X, + y_train=y, + X_test=X, + num_gfr=0, + num_burnin=0, + num_mcmc=5, + general_params={"standardize": False, "random_seed": 1}, + observation_weights_train=weights, + ) + return np.asarray(m.y_hat_test) + + # Non-uniform observation weights must change the fit (same seed). + assert not np.allclose(fit(None), fit(w)) + + def test_bcf_observation_weights_honored(self): + rng = np.random.default_rng(5) + n, p = 200, 5 + X = rng.uniform(size=(n, p)) + Z = rng.binomial(1, 0.5, size=n).astype(float) + propensity = np.full(n, 0.5) + y = X[:, 0] + Z * X[:, 1] + rng.normal(size=n) + w = rng.uniform(0.1, 2.0, size=n) + + def fit(weights): + m = BCFModel() + m.sample( + X_train=X, + Z_train=Z, + y_train=y, + propensity_train=propensity, + X_test=X, + Z_test=Z, + propensity_test=propensity, + num_gfr=0, + num_burnin=0, + num_mcmc=5, + general_params={"standardize": False, "random_seed": 1}, + observation_weights_train=weights, + ) + return np.asarray(m.y_hat_test) + + assert not np.allclose(fit(None), fit(w)) + + def test_bcf_internal_propensity_reproducible(self): + rng = np.random.default_rng(7) + n, p = 200, 5 + X = rng.uniform(size=(n, p)) + Z = rng.binomial(1, 0.5, size=n).astype(float) + y = X[:, 0] + Z * X[:, 1] + rng.normal(size=n) + + def fit(): + m = BCFModel() + # No propensity_train -> BCF estimates it with an internal BART model. + m.sample( + X_train=X, + Z_train=Z, + y_train=y, + num_gfr=0, + num_burnin=0, + num_mcmc=10, + general_params={"random_seed": 99}, + ) + return np.asarray(m.y_hat_train) + + # The internally-estimated propensity (and hence the full fit) must be + # reproducible across runs with a fixed random_seed. + np.testing.assert_allclose(fit(), fit()) + + def test_verbose_prints_progress(self, capfd): + rng = np.random.default_rng(8) + n, p = 100, 3 + X = rng.uniform(size=(n, p)) + y = X[:, 0] + rng.normal(size=n) + kwargs = dict(X_train=X, y_train=y, num_gfr=5, num_burnin=0, num_mcmc=20) + + BARTModel().sample(**kwargs, general_params={"random_seed": 1, "verbose": True}) + out = capfd.readouterr().out + assert "GFR" in out and "MCMC" in out + + BARTModel().sample( + **kwargs, general_params={"random_seed": 1, "verbose": False} + ) + out_silent = capfd.readouterr().out + assert "GFR" not in out_silent and "MCMC" not in out_silent diff --git a/test/python/test_json.py b/test/python/test_json.py index b6f9b36f..9c0463c6 100644 --- a/test/python/test_json.py +++ b/test/python/test_json.py @@ -136,7 +136,7 @@ def outcome_mean(X): # Roundtrip to / from JSON json_test = JSONSerializer() - json_test.add_forest(bart_model.forest_container_mean) + json_test.add_forest(bart_model.extract_forest('mean')) forest_container = json_test.get_forest_container("forest_0") # Predict from the deserialized forest container diff --git a/test/python/test_json_primitives.py b/test/python/test_json_primitives.py new file mode 100644 index 00000000..e8d9e38b --- /dev/null +++ b/test/python/test_json_primitives.py @@ -0,0 +1,38 @@ +"""Unit tests for the JSON key primitives (rename/erase) used by schema migrations.""" + +from stochtree.serialization import JSONSerializer + + +def test_rename_and_erase_top_level(): + s = JSONSerializer() + s.add_scalar("old", 3.5) + + s.rename_field("old", "new") + assert not s.contains("old") + assert s.contains("new") + assert s.get_scalar("new") == 3.5 + + s.erase_field("new") + assert not s.contains("new") + + +def test_rename_and_erase_subfolder(): + s = JSONSerializer() + s.add_scalar("old", 7.0, subfolder_name="sub") + + s.rename_field("old", "new", subfolder_name="sub") + assert not s.contains("old", subfolder_name="sub") + assert s.contains("new", subfolder_name="sub") + assert s.get_scalar("new", subfolder_name="sub") == 7.0 + + s.erase_field("new", subfolder_name="sub") + assert not s.contains("new", subfolder_name="sub") + + +def test_rename_and_erase_missing_is_noop(): + s = JSONSerializer() + # Neither should raise on an absent field. + s.rename_field("nope", "whatever") + s.erase_field("nope") + assert not s.contains("whatever") + assert not s.contains("nope") diff --git a/test/python/test_kernel.py b/test/python/test_kernel.py index 6a6bff09..65d5e5eb 100644 --- a/test/python/test_kernel.py +++ b/test/python/test_kernel.py @@ -2,6 +2,7 @@ import pandas as pd from stochtree import ( + BARTModel, Dataset, Forest, ForestContainer, @@ -80,3 +81,27 @@ def test_forest(self): # Assertion np.testing.assert_almost_equal(computed, expected) assert max_leaf_index == [3] + + def test_bart_model(self): + # Regression: the kernel functions on a BART *model* must reach the forests via + # extract_forest(), not the removed direct-forest properties (which now raise). + rng = np.random.default_rng(42) + X = rng.uniform(0, 1, (100, 3)) + y = X[:, 0] * 2 + rng.normal(0, 0.5, 100) + model = BARTModel() + model.sample( + X_train=X, + y_train=y, + num_gfr=0, + num_burnin=0, + num_mcmc=5, + general_params={"random_seed": 42}, + ) + + leaf_indices = compute_forest_leaf_indices(model, X, forest_type="mean") + max_leaf_index = compute_forest_max_leaf_index(model, forest_type="mean") + + assert leaf_indices.size > 0 + assert np.all(leaf_indices >= 0) + assert np.all(max_leaf_index >= 0) + assert leaf_indices.max() <= np.max(max_leaf_index) diff --git a/test/python/test_observation_weights.py b/test/python/test_observation_weights.py index 6c89216a..445fdadb 100644 --- a/test/python/test_observation_weights.py +++ b/test/python/test_observation_weights.py @@ -39,7 +39,7 @@ def test_uniform_weights_match_no_weights(self): m1.sample(**kwargs) m2 = BARTModel() - m2.sample(**kwargs, observation_weights=np.ones(n_train)) + m2.sample(**kwargs, observation_weights_train=np.ones(n_train)) np.testing.assert_array_equal(m1.y_hat_train, m2.y_hat_train) np.testing.assert_array_equal(m1.y_hat_test, m2.y_hat_test) @@ -54,7 +54,7 @@ def test_nonuniform_weights_output_shape(self): m = BARTModel() m.sample( X_train=X_train, y_train=y_train, X_test=X_test, - observation_weights=weights, + observation_weights_train=weights, num_gfr=0, num_burnin=0, num_mcmc=num_mcmc, ) assert m.y_hat_train.shape == (n_train, num_mcmc) @@ -68,7 +68,7 @@ def test_zero_weights_prior_mode(self): m = BARTModel() m.sample( X_train=X_train, y_train=y_train, - observation_weights=np.zeros(n_train), + observation_weights_train=np.zeros(n_train), num_gfr=0, num_burnin=0, num_mcmc=num_mcmc, ) assert m.y_hat_train.shape == (n_train, num_mcmc) @@ -78,7 +78,7 @@ def test_invalid_type_raises(self): with pytest.raises(ValueError, match="numpy array"): BARTModel().sample( X_train=X_train, y_train=y_train, - observation_weights=list(np.ones(n_train)), + observation_weights_train=list(np.ones(n_train)), num_gfr=0, num_burnin=0, num_mcmc=5, ) @@ -87,7 +87,7 @@ def test_2d_weights_raises(self): with pytest.raises(ValueError, match="1-dimensional"): BARTModel().sample( X_train=X_train, y_train=y_train, - observation_weights=np.ones((n_train, 2)), + observation_weights_train=np.ones((n_train, 2)), num_gfr=0, num_burnin=0, num_mcmc=5, ) @@ -98,7 +98,7 @@ def test_negative_weights_raises(self): with pytest.raises(ValueError, match="negative"): BARTModel().sample( X_train=X_train, y_train=y_train, - observation_weights=weights, + observation_weights_train=weights, num_gfr=0, num_burnin=0, num_mcmc=5, ) @@ -107,7 +107,7 @@ def test_all_zero_with_gfr_raises(self): with pytest.raises(ValueError, match="num_gfr"): BARTModel().sample( X_train=X_train, y_train=y_train, - observation_weights=np.zeros(n_train), + observation_weights_train=np.zeros(n_train), num_gfr=5, num_burnin=0, num_mcmc=10, ) @@ -119,17 +119,17 @@ def test_cloglog_raises(self): with pytest.raises(ValueError, match="cloglog"): BARTModel().sample( X_train=X, y_train=y, - observation_weights=np.ones(n), + observation_weights_train=np.ones(n), num_gfr=0, num_burnin=0, num_mcmc=5, general_params={"outcome_model": OutcomeModel(outcome="ordinal", link="cloglog")}, ) - def test_variance_forest_warns(self): + def test_variance_forest_raises(self): X_train, y_train, _, n_train, _ = make_bart_data() - with pytest.warns(UserWarning, match="variance forest"): + with pytest.raises(ValueError, match="not compatible with a variance forest"): BARTModel().sample( X_train=X_train, y_train=y_train, - observation_weights=np.ones(n_train), + observation_weights_train=np.ones(n_train), num_gfr=0, num_burnin=0, num_mcmc=5, variance_forest_params={"num_trees": 5}, ) @@ -149,7 +149,7 @@ def test_uniform_weights_match_no_weights(self): m1.sample(**kwargs) m2 = BCFModel() - m2.sample(**kwargs, observation_weights=np.ones(n_train)) + m2.sample(**kwargs, observation_weights_train=np.ones(n_train)) np.testing.assert_array_equal(m1.y_hat_train, m2.y_hat_train) np.testing.assert_array_equal(m1.tau_hat_train, m2.tau_hat_train) @@ -164,7 +164,7 @@ def test_nonuniform_weights_output_shape(self): m.sample( X_train=X_train, Z_train=Z_train, y_train=y_train, propensity_train=pi_train, X_test=X_test, Z_test=Z_test, propensity_test=pi_test, - observation_weights=weights, + observation_weights_train=weights, num_gfr=0, num_burnin=0, num_mcmc=num_mcmc, ) assert m.y_hat_train.shape == (n_train, num_mcmc) @@ -180,7 +180,7 @@ def test_negative_weights_raises(self): BCFModel().sample( X_train=X_train, Z_train=Z_train, y_train=y_train, propensity_train=pi_train, - observation_weights=weights, + observation_weights_train=weights, num_gfr=0, num_burnin=0, num_mcmc=5, ) @@ -190,7 +190,7 @@ def test_all_zero_with_gfr_raises(self): BCFModel().sample( X_train=X_train, Z_train=Z_train, y_train=y_train, propensity_train=pi_train, - observation_weights=np.zeros(n_train), + observation_weights_train=np.zeros(n_train), num_gfr=5, num_burnin=0, num_mcmc=10, ) @@ -200,7 +200,28 @@ def test_cloglog_raises(self): BCFModel().sample( X_train=X_train, Z_train=Z_train, y_train=y_train, propensity_train=pi_train, - observation_weights=np.ones(n_train), + observation_weights_train=np.ones(n_train), num_gfr=0, num_burnin=0, num_mcmc=5, general_params={"outcome_model": OutcomeModel(outcome="binary", link="cloglog")}, ) + + +def test_observation_weights_deprecated_alias_warns(): + """The deprecated `observation_weights` alias still works but emits a DeprecationWarning.""" + X_train, y_train, X_test, n_train, n_test = make_bart_data() + rng = np.random.default_rng(0) + weights = rng.uniform(0.5, 2.0, n_train) + kwargs = dict( + X_train=X_train, y_train=y_train, X_test=X_test, + num_gfr=0, num_burnin=0, num_mcmc=10, + general_params={"random_seed": 1}, + ) + m_new = BARTModel() + m_new.sample(**kwargs, observation_weights_train=weights) + + m_old = BARTModel() + with pytest.warns(DeprecationWarning, match="observation_weights"): + m_old.sample(**kwargs, observation_weights=weights) + + # Deprecated alias must produce the same result as the new parameter. + np.testing.assert_allclose(m_new.y_hat_test, m_old.y_hat_test) diff --git a/test/python/test_predict.py b/test/python/test_predict.py index 61eb5c07..23ce4ad8 100644 --- a/test/python/test_predict.py +++ b/test/python/test_predict.py @@ -916,15 +916,19 @@ def test_bart_cloglog_ordinal_posterior_predictive(self): assert set(np.unique(ppd1)).issubset(set(range(1, n_categories + 1))) def test_bart_cloglog_ordinal_probability_transform_k4(self): + # K=4 is the minimal case that exposes cumulative-product bugs; + # K=3 is accidentally correct via the residual-last-class formula. rng = np.random.default_rng(123) - n = 500; p = 3; n_categories = 4 + n = 500 + p = 3 + n_categories = 4 X = rng.uniform(size=(n, p)) beta = np.full(p, 1 / np.sqrt(p)) true_lambda = X @ beta # Balanced cutpoints: ~25% per category so all 4 are reliably observed. - # Unbalanced cutpoints give ~0.3% K=4 probability, often no K=4 training obs. gamma_true = np.array([-2.1, -1.7, -1.2]) + # Sequential DGP: P(Y=k) = prod_{j= 0) - row_sums = p_model.sum(axis=1) + row_sums = p_model.sum(axis=1) # (n_test, num_mcmc) np.testing.assert_allclose(row_sums, np.ones((n_test, num_mcmc)), atol=1e-10) def test_bart_gaussian_interval_and_contrast(self): diff --git a/test/python/test_samples_wrapper.py b/test/python/test_samples_wrapper.py new file mode 100644 index 00000000..72495f2d --- /dev/null +++ b/test/python/test_samples_wrapper.py @@ -0,0 +1,225 @@ +import numpy as np +import pytest + +from stochtree import BARTModel, BCFModel +from stochtree_cpp import BARTSamplesCpp, BCFSamplesCpp + + +class TestBARTSamplesCpp: + """Isolated tests for the thin single-owner wrapper around StochTree::BARTSamples. + + These exercise the wrapper plumbing (FromJson -> accessors / param-vector marshalling / + materialize-on-demand) independently of the model re-point, so the wrapper can be validated + before bart.py routes through it. + """ + + def _fit(self, seed=42, n=100, p=4): + rng = np.random.default_rng(seed) + X = rng.uniform(0, 1, (n, p)) + y = X[:, 0] * 2 + rng.normal(0, 0.5, n) + model = BARTModel() + model.sample( + X_train=X, y_train=y, num_gfr=0, num_burnin=0, num_mcmc=10, + general_params={"random_seed": seed}, + ) + return model + + def test_from_json_accessors_and_materialize(self): + model = self._fit() + samples = BARTSamplesCpp.from_json_string(model.to_json()) + + # Scalars / counts match the model + assert samples.num_samples() == model.num_samples + assert np.isclose(samples.y_bar(), model.y_bar) + assert np.isclose(samples.y_std(), model.y_std) + + # Parameter traces round-trip to the model's arrays (one entry per draw) + if model.sample_sigma2_global: + gv = samples.global_var_samples() + assert gv.shape[0] == model.num_samples + np.testing.assert_allclose(gv, model.global_var_samples) + if model.sample_sigma2_leaf: + np.testing.assert_allclose( + samples.leaf_scale_samples(), model.leaf_scale_samples + ) + + # Materialized mean forest is a faithful deep copy (byte-identical serialization) + assert samples.has_mean_forest() + fc = samples.materialize_mean_forest() + assert fc.NumSamples() == model.num_samples + assert ( + fc.DumpJsonString() + == model.extract_forest('mean').forest_container_cpp.DumpJsonString() + ) + + # This model has no variance forest + assert not samples.has_variance_forest() + assert samples.materialize_variance_forest() is None + + def test_to_json_string_round_trips(self): + # The wrapper's own ToJson -> FromJson preserves the samples subtree. + model = self._fit() + samples = BARTSamplesCpp.from_json_string(model.to_json()) + restored = BARTSamplesCpp.from_json_string(samples.to_json_string()) + assert restored.num_samples() == samples.num_samples() + np.testing.assert_allclose( + restored.global_var_samples(), samples.global_var_samples() + ) + assert ( + restored.materialize_mean_forest().DumpJsonString() + == samples.materialize_mean_forest().DumpJsonString() + ) + + def test_merge_appends_draws(self): + # Build two wrappers from the SAME model's JSON so they share standardization (Merge guards + # against mismatched y_bar/y_std). Merging concatenates their draws. + js = self._fit(seed=1).to_json() + a = BARTSamplesCpp.from_json_string(js) + b = BARTSamplesCpp.from_json_string(js) + n_a, n_b = a.num_samples(), b.num_samples() + gv_a = a.global_var_samples().copy() + gv_b = b.global_var_samples().copy() + + a.merge(b) + + assert a.num_samples() == n_a + n_b + assert a.materialize_mean_forest().NumSamples() == n_a + n_b + np.testing.assert_allclose(a.global_var_samples(), np.concatenate([gv_a, gv_b])) + + def test_from_components(self): + # Assemble the wrapper from a fitted model's existing forest container + parameter arrays + # (the construction path the model re-point will use), and check it matches the model. + model = self._fit() + samples = BARTSamplesCpp.from_components( + model.extract_forest('mean').forest_container_cpp, + None, # no variance forest + model.global_var_samples if model.sample_sigma2_global else None, + model.leaf_scale_samples if model.sample_sigma2_leaf else None, + float(model.y_bar), + float(model.y_std), + int(model.num_samples), + ) + assert samples.num_samples() == model.num_samples + assert np.isclose(samples.y_bar(), model.y_bar) + if model.sample_sigma2_global: + np.testing.assert_allclose(samples.global_var_samples(), model.global_var_samples) + # Deep-copied forest matches the source byte-for-byte + assert ( + samples.materialize_mean_forest().DumpJsonString() + == model.extract_forest('mean').forest_container_cpp.DumpJsonString() + ) + assert not samples.has_variance_forest() + + def test_direct_forest_access_raises(self): + # The public forest_container_* attributes were removed; accessing them raises and the + # message points at the supported extraction call. + model = self._fit() + with pytest.raises(AttributeError, match="extract_forest"): + _ = model.forest_container_mean + with pytest.raises(AttributeError, match="extract_forest"): + _ = model.forest_container_variance + # The documented extraction path works and returns a deep copy. + assert model.extract_forest("mean") is not None + assert model.samples is not None + + +class TestBCFSamplesCpp: + """Isolated tests for the single-owner wrapper around StochTree::BCFSamples.""" + + def _fit(self, seed=42, n=120, p=4): + rng = np.random.default_rng(seed) + X = rng.uniform(0, 1, (n, p)) + pi = 0.3 + 0.4 * X[:, 1] + Z = rng.binomial(1, pi).astype(np.float64).reshape(-1, 1) + mu = 1.0 + 2.0 * X[:, 0] + tau = 1.5 * X[:, 2] + y = mu + tau * Z[:, 0] + 0.5 * rng.standard_normal(n) + model = BCFModel() + model.sample( + X_train=X, Z_train=Z, y_train=y, propensity_train=pi.reshape(-1, 1), + num_gfr=0, num_burnin=0, num_mcmc=10, + general_params={"random_seed": seed}, + ) + return model + + def test_from_json_accessors_and_materialize(self): + model = self._fit() + samples = BCFSamplesCpp.from_json_string(model.to_json()) + + assert samples.num_samples() == model.num_samples + assert samples.treatment_dim() == model.treatment_dim + assert np.isclose(samples.y_bar(), model.y_bar) + assert np.isclose(samples.y_std(), model.y_std) + + if model.sample_sigma2_global: + np.testing.assert_allclose(samples.global_var_samples(), model.global_var_samples) + if model.sample_sigma2_leaf_mu: + np.testing.assert_allclose( + samples.leaf_scale_mu_samples(), model.leaf_scale_mu_samples + ) + + # Both prognostic and treatment forests materialize as faithful deep copies + assert samples.has_mu_forest() + assert samples.has_tau_forest() + assert ( + samples.materialize_mu_forest().DumpJsonString() + == model.extract_forest('prognostic').forest_container_cpp.DumpJsonString() + ) + assert ( + samples.materialize_tau_forest().DumpJsonString() + == model.extract_forest('treatment').forest_container_cpp.DumpJsonString() + ) + assert not samples.has_variance_forest() + assert samples.materialize_variance_forest() is None + + def test_merge_appends_draws(self): + js = self._fit(seed=1).to_json() + a = BCFSamplesCpp.from_json_string(js) + b = BCFSamplesCpp.from_json_string(js) + n_a, n_b = a.num_samples(), b.num_samples() + + a.merge(b) + + assert a.num_samples() == n_a + n_b + assert a.materialize_mu_forest().NumSamples() == n_a + n_b + assert a.materialize_tau_forest().NumSamples() == n_a + n_b + + def test_from_components(self): + model = self._fit() + samples = BCFSamplesCpp.from_components( + model.extract_forest('prognostic').forest_container_cpp, + model.extract_forest('treatment').forest_container_cpp, + None, # no variance forest + model.global_var_samples if model.sample_sigma2_global else None, + model.leaf_scale_mu_samples if model.sample_sigma2_leaf_mu else None, + model.leaf_scale_tau_samples if model.sample_sigma2_leaf_tau else None, + getattr(model, "tau_0_samples", None), + model.b0_samples if model.adaptive_coding else None, + model.b1_samples if model.adaptive_coding else None, + float(model.y_bar), + float(model.y_std), + int(model.num_samples), + int(model.treatment_dim), + ) + assert samples.num_samples() == model.num_samples + assert samples.treatment_dim() == model.treatment_dim + if model.sample_sigma2_global: + np.testing.assert_allclose(samples.global_var_samples(), model.global_var_samples) + assert ( + samples.materialize_mu_forest().DumpJsonString() + == model.extract_forest('prognostic').forest_container_cpp.DumpJsonString() + ) + assert ( + samples.materialize_tau_forest().DumpJsonString() + == model.extract_forest('treatment').forest_container_cpp.DumpJsonString() + ) + + def test_direct_forest_access_raises(self): + model = self._fit() + with pytest.raises(AttributeError, match="extract_forest"): + _ = model.forest_container_mu + with pytest.raises(AttributeError, match="extract_forest"): + _ = model.forest_container_tau + assert model.extract_forest("prognostic") is not None + assert model.extract_forest("treatment") is not None + assert model.samples is not None diff --git a/test/python/test_serialization_compat.py b/test/python/test_serialization_compat.py index cd694f4c..6344461d 100644 --- a/test/python/test_serialization_compat.py +++ b/test/python/test_serialization_compat.py @@ -14,6 +14,7 @@ from pathlib import Path import numpy as np +import pandas as pd import pytest from stochtree import BARTModel, BCFModel @@ -477,3 +478,148 @@ def test_multichain_and_combined_roundtrip(self): combined.from_json_string_list([m_a.to_json(), m_b.to_json()]) assert combined.num_samples == 20 assert isinstance(str(combined), str) + + +def _matrix_covariates(categorical, k, seed=7): + """Predict-time covariates matching the v1 fixture matrix training layout: + 4 numeric columns, plus an unordered categorical column for the categorical + fixtures (mirrors generate_v1_fixtures.py).""" + rng = np.random.default_rng(seed) + X_num = rng.uniform(size=(k, 4)) + if not categorical: + return X_num + X = pd.DataFrame(X_num, columns=[f"x{i}" for i in range(4)]) + X["cat"] = pd.Categorical(rng.choice(["a", "b", "c"], size=k)) + return X + + +class TestV1Snapshot: + """Load the checked-in schema_version=1 fixture matrix (unified envelope). + + The matrix is {bart, bcf} x {numeric, categorical} x {no-rfx, rfx}; these lock + the on-disk v1 format directly -- named forest keys, ``covariate_preprocessor`` + (identity vs categorical encoding), and the ``random_effects`` layout + (including the relocated ``rfx_unique_group_ids``) -- whereas the legacy v0 + fixtures above guard the v0 -> v1 migration path. Regenerate with + ``python test/python/fixtures/generate_v1_fixtures.py``. + """ + + @pytest.mark.parametrize("categorical", [False, True]) + @pytest.mark.parametrize("rfx", [False, True]) + def test_bart_v1_matrix(self, categorical, rfx): + kind = "categorical" if categorical else "numeric" + name = f"bart_{kind}{'_rfx' if rfx else ''}_v1.json" + obj = _load_fixture(name) + assert obj["schema_version"] == 1 + assert set(obj["forests"]) == {"mean_forest"} + + m = BARTModel() + m.from_json(_to_json_string(obj)) + assert m.has_rfx == rfx + + k = 12 + X = _matrix_covariates(categorical, k) + kw = {} + if rfx: + kw["rfx_group_ids"] = np.arange(k) % 3 + kw["rfx_basis"] = np.ones((k, 1)) + preds = m.predict(X, **kw) + assert preds["y_hat"].shape[0] == k + + @pytest.mark.parametrize("categorical", [False, True]) + @pytest.mark.parametrize("rfx", [False, True]) + def test_bcf_v1_matrix(self, categorical, rfx): + kind = "categorical" if categorical else "numeric" + name = f"bcf_{kind}{'_rfx' if rfx else ''}_v1.json" + obj = _load_fixture(name) + assert obj["schema_version"] == 1 + assert {"prognostic_forest", "treatment_forest"} <= set(obj["forests"]) + + m = BCFModel() + m.from_json(_to_json_string(obj)) + assert m.has_rfx == rfx + + k = 12 + X = _matrix_covariates(categorical, k) + rng = np.random.default_rng(3) + Z = rng.binomial(1, 0.5, k).astype(float) + pi = np.full(k, 0.5) + kw = {} + if rfx: + kw["rfx_group_ids"] = np.arange(k) % 3 + kw["rfx_basis"] = np.ones((k, 1)) + preds = m.predict(X, Z, pi, **kw) + assert preds["y_hat"].shape[0] == k + + +def _as_foreign(fixture_name, **overrides): + """Relabel a fixture's ``platform`` as the *other* platform so the + cross-platform gate can be exercised from within Python's own test suite + (the gate peeks generic flags; the foreign preprocessor body is ignored).""" + obj = _load_fixture(fixture_name) + obj["platform"] = "R" + obj.update(overrides) + return _to_json_string(obj) + + +class TestCrossPlatformGate: + """WS-E: cross-platform load succeeds for portable (all-numeric, integer-rfx) + models and is refused with a clear error otherwise. Same-platform loads ignore + the flags.""" + + def test_numeric_bart_cross_platform_loads_and_predicts(self): + m = BARTModel() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + m.from_json(_as_foreign("bart_numeric_v1.json")) + preds = m.predict(np.random.default_rng(0).uniform(size=(8, 4))) + assert preds["y_hat"].shape[0] == 8 + + def test_numeric_bcf_cross_platform_loads_and_predicts(self): + m = BCFModel() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + m.from_json(_as_foreign("bcf_numeric_v1.json")) + rng = np.random.default_rng(0) + X = rng.uniform(size=(8, 4)) + Z = rng.binomial(1, 0.5, 8).astype(float) + preds = m.predict(X, Z, np.full(8, 0.5)) + assert preds["y_hat"].shape[0] == 8 + + def test_categorical_cross_platform_refused(self): + with pytest.raises(ValueError, match="non-numeric covariates"): + BARTModel().from_json(_as_foreign("bart_categorical_v1.json")) + with pytest.raises(ValueError, match="non-numeric covariates"): + BCFModel().from_json(_as_foreign("bcf_categorical_v1.json")) + + def test_string_rfx_cross_platform_refused(self): + # Force the rfx-incompatible flag (Python never writes string rfx itself). + obj = _load_fixture("bart_numeric_rfx_v1.json") + obj["platform"] = "R" + obj["random_effects"]["cross_platform_compatible"] = False + with pytest.raises(ValueError, match="random effects"): + BARTModel().from_json(_to_json_string(obj)) + + def test_same_platform_categorical_not_refused(self): + # Loaded on its own platform, a categorical model must NOT trip the gate. + m = BARTModel() + m.from_json(_to_json_string(_load_fixture("bart_categorical_v1.json"))) + assert m.sampled + + +def test_future_schema_version_raises(): + """A model stamped with a newer schema_version than this install supports must error.""" + from stochtree.serialization import SCHEMA_VERSION + + rng = np.random.default_rng(0) + X = rng.uniform(size=(80, 3)) + y = X[:, 0] + rng.normal(size=80) + m = BARTModel() + m.sample(X_train=X, y_train=y, num_gfr=0, num_burnin=0, num_mcmc=5) + + payload = json.loads(m.to_json()) + assert payload["schema_version"] == SCHEMA_VERSION + payload["schema_version"] = SCHEMA_VERSION + 1 + + with pytest.raises(ValueError, match="schema_version"): + BARTModel().from_json(json.dumps(payload)) diff --git a/tools/debug/debug_cpp_sampler.R b/tools/debug/debug_cpp_sampler.R new file mode 100644 index 00000000..69d0559f --- /dev/null +++ b/tools/debug/debug_cpp_sampler.R @@ -0,0 +1,67 @@ +################################################################################ +# Minimal script for debugging the C++ sampler under lldb. +# +# Usage (from the repo root): +# lldb -- R --vanilla -f tools/debug/debug_cpp_sampler.R +# # then at the (lldb) prompt: +# # run +# # bt (after the crash, to get a backtrace) +# # frame info (to see the crashing frame) +# +# Alternatively, attach to an already-running R process: +# lldb -p $(pgrep -n R) +################################################################################ + +suppressPackageStartupMessages(devtools::load_all(".")) + +# --- Data generation (mirrors debug_cpp_sampler.py) -------------------------- +seed <- 1001 +n <- 10000 +p <- 10 +set.seed(1234) + +X <- matrix(runif(n * p), nrow = n, ncol = p) +f_X <- ifelse( + X[, 1] < 0.25, + -7.5, + ifelse(X[, 1] < 0.50, -2.5, ifelse(X[, 1] < 0.75, 2.5, 7.5)) +) +y <- f_X + rnorm(n, sd = 1.0) + +n_test <- round(0.2 * n) +test_inds <- sort(sample(seq_len(n), n_test, replace = FALSE)) +train_inds <- setdiff(seq_len(n), test_inds) + +X_train <- X[train_inds, ] +X_test <- X[test_inds, ] +y_train <- y[train_inds] +y_test <- y[test_inds] + +cat(sprintf( + "n_train=%d n_test=%d p=%d seed=%d\n", + length(train_inds), + n_test, + p, + seed +)) +cat("Calling bart() with run_cpp=TRUE ...\n") + +# --- Run C++ sampler ---------------------------------------------------------- +m <- bart( + X_train = X_train, + y_train = y_train, + X_test = X_test, + num_gfr = 10, + num_burnin = 0, + num_mcmc = 100, + general_params = list(random_seed = seed), + mean_forest_params = list(num_trees = 200), + run_cpp = TRUE +) + +cat("Completed successfully.\n") +cat(sprintf( + "dim(y_hat_test): %d x %d\n", + nrow(m$y_hat_test), + ncol(m$y_hat_test) +)) diff --git a/tools/debug/debug_cpp_sampler.py b/tools/debug/debug_cpp_sampler.py new file mode 100644 index 00000000..67d22354 --- /dev/null +++ b/tools/debug/debug_cpp_sampler.py @@ -0,0 +1,53 @@ +"""Minimal script for debugging the C++ sampler under lldb. + +Usage: + source venv/bin/activate + lldb -- python debug/debug_cpp_sampler.py + # then at the (lldb) prompt: + # run + # bt (after the crash, to get a backtrace) + # frame info (to see the crashing frame) +""" + +import numpy as np +from stochtree import BARTModel + +seed = 1001 +n = 10000 +p = 10 +rng = np.random.default_rng(1234) + +X = rng.uniform(size=(n, p)) +f_X = ( + np.where((X[:, 0] >= 0.00) & (X[:, 0] < 0.25), -7.5, 0) + + np.where((X[:, 0] >= 0.25) & (X[:, 0] < 0.50), -2.5, 0) + + np.where((X[:, 0] >= 0.50) & (X[:, 0] < 0.75), 2.5, 0) + + np.where((X[:, 0] >= 0.75) & (X[:, 0] < 1.00), 7.5, 0) +) +y = f_X + rng.normal(scale=1.0, size=n) + +n_test = round(0.2 * n) +test_inds = rng.choice(n, size=n_test, replace=False) +train_inds = np.setdiff1d(np.arange(n), test_inds) + +X_train, X_test = X[train_inds], X[test_inds] +y_train, y_test = y[train_inds], y[test_inds] + +print(f"n_train={len(train_inds)} n_test={n_test} p={p} seed={seed}") +print("Calling BARTModel.sample() with run_cpp=True ...") + +m = BARTModel() +m.sample( + X_train=X_train, + y_train=y_train, + X_test=X_test, + num_gfr=10, + num_burnin=0, + num_mcmc=100, + general_params={"random_seed": seed}, + mean_forest_params={"num_trees": 200}, + run_cpp=True, +) + +print("Completed successfully.") +print(f"y_hat_test shape: {m.y_hat_test.shape}") diff --git a/tools/debug/multivariate_bcf_comparison.R b/tools/debug/multivariate_bcf_comparison.R new file mode 100644 index 00000000..cf86d32f --- /dev/null +++ b/tools/debug/multivariate_bcf_comparison.R @@ -0,0 +1,119 @@ +################################################################################ +# Multivariate treatment BCF: R slow path vs R C++ path. +# +# Generates data in R, runs both paths, and prints correlations / RMSEs +# against ground truth for mu, tau1, tau2, and y. +# +# Usage (from repo root): +# NOT_CRAN=true Rscript -e "devtools::load_all('.'); source('tools/debug/multivariate_bcf_comparison.R')" +################################################################################ + +suppressPackageStartupMessages(devtools::load_all(".")) + +# ── DGP ─────────────────────────────────────────────────────────────────────── +set.seed(1234) +n <- 500 +p <- 5 +snr <- 2.0 + +X <- matrix(runif(n * p), ncol = p) +pi_x <- cbind(0.25 + 0.5 * X[, 1], 0.75 - 0.5 * X[, 2]) +mu_x <- pi_x[, 1] * 5 + pi_x[, 2] * 2 + 2 * X[, 3] +tau_x <- cbind(X[, 2], X[, 3]) +Z <- matrix(NA_integer_, nrow = n, ncol = 2) +for (j in 1:2) Z[, j] <- rbinom(n, 1, pi_x[, j]) +E_XZ <- mu_x + rowSums(Z * tau_x) +y <- E_XZ + rnorm(n, 0, 1) * (sd(E_XZ) / snr) + +test_set_pct <- 0.2 +n_test <- round(test_set_pct * n) +test_inds <- sort(sample(seq_len(n), n_test, replace = FALSE)) +train_inds <- setdiff(seq_len(n), test_inds) + +X_train <- X[train_inds, ]; X_test <- X[test_inds, ] +Z_train <- Z[train_inds, ]; Z_test <- Z[test_inds, ] +pi_train <- pi_x[train_inds, ]; pi_test <- pi_x[test_inds, ] +y_train <- y[train_inds]; y_test <- y[test_inds] +mu_test <- mu_x[test_inds] +tau_test <- tau_x[test_inds, ] + +cat(sprintf("n_train=%d n_test=%d treatment_dim=2\n\n", length(train_inds), n_test)) + +num_gfr <- 10 +num_burnin <- 0 +num_mcmc <- 100 +general_params <- list(adaptive_coding = FALSE) +prog_params <- list(sample_sigma2_leaf = FALSE) +trt_params <- list(sample_sigma2_leaf = FALSE) + +# Helper: posterior means from a 3D tau array (n x treatment_dim x num_samples) +tau_mean <- function(tau_arr, dim_idx) apply(tau_arr[, dim_idx, ], 1, mean) + +report <- function(label, mu_hat, tau_hat1, tau_hat2, y_hat) { + cat(sprintf("[%s]\n", label)) + cat(sprintf(" cor(mu_hat, mu_true) = %.4f RMSE = %.4f\n", + cor(mu_hat, mu_test), sqrt(mean((mu_hat - mu_test)^2)))) + cat(sprintf(" cor(tau_hat1, tau_true1) = %.4f RMSE = %.4f\n", + cor(tau_hat1, tau_test[, 1]), sqrt(mean((tau_hat1 - tau_test[, 1])^2)))) + cat(sprintf(" cor(tau_hat2, tau_true2) = %.4f RMSE = %.4f\n", + cor(tau_hat2, tau_test[, 2]), sqrt(mean((tau_hat2 - tau_test[, 2])^2)))) + cat(sprintf(" cor(y_hat, y_test) = %.4f RMSE = %.4f\n", + cor(y_hat, y_test), sqrt(mean((y_hat - y_test)^2)))) +} + +# ── R slow path ─────────────────────────────────────────────────────────────── +cat("Running R slow path ...\n") +t0 <- proc.time() +bcf_slow <- bcf( + X_train = X_train, Z_train = Z_train, y_train = y_train, + propensity_train = pi_train, + X_test = X_test, Z_test = Z_test, propensity_test = pi_test, + num_gfr = num_gfr, num_burnin = num_burnin, num_mcmc = num_mcmc, + general_params = general_params, + prognostic_forest_params = prog_params, + treatment_effect_forest_params = trt_params +) +slow_time <- (proc.time() - t0)[3] +cat(sprintf(" done in %.1f s\n\n", slow_time)) + +# tau_hat_test: (n_test, treatment_dim, num_mcmc) for multivariate slow path +report( + "Slow path", + rowMeans(bcf_slow$mu_hat_test), + tau_mean(bcf_slow$tau_hat_test, 1), + tau_mean(bcf_slow$tau_hat_test, 2), + rowMeans(bcf_slow$y_hat_test) +) + +# ── R C++ path ──────────────────────────────────────────────────────────────── +cat("\nRunning R C++ path ...\n") +t0 <- proc.time() +bcf_cpp <- bcf( + X_train = X_train, Z_train = Z_train, y_train = y_train, + propensity_train = pi_train, + X_test = X_test, Z_test = Z_test, propensity_test = pi_test, + num_gfr = num_gfr, num_burnin = num_burnin, num_mcmc = num_mcmc, + general_params = general_params, + prognostic_forest_params = prog_params, + treatment_effect_forest_params = trt_params, + run_cpp = TRUE +) +cpp_time <- (proc.time() - t0)[3] +cat(sprintf(" done in %.1f s\n\n", cpp_time)) + +# tau_hat_test: (n_test, treatment_dim, num_mcmc) for multivariate C++ path +report( + "C++ path", + rowMeans(bcf_cpp$mu_hat_test), + tau_mean(bcf_cpp$tau_hat_test, 1), + tau_mean(bcf_cpp$tau_hat_test, 2), + rowMeans(bcf_cpp$y_hat_test) +) + +# ── Cross-path agreement ────────────────────────────────────────────────────── +cat("\n[Cross-path agreement: cor(C++, slow)]\n") +cat(sprintf(" mu_hat: %.4f\n", cor(rowMeans(bcf_cpp$mu_hat_test), rowMeans(bcf_slow$mu_hat_test)))) +cat(sprintf(" tau_hat1: %.4f\n", cor(tau_mean(bcf_cpp$tau_hat_test, 1), tau_mean(bcf_slow$tau_hat_test, 1)))) +cat(sprintf(" tau_hat2: %.4f\n", cor(tau_mean(bcf_cpp$tau_hat_test, 2), tau_mean(bcf_slow$tau_hat_test, 2)))) +cat(sprintf(" y_hat: %.4f\n", cor(rowMeans(bcf_cpp$y_hat_test), rowMeans(bcf_slow$y_hat_test)))) +cat(sprintf("\nSpeedup: %.2fx\n", slow_time / cpp_time))