Added ability to update variance weights

Author: andrewherren

R/cpp11.R

@@ -36,6 +36,10 @@ forest_dataset_update_basis_cpp <- function(dataset_ptr, basis) {
   invisible(.Call(`_stochtree_forest_dataset_update_basis_cpp`, dataset_ptr, basis))
 }
 
+forest_dataset_update_var_weights_cpp <- function(dataset_ptr, weights) {
+  invisible(.Call(`_stochtree_forest_dataset_update_var_weights_cpp`, dataset_ptr, weights))
+}
+
 forest_dataset_add_weights_cpp <- function(dataset_ptr, weights) {
   invisible(.Call(`_stochtree_forest_dataset_add_weights_cpp`, dataset_ptr, weights))
 }

man/ForestModel.Rd

@@ -84,6 +84,17 @@ void forest_dataset_update_basis_cpp(cpp11::external_pointer<StochTree::ForestDa
     UNPROTECT(1);
 }
 
+[[cpp11::register]]
+void forest_dataset_update_var_weights_cpp(cpp11::external_pointer<StochTree::ForestDataset> 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);
+}
+
 [[cpp11::register]]
 void forest_dataset_add_weights_cpp(cpp11::external_pointer<StochTree::ForestDataset> dataset_ptr, cpp11::doubles weights) {
     // Add weights

man/bcf.Rd

@@ -72,6 +72,14 @@ extern "C" SEXP _stochtree_forest_dataset_update_basis_cpp(SEXP dataset_ptr, SEX
   END_CPP11
 }
 // R_data.cpp
+void forest_dataset_update_var_weights_cpp(cpp11::external_pointer<StochTree::ForestDataset> dataset_ptr, cpp11::doubles weights);
+extern "C" SEXP _stochtree_forest_dataset_update_var_weights_cpp(SEXP dataset_ptr, SEXP weights) {
+  BEGIN_CPP11
+    forest_dataset_update_var_weights_cpp(cpp11::as_cpp<cpp11::decay_t<cpp11::external_pointer<StochTree::ForestDataset>>>(dataset_ptr), cpp11::as_cpp<cpp11::decay_t<cpp11::doubles>>(weights));
+    return R_NilValue;
+  END_CPP11
+}
+// R_data.cpp
 void forest_dataset_add_weights_cpp(cpp11::external_pointer<StochTree::ForestDataset> dataset_ptr, cpp11::doubles weights);
 extern "C" SEXP _stochtree_forest_dataset_add_weights_cpp(SEXP dataset_ptr, SEXP weights) {
   BEGIN_CPP11
@@ -1541,6 +1549,7 @@ static const R_CallMethodDef CallEntries[] = {
     {"_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_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,                2},
     {"_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},

src/R_data.cpp

@@ -67,6 +67,14 @@ class ForestDatasetCpp {
     dataset_->AddVarianceWeights(data_ptr, num_row);
   }
 
+  void UpdateVarianceWeights(py::array_t<double> weight_vector, data_size_t num_row) {
+    // Extract pointer to contiguous block of memory
+    double* data_ptr = static_cast<double*>(weight_vector.mutable_data());
+    
+    // Load covariates
+    dataset_->AddVarianceWeights(data_ptr, num_row);
+  }
+
   data_size_t NumRows() {
     return dataset_->NumObservations();
   }

src/cpp11.cpp

@@ -61,6 +61,16 @@ def update_basis(self, basis: np.array):
         basis_ = np.expand_dims(basis, 1) if np.ndim(basis) == 1 else basis
         n, p = basis_.shape
         basis_rowmajor = np.ascontiguousarray(basis_)
+        if not self.has_basis():
+            raise ValueError("This dataset does not have a basis to update. Please use `add_basis` to create and initialize the values in the Dataset's basis matrix.")
+        if not isinstance(basis, np.ndarray):
+            raise ValueError("basis must be a numpy array.")
+        if basis.ndim != 2:
+            raise ValueError("basis must be a 2-dimensional numpy array.")
+        if self.num_basis() != p:
+            raise ValueError(f"The number of columns in the new basis ({p}) must match the number of columns in the existing basis ({self.num_basis()}).")
+        if self.num_observations() != n:
+            raise ValueError(f"The number of rows in the new basis ({n}) must match the number of rows in the existing basis ({self.num_observations()}).")
         self.dataset_cpp.UpdateBasis(basis_rowmajor, n, p, True)
 
     def add_variance_weights(self, variance_weights: np.array):
@@ -74,6 +84,27 @@ def add_variance_weights(self, variance_weights: np.array):
         """
         n = variance_weights.size
         self.dataset_cpp.AddVarianceWeights(variance_weights, n)
+    
+    def update_variance_weights(self, variance_weights: np.array):
+        """
+        Update variance weights in a dataset. Allows users to build an ensemble that depends on 
+        variance weights that are updated throughout the sampler.
+
+        Parameters
+        ----------
+        variance_weights : np.array
+            Univariate numpy array of variance weights.
+        """
+        n = variance_weights.size
+        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.")
+        if not isinstance(variance_weights, np.ndarray):
+            raise ValueError("variance_weights must be a numpy array.")
+        if variance_weights.ndim != 1:
+            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.AddVarianceWeights(variance_weights, n)
 
     def num_observations(self) -> int:
         """