Merge pull request #391 from StochasticTree/worktree-float32-fix-backport-0.4.3

Add predict() coverage to float32 test suite

Author: andrewherren

test/python/test_bart.py

@@ -1604,6 +1604,8 @@ def test_bart_float32_runs(self):
         )
         assert bart_model.y_hat_train.shape == (self.n_train, self.num_mcmc)
         assert bart_model.y_hat_test.shape == (self.n_test, self.num_mcmc)
+        preds = bart_model.predict(X=self.X_test)
+        assert preds["y_hat"].shape == (self.n_test, self.num_mcmc)
 
     def test_bart_float32_matches_float64(self):
         """float32 and float64 inputs with the same seed should produce close results."""
@@ -1628,6 +1630,9 @@ def test_bart_float32_matches_float64(self):
             general_params={"random_seed": 1},
         )
         np.testing.assert_allclose(bart32.y_hat_train, bart64.y_hat_train, rtol=1e-5)
+        pred32 = bart32.predict(X=self.X_test)
+        pred64 = bart32.predict(X=self.X_test.astype(np.float64))
+        np.testing.assert_allclose(pred32["y_hat"], pred64["y_hat"], rtol=1e-5)
 
     def test_bart_float32_leaf_basis(self):
         rng = np.random.default_rng(7)
@@ -1646,6 +1651,8 @@ def test_bart_float32_leaf_basis(self):
         )
         assert bart_model.y_hat_train.shape == (self.n_train, self.num_mcmc)
         assert bart_model.y_hat_test.shape == (self.n_test, self.num_mcmc)
+        preds = bart_model.predict(X=self.X_test, leaf_basis=basis_test)
+        assert preds["y_hat"].shape == (self.n_test, self.num_mcmc)
 
     def test_bart_float32_leaf_basis_matches_float64(self):
         rng = np.random.default_rng(7)
@@ -1663,6 +1670,9 @@ def test_bart_float32_leaf_basis_matches_float64(self):
                       X_test=self.X_test.astype(np.float64),
                       leaf_basis_test=basis_test.astype(np.float64), **common)
         np.testing.assert_allclose(bart32.y_hat_train, bart64.y_hat_train, rtol=1e-5)
+        pred32 = bart32.predict(X=self.X_test, leaf_basis=basis_test)
+        pred64 = bart32.predict(X=self.X_test.astype(np.float64), leaf_basis=basis_test.astype(np.float64))
+        np.testing.assert_allclose(pred32["y_hat"], pred64["y_hat"], rtol=1e-5)
 
     def test_bart_float32_rfx(self):
         rng = np.random.default_rng(7)
@@ -1686,6 +1696,8 @@ def test_bart_float32_rfx(self):
         )
         assert bart_model.y_hat_train.shape == (self.n_train, self.num_mcmc)
         assert bart_model.y_hat_test.shape == (self.n_test, self.num_mcmc)
+        preds = bart_model.predict(X=self.X_test, rfx_group_ids=group_ids_test, rfx_basis=rfx_basis_test)
+        assert preds["y_hat"].shape == (self.n_test, self.num_mcmc)
 
     def test_bart_float32_rfx_matches_float64(self):
         rng = np.random.default_rng(7)
@@ -1706,3 +1718,6 @@ def test_bart_float32_rfx_matches_float64(self):
                       rfx_basis_train=rfx_basis_train.astype(np.float64),
                       rfx_basis_test=rfx_basis_test.astype(np.float64), **common)
         np.testing.assert_allclose(bart32.y_hat_train, bart64.y_hat_train, rtol=1e-4)
+        pred32 = bart32.predict(X=self.X_test, rfx_group_ids=group_ids_test, rfx_basis=rfx_basis_test)
+        pred64 = bart32.predict(X=self.X_test.astype(np.float64), rfx_group_ids=group_ids_test, rfx_basis=rfx_basis_test.astype(np.float64))
+        np.testing.assert_allclose(pred32["y_hat"], pred64["y_hat"], rtol=1e-4)

test/python/test_bcf.py

@@ -1159,6 +1159,9 @@ def test_bcf_float32_with_propensity(self):
         assert bcf_model.y_hat_test.shape == (self.n_test, self.num_mcmc)
         assert bcf_model.tau_hat_train.shape == (self.n_train, self.num_mcmc)
         assert bcf_model.tau_hat_test.shape == (self.n_test, self.num_mcmc)
