Merge pull request #161 from StochasticTree/numpy-dtype-hotfix

Specify float dtype in numpy ones and zeros

Author: andrewherren

stochtree/bart.py

@@ -734,13 +734,13 @@ def sample(
             )
             if self.has_basis:
                 if sigma_leaf is None:
-                    current_leaf_scale = np.zeros((self.num_basis, self.num_basis))
+                    current_leaf_scale = np.zeros((self.num_basis, self.num_basis), dtype=float)
                     np.fill_diagonal(
                         current_leaf_scale,
                         np.squeeze(np.var(resid_train)) / num_trees_mean,
                     )
                 elif isinstance(sigma_leaf, float):
-                    current_leaf_scale = np.zeros((self.num_basis, self.num_basis))
+                    current_leaf_scale = np.zeros((self.num_basis, self.num_basis), dtype=float)
                     np.fill_diagonal(current_leaf_scale, sigma_leaf)
                 elif isinstance(sigma_leaf, np.ndarray):
                     if sigma_leaf.ndim != 2:
@@ -834,7 +834,7 @@ def sample(
                 alpha_init = np.array([1])
             elif num_rfx_components > 1:
                 alpha_init = np.concatenate(
-                    (np.ones(1), np.zeros(num_rfx_components - 1))
+                    (np.ones(1, dtype=float), np.zeros(num_rfx_components - 1, dtype=float))
                 )
             else:
                 raise ValueError("There must be at least 1 random effect component")

stochtree/bcf.py

@@ -1163,7 +1163,7 @@ def sample(
             raise ValueError("sigma_leaf_mu must be a scalar")
         if isinstance(sigma_leaf_tau, float):
             if Z_train.shape[1] > 1:
-                current_leaf_scale_tau = np.zeros((Z_train.shape[1], Z_train.shape[1]))
+                current_leaf_scale_tau = np.zeros((Z_train.shape[1], Z_train.shape[1]), dtype=float)
                 np.fill_diagonal(current_leaf_scale_tau, sigma_leaf_tau)
             else:
                 current_leaf_scale_tau = np.array([[sigma_leaf_tau]])
@@ -1230,7 +1230,7 @@ def sample(
                 alpha_init = np.array([1])
             elif num_rfx_components > 1:
                 alpha_init = np.concatenate(
-                    (np.ones(1), np.zeros(num_rfx_components - 1))
+                    (np.ones(1, dtype=float), np.zeros(num_rfx_components - 1, dtype=float))
                 )
             else:
                 raise ValueError("There must be at least 1 random effect component")
@@ -1488,7 +1488,7 @@ def sample(
 
         # Initialize the leaves of each tree in the treatment forest
         if self.multivariate_treatment:
-            init_tau = np.zeros(Z_train.shape[1])
+            init_tau = np.zeros(Z_train.shape[1], dtype=float)
         else:
             init_tau = np.array([0.0])
         forest_sampler_tau.prepare_for_sampler(

stochtree/config.py

@@ -113,7 +113,7 @@ def __init__(
             raise ValueError("`leaf_dimension` must be an integer greater than 0")
         if leaf_model_scale is None:
             diag_value = 1.0 / num_trees
-            leaf_model_scale_array = np.zeros((leaf_dimension, leaf_dimension), float)
+            leaf_model_scale_array = np.zeros((leaf_dimension, leaf_dimension), dtype=float)
             np.fill_diagonal(leaf_model_scale_array, diag_value)
         else:
             if isinstance(leaf_model_scale, np.ndarray):
@@ -128,7 +128,7 @@ def __init__(
                         "`leaf_model_scale` must be positive, if provided as scalar"
                     )
                 leaf_model_scale_array = np.zeros(
-                    (leaf_dimension, leaf_dimension), float
+                    (leaf_dimension, leaf_dimension), dtype=float
                 )
                 np.fill_diagonal(leaf_model_scale_array, leaf_model_scale)
             else:
@@ -278,7 +278,7 @@ def update_leaf_model_scale(
                     "`leaf_model_scale` must be positive, if provided as scalar"
                 )
             leaf_model_scale_array = np.zeros(
-                (self.leaf_dimension, self.leaf_dimension), float
+                (self.leaf_dimension, self.leaf_dimension), dtype=float
             )
             np.fill_diagonal(leaf_model_scale_array, leaf_model_scale)
         else: