Updated interfaces and added unit tests

andrewherren · andrewherren · commit 69d60aa42875 · 2025-09-22T18:01:35.000-05:00
diff --git a/R/cpp11.R b/R/cpp11.R
@@ -80,6 +80,10 @@ rfx_dataset_update_var_weights_cpp <- function(dataset_ptr, weights, exponentiat
   invisible(.Call(`_stochtree_rfx_dataset_update_var_weights_cpp`, dataset_ptr, weights, exponentiate))
 }
 
+rfx_dataset_update_group_labels_cpp <- function(dataset_ptr, group_labels) {
+  invisible(.Call(`_stochtree_rfx_dataset_update_group_labels_cpp`, dataset_ptr, group_labels))
+}
+
 rfx_dataset_num_basis_cpp <- function(dataset) {
   .Call(`_stochtree_rfx_dataset_num_basis_cpp`, dataset)
 }
diff --git a/src/cpp11.cpp b/src/cpp11.cpp
@@ -157,6 +157,14 @@ extern "C" SEXP _stochtree_rfx_dataset_update_var_weights_cpp(SEXP dataset_ptr,
   END_CPP11
 }
 // R_data.cpp
+void rfx_dataset_update_group_labels_cpp(cpp11::external_pointer<StochTree::RandomEffectsDataset> dataset_ptr, cpp11::integers group_labels);
+extern "C" SEXP _stochtree_rfx_dataset_update_group_labels_cpp(SEXP dataset_ptr, SEXP group_labels) {
+  BEGIN_CPP11
+    rfx_dataset_update_group_labels_cpp(cpp11::as_cpp<cpp11::decay_t<cpp11::external_pointer<StochTree::RandomEffectsDataset>>>(dataset_ptr), cpp11::as_cpp<cpp11::decay_t<cpp11::integers>>(group_labels));
+    return R_NilValue;
+  END_CPP11
+}
+// R_data.cpp
 int rfx_dataset_num_basis_cpp(cpp11::external_pointer<StochTree::RandomEffectsDataset> dataset);
 extern "C" SEXP _stochtree_rfx_dataset_num_basis_cpp(SEXP dataset) {
   BEGIN_CPP11
@@ -1697,6 +1705,7 @@ static const R_CallMethodDef CallEntries[] = {
     {"_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},
diff --git a/src/py_stochtree.cpp b/src/py_stochtree.cpp
@@ -67,12 +67,12 @@ class ForestDatasetCpp {
     dataset_->AddVarianceWeights(data_ptr, num_row);
   }
 
-  void UpdateVarianceWeights(py::array_t<double> weight_vector, data_size_t num_row) {
+  void UpdateVarianceWeights(py::array_t<double> weight_vector, data_size_t num_row, bool exponentiate) {
     // Extract pointer to contiguous block of memory
     double* data_ptr = static_cast<double*>(weight_vector.mutable_data());
     
     // Load covariates
-    dataset_->UpdateVarWeights(data_ptr, num_row);
+    dataset_->UpdateVarWeights(data_ptr, num_row, exponentiate);
   }
 
   data_size_t NumRows() {
@@ -2067,6 +2067,7 @@ PYBIND11_MODULE(stochtree_cpp, m) {
     .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)
@@ -2200,6 +2201,9 @@ PYBIND11_MODULE(stochtree_cpp, m) {
       .def("AddGroupLabels", &RandomEffectsDatasetCpp::AddGroupLabels)
       .def("AddBasis", &RandomEffectsDatasetCpp::AddBasis)
       .def("AddVarianceWeights", &RandomEffectsDatasetCpp::AddVarianceWeights)
+      .def("UpdateGroupLabels", &RandomEffectsDatasetCpp::UpdateGroupLabels)
+      .def("UpdateBasis", &RandomEffectsDatasetCpp::UpdateBasis)
+      .def("UpdateVarianceWeights", &RandomEffectsDatasetCpp::UpdateVarianceWeights)
       .def("HasGroupLabels", &RandomEffectsDatasetCpp::HasGroupLabels)
       .def("HasBasis", &RandomEffectsDatasetCpp::HasBasis)
       .def("HasVarianceWeights", &RandomEffectsDatasetCpp::HasVarianceWeights);
diff --git a/stochtree/data.py b/stochtree/data.py
@@ -94,7 +94,7 @@ def add_variance_weights(self, variance_weights: np.array):
         
         self.dataset_cpp.AddVarianceWeights(variance_weights_, n)
     
-    def update_variance_weights(self, variance_weights: np.array):
+    def update_variance_weights(self, variance_weights: np.array, exponentiate: bool = False):
         """
         Update variance weights in a dataset. Allows users to build an ensemble that depends on 
         variance weights that are updated throughout the sampler.
@@ -103,6 +103,8 @@ def update_variance_weights(self, variance_weights: np.array):
         ----------
         variance_weights : np.array
             Univariate numpy array of variance weights.
+        exponentiate : bool
+            Whether to exponentiate the variance weights before storing them in the dataset.
         """
         if not self.has_variance_weights():
             raise ValueError("This dataset does not have variance weights to update. Please use `add_variance_weights` to create and initialize the values in the Dataset's variance weight vector.")
@@ -114,7 +116,7 @@ def update_variance_weights(self, variance_weights: np.array):
             raise ValueError("variance_weights must be a 1-dimensional numpy array.")
         if self.num_observations() != n:
             raise ValueError(f"The number of rows in the new variance_weights vector ({n}) must match the number of rows in the existing vector ({self.num_observations()}).")
-        self.dataset_cpp.UpdateVarianceWeights(variance_weights_, n)
+        self.dataset_cpp.UpdateVarianceWeights(variance_weights_, n, exponentiate)
 
     def num_observations(self) -> int:
         """
diff --git a/stochtree/random_effects.py b/stochtree/random_effects.py
@@ -111,7 +111,7 @@ def add_variance_weights(self, variance_weights: np.array):
         n = variance_weights_.shape[0]
         self.rfx_dataset_cpp.AddVarianceWeights(variance_weights_, n)
     
-    def update_variance_weights(self, variance_weights: np.array):
+    def update_variance_weights(self, variance_weights: np.array, exponentiate: bool = False):
         """
         Update variance weights in a dataset. Allows users to build an ensemble that depends on 
         variance weights that are updated throughout the sampler.
@@ -120,6 +120,8 @@ def update_variance_weights(self, variance_weights: np.array):
         ----------
         variance_weights : np.array
             Univariate numpy array of variance weights.
+        exponentiate : bool
+            Whether to exponentiate the variance weights before storing them in the dataset.
         """
         if not self.has_variance_weights():
             raise ValueError("This dataset does not have variance weights to update. Please use `add_variance_weights` to create and initialize the values in the Dataset's variance weight vector.")
@@ -133,7 +135,7 @@ def update_variance_weights(self, variance_weights: np.array):
         n = variance_weights_.shape[0]
         if self.num_observations() != n:
             raise ValueError(f"The number of rows in the new variance_weights vector ({n}) must match the number of rows in the existing vector ({self.num_observations()}).")
-        self.rfx_dataset_cpp.UpdateVarianceWeights(variance_weights, n)
+        self.rfx_dataset_cpp.UpdateVarianceWeights(variance_weights, n, exponentiate)
 
     def num_observations(self) -> int:
         """
diff --git a/test/R/testthat/test-dataset.R b/test/R/testthat/test-dataset.R
@@ -0,0 +1,56 @@
+test_that("ForestDataset can be constructed and updated", {
+    # Generate data
+    n <- 20
+    num_covariates <- 10
+    num_basis <- 5
+    covariates <- matrix(runif(n * num_covariates), ncol = num_covariates)
+    basis <- matrix(runif(n * num_basis), ncol = num_basis)
+    variance_weights <- runif(n)
+    
+    # Copy data to a ForestDataset object
+    forest_dataset <- createForestDataset(covariates, basis, variance_weights)
+    
+    # Run first round of expectations
+    expect_equal(forest_dataset$num_observations(), n)
+    expect_equal(forest_dataset$num_covariates(), num_covariates)
+    expect_equal(forest_dataset$num_basis(), num_basis)
+    expect_equal(forest_dataset$has_variance_weights(), T)
+    
+    # Update data
+    new_basis <- matrix(runif(n * num_basis), ncol = num_basis)
+    new_variance_weights <- runif(n)
+    expect_no_error(
+        forest_dataset$update_basis(new_basis)
+    )
+    expect_no_error(
+        forest_dataset$update_variance_weights(new_variance_weights)
+    )
+})
+
+test_that("RandomEffectsDataset can be constructed and updated", {
+    # Generate data
+    n <- 20
+    num_groups <- 4
+    num_basis <- 5
+    group_ids <- sample(as.integer(1:num_groups), size = n, replace = T)
+    rfx_basis <- cbind(1, matrix(runif(n*(num_basis-1)), ncol = (num_basis-1)))
+    variance_weights <- runif(n)
+    
+    # Copy data to a RandomEffectsDataset object
+    rfx_dataset <- createRandomEffectsDataset(group_ids, rfx_basis, variance_weights)
+    
+    # Run first round of expectations
+    expect_equal(rfx_dataset$num_observations(), n)
+    expect_equal(rfx_dataset$num_basis(), num_basis)
+    expect_equal(rfx_dataset$has_variance_weights(), T)
+    
+    # Update data
+    new_rfx_basis <- matrix(runif(n * num_basis), ncol = num_basis)
+    new_variance_weights <- runif(n)
+    expect_no_error(
+        rfx_dataset$update_basis(new_basis)
+    )
+    expect_no_error(
+        rfx_dataset$update_variance_weights(new_variance_weights)
+    )
+})
diff --git a/test/python/test_data.py b/test/python/test_data.py
@@ -1,6 +1,6 @@
 import numpy as np
 
-from stochtree import Dataset
+from stochtree import Dataset, RandomEffectsDataset
 
 class TestDataset:
     def test_dataset_update(self):
@@ -29,3 +29,34 @@ def test_dataset_update(self):
         with np.testing.assert_no_warnings():
             forest_dataset.update_basis(new_basis)
             forest_dataset.update_variance_weights(new_variance_weights)
+
+class TestRFXDataset:
+    def test_rfx_dataset_update(self):
+        # Generate data
+        n = 20
+        num_groups = 4
+        num_basis = 5
+        rng = np.random.default_rng()
+        group_labels = rng.choice(num_groups, size=n)
+        basis = np.empty((n, num_basis))
+        basis[:, 0] = 1.0
+        if num_basis > 1:
+            basis[:, 1:] = rng.uniform(-1, 1, (n, num_basis - 1))
+        variance_weights = rng.uniform(0, 1, size=n)
+        
+        # Construct dataset
+        rfx_dataset = RandomEffectsDataset()
+        rfx_dataset.add_group_labels(group_labels)
+        rfx_dataset.add_basis(basis)
+        rfx_dataset.add_variance_weights(variance_weights)
+        assert rfx_dataset.num_observations() == n
+        assert rfx_dataset.num_basis() == num_basis
+        assert rfx_dataset.has_variance_weights()
+        
+        # Update dataset
+        new_basis = rng.uniform(0, 1, size=(n, num_basis))
+        new_variance_weights = rng.uniform(0, 1, size=n)
+        with np.testing.assert_no_warnings():
+            rfx_dataset.update_basis(new_basis)
+            rfx_dataset.update_variance_weights(new_variance_weights)
+

Original file line number	Diff line number	Diff line change
`@@ -80,6 +80,10 @@ rfx_dataset_update_var_weights_cpp <- function(dataset_ptr, weights, exponentiat`
`80`	`80`	invisible(.Call(`_stochtree_rfx_dataset_update_var_weights_cpp`, dataset_ptr, weights, exponentiate))
`81`	`81`	`}`
`82`	`82`
	`83`	`+rfx_dataset_update_group_labels_cpp <- function(dataset_ptr, group_labels) {`
	`84`	+ invisible(.Call(`_stochtree_rfx_dataset_update_group_labels_cpp`, dataset_ptr, group_labels))
	`85`	`+}`
	`86`	`+`
`83`	`87`	`rfx_dataset_num_basis_cpp <- function(dataset) {`
`84`	`88`	.Call(`_stochtree_rfx_dataset_num_basis_cpp`, dataset)
`85`	`89`	`}`