cluster_tree.hpp

#ifndef INFRARED_CLUSTER_TREE_HPP
#define INFRARED_CLUSTER_TREE_HPP
 
/*
 * InfraRed ---  A generic engine for Boltzmann sampling over constraint networks
 * (C) Sebastian Will, 2018
 *
 * This file is part of the InfraRed source code.
 *
 * InfraRed provides a generic framework for tree decomposition-based
 * Boltzmann sampling over constraint networks
 */
 
#include <set>
#include <optional>
 
#include "graph.hpp"
#include "feature_network.hpp"
 
namespace ired {
    template<class FunValue=double, class EvaluationPolicy=StdEvaluationPolicy<FunValue>>
    class ClusterTree {
 
    public:
        using feature_network_t = FeatureNetwork<FunValue, EvaluationPolicy>;
 
        using evaluation_policy_t = typename feature_network_t::evaluation_policy_t;
 
        using var_idx_t = typename feature_network_t::var_idx_t;
        using cluster_t = typename feature_network_t::cluster_t;
        using assignment_t = typename feature_network_t::assignment_t;
        using fun_value_t = typename feature_network_t::fun_value_t;
        using constraint_t = typename feature_network_t::constraint_t;
        using function_t = typename feature_network_t::function_t;
 
        using message_t = MaterializedFunction<fun_value_t>;
 
        struct vertex_info_t {
            vertex_info_t()
                : cluster() {}
            vertex_info_t(const cluster_t &cluster)
                : cluster(cluster) {}
            cluster_t cluster;
        };
 
        struct edge_info_t {
            edge_info_t(function_t *message=nullptr)
                : message(message) {}
            function_t *message;
        };
 
        using tree_t = graph::adjacency_list<vertex_info_t, edge_info_t>;
 
        using vertex_descriptor_t = typename tree_t::vertex_descriptor_t;
 
        ClusterTree() {
        }
 
        explicit
        ClusterTree(const FiniteDomainVector &domains)
            : fn_( domains ) {
        }
 
        explicit
        ClusterTree(std::vector<int> domsizes)
            : fn_( domains_from_domsizes( domsizes ) ) {
        }
 
        ClusterTree(int num_vars, const FiniteDomain &domain)
            : fn_(num_vars, domain) {
        }
 
        ClusterTree(int num_vars, int domsize)
            : fn_( num_vars, FiniteDomain(0, domsize-1) ) {
        }
 
        ~ClusterTree() {}
 
        const auto & feature_network() const {
            return fn_;
        }
 
        auto
        add_root_cluster(const std::vector<var_idx_t> &vars) {
            return tree_.add_vertex(cluster_t(vars));
        }
 
        auto
        add_child_cluster( vertex_descriptor_t parent, const std::vector<int> &vars) {
            auto child = tree_.add_vertex(cluster_t(vars));
            tree_.add_edge(parent, child, edge_info_t());
            return child;
        }
 
        void
        add_constraint( vertex_descriptor_t node, const std::shared_ptr<constraint_t> &x ) {
            tree_[node].cluster.add_constraint( fn_.add_constraint(x) );
        }
 
        void
        add_function( vertex_descriptor_t node, const std::shared_ptr<function_t> &x ) {
            tree_[node].cluster.add_function( fn_.add_function(x) );
        }
 
        auto
        evaluate();
 
        bool is_consistent() {
            return evaluate() != evaluation_policy_t::zero();
        }
 
        auto traceback();
 
        auto restricted_traceback(const std::set<var_idx_t> &variables,
            const assignment_t &assignment);
 
    private:
        feature_network_t fn_;
        tree_t tree_;
 
        bool evaluated_ = false;
        fun_value_t evaluation_result_;
        bool single_empty_rooted_ = false;
 
        vertex_descriptor_t root_;
 
        std::optional<std::set<var_idx_t>> trace_variables;
 
        // insert pseudo root to connect trees of the forest (unless already done)
        // @returns new root
        auto
        single_empty_root();
 
        auto domains_from_domsizes( std::vector<int> &domsizes ) {
            auto domains = FiniteDomainVector();
            for (auto x: domsizes) {
                domains.push_back( FiniteDomain( 0, x - 1 ) );
            }
            return domains;
        }
 
        void
        dfs_evaluate(vertex_descriptor_t v) {
            // notes: this terminates if there are no children
            // as well, we can know, that tree_ is a tree, without any
            // cycles
 
            for(auto &e: tree_.adj_edges(v)) {
                // recursively evaluate subtree at e.target()
                dfs_evaluate(e.target());
 
                // then, compute the message for the current edge
                // from cluster child to cluster parent
 
                const auto &parent = tree_[ e.source() ].cluster;
                const auto &child =  tree_[ e.target() ].cluster;
 
                auto sep  = child.sep_vars(parent);
                auto diff = child.diff_vars(parent);
 
                // concat sep + diff
                auto sep_diff = sep;
                sep_diff.insert(sep_diff.end(),diff.begin(),diff.end());
 
                auto message = std::make_unique<message_t>(sep, fn_);
 
                auto a = assignment_t(fn_.domains());
 
                auto it = a.make_iterator
                    (sep_diff,
                     fn_,
                     child.constraints(),
                     child.functions(),
                     //evaluate 0-ary functions
                     a.eval_determined(child.functions(), evaluation_policy_t())
                     );
 
                fun_value_t x = evaluation_policy_t::zero();
 
                it.register_finish_stage2_hook
                    (sep.size(),
                     [&message,&a,&x] () {
                        message->set(a, x);
                        x = evaluation_policy_t::zero();
                    } );
 
                for(; ! it.finished() ; ++it ) {
                    x = evaluation_policy_t::plus( x, it.value() );
                }
 
                // register message in cn, such that it persists!
                auto msg = fn_.add_function(std::move(message));
                // then, register in cluster parent
                tree_[ e.source() ].cluster.add_function(msg);
 
                // ... and as edge property
                e.message = msg;
            }
        }
 
