Function GetTreeSplitsDescriptions

catboost/private/libs/algo/tree_print.cpp:180–224

utility function for python_package/catboost/core.py plot_tree function

// utility function for python_package/catboost/core.py plot_tree function
TVector<TString> GetTreeSplitsDescriptions(const TFullModel& model, size_t treeIdx, const NCB::TDataProviderPtr pool) {
    CB_ENSURE(treeIdx < model.GetTreeCount(),
        "Requested tree splits description for tree " << treeIdx << ", but model has " << model.GetTreeCount());

    if (pool) {
        CheckModelAndDatasetCompatibility(model, *pool->ObjectsData.Get());
    }

    TVector<TString> splits;

    const auto binFeatures = model.ModelTrees->GetBinFeatures();

    size_t treeSplitEnd = (treeIdx + 1 < model.GetTreeCount())
        ? model.ModelTrees->GetModelTreeData()->GetTreeStartOffsets()[treeIdx + 1]
        : model.ModelTrees->GetModelTreeData()->GetTreeSplits().size();

    THashMap<ui32, TString> catFeaturesHash;
    NCB::TFeaturesLayout featuresLayout;
    if (pool) {
        catFeaturesHash = MergeCatFeaturesHashToString(pool.Get()->ObjectsData.Get()[0]);
        featuresLayout = *(pool.Get()->MetaInfo.FeaturesLayout.Get());
    } else {
        TVector<ui32> catFeaturesExternalIndexes;
        for (const auto& feature: model.ModelTrees->GetCatFeatures()) {
            catFeaturesExternalIndexes.push_back(feature.Position.FlatIndex);
        }
        featuresLayout = NCB::TFeaturesLayout(model.GetNumFloatFeatures() + model.GetNumCatFeatures(), catFeaturesExternalIndexes, {}, {}, {}, false);
    }

    for (size_t splitIdx = model.ModelTrees->GetModelTreeData()->GetTreeStartOffsets()[treeIdx]; splitIdx < treeSplitEnd; ++splitIdx) {
        TModelSplit binFeature = binFeatures[model.ModelTrees->GetModelTreeData()->GetTreeSplits()[splitIdx]];
        TString featureDescription = BuildDescription(featuresLayout, binFeature);

        if (binFeature.Type == ESplitType::OneHotFeature) {
            CB_ENSURE(pool,
                "Please pass training dataset to plot_tree function, "
                "training dataset is required if categorical features are present in the model.");
            featureDescription += catFeaturesHash[(ui32)binFeature.OneHotFeature.Value];
        }

        splits.push_back(featureDescription);
    }

    return splits;
}

TVector<TString> GetTreeLeafValuesDescriptions(const TFullModel& model, size_t treeIdx) {
    CB_ENSURE(treeIdx < model.GetTreeCount(),