Added methods to query data from a C++ object back to python as a numpy array

Author: andrewherren

src/py_stochtree.cpp

@@ -75,6 +75,48 @@ class ForestDatasetCpp {
     dataset_->UpdateVarWeights(data_ptr, num_row, exponentiate);
   }
 
+  py::array_t<double> GetCovariates() {
+    // Initialize n x p numpy array to store the covariates
+    data_size_t n = dataset_->NumObservations();
+    int num_covariates = dataset_->NumCovariates();
+    auto result = py::array_t<double>(py::detail::any_container<py::ssize_t>({n, num_covariates}));
+    auto accessor = result.mutable_unchecked<2>();
+    for (size_t i = 0; i < n; i++) {
+      for (int j = 0; j < num_covariates; j++) {
+        accessor(i,j) = dataset_->CovariateValue(i,j);
+      }
+    }
+
+    return result;
+  }
+
+  py::array_t<double> GetBasis() {
+    // Initialize n x k numpy array to store the basis
+    data_size_t n = dataset_->NumObservations();
+    int num_basis = dataset_->NumBasis();
+    auto result = py::array_t<double>(py::detail::any_container<py::ssize_t>({n, num_basis}));
+    auto accessor = result.mutable_unchecked<2>();
+    for (size_t i = 0; i < n; i++) {
+      for (int j = 0; j < num_basis; j++) {
+        accessor(i,j) = dataset_->BasisValue(i,j);
+      }
+    }
+
+    return result;
+  }
+
+  py::array_t<double> GetVarianceWeights() {
+    // Initialize n x 1 numpy array to store the variance weights
+    data_size_t n = dataset_->NumObservations();
+    auto result = py::array_t<double>(py::detail::any_container<py::ssize_t>({n}));
+    auto accessor = result.mutable_unchecked<1>();
+    for (size_t i = 0; i < n; i++) {
+      accessor(i) = dataset_->VarWeightValue(i);
+    }
+
+    return result;
+  }
+
   data_size_t NumRows() {
     return dataset_->NumObservations();
   }
@@ -1313,6 +1355,36 @@ class RandomEffectsDatasetCpp {
     }
     rfx_dataset_->UpdateGroupLabels(group_labels_vec, num_row);
   }
+  py::array_t<double> GetBasis() {
+    int num_row = rfx_dataset_->NumObservations();
+    int num_col = rfx_dataset_->NumBases();
+    auto result = py::array_t<double>(py::detail::any_container<py::ssize_t>({num_row, num_col}));
+    auto accessor = result.mutable_unchecked<2>();
+    for (py::ssize_t i = 0; i < num_row; i++) {
+      for (int j = 0; j < num_col; j++) {
+        accessor(i,j) = rfx_dataset_->BasisValue(i,j);
+      }
+    }
+    return result;
+  }
+  py::array_t<double> GetVarianceWeights() {
+    int num_row = rfx_dataset_->NumObservations();
+    auto result = py::array_t<double>(py::detail::any_container<py::ssize_t>({num_row}));
+    auto accessor = result.mutable_unchecked<1>();
+    for (py::ssize_t i = 0; i < num_row; i++) {
+      accessor(i) = rfx_dataset_->VarWeightValue(i);
+    }
+    return result;
+  }
+  py::array_t<int> GetGroupLabels() {
+    int num_row = rfx_dataset_->NumObservations();
+    auto result = py::array_t<int>(py::detail::any_container<py::ssize_t>({num_row}));
+    auto accessor = result.mutable_unchecked<1>();
+    for (py::ssize_t i = 0; i < num_row; i++) {
+      accessor(i) = rfx_dataset_->GroupId(i);
+    }
+    return result;
+  }
   bool HasGroupLabels() {return rfx_dataset_->HasGroupLabels();}
   bool HasBasis() {return rfx_dataset_->HasBasis();}
   bool HasVarianceWeights() {return rfx_dataset_->HasVarWeights();}