        void
        dfs_traceback(vertex_descriptor_t v, assignment_t &a) {
 
            for(const auto &e: tree_.adj_edges(v)) {
 
                const auto &parent = tree_[ e.source() ].cluster;
                const auto &child =  tree_[ e.target() ].cluster;
 
                auto diff = child.diff_vars(parent);
                bool diff_changed = false;
 
                if (trace_variables) { // if the trace is restricted,
                                       // then remove other vars from diff
                    auto new_end = remove_if(diff.begin(), diff.end(),
                            [&](var_idx_t i) {
                                return trace_variables->find(i) == trace_variables->end();
                            });
                    diff_changed = new_end != diff.end();
                    diff.erase(new_end, diff.end());
                }
 
                if (! diff.empty()) {
 
                    a.set_undet(diff);
 
                    assert(a.eval_determined(child.constraints(),
                                             StdEvaluationPolicy<bool>()));
 
                    auto it = a.make_iterator(diff, fn_,
                                              child.constraints(),
                                              child.functions(),
                                              a.eval_determined(child.functions(),
                                                                evaluation_policy_t())
                                              );
 
                    // -----
                    // determine value for selector
                    auto value = evaluation_policy_t::zero();
                    if (diff_changed) {
                        // recompute value for selector
                        for( ; ! it.finished(); ++it ) {
                            value = evaluation_policy_t::plus( value, it.value() );
                        }
                        it.reset();
                    } else {
                        // get value by evaluating message:
                        // evaluate message at partial assignment a;
                        value = (*e.message)(a);
                    }
                    // initialize the Selector
                    auto selector = typename evaluation_policy_t::selector(value);
 
                    // enumerate all combinations of values assigned to diff variables;
                    // recompute SUM_a(PROD_f(f(a))), where a runs over these
                    // combinations
                    auto x = evaluation_policy_t::zero();
                    for( ; ! it.finished(); ++it ) {
                        x = evaluation_policy_t::plus( x, it.value() );
 
                        if (selector.select(x)) {
                            break;
                        }
                    }
                }
                // trace back from target
                dfs_traceback(e.target(), a);
            }
        }
 
    }; // end class ClusterTree
 
    // return single empty cluster that roots the tree; if such a
    // cluster exists or was generated before, simply return it;
    // otherwise, construct a new empty cluster, connect it as parent
    // to all existing roots, and return it.
    template<class FunValue, class EvaluationPolicy>
    auto
    ClusterTree<FunValue,EvaluationPolicy>::single_empty_root() {
        if (single_empty_rooted_) {return root_;}
 
        // find all root nodes of the tree
 
        std::set<vertex_descriptor_t> old_roots;
 
        for (int i=0; i < int(tree_.size()); ++i) {
            old_roots.insert(i);
        }
 
        for (int i=0; i < int(tree_.size()); ++i) {
            for (auto &e: tree_.adj_edges(i)) {
                old_roots.erase( e.target() );
            }
        }
 
        if ( old_roots.size()==1 && tree_[ *old_roots.begin() ].cluster.empty() ) {
            root_ = *old_roots.begin();
        } else {
            // insert new root and point to all old roots
            root_ = tree_.add_vertex();
 
            for (auto old_root : old_roots) {
                tree_.add_edge(root_, old_root);
            }
        }
 
        single_empty_rooted_ = true;
 
        return root_;
    }
 
    template<class FunValue, class EvaluationPolicy>
    auto
    ClusterTree<FunValue,EvaluationPolicy>
    ::evaluate() {
        auto root = single_empty_root();
 
        if (!evaluated_) {
            dfs_evaluate(root);
 
            auto a = assignment_t(fn_.domains());
            evaluation_result_ = a.eval_determined(tree_[root].cluster.functions(), evaluation_policy_t());
 
            evaluated_ = true;
        }
        return evaluation_result_;
    }
 
    template<class FunValue, class EvaluationPolicy>
    auto
    ClusterTree<FunValue,EvaluationPolicy>::traceback() {
 
        assert(evaluated_);
        //assert(is_consistent());
 
        auto a = assignment_t(fn_.domains());
 
        dfs_traceback(single_empty_root(), a);
 
        return a;
    }
 
    template<class FunValue, class EvaluationPolicy>
    auto
    ClusterTree<FunValue,EvaluationPolicy>
    ::restricted_traceback(const std::set<var_idx_t> &variables,
        const assignment_t &assignment) {
 
        assert(evaluated_);
 
        //trace_variables = std::set<var_idx_t>(variables.begin(),variables.end());
        trace_variables = variables;
 
        auto a = assignment_t(assignment);
 
        dfs_traceback(single_empty_root(), a);
 
        trace_variables.reset();
 
        return a;
    }
}
 
#endif