Easiest Way to Add Variable Sparsity Prior

[kevinli1324]

Hi! The use case I'm working on would benefit from some sparsity inducing prior for variable selection like the type used by Linero 2018. What would be the easiest way to add this feature into the library (if just for my own use)

Thanks!

[andrewherren]

Hi there! Thanks for reaching out!

The quickest way to get this up and running is to do it in R directly (or Python, it is also feasible there), using what we sometimes call the "low-level" or "researcher" interface to stochtree. See here for several demos of that interface. Broadly, you can recreate the functionality of bart()or bcf() by instantiating R wrappers to all of the C++ classes and then writing a for loop that updates the forests and any other parameters in the model. What you'd need to do specifically here is implement the Dirichlet sampler in R and then reflect those new variable sampling probabilities through to the ForestModelConfig object via the update_variable_weights method. You should be able to extract split counts from a Forest object with the get_forest_split_counts method (which you'd use for the Dirichlet update).

To make it faster (and to get it incorporated back into the stochtree project), you'd likely need to implement the Dirichlet sampler directly in C++. That implementation should be fairly independent of the rest of the stochtree codebase, though I'd recommend mirroring the patterns of some of our other non-tree samplers, for example inverse gamma and discrete sampling without replacement.