[BUG] Some params not reloaded in createBCFModelFromJsonString()

[evan-magnusson]

Description

In the R version of stochtree to save and reload a BCF model, when using saveBCFModelToJsonString() and then createBCFModelFromJsonString(), I've found that 1) printing the object results in an error, as does attempting to use predict() on the loaded model. It appears that certain required parameters are not being reloaded by createBCFModelFromJsonString() (but are in the json string) - when I add them manually, everything works as intended. There may be other important parameters that aren't loaded as well, but my code caught these.

Reproducing

Given an object bcf_fit:

json_string_to_export <- saveBCFModelToJsonString(bcf_fit) 
# Save string to file location, then read that string as string_bcf (omitted code here)
reloaded_bcf_model <- createBCFModelFromJsonString(string_bcf)

Error 1:

reloaded_bcf_model

ERROR while rich displaying an object: Error in if (x$model_params$binary_treatment) {: argument is of length zero

Traceback:
1. sapply(x, f, simplify = simplify)
2. lapply(X = X, FUN = FUN, ...)
3. FUN(X[[i]], ...)
4. tryCatch(withCallingHandlers({
 .     if (!mime %in% names(repr::mime2repr)) 
 .         stop("No repr_* for mimetype ", mime, " in repr::mime2repr")
 .     rpr <- repr::mime2repr[[mime]](obj)
 .     if (is.null(rpr)) 
 .         return(NULL)
 .     prepare_content(is.raw(rpr), rpr)
 . }, error = error_handler), error = outer_handler)
5. tryCatchList(expr, classes, parentenv, handlers)
6. tryCatchOne(expr, names, parentenv, handlers[[1L]])
7. doTryCatch(return(expr), name, parentenv, handler)
8. withCallingHandlers({
 .     if (!mime %in% names(repr::mime2repr)) 
 .         stop("No repr_* for mimetype ", mime, " in repr::mime2repr")
 .     rpr <- repr::mime2repr[[mime]](obj)
 .     if (is.null(rpr)) 
 .         return(NULL)
 .     prepare_content(is.raw(rpr), rpr)
 . }, error = error_handler)
9. repr::mime2repr[[mime]](obj)
10. repr_text.default(obj)
11. paste(capture.output(print(obj)), collapse = "\n")
12. capture.output(print(obj))
13. withVisible(...elt(i))
14. print(obj)
15. print.bcfmodel(obj)

Error 2:

predict(reloaded_bcf_model, data, treatment)

Also results in an error (don't have the text at the moment).

The following code fixes the issue:

reloaded_bcf_model$model_params$binary_treatment <- TRUE
reloaded_bcf_model$bart_propensity_model$fit$numeric_vars <- vars_I_used
reloaded_bcf_model$bart_propensity_model$train_set_metadata$numeric_vars <- vars_I_used

Version: 0.4.1

[andrewherren]

Hi @evan-magnusson, thanks for raising this! It's a very timely issue as I'm working through an RFC on serialization / deserialization. This portion of the codebase definitely needs a lot more standardization and testing, along with shims for previous versions, to make it bulletproof (hence the focus of the RFC).

I'm sorry this workflow broke down for your model, and I definitely want to get it back on track.

We're about to release 0.4.3, but I'll postpone that to ensure a fix for this is included. I'll follow up when a fix is available on main, shouldn't be any later than tomorrow EOD.

[andrewherren]

@evan-magnusson If you install from main, this problem should be fixed. I've added a lot more unit tests for the JSON roundtrip (in both R and Python) and fixed several reload issues like the one you raised.

As a quick simulation-based reprex, this script should run now

# Load library
library(stochtree)

# Generate data
n <- 100
p <- 10
X <- matrix(runif(n * p), ncol = p)
pi_x <- 0.2 + 0.6 * X[, 1]
Z <- rbinom(n, 1, pi_x)
mu_x <- 3 * X[, 1]
tau_x <- 0.5 * X[, 2]
y <- mu_x + tau_x * Z + rnorm(n)

# Model fit and predictions
bcf_model <- bcf(X_train = X, Z_train = Z, y_train = y)
preds_orig <- predict(bcf_model, X = X, Z = Z)

# Inspect model
print(bcf_model)
summary(bcf_model)

# JSON roundtrip
bcf_json_string <- saveBCFModelToJsonString(bcf_model)
reloaded_bcf_model <- createBCFModelFromJsonString(bcf_json_string)

# Inspect reloaded model
print(reloaded_bcf_model)
summary(reloaded_bcf_model)

# Predict from reloaded model
preds_reloaded <- predict(reloaded_bcf_model, X = X, Z = Z)

# Check agreement
isTRUE(all.equal(preds_orig$y_hat, preds_reloaded$y_hat))
isTRUE(all.equal(preds_orig$tau_hat, preds_reloaded$tau_hat))
isTRUE(all.equal(preds_orig$cate, preds_reloaded$cate))
isTRUE(all.equal(preds_orig$mu_hat, preds_reloaded$mu_hat))
isTRUE(all.equal(preds_orig$prognostic_function, preds_reloaded$prognostic_function))

[evan-magnusson]

Thank you so much!