R: use rvars from posterior as output format (#9)

* R: use rvars from posterior as output format * Install posterior in CI
WardBrian · Nov 17, 2023 · 861d5b5 · 861d5b5
1 parent 3dc8634
commit 861d5b5
Show file tree

Hide file tree

Showing 9 changed files with 85 additions and 47 deletions.
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
@@ -136,6 +136,7 @@ jobs:
             any::R6
             any::testthat
             any::devtools
+            any::posterior
 
       - name: Restore Stan
         uses: actions/cache@v3

diff --git a/clients/R/DESCRIPTION b/clients/R/DESCRIPTION
@@ -14,5 +14,6 @@ Suggests:
     testthat (>= 3.0.0)
 Imports:
     R6 (>= 2.4.0),
+    posterior (>= 1.4.0),
     rlang (>= 0.3.0)
 Config/testthat/edition: 3
diff --git a/clients/R/R/ffistan.R b/clients/R/R/ffistan.R
@@ -109,8 +109,7 @@ FFIStanModel <- R6::R6Class("FFIStanModel", public = list(initialize = function(
             err = raw(8), PACKAGE = private$lib_name)
         handle_error(vars$return_code, private$lib_name, vars$err)
         # reshape the output matrix
-        out <- aperm(array(vars$out, dim = c(num_params, num_draws, num_chains),
-            dimnames = list(params, NULL, NULL)), c(3, 2, 1))
+        out <- output_as_rvars(params, num_draws, num_chains, vars$out)
 
         if (save_metric) {
             if (metric == HMCMetric$DENSE) {
@@ -120,10 +119,10 @@ FFIStanModel <- R6::R6Class("FFIStanModel", public = list(initialize = function(
                 metric <- aperm(array(vars$metric, dim = c(free_params, num_chains)),
                   c(2, 1))
             }
-            return(list(params = params, draws = out, metric = metric))
+            return(list(draws = out, metric = metric))
         }
 
-        list(params = params, draws = out)
+        out
     })
 }, pathfinder = function(data = "", num_paths = 4, inits = NULL, seed = NULL, id = 1,
     init_radius = 2, num_draws = 1000, max_history_size = 5, init_alpha = 0.001,
@@ -167,10 +166,8 @@ FFIStanModel <- R6::R6Class("FFIStanModel", public = list(initialize = function(
             as.integer(num_elbo_draws), as.integer(num_multi_draws), as.integer(refresh),
             as.integer(num_threads), out = double(output_size), err = raw(8), PACKAGE = private$lib_name)
         handle_error(vars$return_code, private$lib_name, vars$err)
-        # reshape the output matrix
-        out <- t(array(vars$out, dim = c(num_params, num_output), dimnames = list(params,
-            NULL)))
-        list(params = params, draws = out)
+
+        output_as_rvars(params, num_output, 1, vars$out)
     })
 }, optimize = function(data = "", init = NULL, seed = NULL, id = 1, init_radius = 2,
     algorithm = OptimizationAlgorithm$LBFGS, jacobian = FALSE, num_iterations = 2000,
@@ -194,10 +191,9 @@ FFIStanModel <- R6::R6Class("FFIStanModel", public = list(initialize = function(
             as.double(tol_param), as.integer(refresh), as.integer(num_threads), out = double(output_size),
             err = raw(8), PACKAGE = private$lib_name)
         handle_error(vars$return_code, private$lib_name, vars$err)
-        out <- array(vars$out, dim = c(num_params), dimnames = list(params))
-        list(params = params, optimum = out)
-    })
 
+        output_as_rvars(params, 1, 1, vars$out)
+    })
 }), private = list(lib = NA, lib_name = NA, sep = NA, with_model = function(data,
     seed, block) {
     ffi_ret <- .C("ffistan_create_model_R", model = raw(8), as.character(data), as.integer(seed),

diff --git a/clients/R/R/output.R b/clients/R/R/output.R
@@ -0,0 +1,23 @@
+
+# copied from cmdstanr, definitely doesn't handle tuples, but then neither does
+# posterior
+repair_variable_names <- function(names) {
+    names <- sub("\\.", "[", names)
+    names <- gsub("\\.", ",", names)
+    names[grep("\\[", names)] <- paste0(names[grep("\\[", names)], "]")
+    names
+}
+
+
+output_as_rvars <- function(names, num_draws, num_chains, draws) {
+    names <- repair_variable_names(names)
+    num_params <- length(names)
+    dims <- c(num_params, num_draws, num_chains)
+
+    # all our outputs are row-major
+    draws <- array(draws, dim = dims, dimnames = list(names, NULL, NULL))
+    # so we need to rearrange. posterior likes draws x chains x params
+    draws <- aperm(draws, c(2, 3, 1))
+
+    posterior::as_draws_rvars(draws)
+}
diff --git a/clients/R/example.R b/clients/R/example.R
@@ -4,10 +4,10 @@ model <- FFIStanModel$new("./test_models/bernoulli/bernoulli_model.so")
 data <- "./test_models/bernoulli/bernoulli.data.json"
 
 fit <- model$sample(data)
-print(colMeans(fit$draws, dims = 2)[8])
+print(fit$theta)
 
-pf = model$pathfinder(data)
-print(colMeans(pf$draws)[3])
+pf <- model$pathfinder(data)
+print(pf$theta)
 
-o = model$optimize(data)
-print(o$optimum[2])
+o <- model$optimize(data)
+print(o$theta)
diff --git a/clients/R/tests/testthat/setup.R b/clients/R/tests/testthat/setup.R
@@ -1,3 +1,7 @@
+options(warnPartialMatchDollar = TRUE, warn = 2)
+options(warnPartialMatchArgs = TRUE, warn = 2)
+options(warnPartialMatchAttr = TRUE, warn = 2)
+
 stan_folder <- file.path("..", "..", "..", "..", "test_models")
 
 bernoulli_model <- ffistan::FFIStanModel$new(file.path(stan_folder, "bernoulli",
@@ -11,3 +15,12 @@ multimodal_model <- ffistan::FFIStanModel$new(file.path(stan_folder, "multimodal
     "multimodal_model.so"))
 simple_jacobian_model <- ffistan::FFIStanModel$new(file.path(stan_folder, "simple_jacobian",
     "simple_jacobian_model.so"))
+
+# hack around the fact that comparisons to NULL result in logical(0) and
+# all(logical(0)) is TRUE, for some reason.
+.builtin_all <- all
+all <- function(l) {
+    if (length(l) == 0)
+        return(isTRUE(l))
+    .builtin_all(l)
+}
diff --git a/clients/R/tests/testthat/test_optimize.R b/clients/R/tests/testthat/test_optimize.R
@@ -4,11 +4,11 @@ ALGORITHMS <- c(ffistan::OptimizationAlgorithm$NEWTON, ffistan::OptimizationAlgo
 test_that("data args work", {
 
     out1 <- bernoulli_model$optimize(BERNOULLI_DATA)
-    expect_true((out1$optimum[2]) > 0.19 && (out1$optimum[2]) < 0.21)
+    expect_true(mean(out1$theta) > 0.19 && mean(out1$theta) < 0.21)
 
     data_file <- file.path(stan_folder, "bernoulli", "bernoulli.data.json")
     out2 <- bernoulli_model$optimize(data = data_file)
-    expect_true(mean(out2$optimum[2]) > 0.19 && mean(out2$optimum[2]) < 0.21)
+    expect_true(mean(out2$theta) > 0.19 && mean(out2$theta) < 0.21)
 
 })
 
@@ -20,10 +20,12 @@ test_that("algorithm and jacobian args work", {
             out <- simple_jacobian_model$optimize(algorithm = algorithm, jacobian = jacobian,
                 seed = 1234)
 
+            sigma <- posterior::extract_variable(out, "sigma")
+
             if (jacobian) {
-                expect_equal(out$optimum[2], 3.3, tolerance = 0.01, ignore_attr = TRUE)
+                expect_equal(sigma, 3.3, tolerance = 0.01, ignore_attr = TRUE)
             } else {
-                expect_equal(out$optimum[2], 3, tolerance = 0.01, ignore_attr = TRUE)
+                expect_equal(sigma, 3, tolerance = 0.01, ignore_attr = TRUE)
             }
 
         }
@@ -36,10 +38,10 @@ test_that("seed works", {
     out1 <- bernoulli_model$optimize(BERNOULLI_DATA, seed = 123)
     out2 <- bernoulli_model$optimize(BERNOULLI_DATA, seed = 123)
 
-    expect_equal(out1$optimum, out2$optimum)
+    expect_equal(out1, out2)
 
     out3 <- bernoulli_model$optimize(BERNOULLI_DATA, seed = 456)
-    expect_error(expect_equal(out1$optimum, out3$optimum))
+    expect_error(expect_equal(out1, out3))
 
 })
 

diff --git a/clients/R/tests/testthat/test_pathfinder.R b/clients/R/tests/testthat/test_pathfinder.R
@@ -1,11 +1,11 @@
 test_that("data arguments work", {
 
     out1 <- bernoulli_model$pathfinder(BERNOULLI_DATA)
-    expect_true(mean(out1$draws[, 3]) > 0.2 && mean(out1$draws[, 3]) < 0.3)
+    expect_true(mean(out1$theta) > 0.2 && mean(out1$theta) < 0.3)
 
     data_file <- file.path(stan_folder, "bernoulli", "bernoulli.data.json")
     out2 <- bernoulli_model$pathfinder(data = data_file)
-    expect_true(mean(out2$draws[, 3]) > 0.2 && mean(out2$draws[, 3]) < 0.3)
+    expect_true(mean(out2$theta) > 0.2 && mean(out2$theta) < 0.3)
 
 })
 
@@ -14,11 +14,11 @@ test_that("output sizes are correct", {
 
     out1 <- bernoulli_model$pathfinder(BERNOULLI_DATA, num_paths = 4, num_draws = 101,
         num_multi_draws = 99)
-    expect_equal(dim(out1$draws)[1], 99)
+    expect_equal(posterior::ndraws(out1), 99)
 
     out2 <- bernoulli_model$pathfinder(BERNOULLI_DATA, num_paths = 1, num_draws = 101,
         num_multi_draws = 99)
-    expect_equal(dim(out2$draws)[1], 101)
+    expect_equal(posterior::ndraws(out2), 101)
 
 })
 
@@ -27,27 +27,26 @@ test_that("seed works", {
     out1 <- bernoulli_model$pathfinder(BERNOULLI_DATA, seed = 123)
     out2 <- bernoulli_model$pathfinder(BERNOULLI_DATA, seed = 123)
 
-    expect_equal(out1$draws, out2$draws)
+    expect_equal(out1$theta, out2$theta)
 
     out3 <- bernoulli_model$pathfinder(BERNOULLI_DATA, seed = 456)
-    expect_error(expect_equal(out1$draws, out3$draws))
+    expect_error(expect_equal(out1$theta, out3$theta))
 
 })
 test_that("inits work", {
 
     init1 <- "{\"mu\": -1000}"
     out1 <- multimodal_model$pathfinder(inits = init1)
-    expect_true(all(out1$draws[, 3] < 0))
+    expect_true(all(out1$mu < 0))
 
     init2 <- "{\"mu\": 1000}"
     out2 <- multimodal_model$pathfinder(inits = init2)
-    expect_true(all(out2$draws[, 3] > 0))
+    expect_true(all(out2$mu > 0))
 
     temp_file <- tempfile(fileext = ".json")
     write(init1, temp_file)
     out3 <- multimodal_model$pathfinder(num_paths = 2, inits = c(temp_file, init1))
-    expect_true(all(out3$draws[, 3] < 0))
-
+    expect_true(all(out3$mu < 0))
 })
 
 

diff --git a/clients/R/tests/testthat/test_sample.R b/clients/R/tests/testthat/test_sample.R
@@ -2,23 +2,23 @@ test_that("data arguments work", {
 
     out1 <- bernoulli_model$sample(BERNOULLI_DATA, num_warmup = 100, num_samples = 100)
 
-    expect_true(mean(out1$draws[, , 8]) > 0.2 && mean(out1$draws[, , 8]) < 0.3)
+    expect_true(mean(out1$theta) > 0.2 && mean(out1$theta) < 0.3)
 
     data_file <- file.path(stan_folder, "bernoulli", "bernoulli.data.json")
     out2 <- bernoulli_model$sample(data = data_file, num_warmup = 100, num_samples = 100)
-    expect_true(mean(out2$draws[, , 8]) > 0.2 && mean(out2$draws[, , 8]) < 0.3)
+    expect_true(mean(out2$theta) > 0.2 && mean(out2$theta) < 0.3)
 
 })
 
 test_that("save_warmup works", {
 
     out <- bernoulli_model$sample(BERNOULLI_DATA, num_warmup = 12, num_samples = 34,
         save_warmup = FALSE)
-    expect_equal(dim(out$draws)[2], 34)
+    expect_equal(posterior::niterations(out), 34)
 
     out <- bernoulli_model$sample(BERNOULLI_DATA, num_warmup = 12, num_samples = 34,
         save_warmup = TRUE)
-    expect_equal(dim(out$draws)[2], 12 + 34)
+    expect_equal(posterior::niterations(out), 12 + 34)
 
 })
 
@@ -29,12 +29,12 @@ test_that("seed works", {
     out2 <- bernoulli_model$sample(BERNOULLI_DATA, seed = 123, num_warmup = 100,
         num_samples = 100)
 
-    expect_equal(out1$draws, out2$draws)
+    expect_equal(out1, out2)
 
     out3 <- bernoulli_model$sample(BERNOULLI_DATA, seed = 456, num_warmup = 100,
         num_samples = 100)
 
-    expect_error(expect_equal(out1$draws, out3$draws))
+    expect_error(expect_equal(out1, out3))
 
 })
 
@@ -75,9 +75,10 @@ test_that("init_inv_metric is used", {
             adapt = adapt, metric = HMCMetric$DIAGONAL, init_inv_metric = diag_metric,
             save_metric = TRUE, seed = 1234)
 
-        chain_one_divergences <- sum(out_diag$draws[1, , 6])
+        divergent <- out_diag$draws$divergent__
+        chain_one_divergences <- sum(posterior::subset_draws(divergent, chain = 1))
         expect_true(chain_one_divergences > ifelse(adapt, 10, 500))
-        chain_two_divergences <- sum(out_diag$draws[2, , 6])
+        chain_two_divergences <- sum(posterior::subset_draws(divergent, chain = 2))
         expect_true(chain_two_divergences < 10)
         expect_true(chain_two_divergences < chain_one_divergences)
         expect_false(all(diag_metric == t(out_diag$metric)))
@@ -89,9 +90,10 @@ test_that("init_inv_metric is used", {
             adapt = adapt, metric = HMCMetric$DENSE, init_inv_metric = dense_metric,
             save_metric = TRUE, seed = 1234)
 
-        chain_one_divergences <- sum(out_dense$draws[1, , 6])
+        divergent <- out_dense$draws$divergent__
+        chain_one_divergences <- sum(posterior::subset_draws(divergent, chain = 1))
         expect_true(chain_one_divergences > ifelse(adapt, 10, 500))
-        chain_two_divergences <- sum(out_dense$draws[2, , 6])
+        chain_two_divergences <- sum(posterior::subset_draws(divergent, chain = 2))
         expect_true(chain_two_divergences < 10)
         expect_true(chain_two_divergences < chain_one_divergences)
         expect_false(all(dense_metric == aperm(out_dense$metric, c(3, 2, 1))))
@@ -103,20 +105,21 @@ test_that("multiple inits work", {
     init1 <- "{\"mu\": -100}"
     out1 <- multimodal_model$sample(num_chains = 2, num_warmup = 100, num_samples = 100,
         inits = init1)
-    expect_true(all(out1$draws[, , 8] < 0))
+    expect_true(all(out1$mu < 0))
 
     init2 <- "{\"mu\": 100}"
     out2 <- multimodal_model$sample(num_chains = 2, num_warmup = 100, num_samples = 100,
         inits = list(init1, init2))
-    expect_true(all(out2$draws[1, , 8] < 0))
-    expect_true(all(out2$draws[2, , 8] > 0))
+
+    expect_true(all(posterior::subset_draws(out2$mu, chain = 1) < 0))
+    expect_true(all(posterior::subset_draws(out2$mu, chain = 2) > 0))
 
     temp_file <- tempfile(fileext = ".json")
     write(init1, temp_file)
     out3 <- multimodal_model$sample(num_chains = 2, num_warmup = 100, num_samples = 100,
         inits = c(temp_file, init2))
-    expect_true(all(out3$draws[1, , 8] < 0))
-    expect_true(all(out3$draws[2, , 8] > 0))
+    expect_true(all(posterior::subset_draws(out3$mu, chain = 1) < 0))
+    expect_true(all(posterior::subset_draws(out3$mu, chain = 2) > 0))
 
 })