Updated unit tests and R package

Author: andrewherren

NAMESPACE

@@ -38,6 +38,9 @@ export(createRandomEffectSamples)
 export(createRandomEffectsDataset)
 export(createRandomEffectsModel)
 export(createRandomEffectsTracker)
+export(expand_dims_1d)
+export(expand_dims_2d)
+export(expand_dims_2d_diag)
 export(getRandomEffectSamples)
 export(loadForestContainerCombinedJson)
 export(loadForestContainerCombinedJsonString)

R/utils.R

@@ -892,12 +892,12 @@ expand_dims_1d <- function(input, output_size) {
 #' @export
 expand_dims_2d <- function(input, output_rows, output_cols) {
     if (length(input) == 1) {
-        output <- as.matrix(rep(input, output_rows * output_cols), ncol = output_cols)
+        output <- matrix(rep(input, output_rows * output_cols), ncol = output_cols)
     } else if (is.numeric(input)) {
         if (length(input) == output_cols) {
-            output <- matrix(rep(x, output_rows), nrow=output_rows, byrow = T)
+            output <- matrix(rep(input, output_rows), nrow=output_rows, byrow = T)
         } else if (length(input) == output_rows) {
-            output <- matrix(rep(x, output_cols), ncol=output_cols, byrow = F)
+            output <- matrix(rep(input, output_cols), ncol=output_cols, byrow = F)
         } else {
             stop("If `input` is a vector, it must either contain `output_rows` or `output_cols` elements")
         }

man/bart.Rd

@@ -192,6 +192,7 @@ def _expand_dims_1d(input: Union[int, float, np.array], output_size: int) -> np.
     """
     Convert scalar input to 1D numpy array of dimension `output_size`, 
     or check that input array is equivalent to a 1D array of dimension `output_size`.
+    Single element numpy arrays (i.e. `np.array([2.5])`) are treated as scalars.
     
     Parameters
     ----------
@@ -207,11 +208,14 @@ def _expand_dims_1d(input: Union[int, float, np.array], output_size: int) -> np.
     """
     if isinstance(input, np.ndarray):
         input = np.squeeze(input)
-        if input.ndim != 1:
-            raise ValueError("`input` must be convertible to a 1D numpy array")
-        if input.shape[0] != output_size:
-            raise ValueError("`input` must be a 1D numpy array with `output_size` elements")
-        output = input
+        if input.ndim > 1:
+            raise ValueError("`input` must be convertible to a 1D numpy array or scalar")
+        if input.ndim == 0:
+            output = np.repeat(input, output_size)
+        else:
+            if input.shape[0] != output_size:
+                raise ValueError("`input` must be a 1D numpy array with `output_size` elements")
+            output = input
     elif isinstance(input, (int, float)):
         output = np.repeat(input, output_size)
     else:
@@ -227,7 +231,7 @@ def _expand_dims_2d(input: Union[int, float, np.array], output_rows: int, output
         2. `input` is a 1D array of length `output_rows`: output is a (`output_rows`, `output_cols`) array with `input` broadcast across each of `output_cols` columns
         3. `input` is a 1D array of length `output_cols`: output is a (`output_rows`, `output_cols`) array with `input` broadcast across each of `output_rows` rows
         4. `input` is a 2D array of dimension (`output_rows`, `output_cols`): input is passed through as-is
-    All other cases raise a `ValueError`.
+    All other cases raise a `ValueError`. Single element numpy arrays (i.e. `np.array([2.5])`) are treated as scalars.
     
     Parameters
     ----------
@@ -273,6 +277,7 @@ def _expand_dims_2d_diag(input: Union[int, float, np.array], output_size: int) -
     """
     Convert scalar input to 2D square numpy array of dimension `output_size` x `output_size` with `input` along the diagonal, 
     or check that input array is equivalent to a 2D square array of dimension `output_size` x `output_size`.
+    Single element numpy arrays (i.e. `np.array([2.5])`) are treated as scalars.
     
     Parameters
     ----------
@@ -288,13 +293,19 @@ def _expand_dims_2d_diag(input: Union[int, float, np.array], output_size: int) -
     """
     if isinstance(input, np.ndarray):
         input = np.squeeze(input)
-        if input.ndim != 2:
-            raise ValueError("`input` must be convertible to a 2D numpy array")
-        if input.shape[0] != input.shape[1]:
-            raise ValueError("`input` must be a 2D square numpy array")
-        if input.shape[0] != output_size:
-            raise ValueError("`input` must be a 2D square numpy array with exactly `output_size` rows and columns")
-        output = input
+        if (input.ndim != 2) and (input.ndim != 0):
+            raise ValueError("`input` must be convertible to a 2D numpy array or scalar")
+        if input.ndim == 0:
+            output = np.zeros(
+                (output_size, output_size), dtype=float
+            )
+            np.fill_diagonal(output, input)
+        else:
+            if input.shape[0] != input.shape[1]:
+                raise ValueError("`input` must be a 2D square numpy array")
+            if input.shape[0] != output_size:
+                raise ValueError("`input` must be a 2D square numpy array with exactly `output_size` rows and columns")
+            output = input
     elif isinstance(input, (int, float)):
         output = np.zeros(
             (output_size, output_size), dtype=float

man/expand_dims_1d.Rd

@@ -26,19 +26,45 @@ test_that("Array conversion", {
     )
     array_square_3 <- matrix(
         c(2.5,0.0,0.0,0.0,2.9,0.0,0.0,0.0,5.6), 
-        nrow = 2, ncol = 2, byrow = T
+        nrow = 3, ncol = 3, byrow = T
     )
 
     # Error cases
     expect_error(expand_dims_1d(array_1d_1, 5))
     expect_error(expand_dims_1d(array_1d_2, 4))
-    expect_error(expand_dims_1d(array_1d_3, 3))
     expect_error(expand_dims_2d(array_2d_1, 2, 4))
     expect_error(expand_dims_2d(array_2d_2, 3, 4))
     expect_error(expand_dims_2d_diag(array_square_1, 4))
     expect_error(expand_dims_2d_diag(array_square_2, 3))
     expect_error(expand_dims_2d_diag(array_square_3, 2))
 
-    # # Assertion
-    # expect_equal(y_hat_orig, y_hat_reloaded)
+    # Working cases
+    expect_equal(c(scalar_1,scalar_1,scalar_1), expand_dims_1d(scalar_1, 3))
+    expect_equal(c(scalar_2,scalar_2,scalar_2,scalar_2), expand_dims_1d(scalar_2, 4))
+    expect_equal(c(scalar_3,scalar_3), expand_dims_1d(scalar_3, 2))
+    expect_equal(c(array_1d_3,array_1d_3,array_1d_3), expand_dims_1d(array_1d_3, 3))
+    
+    output_exp <- matrix(rep(scalar_1, 6), nrow = 2, byrow = T)
+    expect_equal(output_exp, expand_dims_2d(scalar_1, 2, 3))
+    output_exp <- matrix(rep(scalar_2, 8), nrow = 2, byrow = T)
+    expect_equal(output_exp, expand_dims_2d(scalar_2, 2, 4))
+    output_exp <- matrix(rep(scalar_3, 6), nrow = 3, byrow = T)
+    expect_equal(output_exp, expand_dims_2d(scalar_3, 3, 2))
+    output_exp <- matrix(rep(array_1d_3, 6), nrow = 3, byrow = T)
+    expect_equal(output_exp, expand_dims_2d(array_1d_3, 3, 2))
+    output_exp <- unname(rbind(array_1d_1, array_1d_1))
+    expect_equal(output_exp, expand_dims_2d(array_1d_1, 2, 4))
+    output_exp <- unname(rbind(array_1d_2, array_1d_2, array_1d_2))
+    expect_equal(output_exp, expand_dims_2d(array_1d_2, 3, 3))
+    output_exp <- unname(cbind(array_1d_2, array_1d_2, array_1d_2, array_1d_2))
+    expect_equal(output_exp, expand_dims_2d(array_1d_2, 3, 4))
+    output_exp <- unname(cbind(array_1d_3, array_1d_3, array_1d_3, array_1d_3))
+    expect_equal(output_exp, expand_dims_2d(array_1d_3, 1, 4))
+    output_exp <- unname(rbind(array_1d_3, array_1d_3, array_1d_3, array_1d_3))
+    expect_equal(output_exp, expand_dims_2d(array_1d_3, 4, 1))
+    
+    expect_equal(diag(scalar_1, 3), expand_dims_2d_diag(scalar_1, 3))
+    expect_equal(diag(scalar_2, 2), expand_dims_2d_diag(scalar_2, 2))
+    expect_equal(diag(scalar_3, 4), expand_dims_2d_diag(scalar_3, 4))
+    expect_equal(diag(array_1d_3, 2), expand_dims_2d_diag(array_1d_3, 2))
 })

man/expand_dims_2d.Rd

@@ -143,10 +143,12 @@ def test_array_conversion(self):
         np.testing.assert_array_equal(np.array([scalar_1,scalar_1,scalar_1]), _expand_dims_1d(scalar_1, 3))
         np.testing.assert_array_equal(np.array([scalar_2,scalar_2,scalar_2,scalar_2]), _expand_dims_1d(scalar_2, 4))
         np.testing.assert_array_equal(np.array([scalar_3,scalar_3]), _expand_dims_1d(scalar_3, 2))
+        np.testing.assert_array_equal(np.array([array_1d_3[0],array_1d_3[0],array_1d_3[0]]), _expand_dims_1d(array_1d_3, 3))
 
         np.testing.assert_array_equal(np.array([[scalar_1,scalar_1,scalar_1],[scalar_1,scalar_1,scalar_1]]), _expand_dims_2d(scalar_1, 2, 3))
         np.testing.assert_array_equal(np.array([[scalar_2,scalar_2,scalar_2,scalar_2],[scalar_2,scalar_2,scalar_2,scalar_2]]), _expand_dims_2d(scalar_2, 2, 4))
         np.testing.assert_array_equal(np.array([[scalar_3,scalar_3],[scalar_3,scalar_3],[scalar_3,scalar_3]]), _expand_dims_2d(scalar_3, 3, 2))
+        np.testing.assert_array_equal(np.array([[array_1d_3[0],array_1d_3[0]],[array_1d_3[0],array_1d_3[0]],[array_1d_3[0],array_1d_3[0]]]), _expand_dims_2d(array_1d_3, 3, 2))
         np.testing.assert_array_equal(np.vstack((array_1d_1, array_1d_1)), _expand_dims_2d(array_1d_1, 2, 4))
         np.testing.assert_array_equal(np.vstack((array_1d_2, array_1d_2, array_1d_2)), _expand_dims_2d(array_1d_2, 3, 3))
         np.testing.assert_array_equal(np.column_stack((array_1d_2, array_1d_2, array_1d_2, array_1d_2)), _expand_dims_2d(array_1d_2, 3, 4))
@@ -156,3 +158,4 @@ def test_array_conversion(self):
         np.testing.assert_array_equal(np.array([[scalar_1,0.0,0.0],[0.0,scalar_1,0.0],[0.0,0.0,scalar_1]]), _expand_dims_2d_diag(scalar_1, 3))
         np.testing.assert_array_equal(np.array([[scalar_2,0.0],[0.0,scalar_2]]), _expand_dims_2d_diag(scalar_2, 2))
         np.testing.assert_array_equal(np.array([[scalar_3,0.0,0.0,0.0],[0.0,scalar_3,0.0,0.0],[0.0,0.0,scalar_3,0.0],[0.0,0.0,0.0,scalar_3]]), _expand_dims_2d_diag(scalar_3, 4))
+        np.testing.assert_array_equal(np.array([[array_1d_3[0],0.0],[0.0,array_1d_3[0]]]), _expand_dims_2d_diag(array_1d_3, 2))