+        preds = bcf_model.predict(X=self.X_test, Z=self.Z_test, propensity=self.pi_test)
+        assert preds["y_hat"].shape == (self.n_test, self.num_mcmc)
+        assert preds["tau_hat"].shape == (self.n_test, self.num_mcmc)
 
     def test_bcf_float32_with_propensity_matches_float64(self):
         common = dict(num_gfr=5, num_burnin=0, num_mcmc=self.num_mcmc, general_params={"random_seed": 1})
@@ -1175,6 +1178,10 @@ def test_bcf_float32_with_propensity_matches_float64(self):
                      Z_test=self.Z_test.astype(np.float64),
                      propensity_test=self.pi_test.astype(np.float64), **common)
         np.testing.assert_allclose(bcf32.y_hat_train, bcf64.y_hat_train, rtol=1e-4)
+        pred32 = bcf32.predict(X=self.X_test, Z=self.Z_test, propensity=self.pi_test)
+        pred64 = bcf32.predict(X=self.X_test.astype(np.float64), Z=self.Z_test.astype(np.float64),
+                               propensity=self.pi_test.astype(np.float64))
+        np.testing.assert_allclose(pred32["y_hat"], pred64["y_hat"], rtol=1e-4)
 
     def test_bcf_float32_no_propensity(self):
         """float32 Z, y, X with internal propensity estimation."""
@@ -1191,6 +1198,8 @@ def test_bcf_float32_no_propensity(self):
         )
         assert bcf_model.y_hat_train.shape == (self.n_train, self.num_mcmc)
         assert bcf_model.y_hat_test.shape == (self.n_test, self.num_mcmc)
+        preds = bcf_model.predict(X=self.X_test, Z=self.Z_test)
+        assert preds["y_hat"].shape == (self.n_test, self.num_mcmc)
 
     def test_bcf_float32_no_propensity_matches_float64(self):
         common = dict(num_gfr=5, num_burnin=0, num_mcmc=self.num_mcmc, general_params={"random_seed": 1})
@@ -1204,6 +1213,9 @@ def test_bcf_float32_no_propensity_matches_float64(self):
                      X_test=self.X_test.astype(np.float64),
                      Z_test=self.Z_test.astype(np.float64), **common)
         np.testing.assert_allclose(bcf32.y_hat_train, bcf64.y_hat_train, rtol=1e-4)
+        pred32 = bcf32.predict(X=self.X_test, Z=self.Z_test)
+        pred64 = bcf32.predict(X=self.X_test.astype(np.float64), Z=self.Z_test.astype(np.float64))
+        np.testing.assert_allclose(pred32["y_hat"], pred64["y_hat"], rtol=1e-4)
 
     def test_bcf_float32_rfx(self):
         """float32 rfx_basis_train and rfx_basis_test."""
@@ -1232,6 +1244,9 @@ def test_bcf_float32_rfx(self):
         )
         assert bcf_model.y_hat_train.shape == (self.n_train, self.num_mcmc)
         assert bcf_model.y_hat_test.shape == (self.n_test, self.num_mcmc)
+        preds = bcf_model.predict(X=self.X_test, Z=self.Z_test, propensity=self.pi_test,
+                                   rfx_group_ids=group_ids_test, rfx_basis=rfx_basis_test)
+        assert preds["y_hat"].shape == (self.n_test, self.num_mcmc)
 
     def test_bcf_float32_rfx_matches_float64(self):
         rng = np.random.default_rng(7)
@@ -1256,3 +1271,9 @@ def test_bcf_float32_rfx_matches_float64(self):
                      rfx_basis_train=rfx_basis_train.astype(np.float64),
                      rfx_basis_test=rfx_basis_test.astype(np.float64), **common)
         np.testing.assert_allclose(bcf32.y_hat_train, bcf64.y_hat_train, rtol=1e-4)
+        pred32 = bcf32.predict(X=self.X_test, Z=self.Z_test, propensity=self.pi_test,
+                               rfx_group_ids=group_ids_test, rfx_basis=rfx_basis_test)
+        pred64 = bcf32.predict(X=self.X_test.astype(np.float64), Z=self.Z_test.astype(np.float64),
+                               propensity=self.pi_test.astype(np.float64),
+                               rfx_group_ids=group_ids_test, rfx_basis=rfx_basis_test.astype(np.float64))
+        np.testing.assert_allclose(pred32["y_hat"], pred64["y_hat"], rtol=1e-4)