Boolean guards added for constrained learning

include/LightGBM/tree.h

@@ -157,8 +157,9 @@ class Tree {
   inline int split_feature_inner(int split_idx) const { return split_feature_inner_[split_idx]; }
 
   /*! \brief Get features on leaf's branch*/
-  std::vector<int> branch_features(int leaf) const { return branch_features_[leaf]; }
+  inline std::vector<int> branch_features(int leaf) const { return branch_features_[leaf]; }
 
+  /*! \brief Get unique features used by the current tree*/
   std::set<int> tree_features() const {
      return tree_features_;
   }
@@ -324,6 +325,8 @@ class Tree {
 
   inline bool is_linear() const { return is_linear_; }
 
+  inline bool is_tracking_branch_features() const { return track_branch_features_; }
+
   #ifdef USE_CUDA
   inline bool is_cuda_tree() const { return is_cuda_tree_; }
   #endif  // USE_CUDA

src/boosting/gbdt.cpp

@@ -244,17 +244,18 @@ void GBDT::Train(int snapshot_freq, const std::string& model_output_path) {
       std::string snapshot_out = model_output_path + ".snapshot_iter_" + std::to_string(iter + 1);
       SaveModelToFile(0, -1, config_->saved_feature_importance_type, snapshot_out.c_str());
     }
-    interactions_used.insert(models_[models_.size() - 1]->tree_features());
+
+    if (config_->max_interactions != 0) {
+      interactions_used.insert(models_[models_.size() - 1]->tree_features());
+    }
 
     if (config_->max_interactions != 0 && static_cast<int>(interactions_used.size()) >= config_->max_interactions) {
       auto new_config = std::unique_ptr<Config>(new Config(*config_));
       new_config->tree_interaction_constraints_vector.clear();
       for (auto &inter_set : interactions_used) {
-        new_config->tree_interaction_constraints_vector.emplace_back(
-            inter_set.begin(), inter_set.end());
+        new_config->tree_interaction_constraints_vector.emplace_back(inter_set.begin(), inter_set.end());
       }
       new_config->max_interactions = 0;
-
       ResetConfig(new_config.release());
     }
   }

src/boosting/gbdt.h

@@ -543,7 +543,7 @@ class GBDT : public GBDTBase {
   std::vector<std::vector<std::string>> best_msg_;
   /*! \brief Trained models(trees) */
   std::vector<std::unique_ptr<Tree>> models_;
-  /*! \brief Trained models(trees) */
+  /*! \brief Set of set of features used in all the models */
   std::set<std::set<int>> interactions_used;
   /*! \brief Max feature index of training data*/
   int max_feature_idx_;

src/treelearner/col_sampler.hpp

@@ -163,7 +163,7 @@ class ColSampler {
     ComputeBranchAllowedFeatures(tree, leaf, &branch_allowed_features);
 
 
-        // intersect allowed features for branch and tree
+    // intersect allowed features for branch and tree
     std::unordered_set<int> allowed_features;
 
     if ((tree_interaction_constraints_.empty() && n_tree_interaction_constraints_ == 0) && !interaction_constraints_.empty()) {

src/treelearner/serial_tree_learner.cpp

@@ -335,16 +335,19 @@ void SerialTreeLearner::BeforeTrain() {
 
 bool SerialTreeLearner::BeforeFindBestSplit(const Tree* tree, int left_leaf, int right_leaf) {
   Common::FunctionTimer fun_timer("SerialTreeLearner::BeforeFindBestSplit", global_timer);
-  #pragma omp parallel for schedule(static) num_threads(OMP_NUM_THREADS())
-  for (int i = 0; i < config_->num_leaves; ++i) {
-    int feat_index = best_split_per_leaf_[i].feature;
-    if (feat_index == -1) continue;
 
-    int inner_feat_index = train_data_->InnerFeatureIndex(feat_index);
-    auto allowed_feature = col_sampler_.GetByNode(tree, i);
-    if (!allowed_feature[inner_feat_index]) {
-      RecomputeBestSplitForLeaf(tree, i, &best_split_per_leaf_[i]);
-    }
+  if (tree->is_tracking_branch_features()) {
+      #pragma omp parallel for schedule(static) num_threads(OMP_NUM_THREADS())
+      for (int i = 0; i < config_->num_leaves; ++i) {
+          int feat_index = best_split_per_leaf_[i].feature;
+          if (feat_index == -1) continue;
+
+          int inner_feat_index = train_data_->InnerFeatureIndex(feat_index);
+          auto allowed_feature = col_sampler_.GetByNode(tree, i);
+          if (!allowed_feature[inner_feat_index]) {
+              RecomputeBestSplitForLeaf(tree, i, &best_split_per_leaf_[i]);
+          }
+      }
   }
 
   // check depth of current leaf

tests/python_package_test/test_engine.py

@@ -3665,6 +3665,66 @@ def test_interaction_constraints():
                     train_data, num_boost_round=10)
 
 
+'''
+@pytest.mark.skipif(getenv('TASK', '') == 'cuda_exp', reason='Interaction constraints are not yet supported by CUDA Experimental version')
+def test_tree_interaction_constraints():
+    def check_consistency(est, tree_interaction_constraints):
+        feat_to_index = {feat: i for i, feat in enumerate(est.feature_name())}
+        tree_df = est.trees_to_dataframe()
+        inter_found = set()
+        for tree_index in tree_df["tree_index"].unique():
+            tree_df_per_index = tree_df[tree_df["tree_index"] == tree_index]
+            feat_used = [feat_to_index[feat] for feat in tree_df_per_index["split_feature"].unique() if feat is not None]
+            inter_found.add(tuple(sorted(feat_used)))
+        print(inter_found)
+        for feats_found in inter_found:
+            found = False
+            for real_contraints in tree_interaction_constraints:
+                if set(feats_found) <= set(real_contraints):
+                    found = True
+                    break
+            assert found is True
+    X, y = make_synthetic_regression(n_samples=200)
+    num_features = X.shape[1]
+    train_data = lgb.Dataset(X, label=y)
+    # check that constraint containing all features is equivalent to no constraint
+    params = {'verbose': -1,
+              'seed': 0}
+    est = lgb.train(params, train_data, num_boost_round=10)
+    pred1 = est.predict(X)
+    est = lgb.train(dict(params, interaction_constraints=[list(range(num_features))]), train_data,
+                    num_boost_round=10)
+    num_features = X.shape[1]
+    train_data = lgb.Dataset(X, label=y)
+    # check that tree constraint containing all features is equivalent to no constraint
+    params = {'verbose': -1,
+              'seed': 0}
+    est = lgb.train(params, train_data, num_boost_round=10)
+    pred1 = est.predict(X)
+    est = lgb.train(dict(params, tree_interaction_constraints=[list(range(num_features))]), train_data,
+                    num_boost_round=10)
+    pred2 = est.predict(X)
+    np.testing.assert_allclose(pred1, pred2)
+    # check that each tree is composed exactly of 2 features contained in the contrained set
+    tree_interaction_constraints = [[i, i + 1] for i in range(0, num_features - 1, 2)]
+    print(tree_interaction_constraints)
+    new_params = dict(params, tree_interaction_constraints=tree_interaction_constraints)
+    est = lgb.train(new_params, train_data, num_boost_round=100)
+    check_consistency(est, tree_interaction_constraints)
+    # check if tree features interaction constraints works with multiple set of features
+    tree_interaction_constraints = [[i for i in range(i, i + 5)] for i in range(0, num_features - 5, 5)]
+    print(tree_interaction_constraints)
+    new_params = dict(params, tree_interaction_constraints=tree_interaction_constraints)
+    est = lgb.train(new_params, train_data, num_boost_round=100)
+    check_consistency(est, tree_interaction_constraints)
+    # check if tree features interaction constraints works with multiple set of features
+    tree_interaction_constraints = [[i] for i in range(num_features)]
+    print(tree_interaction_constraints)
+    new_params = dict(params, tree_interaction_constraints=tree_interaction_constraints)
+    est = lgb.train(new_params, train_data, num_boost_round=100)
+    check_consistency(est, tree_interaction_constraints)
+'''
+
 def test_linear_trees_num_threads():
     # check that number of threads does not affect result
     np.random.seed(0)