Updated python and R data interfaces

Author: andrewherren

include/stochtree/data.h

@@ -553,6 +553,21 @@ class RandomEffectsDataset {
       var_weights_.SetElement(i, temp_value);
     }
   }
+  /*!
+   * \brief Update a RandomEffectsDataset's group indices
+   *
+   * \param data_ptr Pointer to first element of a contiguous array of data storing a weight vector
+   * \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<int32_t>& 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
+    for (data_size_t i = 0; i < num_row; ++i) {
+      group_labels_[i] = group_labels[i];
+    }
+  }
   /*!
    * \brief Copy / load group indices for random effects
    * 

src/R_data.cpp

@@ -217,6 +217,14 @@ void rfx_dataset_update_var_weights_cpp(cpp11::external_pointer<StochTree::Rando
     UNPROTECT(1);
 }
 
+[[cpp11::register]]
+void rfx_dataset_update_group_labels_cpp(cpp11::external_pointer<StochTree::RandomEffectsDataset> dataset_ptr, cpp11::integers group_labels) {
+    // Update group labels
+    int n = group_labels.size();
+    std::vector<int32_t> 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<StochTree::RandomEffectsDataset> dataset) {
     return dataset->NumBases();

src/py_stochtree.cpp

@@ -72,7 +72,7 @@ class ForestDatasetCpp {
     double* data_ptr = static_cast<double*>(weight_vector.mutable_data());
 
     // Load covariates
-    dataset_->AddVarianceWeights(data_ptr, num_row);
+    dataset_->UpdateVarWeights(data_ptr, num_row);
   }
 
   data_size_t NumRows() {
@@ -1297,6 +1297,22 @@ class RandomEffectsDatasetCpp {
     double* weight_data_ptr = static_cast<double*>(weights.mutable_data());
     rfx_dataset_->AddVarianceWeights(weight_data_ptr, num_row);
   }
+  void UpdateBasis(py::array_t<double> basis, data_size_t num_row, int num_col, bool row_major) {
+    double* basis_data_ptr = static_cast<double*>(basis.mutable_data());
+    rfx_dataset_->UpdateBasis(basis_data_ptr, num_row, num_col, row_major);
+  }
+  void UpdateVarianceWeights(py::array_t<double> weights, data_size_t num_row, bool exponentiate) {
+    double* weight_data_ptr = static_cast<double*>(weights.mutable_data());
+    rfx_dataset_->UpdateVarWeights(weight_data_ptr, num_row, exponentiate);
+  }
+  void UpdateGroupLabels(py::array_t<int> group_labels, data_size_t num_row) {
+    std::vector<int> group_labels_vec(num_row);
+    auto accessor = group_labels.mutable_unchecked<1>();
+    for (py::ssize_t i = 0; i < num_row; i++) {
+      group_labels_vec[i] = accessor(i);
+    }
+    rfx_dataset_->UpdateGroupLabels(group_labels_vec, num_row);
+  }
   bool HasGroupLabels() {return rfx_dataset_->HasGroupLabels();}
   bool HasBasis() {return rfx_dataset_->HasBasis();}
   bool HasVarianceWeights() {return rfx_dataset_->HasVarWeights();}

stochtree/data.py

@@ -85,8 +85,14 @@ def add_variance_weights(self, variance_weights: np.array):
         variance_weights : np.array
             Univariate numpy array of variance weights.
         """
-        n = variance_weights.size
-        self.dataset_cpp.AddVarianceWeights(variance_weights, n)
+        if not isinstance(variance_weights, np.ndarray):
+            raise ValueError("variance_weights must be a numpy array.")
+        variance_weights_ = np.squeeze(variance_weights)
+        n = variance_weights_.size
+        if variance_weights_.ndim != 1:
+            raise ValueError("variance_weights must be a 1-dimensional numpy array.")
+        
+        self.dataset_cpp.AddVarianceWeights(variance_weights_, n)
 
     def update_variance_weights(self, variance_weights: np.array):
         """
@@ -98,16 +104,17 @@ def update_variance_weights(self, variance_weights: np.array):
         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:
+        variance_weights_ = np.squeeze(variance_weights)
+        n = variance_weights_.size
+        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)
+        self.dataset_cpp.UpdateVarianceWeights(variance_weights_, n)
 
     def num_observations(self) -> int:
         """

stochtree/random_effects.py

@@ -40,6 +40,23 @@ def add_group_labels(self, group_labels: np.array):
         n = group_labels_.shape[0]
         self.rfx_dataset_cpp.AddGroupLabels(group_labels_, n)
 
+    def update_group_labels(self, group_labels: np.array):
+        """
+        Update group labels in a dataset
+
+        Parameters
+        ----------
+        group_labels : np.array
+            One-dimensional numpy array of group labels.
+        """
+        group_labels_ = np.squeeze(group_labels)
+        if group_labels_.ndim > 1:
+            raise ValueError(
+                "group_labels must be a one-dimensional numpy array of group indices"
+            )
+        n = group_labels_.shape[0]
+        self.rfx_dataset_cpp.UpdateGroupLabels(group_labels_, n)
+
     def add_basis(self, basis: np.array):
         """
         Add basis matrix to a dataset
@@ -93,6 +110,30 @@ 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):
+        """
+        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.
+        """
+        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.")
+        variance_weights_ = np.squeeze(variance_weights)
+        if variance_weights_.ndim > 1:
+            raise ValueError(
+                "variance_weights must be a one-dimensional numpy array of group indices"
+            )
+        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)
 
     def num_observations(self) -> int:
         """

test/python/test_data.py

@@ -0,0 +1,31 @@
+import numpy as np
+
+from stochtree import Dataset
+
+class TestDataset:
+    def test_dataset_update(self):
+        # Generate data
+        n = 20
+        num_covariates = 10
+        num_basis = 5
+        rng = np.random.default_rng()
+        covariates = rng.uniform(0, 1, size=(n, num_covariates))
+        basis = rng.uniform(0, 1, size=(n, num_basis))
+        variance_weights = rng.uniform(0, 1, size=n)
+        
+        # Construct dataset
+        forest_dataset = Dataset()
+        forest_dataset.add_covariates(covariates)
+        forest_dataset.add_basis(basis)
+        forest_dataset.add_variance_weights(variance_weights)
+        assert forest_dataset.num_observations() == n
+        assert forest_dataset.num_covariates() == num_covariates
+        assert forest_dataset.num_basis() == num_basis
+        assert forest_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():
+            forest_dataset.update_basis(new_basis)
+            forest_dataset.update_variance_weights(new_variance_weights)