operator.h

Go to the documentation of this file.

#ifndef OPERATOR_H
#define OPERATOR_H
 
#include "../init.h"
#include "tree_node.h"
#include "../util/utils.h"
 
namespace Brush{
namespace util{
    template<typename T, typename Scalar, typename W> 
        requires (!is_one_of_v<Scalar,bool,bJet>)
    Scalar get_weight(const TreeNode& tn, const W** weights=nullptr)
    { 
        Scalar w;
        // Prediction case: weight is stored in the node data.
        if (weights == nullptr)
        {
            w = Scalar(tn.data.W);
        }
        else
        {
            // NLS case 1: floating point weight is stored in weights
            if constexpr (is_same_v<Scalar, W>) 
                w = **weights;
            // NLS case 2: a Jet/Dual weight is stored in weights, but this constant is a 
            // integer type. We need to do some casting
            else if constexpr (is_same_v<Scalar, iJet> && is_same_v<W, fJet>) {
                using WScalar = typename Scalar::Scalar;
                WScalar tmp = WScalar((**weights).a);    
                w = Scalar(tmp);
            }
            // NLS case 3: a Jet/Dual weight is stored in weights, matching Scalar type
            else            
                w = Scalar(**weights);
            *weights = *weights+1;
            
        }
        return w;
    };
    template<typename T, typename Scalar, typename W> 
        requires (is_one_of_v<Scalar,bool,bJet>)
    Scalar get_weight(const TreeNode& tn, const W** weights=nullptr)
    {
        // we cannot weight a boolean feature. Nevertheless, we need to provide
        // an implementation for get_weight behavior, so the metaprogramming
        // doesn't fail to get a matching signature. 
 
        if (tn.data.get_is_weighted())
            // Node's init() function avoids the creation of weighted nodes, 
            // and the setter for `is_weighted` prevent enabling weight on 
            // boolean values. 
            HANDLE_ERROR_THROW(fmt::format("boolean terminal is weighted, but "
            "it should not\n"));
 
        return Scalar(true);
    };
}
// Operator class
 
template<NodeType NT, typename S, bool Fit, typename E=void> 
struct Operator 
{
    using ArgTypes = conditional_t<
        ((UnaryOp<NT> || NaryOp<NT>) && S::ArgCount > 1),
        Array<typename S::FirstArg::Scalar, -1, S::ArgCount>,
        typename S::ArgTypes>;
 
    using RetType = typename S::RetType;
 
    static constexpr size_t ArgCount = S::ArgCount;
 
    template <std::size_t N>
    using NthType = typename S::NthType<N>; 
 
    using W = typename S::WeightType; 
    
    static constexpr auto F = [](const auto& ...args) { 
        Function<NT> f; 
        return f(args...); 
    }; 
 
    Operator() = default;
    // Utilities to grab child outputs.
 
    template<typename T=ArgTypes> requires(is_std_array_v<T> || is_eigen_array_v<T>) 
    T get_kids(const Dataset& d, TreeNode& tn, const W** weights=nullptr) const
    {
        T child_outputs;
        using arg_type = std::conditional_t<is_std_array_v<T>,
            typename T::value_type, Array<typename S::FirstArg::Scalar, -1, 1>>;
        if constexpr (is_eigen_array_v<T>)
            child_outputs.resize(d.get_n_samples(), Eigen::NoChange);
 
        TreeNode* sib = tn.first_child;
        for (int i = 0; i < ArgCount; ++i)
        {
            if (sib == nullptr)
                HANDLE_ERROR_THROW("bad sibling ptr in get kids");
            if constexpr (Fit){
                if constexpr(is_std_array_v<T>)
                    child_outputs.at(i) = sib->fit<arg_type>(d);
                else
                    child_outputs.col(i) = sib->fit<arg_type>(d);
            }
            else{
                if constexpr(is_std_array_v<T>)
                    child_outputs.at(i) = sib->predict<arg_type>(d, weights);
                else
                    child_outputs.col(i) = sib->predict<arg_type>(d, weights);
            }
            sib = sib->next_sibling;
        }
        return child_outputs;
    };
 
    template<int I>
    NthType<I> get_kid(const Dataset& d,TreeNode& tn, const W** weights ) const
    {
        auto sib = tree<TreeNode>::sibling_iterator(tn.first_child) ;
        sib += I;
        if constexpr(Fit)
            return sib->fit<NthType<I>>(d);
        else
            return sib->predict<NthType<I>>(d,weights);
    };
 
    template<typename T, size_t ...Is> requires(is_tuple_v<T>) 
    T get_kids_seq(const Dataset& d, TreeNode& tn, const W** weights, std::index_sequence<Is...>) const 
    { 
        return std::make_tuple(get_kid<Is>(d,tn,weights)...);
    };
 
    template<typename T=ArgTypes> requires(is_tuple_v<T>) 
    T get_kids(const Dataset& d, TreeNode& tn, const W** weights=nullptr) const
    {
        return get_kids_seq<T>(d, tn, weights, std::make_index_sequence<ArgCount>{});
    };
 
 
    template<typename T=ArgTypes> requires ( is_std_array_v<T> || is_tuple_v<T>)
    RetType apply(const T& inputs) const
    {
        return std::apply(F, inputs);
    }
 
    template<typename T=ArgTypes> requires ( is_eigen_array_v<T> && !is_std_array_v<T>)
    RetType apply(const T& inputs) const
    {
        return F(inputs);
    }
 
