#include "xgboost/multi_target_tree_model.h"
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <string_view>
#include <utility>
#include <vector>
#include "io_utils.h"
#include "xgboost/base.h"
#include "xgboost/json.h"
#include "xgboost/logging.h"
#include "xgboost/tree_model.h"
namespace xgboost {
MultiTargetTree::MultiTargetTree(TreeParam const* param)
: param_{param},
left_(1ul, InvalidNodeId()),
right_(1ul, InvalidNodeId()),
parent_(1ul, InvalidNodeId()),
split_index_(1ul, 0),
default_left_(1ul, 0),
split_conds_(1ul, std::numeric_limits<float>::quiet_NaN()),
weights_(param->size_leaf_vector, std::numeric_limits<float>::quiet_NaN()) {
CHECK_GT(param_->size_leaf_vector, 1);
}
template <bool typed, bool feature_is_64>
void LoadModelImpl(Json const& in, std::vector<float>* p_weights, std::vector<bst_node_t>* p_lefts,
std::vector<bst_node_t>* p_rights, std::vector<bst_node_t>* p_parents,
std::vector<float>* p_conds, std::vector<bst_feature_t>* p_fidx,
std::vector<std::uint8_t>* p_dft_left) {
namespace tf = tree_field;
auto get_float = [&](std::string_view name, std::vector<float>* p_out) {
auto& values = get<FloatArrayT<typed>>(get<Object const>(in).find(name)->second);
auto& out = *p_out;
out.resize(values.size());
for (std::size_t i = 0; i < values.size(); ++i) {
out[i] = GetElem<Number>(values, i);
}
};
get_float(tf::kBaseWeight, p_weights);
get_float(tf::kSplitCond, p_conds);
auto get_nidx = [&](std::string_view name, std::vector<bst_node_t>* p_nidx) {
auto& nidx = get<I32ArrayT<typed>>(get<Object const>(in).find(name)->second);
auto& out_nidx = *p_nidx;
out_nidx.resize(nidx.size());
for (std::size_t i = 0; i < nidx.size(); ++i) {
out_nidx[i] = GetElem<Integer>(nidx, i);
}
};
get_nidx(tf::kLeft, p_lefts);
get_nidx(tf::kRight, p_rights);
get_nidx(tf::kParent, p_parents);
auto const& splits = get<IndexArrayT<typed, feature_is_64> const>(in[tf::kSplitIdx]);
p_fidx->resize(splits.size());
auto& out_fidx = *p_fidx;
for (std::size_t i = 0; i < splits.size(); ++i) {
out_fidx[i] = GetElem<Integer>(splits, i);
}
auto const& dft_left = get<U8ArrayT<typed> const>(in[tf::kDftLeft]);
auto& out_dft_l = *p_dft_left;
out_dft_l.resize(dft_left.size());
for (std::size_t i = 0; i < dft_left.size(); ++i) {
out_dft_l[i] = GetElem<Boolean>(dft_left, i);
}
}
void MultiTargetTree::LoadModel(Json const& in) {
namespace tf = tree_field;
bool typed = IsA<F32Array>(in[tf::kBaseWeight]);
bool feature_is_64 = IsA<I64Array>(in[tf::kSplitIdx]);
if (typed && feature_is_64) {
LoadModelImpl<true, true>(in, &weights_, &left_, &right_, &parent_, &split_conds_,
&split_index_, &default_left_);
} else if (typed && !feature_is_64) {
LoadModelImpl<true, false>(in, &weights_, &left_, &right_, &parent_, &split_conds_,
&split_index_, &default_left_);
} else if (!typed && feature_is_64) {
LoadModelImpl<false, true>(in, &weights_, &left_, &right_, &parent_, &split_conds_,
&split_index_, &default_left_);
} else {
LoadModelImpl<false, false>(in, &weights_, &left_, &right_, &parent_, &split_conds_,
&split_index_, &default_left_);
}
}
void MultiTargetTree::SaveModel(Json* p_out) const {
CHECK(p_out);
auto& out = *p_out;
auto n_nodes = param_->num_nodes;
I32Array lefts(n_nodes);
I32Array rights(n_nodes);
I32Array parents(n_nodes);
F32Array conds(n_nodes);
U8Array default_left(n_nodes);
F32Array weights(n_nodes * this->NumTarget());
auto save_tree = [&](auto* p_indices_array) {
auto& indices_array = *p_indices_array;
for (bst_node_t nidx = 0; nidx < n_nodes; ++nidx) {
CHECK_LT(nidx, left_.size());
lefts.Set(nidx, left_[nidx]);
CHECK_LT(nidx, right_.size());
rights.Set(nidx, right_[nidx]);
CHECK_LT(nidx, parent_.size());
parents.Set(nidx, parent_[nidx]);
CHECK_LT(nidx, split_index_.size());
indices_array.Set(nidx, split_index_[nidx]);
conds.Set(nidx, split_conds_[nidx]);
default_left.Set(nidx, default_left_[nidx]);
auto in_weight = this->NodeWeight(nidx);
auto weight_out = common::Span<float>(weights.GetArray())
.subspan(nidx * this->NumTarget(), this->NumTarget());
CHECK_EQ(in_weight.Size(), weight_out.size());
std::copy_n(in_weight.Values().data(), in_weight.Size(), weight_out.data());
}
};
namespace tf = tree_field;
if (this->param_->num_feature >
static_cast<bst_feature_t>(std::numeric_limits<std::int32_t>::max())) {
I64Array indices_64(n_nodes);
save_tree(&indices_64);
out[tf::kSplitIdx] = std::move(indices_64);
} else {
I32Array indices_32(n_nodes);
save_tree(&indices_32);
out[tf::kSplitIdx] = std::move(indices_32);
}
out[tf::kBaseWeight] = std::move(weights);
out[tf::kLeft] = std::move(lefts);
out[tf::kRight] = std::move(rights);
out[tf::kParent] = std::move(parents);
out[tf::kSplitCond] = std::move(conds);
out[tf::kDftLeft] = std::move(default_left);
}
void MultiTargetTree::SetLeaf(bst_node_t nidx, linalg::VectorView<float const> weight) {
CHECK(this->IsLeaf(nidx)) << "Collapsing a split node to leaf " << MTNotImplemented();
auto const next_nidx = nidx + 1;
CHECK_EQ(weight.Size(), this->NumTarget());
CHECK_GE(weights_.size(), next_nidx * weight.Size());
auto out_weight = common::Span<float>(weights_).subspan(nidx * weight.Size(), weight.Size());
for (std::size_t i = 0; i < weight.Size(); ++i) {
out_weight[i] = weight(i);
}
}
void MultiTargetTree::Expand(bst_node_t nidx, bst_feature_t split_idx, float split_cond,
bool default_left, linalg::VectorView<float const> base_weight,
linalg::VectorView<float const> left_weight,
linalg::VectorView<float const> right_weight) {
CHECK(this->IsLeaf(nidx));
CHECK_GE(parent_.size(), 1);
CHECK_EQ(parent_.size(), left_.size());
CHECK_EQ(left_.size(), right_.size());
std::size_t n = param_->num_nodes + 2;
CHECK_LT(split_idx, this->param_->num_feature);
left_.resize(n, InvalidNodeId());
right_.resize(n, InvalidNodeId());
parent_.resize(n, InvalidNodeId());
auto left_child = parent_.size() - 2;
auto right_child = parent_.size() - 1;
left_[nidx] = left_child;
right_[nidx] = right_child;
if (nidx != 0) {
CHECK_NE(parent_[nidx], InvalidNodeId());
}
parent_[left_child] = nidx;
parent_[right_child] = nidx;
split_index_.resize(n);
split_index_[nidx] = split_idx;
split_conds_.resize(n, std::numeric_limits<float>::quiet_NaN());
split_conds_[nidx] = split_cond;
default_left_.resize(n);
default_left_[nidx] = static_cast<std::uint8_t>(default_left);
weights_.resize(n * this->NumTarget());
auto p_weight = this->NodeWeight(nidx);
CHECK_EQ(p_weight.Size(), base_weight.Size());
auto l_weight = this->NodeWeight(left_child);
CHECK_EQ(l_weight.Size(), left_weight.Size());
auto r_weight = this->NodeWeight(right_child);
CHECK_EQ(r_weight.Size(), right_weight.Size());
for (std::size_t i = 0; i < base_weight.Size(); ++i) {
p_weight(i) = base_weight(i);
l_weight(i) = left_weight(i);
r_weight(i) = right_weight(i);
}
}
bst_target_t MultiTargetTree::NumTarget() const { return param_->size_leaf_vector; }
std::size_t MultiTargetTree::Size() const { return parent_.size(); }
}