MAINT: move tree sampling code out of a test module and into the main sampling module