    template<typename T=ArgTypes, typename Scalar=RetType::Scalar>
    RetType eval(const Dataset& d, TreeNode& tn, const W** weights=nullptr) const
    {
        auto inputs = get_kids(d, tn, weights);
        if constexpr (is_one_of_v<Scalar,float,fJet>)
        {
            if (tn.data.get_is_weighted())
            {
                auto w = util::get_weight<RetType,Scalar,W>(tn, weights);
                return this->apply(inputs)*w;
            }
        }
        return this->apply(inputs);
    };
 
    // overloaded version for offset sum
    template<typename T=ArgTypes, typename Scalar=RetType::Scalar>
    requires is_in_v<NT, NodeType::OffsetSum>
    RetType eval(const Dataset& d, TreeNode& tn, const W** weights=nullptr) const
    {
        auto inputs = get_kids(d, tn, weights);
        if constexpr (is_one_of_v<Scalar,float,fJet>)
        {
            if (tn.data.get_is_weighted())
            {
                auto w = util::get_weight<RetType,Scalar,W>(tn, weights);
                return this->apply(inputs) + w;
            }
        }
        return this->apply(inputs);
    };
};
 
template<typename S, bool Fit>
struct Operator<NodeType::Terminal, S, Fit>
{
    using RetType = typename S::RetType;
    using W = typename S::WeightType; 
 
    // Standard C++ types
    template<typename T=RetType, typename Scalar=typename T::Scalar> 
        requires (is_one_of_v<Scalar,bool,int,float>)
    RetType eval(const Dataset& d, const TreeNode& tn, const W** weights=nullptr) const 
    { 
        if constexpr (is_one_of_v<Scalar,float,fJet>)
        {
            if (tn.data.get_is_weighted())
            {
                auto w = util::get_weight<RetType,Scalar,W>(tn, weights);
                return this->get<RetType>(d, tn.data.get_feature())*w;
            }
        }
        return this->get<RetType>(d,tn.data.get_feature());
    };
 
    // Jet types
    template <typename T = RetType, typename Scalar=typename T::Scalar>
        requires( is_one_of_v<Scalar, bJet, iJet, fJet>)
    RetType eval(const Dataset &d, const TreeNode &tn, const W **weights = nullptr) const
    {
        using nonJetType = UnJetify_t<RetType>; 
        if constexpr (is_one_of_v<Scalar,float,fJet>)
        {
            if (tn.data.get_is_weighted())
            {
                auto w = util::get_weight<RetType,Scalar,W>(tn, weights);
                return this->get<nonJetType>(d, tn.data.get_feature()).template cast<Scalar>()*w;
            }
        }
        return this->get<nonJetType>(d, tn.data.get_feature()).template cast<Scalar>();
    };
 
    // Accessing dataset directly
    template<typename T>
    auto get(const Dataset& d, const string& feature) const
    {
        if (std::holds_alternative<T>(d[feature]))
            return std::get<T>(d[feature]); 
 
        HANDLE_ERROR_THROW(fmt::format("Failed to return type {} for '{}'\n",
            DataTypeEnum<RetType>::value,
            feature
        ));
 
        return T(); 
    }
};
 
// Constant Overloads
template<typename S, bool Fit> 
struct Operator<NodeType::Constant, S, Fit>
{
    using RetType = typename S::RetType;
    using W = typename S::WeightType;
 
    template<typename T=RetType, typename Scalar=T::Scalar, int N=T::NumDimensions> 
    RetType eval(const Dataset& d, TreeNode& tn, const W** weights=nullptr) const 
    { 
        // Scalar w = get_weight(tn, weights);
        Scalar w = util::get_weight<RetType,Scalar,W>(tn, weights);
        if constexpr (N == 1)
            return RetType::Constant(d.get_n_samples(), w); 
        else
            return RetType::Constant(d.get_n_samples(), d.get_n_features(), w); 
    };
    
};
 
// MeanLabel overload
template<typename S, bool Fit> 
struct Operator<NodeType::MeanLabel, S, Fit>
{
    using RetType = typename S::RetType;
    using W = typename S::WeightType; 
 
    RetType fit(const Dataset& d, TreeNode& tn) const {
        tn.data.W = d.y.mean();
        return predict(d, tn);
    };
 
    template<typename T=RetType, typename Scalar=T::Scalar, int N=T::NumDimensions> 
    RetType predict(const Dataset& d, TreeNode& tn, const W** weights=nullptr) const 
    { 
        Scalar w = util::get_weight<RetType,Scalar,W>(tn, weights);
        if constexpr (N == 1)
            return RetType::Constant(d.get_n_samples(), w); 
        else
            return RetType::Constant(d.get_n_samples(), d.get_n_features(), w); 
    };
 
    RetType eval(const Dataset& d, TreeNode& tn, const W** weights=nullptr) const {
        if constexpr (Fit)
            return fit(d,tn); 
        else
            return predict(d,tn,weights);
    };
};
 
// Operator overloads 
//  Split
#include "split.h"
// Dispatch functions
template<typename R, NodeType NT, typename S, bool Fit, typename W> 
inline R DispatchOp(const Dataset& d, TreeNode& tn, const W** weights) 
{
    const auto op = Operator<NT,S,Fit>{};
    return op.eval(d, tn, weights);
};
 
template<typename R, NodeType NT, typename S, bool Fit> 
inline R DispatchOp(const Dataset& d, TreeNode& tn) 
{
    const auto op = Operator<NT,S,Fit>{};
    return op.eval(d, tn);
};
 
} // Brush
 
#endif