#include <algorithm>
#include <unordered_set>
#include <vector>
#include "xgboost/span.h"
#include "xgboost/json.h"
#include "constraints.h"
#include "param.h"
namespace xgboost {
void FeatureInteractionConstraintHost::Configure(tree::TrainParam const& param,
bst_feature_t const n_features) {
if (param.interaction_constraints.empty()) {
enabled_ = !param.interaction_constraints.empty();
return; }
enabled_ = true;
this->interaction_constraint_str_ = param.interaction_constraints;
this->n_features_ = n_features;
this->Reset();
}
void FeatureInteractionConstraintHost::Reset() {
if (!enabled_) {
return;
}
std::vector<std::vector<bst_feature_t>> tmp;
try {
ParseInteractionConstraint(this->interaction_constraint_str_, &tmp);
} catch (dmlc::Error const &e) {
LOG(FATAL) << "Failed to parse feature interaction constraint:\n"
<< this->interaction_constraint_str_ << "\n"
<< "With error:\n" << e.what();
}
for (const auto& e : tmp) {
interaction_constraints_.emplace_back(e.begin(), e.end());
}
node_constraints_.clear();
node_constraints_.resize(1, std::unordered_set<bst_feature_t>());
node_constraints_[0].reserve(n_features_);
for (bst_feature_t i = 0; i < n_features_; ++i) {
node_constraints_[0].insert(i);
}
splits_.clear();
splits_.resize(1, std::unordered_set<bst_feature_t>());
}
void FeatureInteractionConstraintHost::SplitImpl(
bst_node_t node_id, bst_feature_t feature_id, bst_node_t left_id, bst_node_t right_id) {
bst_node_t newsize = std::max(left_id, right_id) + 1;
auto feature_splits = splits_[node_id]; feature_splits.insert(feature_id); splits_.resize(newsize);
splits_[left_id] = feature_splits;
splits_[right_id] = feature_splits;
CHECK_NE(newsize, 0);
node_constraints_.resize(newsize, std::unordered_set<bst_feature_t>());
for (bst_feature_t fid : feature_splits) {
node_constraints_[left_id].insert(fid);
node_constraints_[right_id].insert(fid);
}
for (const auto &constraint : interaction_constraints_) {
bst_uint flag = 1;
for (bst_uint checkvar : feature_splits) {
if (constraint.count(checkvar) == 0) {
flag = 0;
break; }
}
if (flag == 1) {
for (bst_uint k : constraint) {
node_constraints_[left_id].insert(k);
node_constraints_[right_id].insert(k);
}
}
}
}
}