MAINT: add naive code to get distribution over states at nodes in a tree, and add failing test