diff --git a/include/stochtree/tree_sampler.h b/include/stochtree/tree_sampler.h index 0796ec07..a47a53d4 100644 --- a/include/stochtree/tree_sampler.h +++ b/include/stochtree/tree_sampler.h @@ -863,6 +863,70 @@ static inline void GFRSampleOneIter(TreeEnsemble& active_forest, ForestTracker& } } +// Proposal availability uses observation counts only. Selecting a constant +// variable or a leaf at maximum depth still yields a rejected grow proposal. +struct MCMCProposalState { + int leaves; + int leaf_parents; + int grow_candidates; + double GrowProbability() const { + return grow_candidates > 0 ? (leaf_parents > 0 ? 0.5 : 1.0) : 0.0; + } + double PruneProbability() const { + return leaf_parents > 0 ? (grow_candidates > 0 ? 0.5 : 1.0) : 0.0; + } +}; + +static inline bool MCMCSizeAllowsGrow(data_size_t n, int minimum) { + // Equality permits two children with exactly minimum observations each. + return n >= 2 * minimum; +} + +static inline MCMCProposalState MCMCGetProposalState(Tree* tree, ForestTracker& tracker, int tree_num, int minimum) { + MCMCProposalState state{tree->NumLeaves(), tree->NumLeafParents(), 0}; + for (int leaf : tree->GetLeaves()) { + state.grow_candidates += MCMCSizeAllowsGrow(tracker.UnsortedNodeSize(tree_num, leaf), minimum); + } + return state; +} + +static inline MCMCProposalState MCMCStateAfterGrow(Tree* tree, MCMCProposalState state, int leaf, + data_size_t left_n, data_size_t right_n, int minimum) { + state.leaves += 1; + state.leaf_parents += 1; + if (!tree->IsRoot(leaf) && tree->IsLeafParent(tree->Parent(leaf))) { + state.leaf_parents -= 1; + } + state.grow_candidates += MCMCSizeAllowsGrow(left_n, minimum) + MCMCSizeAllowsGrow(right_n, minimum) + - MCMCSizeAllowsGrow(left_n + right_n, minimum); + return state; +} + +static inline MCMCProposalState MCMCStateAfterPrune(Tree* tree, MCMCProposalState state, int node, + data_size_t left_n, data_size_t right_n, int minimum) { + state.leaves -= 1; + state.leaf_parents -= 1; + if (!tree->IsRoot(node)) { + int parent = tree->Parent(node); + int sibling = tree->LeftChild(parent) == node ? tree->RightChild(parent) : tree->LeftChild(parent); + if (tree->IsLeaf(sibling)) state.leaf_parents += 1; + } + state.grow_candidates += MCMCSizeAllowsGrow(left_n + right_n, minimum) + - MCMCSizeAllowsGrow(left_n, minimum) - MCMCSizeAllowsGrow(right_n, minimum); + return state; +} + +// The split-rule density cancels with the corresponding conditional tree prior. +// These terms cover move type and uniform leaf / leaf-parent selection only. +static inline double MCMCLogProposalRatio(const MCMCProposalState& from, const MCMCProposalState& to, bool grow) { + if (grow) { + return std::log(to.PruneProbability()) - std::log(static_cast(to.leaf_parents)) + - std::log(from.GrowProbability()) + std::log(static_cast(from.leaves)); + } + return std::log(to.GrowProbability()) - std::log(static_cast(to.leaves)) + - std::log(from.PruneProbability()) + std::log(static_cast(from.leaf_parents)); +} + template static inline void MCMCGrowTreeOneIter(Tree* tree, ForestTracker& tracker, LeafModel& leaf_model, ForestDataset& dataset, ColumnVector& residual, TreePrior& tree_prior, std::mt19937& gen, int tree_num, std::vector& variable_weights, @@ -927,28 +991,13 @@ static inline void MCMCGrowTreeOneIter(Tree* tree, ForestTracker& tracker, LeafM double pgl = tree_prior.GetAlpha() * std::pow(1+leaf_depth+1, -tree_prior.GetBeta()); double pgr = tree_prior.GetAlpha() * std::pow(1+leaf_depth+1, -tree_prior.GetBeta()); - // Determine whether a "grow" move is possible from the newly formed tree - // in order to compute the probability of choosing "prune" from the new tree - // (which is always possible by construction) - bool non_constant = NodesNonConstantAfterSplit(dataset, tracker, split, tree_num, leaf_chosen, var_chosen); - bool min_samples_left_check = left_n >= 2*tree_prior.GetMinSamplesLeaf(); - bool min_samples_right_check = right_n >= 2*tree_prior.GetMinSamplesLeaf(); - double prob_prune_new; - if (non_constant && (min_samples_left_check || min_samples_right_check)) { - prob_prune_new = 0.5; - } else { - prob_prune_new = 1.0; - } - - // Determine the number of leaves in the current tree and leaf parents in the proposed tree - int num_leaf_parents = tree->NumLeafParents(); - double p_leaf = 1/static_cast(num_leaves); - double p_leaf_parent = 1/static_cast(num_leaf_parents+1); + auto old_state = MCMCGetProposalState(tree, tracker, tree_num, tree_prior.GetMinSamplesLeaf()); + auto new_state = MCMCStateAfterGrow(tree, old_state, leaf_chosen, left_n, right_n, tree_prior.GetMinSamplesLeaf()); // Compute the final MH ratio double log_mh_ratio = ( - std::log(pg) + std::log(1-pgl) + std::log(1-pgr) - std::log(1-pg) + std::log(prob_prune_new) + - std::log(p_leaf_parent) - std::log(prob_grow_old) - std::log(p_leaf) - no_split_log_marginal_likelihood + split_log_marginal_likelihood + std::log(pg) + std::log(1-pgl) + std::log(1-pgr) - std::log(1-pg) + + MCMCLogProposalRatio(old_state, new_state, true) - no_split_log_marginal_likelihood + split_log_marginal_likelihood ); // Threshold at 0 if (log_mh_ratio > 0) { @@ -1001,36 +1050,13 @@ static inline void MCMCPruneTreeOneIter(Tree* tree, ForestTracker& tracker, Leaf double pgl = tree_prior.GetAlpha() * std::pow(1+leaf_parent_depth+1, -tree_prior.GetBeta()); double pgr = tree_prior.GetAlpha() * std::pow(1+leaf_parent_depth+1, -tree_prior.GetBeta()); - // Determine whether a "prune" move is possible from the new tree, - // in order to compute the probability of choosing "grow" from the new tree - // (which is always possible by construction) - bool non_root_tree = tree->NumNodes() > 1; - double prob_grow_new; - if (non_root_tree) { - prob_grow_new = 0.5; - } else { - prob_grow_new = 1.0; - } - - // Determine whether a "grow" move was possible from the old tree, - // in order to compute the probability of choosing "prune" from the old tree - bool non_constant_left = NodeNonConstant(dataset, tracker, tree_num, left_node); - bool non_constant_right = NodeNonConstant(dataset, tracker, tree_num, right_node); - double prob_prune_old; - if (non_constant_left && non_constant_right) { - prob_prune_old = 0.5; - } else { - prob_prune_old = 1.0; - } - - // Determine the number of leaves in the current tree and leaf parents in the proposed tree - double p_leaf = 1/static_cast(num_leaves-1); - double p_leaf_parent = 1/static_cast(num_leaf_parents); + auto old_state = MCMCGetProposalState(tree, tracker, tree_num, tree_prior.GetMinSamplesLeaf()); + auto new_state = MCMCStateAfterPrune(tree, old_state, leaf_parent_chosen, left_n, right_n, tree_prior.GetMinSamplesLeaf()); // Compute the final MH ratio double log_mh_ratio = ( - std::log(1-pg) - std::log(pg) - std::log(1-pgl) - std::log(1-pgr) + std::log(prob_prune_old) + - std::log(p_leaf) - std::log(prob_grow_new) - std::log(p_leaf_parent) + no_split_log_marginal_likelihood - split_log_marginal_likelihood + std::log(1-pg) - std::log(pg) - std::log(1-pgl) - std::log(1-pgr) + + MCMCLogProposalRatio(old_state, new_state, false) + no_split_log_marginal_likelihood - split_log_marginal_likelihood ); // Threshold at 0 if (log_mh_ratio > 0) { @@ -1052,37 +1078,12 @@ template & variable_weights, int tree_num, double global_variance, int num_threads, LeafSuffStatConstructorArgs&... leaf_suff_stat_args) { - // Determine whether it is possible to grow any of the leaves - bool grow_possible = false; - std::vector leaves = tree->GetLeaves(); - for (auto& leaf: leaves) { - if (tracker.UnsortedNodeSize(tree_num, leaf) > 2 * tree_prior.GetMinSamplesLeaf()) { - grow_possible = true; - break; - } - } - - // Determine whether it is possible to prune the tree - bool prune_possible = false; - if (tree->NumValidNodes() > 1) { - prune_possible = true; - } - - // Determine the relative probability of grow vs prune (0 = grow, 1 = prune) - double prob_grow; - std::vector step_probs(2); - if (grow_possible && prune_possible) { - step_probs = {0.5, 0.5}; - prob_grow = 0.5; - } else if (!grow_possible && prune_possible) { - step_probs = {0.0, 1.0}; - prob_grow = 0.0; - } else if (grow_possible && !prune_possible) { - step_probs = {1.0, 0.0}; - prob_grow = 1.0; - } else { - Log::Fatal("In this tree, neither grow nor prune is possible"); - } + auto state = MCMCGetProposalState(tree, tracker, tree_num, tree_prior.GetMinSamplesLeaf()); + double prob_grow = state.GrowProbability(); + double prob_prune = state.PruneProbability(); + // An undersized root has no legal move; retain the current tree. + if (prob_grow == 0.0 && prob_prune == 0.0) return; + std::vector step_probs{prob_grow, prob_prune}; walker_vose step_dist(step_probs.begin(), step_probs.end()); // Draw a split rule at random diff --git a/test/cpp/test_sampler_acceptance.cpp b/test/cpp/test_sampler_acceptance.cpp new file mode 100644 index 00000000..4b3bc6ab --- /dev/null +++ b/test/cpp/test_sampler_acceptance.cpp @@ -0,0 +1,88 @@ +#include +#include +#include + +namespace { +struct ProposalFixture { + StochTree::ForestDataset data; + std::vector types{StochTree::FeatureType::kNumeric}; + StochTree::Tree tree; + std::unique_ptr tracker; + StochTree::TreePrior prior{0.6, 2., 2}; + std::mt19937 gen{19}; + ProposalFixture(int n) { + std::vector x(n); std::iota(x.begin(), x.end(), 0.); + data.AddCovariates(x.data(), n, 1, true); + tracker = std::make_unique(data.GetCovariates(), types, 1, n); + tree.Init(1); + } + auto State() { return StochTree::MCMCGetProposalState(&tree, *tracker, 0, prior.GetMinSamplesLeaf()); } + void Grow(int leaf, double cut) { + StochTree::TreeSplit split(cut); + StochTree::AddSplitToModel(*tracker, data, prior, split, gen, &tree, 0, leaf, 0); + } +}; +void Same(const StochTree::MCMCProposalState& a, const StochTree::MCMCProposalState& b) { + EXPECT_EQ(a.leaves, b.leaves); EXPECT_EQ(a.leaf_parents, b.leaf_parents); + EXPECT_EQ(a.grow_candidates, b.grow_candidates); +} +} + +TEST(SamplerAcceptance, GrowThenPruneUsesActualWholeTreeState) { + ProposalFixture f(12); + for (auto proposal : std::vector>{{0,3.5},{2,7.5},{1,1.5}}) { + int leaf=proposal.first; + int n=f.tracker->UnsortedNodeSize(0,leaf); + int left_n=0; + for (auto i=f.tracker->UnsortedNodeBeginIterator(0,leaf);i!=f.tracker->UnsortedNodeEndIterator(0,leaf);++i) + left_n += f.data.CovariateValue(*i,0)<=proposal.second; + auto old=f.State(); + auto predicted=StochTree::MCMCStateAfterGrow(&f.tree,old,leaf,left_n,n-left_n,2); + f.Grow(leaf,proposal.second); + auto actual=f.State(); Same(predicted,actual); + auto reverse=StochTree::MCMCStateAfterPrune(&f.tree,actual,leaf,left_n,n-left_n,2); + Same(reverse,old); + double q_forward=old.GrowProbability()/old.leaves; + double q_reverse=actual.PruneProbability()/actual.leaf_parents; + EXPECT_NEAR(StochTree::MCMCLogProposalRatio(old,actual,true),std::log(q_reverse/q_forward),1e-14); + EXPECT_NEAR(StochTree::MCMCLogProposalRatio(actual,old,false),std::log(q_forward/q_reverse),1e-14); + } +} + +TEST(SamplerAcceptance, OtherLeavesAndLeafParentReplacementMatter) { + ProposalFixture f(20); f.Grow(0,3.5); + auto old=f.State(); + auto next=StochTree::MCMCStateAfterGrow(&f.tree,old,1,2,2,2); + f.Grow(1,1.5); Same(next,f.State()); + EXPECT_EQ(next.leaf_parents,1); // root stops being a leaf parent + EXPECT_EQ(next.grow_candidates,1); // the unaffected 16-observation leaf + EXPECT_DOUBLE_EQ(next.PruneProbability(),0.5); +} + +TEST(SamplerAcceptance, EqualityBoundaryAndPruningToRoot) { + ProposalFixture f(4); + auto root=f.State(); EXPECT_DOUBLE_EQ(root.GrowProbability(),1.); + f.Grow(0,1.5); auto split=f.State(); + EXPECT_DOUBLE_EQ(split.PruneProbability(),1.); + auto back=StochTree::MCMCStateAfterPrune(&f.tree,split,0,2,2,2); + Same(back,root); EXPECT_DOUBLE_EQ(back.GrowProbability(),1.); + EXPECT_DOUBLE_EQ(StochTree::MCMCLogProposalRatio(root,split,true),0.); + EXPECT_DOUBLE_EQ(StochTree::MCMCLogProposalRatio(split,root,false),0.); +} + +TEST(SamplerAcceptance, InactiveNodesDoNotPreventRootProbability) { + ProposalFixture f(8); f.Grow(0,3.5); + StochTree::RemoveSplitFromModel(*f.tracker,f.data,f.prior,f.gen,&f.tree,0,0,1,2); + EXPECT_DOUBLE_EQ(f.State().GrowProbability(),1.); + EXPECT_DOUBLE_EQ(f.State().PruneProbability(),0.); +} + +TEST(SamplerAcceptance, UndersizedRootRetainsTree) { + ProposalFixture f(3); + std::vector y(3,0.); StochTree::ColumnVector residual(y.data(),3); + StochTree::GaussianConstantLeafModel model(1.); + StochTree::ForestContainer forests(1); std::vector weights{1.}; + StochTree::MCMCSampleTreeOneIter( + &f.tree,*f.tracker,forests,model,f.data,residual,f.prior,f.gen,weights,0,1.,1); + EXPECT_EQ(f.tree.NumLeaves(),1); +}