program.h

Go to the documentation of this file.

/* Brush
 
copyright 2020 William La Cava
license: GNU/GPL v3
*/
#ifndef PROGRAM_H
#define PROGRAM_H
//external includes
 
 
#include <string>
#include "assert.h"
 
// internal includes
 
// #include "data/data.h"
#include "../init.h"
#include "tree_node.h"
#include "node.h"
#include "../vary/search_space.h"
#include "../params.h"
#include "../util/utils.h"
#include "functions.h"
// #include "../variation.h"
// #include "weight_optimizer.h"
 
 
using std::cout;
using std::string;
using Brush::Data::Dataset;
using Brush::SearchSpace;
 
namespace Brush {
 
 
typedef tree<Node>::pre_order_iterator Iter; 
typedef tree<Node>::post_order_iterator PostIter; 
 
using PT = ProgramType;
 
// for unsupervised learning, classification and regression. 

template<PT PType> struct Program 
{
    static constexpr PT program_type = PType;

    using RetType = typename std::conditional_t<PType == PT::Regressor, ArrayXf,
        std::conditional_t<PType == PT::BinaryClassifier, ArrayXb,
        std::conditional_t<PType == PT::MulticlassClassifier, ArrayXi,
        std::conditional_t<PType == PT::Representer, ArrayXXf, ArrayXf
        >>>>;

    using TreeType = std::conditional_t<PType == PT::BinaryClassifier, ArrayXf,
        std::conditional_t<PType == PT::MulticlassClassifier, ArrayXXf, 
        RetType>>;

    bool is_fitted_ = false;

    // Fitness fitness;
    
    tree<Node> Tree; 
    
    std::optional<std::reference_wrapper<SearchSpace>> SSref;
 
    Program() = default;
    Program(const std::reference_wrapper<SearchSpace> s, const tree<Node> t)
        : Tree(t), is_fitted_(false)
    {
        SSref = std::optional<std::reference_wrapper<SearchSpace>>{s};
    }
 
    Program<PType> copy() { return Program<PType>(*this); }
 
    inline void set_search_space(const std::reference_wrapper<SearchSpace> s)
    {
        SSref = std::optional<std::reference_wrapper<SearchSpace>>{s};
    }

    int complexity() const{
        auto head = Tree.begin(); 
        
        return head.node->get_complexity();
    }

    int linear_complexity() const{
        auto head = Tree.begin(); 
        
        return head.node->get_linear_complexity();
    }

    int size(bool include_weight=true) const{
        auto head = Tree.begin(); 
        
        return head.node->get_size(include_weight);
    }

    int size_at(Iter& top, bool include_weight=true) const{
 
        return top.node->get_size(include_weight);
    }

    int depth() const{
        //tree.hh count the number of edges. We need to ensure that a single-node
        //tree has depth>0
        return 1+Tree.max_depth();
    }

    int depth_at(Iter& top) const{
        return 1+Tree.max_depth(top);
    }

    int depth_to_reach(Iter& top) const{
        return 1+Tree.depth(top);
    }
 
    Program<PType>& fit(const Dataset& d)
    {
        return fit(d, {});
    }

    Program<PType>& fit(const Dataset& d,
                        const vector<float>& class_weights)
    {
        TreeType out =  Tree.begin().node->fit<TreeType>(d);
        this->is_fitted_ = true;
        update_weights(d, class_weights);
        // this->valid = true;
        return *this;
    };

    Program<PType>& replace_program(const Program<PType>& new_program)
    {
        this->Tree = new_program.Tree;
        this->is_fitted_ = false;
        
        // Update search space reference if the new program has one, otherwise keep current
        if (new_program.SSref.has_value())
            this->SSref = new_program.SSref;
        // If new_program doesn't have a search space, keep this->SSref as is
        
        return *this;
    };

    Program<PType>& replace_program(const json& j)
    {
        Program<PType> new_program = j;
        return replace_program(new_program);
    };
 
    template <typename R, typename W>
    R predict_with_weights(const Dataset &d, const W** weights)
    {
        if (!is_fitted_)
            HANDLE_ERROR_THROW("Program is not fitted. Call 'fit' first.\n");
 
        return Tree.begin().node->predict<R>(d, weights);
    };
 
    auto predict_with_weights(const Dataset &d, const ArrayXf& weights)
    {
        float const * wptr = weights.data(); 
        return this->predict_with_weights<RetType>(d, &wptr);
    };

    template <typename R = RetType>
    TreeType predict(const Dataset &d) requires(is_same_v<R, TreeType>)
    {
        if (!is_fitted_)
            HANDLE_ERROR_THROW("Program is not fitted. Call 'fit' first.\n");
 
        return Tree.begin().node->predict<TreeType>(d);
    };

    template <typename R = RetType>
    ArrayXb predict(const Dataset &d)   requires(is_same_v<R, ArrayXb>)
    {
        if (!is_fitted_)
            HANDLE_ERROR_THROW("Program is not fitted. Call 'fit' first.\n");
            
        return (Tree.begin().node->predict<TreeType>(d) > 0.5);
    };

    template <typename R = RetType>
    ArrayXi predict(const Dataset &d)   requires(is_same_v<R, ArrayXi>)
    {
        if (!is_fitted_)
            HANDLE_ERROR_THROW("Program is not fitted. Call 'fit' first.\n");
 
        TreeType out = Tree.begin().node->predict<TreeType>(d);
        auto argmax = Function<NodeType::ArgMax>{};
        return argmax(out);
    };
 
    // template <typename R = RetType>
    template <PT P = PType>
        requires((P == PT::BinaryClassifier) || (P == PT::MulticlassClassifier))
    TreeType predict_proba(const Dataset &d) 
    { 
        return predict<TreeType>(d); 
    };

    Program<PType>& fit(const Ref<const ArrayXXf>& X, const Ref<const ArrayXf>& y)
    {
        Dataset d(X,y);
        return fit(d);
    };

    RetType predict(const Ref<const ArrayXXf>& X)
    {
        Dataset d(X);
        return predict(d);
    };

    template <PT P = PType>
        requires((P == PT::BinaryClassifier) || (P == PT::MulticlassClassifier))
    TreeType predict_proba(const Ref<const ArrayXXf>& X) 
    {
        Dataset d(X);
        return predict_proba(d);
    };

    void update_weights(const Dataset& d,
                        const vector<float>& class_weights);

    int get_n_weights() const
    {
        int count=0;
        // check tree nodes for weights
        for (PostIter i = Tree.begin_post(); i != Tree.end_post(); ++i)
        {
            const auto& node = i.node->data;
 
            // we do not want to include fixed weights, because they should not be changed.
            // get_n_weights, get_weights, and set_weights, are functions used in 
            // parameter optimization --- we can simply make weights invisible to parameter
            // optimization, and they will not be changed (so they remain fixed).
            if (node.weight_is_fixed)
                continue;
 
            // some nodes cannot have their weights optimized, others must have.
            // It is important that this condition also matches the condition in 
            // the methods get_weights and set_weights.
            if (Is<NodeType::OffsetSum>(node.node_type)
            || (node.get_is_weighted() && IsWeighable(node.ret_type)
                && IsWeighable(node.node_type)) )
                ++count;
        }
        return count;
    }

    ArrayXf get_weights()
    {
        ArrayXf weights(get_n_weights()); 
        int i = 0;
        for (PostIter t = Tree.begin_post(); t != Tree.end_post(); ++t)
        {
            const auto& node = t.node->data; 
 
            // skip fixed weights (this also avoid changing offsetSum weight if is locked)
            if (node.weight_is_fixed)
                continue;
 
            if ( Is<NodeType::OffsetSum>(node.node_type)
            ||   (node.get_is_weighted() && IsWeighable(node.ret_type)
                && IsWeighable(node.node_type)) )
            {
                weights(i) = node.W;
                ++i;
            }
        }
        return weights;
    }

    void set_weights(const ArrayXf& weights)
    {
        // take the weights set them in the tree. 
        // return the weights of the tree as an array 
        if (weights.size() != get_n_weights())
            HANDLE_ERROR_THROW("Tried to set_weights of incorrect size");
 
        int j = 0;
        for (PostIter i = Tree.begin_post(); i != Tree.end_post(); ++i)
        {
            auto& node = i.node->data; 
 
            // skip fixed weights (this also avoid changing offsetSum weight if is locked)
            if (node.weight_is_fixed)
                continue;
 
            if ( Is<NodeType::OffsetSum>(node.node_type)
            ||   (node.get_is_weighted() && IsWeighable(node.ret_type)
                && IsWeighable(node.node_type)) )
            {
                node.W = weights(j);
                ++j;
            }
        }
    }

    void lock_nodes(int end_depth=0, bool keep_leaves_unlocked=true, bool keep_current_weights=false)
    {
        // This also support unlocking the ndoes by setting 0, true, false.
        // OBS for unlocking the node: Do not change prob_change here, because some
        // nodes are meant to never be replaced (e.g. logistic root).
        // unlocking and changing probabilities are different things.
        // every `if` performing a lock should have its counterpart `else` performing
        // unlocking.
        
        // iterate over the nodes, locking them if their depth does not exceed end_depth.
        if (end_depth<0) {
            return;
        }
 
        // we need the iterator to calculate the depth, but 
        // the lambda below iterate using nodes. So we are creating an iterator
        // and using it to access depth.
        auto tree_iter = Tree.begin();
 
        std::for_each(Tree.begin(), Tree.end(),
            [&](auto& n){ 
                auto d = Tree.depth(tree_iter);
                std::advance(tree_iter, 1);
 
                // weights (this will work for all nodes)
                if (n.get_is_weighted() && d<end_depth) // this will lock and unlock weights!
                    n.weight_is_fixed = keep_current_weights;
                else{
                    n.weight_is_fixed = false;
                }
 
                // leaves (terminals, constants, and splitbest variables)
                if (IsLeaf(n.node_type)) {
                    if (d<end_depth && !keep_leaves_unlocked)
                        n.node_is_fixed = true; // leaves should be locked based on depth 
                    else {
                        // either we are outside end_depth or leaves should be unlocked.
                        // this `else` also helps unlocking th enode
                        n.node_is_fixed = false;
                    }
 
                    return; // stop here
                }
 
                // non-terminal nodes
                if (d<end_depth) // This if-else could be a single line, but im trying to make it readable
                    n.node_is_fixed = true;
                else // this else clausure helps unlocking the node
                    n.node_is_fixed = false;
 
                // special case - split best nodes.
                // If we are skipping leaves, then the split feature is unlocked;
                // Otherwise, then we lock based on depth.
                if (n.node_type==NodeType::SplitBest)
                {
                    if (keep_leaves_unlocked)
                    {
                        n.set_keep_split_feature(false);
                    }
                    else // leaves can be locked
                    {
                        // check if we should lock based on depth
                        n.set_keep_split_feature(d+1<end_depth);
                    }
                }
            }
        );
    }

    string get_model(string fmt="compact", bool pretty=false) const
    {
        auto head = Tree.begin(); 
        if (fmt=="tree")
            return head.node->get_tree_model(pretty);
        else if (fmt=="dot")
            return get_dot_model(); ;
        return head.node->get_model(pretty);
    }

    string get_dot_model(string extras="") const
    {
        // TODO: make node IDs stable (hash or index) and labels reflect nodetype names. 
        // ref: https://stackoverflow.com/questions/10579041/graphviz-create-new-node-with-this-same-label#10579155
        string out = "digraph G {\n";
        if (! extras.empty())
            out += fmt::format("{}\n", extras);
 
        auto get_id = [](const auto& n){
            if (Is<NodeType::Terminal>(n->data.node_type)) 
                return n->data.get_name(false);
 
            return fmt::format("{}",fmt::ptr(n)).substr(2);
        };
        // bool first = true;
        std::map<string, unsigned int> node_count;
        int i = 0;
        for (Iter iter = Tree.begin(); iter!=Tree.end(); iter++)
        {
            const auto& parent = iter.node;
            // const auto& parent_data = iter.node->data;
 
            string parent_id = get_id(parent);
            // if (Is<NodeType::Terminal>(parent_data.node_type)) 
            //     parent_id = parent_data.get_name(false);
            // else{
            //     parent_id = fmt::format("{}",fmt::ptr(iter.node)).substr(2);
            // }
            // // parent_id = parent_id.substr(2);
 
            // This is for the root --------------------------------------------
            // if the first node is weighted, make a dummy output node so that the 
            // first node's weight can be shown
            if (i==0 && parent->data.get_is_weighted())
            {
                // making the weight red if fixed
                string font_color = "";
                if (parent->data.weight_is_fixed) {
                    font_color = ", fontcolor=lightcoral";
                }
 
                out += "y [shape=box];\n";
                out += fmt::format("y -> \"{}\" [label=\"{:.2f}\"{}];\n", 
                        // parent_data.get_name(false),
                        parent_id,
                        parent->data.W,
                        font_color
                        );
            }
 
            // add the node
            bool is_constant = Is<NodeType::Constant, NodeType::MeanLabel>(parent->data.node_type);
            string node_label = parent->data.get_name(is_constant);
 
            if (Is<NodeType::SplitBest>(parent->data.node_type)) {
                std::string feature = parent->data.get_feature();
                std::string threshold = fmt::format("{:.2f}", parent->data.W);
 
                // Append markers for fixed flags
                if (parent->data.keep_split_feature)
                    feature += "^";  // split feature fixed
                if (parent->data.weight_is_fixed)
                    threshold += "*"; // split weight fixed
 
                node_label = fmt::format("{} >= {}?", feature, threshold);
            }
            if (Is<NodeType::OffsetSum>(parent->data.node_type)){
                node_label = fmt::format("Add"); 
            }
            
            string node_style = parent->data.get_prob_change() >0.0 ? "" : ", style=filled, fillcolor=lightcoral";
            out += fmt::format("\"{}\" [label=\"{}\"{}];\n", parent_id, node_label, node_style);
 
            // add edges to the node's children
            auto kid = iter.node->first_child;
            for (int j = 0; j < iter.number_of_children(); ++j)
            {
                string edge_label="";
                string head_label="";
                string tail_label="";
                bool use_head_tail_labels = false;
                
                string kid_id = get_id(kid);
                // string kid_id = fmt::format("{}",fmt::ptr(kid));
                // kid_id = kid_id.substr(2);
 
                if (kid->data.get_is_weighted()
                && Isnt<NodeType::Constant, NodeType::MeanLabel, 
                        NodeType::OffsetSum, NodeType::SplitBest>(kid->data.node_type))
                {
                    edge_label = fmt::format("{:.2f}",kid->data.W);
                }
 
                if (Is<NodeType::SplitOn>(parent->data.node_type)){
                    use_head_tail_labels=true;
                    if (j == 0)
                        tail_label = fmt::format(">= {:.2f}",parent->data.W); 
                    else if (j==1)
                        tail_label = "Y"; 
                    else
                        tail_label = "N";
 
                    head_label=edge_label;
                }
                else if (Is<NodeType::SplitBest>(parent->data.node_type)){
                    use_head_tail_labels=true;
                    if (j == 0){
                        tail_label = "Y"; 
                    }
                    else
                        tail_label = "N";
 
                    head_label = edge_label;
                }
                else if (Is<NodeType::Softmax>(parent->data.node_type)){
                    // Each child supplies logits for one output class.
                    edge_label = fmt::format("class {}", j);
                }
 
                // drawing the edges
                string font_color = "";
                if (kid->data.weight_is_fixed) {
                    font_color = ", fontcolor=lightcoral";
                }
 
                if (use_head_tail_labels){
                    out += fmt::format("\"{}\" -> \"{}\" [headlabel=\"{}\",taillabel=\"{}\"{}];\n", 
                            parent_id,
                            kid_id,
                            head_label,
                            tail_label,
                            font_color
                            );
                }
                else{
                    out += fmt::format("\"{}\" -> \"{}\" [label=\"{}\"{}];\n", 
                            parent_id,
                            kid_id,
                            edge_label,
                            font_color
                            );
                }
                kid = kid->next_sibling;
            }
        
            // adding the offset as the last child
            if (Is<NodeType::OffsetSum>(parent->data.node_type)){
                // drawing the edge
                out += fmt::format("\"{}\" -> \"{}\" [label=\"\"];\n", 
                        parent_id,
                        parent_id+"Offset"
                        );
                        
                // drawing the node
                out += fmt::format("\"{}\" [label=\"{:.2f}\"{}];\n",
                        parent_id+"Offset",
                        parent->data.W,
                        node_style
                        ); 
            }
                 
            ++i;
        }
 
        out += "label=\"^ split feature fixed, * split threshold fixed\";\nlabelloc=bottom;\nfontsize=10;";
        out += "}\n";
 
        return out;
    }

    vector<Node> linearize() const {
        vector<Node> linear_program;
        for (PostIter i = Tree.begin_post(); i != Tree.end_post(); ++i)
            linear_program.push_back(i.node->data);
        return linear_program;
    }
}; // Program
} // Brush

// weight optimization
#include "optimizer/weight_optimizer.h"
// #include "../variation.h"
namespace Brush{
 
template<ProgramType PType> 
void Program<PType>::update_weights(const Dataset& d,
                                    const vector<float>& class_weights)
{
    // Updates the weights within a tree. 
    // make an optimizer
    auto WO = WeightOptimizer(); 
    // get new weights from optimization.
    WO.update((*this), d, class_weights);
};
 

// serialization
// serialization for program
template<ProgramType PType>
void to_json(json &j, const Program<PType> &p)
{
    j = json{{"Tree",p.Tree}, {"is_fitted_", p.is_fitted_}}; 
}
 
template<ProgramType PType>
void from_json(const json &j, Program<PType>& p)
{
    j.at("Tree").get_to(p.Tree);
    j.at("is_fitted_").get_to(p.is_fitted_);
}
 
}//namespace Brush
 
 
 
#endif