@@ -2071,6 +2143,9 @@ PYBIND11_MODULE(stochtree_cpp, m) {
     .def("NumRows", &ForestDatasetCpp::NumRows)
     .def("NumCovariates", &ForestDatasetCpp::NumCovariates)
     .def("NumBasis", &ForestDatasetCpp::NumBasis)
+    .def("GetCovariates", &ForestDatasetCpp::GetCovariates)
+    .def("GetBasis", &ForestDatasetCpp::GetBasis)
+    .def("GetVarianceWeights", &ForestDatasetCpp::GetVarianceWeights)
     .def("HasBasis", &ForestDatasetCpp::HasBasis)
     .def("HasVarianceWeights", &ForestDatasetCpp::HasVarianceWeights);
 
@@ -2204,6 +2279,9 @@ PYBIND11_MODULE(stochtree_cpp, m) {
       .def("UpdateGroupLabels", &RandomEffectsDatasetCpp::UpdateGroupLabels)
       .def("UpdateBasis", &RandomEffectsDatasetCpp::UpdateBasis)
       .def("UpdateVarianceWeights", &RandomEffectsDatasetCpp::UpdateVarianceWeights)
+      .def("GetGroupLabels", &RandomEffectsDatasetCpp::GetGroupLabels)
+      .def("GetBasis", &RandomEffectsDatasetCpp::GetBasis)
+      .def("GetVarianceWeights", &RandomEffectsDatasetCpp::GetVarianceWeights)
       .def("HasGroupLabels", &RandomEffectsDatasetCpp::HasGroupLabels)
       .def("HasBasis", &RandomEffectsDatasetCpp::HasBasis)
       .def("HasVarianceWeights", &RandomEffectsDatasetCpp::HasVarianceWeights);

stochtree/data.py

@@ -150,6 +150,39 @@ def num_basis(self) -> int:
             Dimension of the basis vector in the dataset, returning 0 if the dataset does not have a basis
         """
         return self.dataset_cpp.NumBasis()
+    
+    def get_covariates(self) -> np.array:
+        """
+        Return the covariates in a Dataset as a numpy array
+
+        Returns
+        -------
+        np.array
+            Covariate data
+        """
+        return self.dataset_cpp.GetCovariates()
+    
+    def get_basis(self) -> np.array:
+        """
+        Return the bases in a Dataset as a numpy array
+
+        Returns
+        -------
+        np.array
+            Basis data
+        """
+        return self.dataset_cpp.GetBasis()
+    
+    def get_variance_weights(self) -> np.array:
+        """
+        Return the variance weights in a Dataset as a numpy array
+
+        Returns
+        -------
+        np.array
+            Variance weights data
+        """
+        return self.dataset_cpp.GetVarianceWeights()
 
     def has_basis(self) -> bool:
         """

stochtree/random_effects.py

@@ -136,6 +136,39 @@ def update_variance_weights(self, variance_weights: np.array, exponentiate: bool
         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, exponentiate)
+    
+    def get_group_labels(self) -> np.array:
+        """
+        Return the group labels in a RandomEffectsDataset as a numpy array
+
+        Returns
+        -------
+        np.array
+            One-dimensional numpy array of group labels.
+        """
+        return self.rfx_dataset_cpp.GetGroupLabels()
+    
+    def get_basis(self) -> np.array:
+        """
+        Return the bases in a RandomEffectsDataset as a numpy array
+
+        Returns
+        -------
+        np.array
+            Two-dimensional numpy array of basis vectors.
+        """
+        return self.rfx_dataset_cpp.GetBasis()
+    
+    def get_variance_weights(self) -> np.array:
+        """
+        Return the variance weights in a RandomEffectsDataset as a numpy array
+
+        Returns
+        -------
+        np.array
+            One-dimensional numpy array of variance weights.
+        """
+        return self.rfx_dataset_cpp.GetVarianceWeights()
 
     def num_observations(self) -> int:
         """

test/python/test_data.py

@@ -29,6 +29,11 @@ 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)
+        
+        # Check that we recover the correct data through get_covariates, get_basis, and get_variance_weights
+        np.testing.assert_array_equal(forest_dataset.get_covariates(), covariates)
+        np.testing.assert_array_equal(forest_dataset.get_basis(), new_basis)
+        np.testing.assert_array_equal(forest_dataset.get_variance_weights(), new_variance_weights)
 
 class TestRFXDataset:
     def test_rfx_dataset_update(self):
@@ -59,4 +64,9 @@ def test_rfx_dataset_update(self):
         with np.testing.assert_no_warnings():
             rfx_dataset.update_basis(new_basis)
             rfx_dataset.update_variance_weights(new_variance_weights)
+        
+        # Check that we recover the correct data through get_group_labels, get_basis, and get_variance_weights
+        np.testing.assert_array_equal(rfx_dataset.get_group_labels(), group_labels)
+        np.testing.assert_array_equal(rfx_dataset.get_basis(), new_basis)
+        np.testing.assert_array_equal(rfx_dataset.get_variance_weights(), new_variance_weights)