Is using the BART's predict function a correct way to predict new data?

If you run

μ_pred = pmb.predict(idata, rng, X.values, 100)
y_pred = norm.rvs(μ_pred, idata.posterior["σ"].mean())

y_pred will have shape (100, len(X)), that is 100 predictions, each one having the same size as the observed data.

Then you can work with that as you need, for example if you want the mean over the 100 predictions for each datapoint you can do y_pred.mean(0), or if you want the overall mean just y_pred.mean(). The same goes for the standard deviation the HDI, or whatever quantity you want to compute.

Alternatively, instead of using the mean of sigma, you can take samples from the posterior of sigma, something like this.

μ_pred = pmb.predict(idata, rng, X.values, 100).squeeze().T
σ_ = az.extract(idata, var_names="σ" ,num_samples=100)
y_pred = norm.rvs(μ_pred.T, σ_).T
1 Like