WeightedFastHoughTree.h

/**************************************************************************
 * basf2 (Belle II Analysis Software Framework)                           *
 * Author: The Belle II Collaboration                                     *
 *                                                                        *
 * See git log for contributors and copyright holders.                    *
 * This file is licensed under LGPL-3.0, see LICENSE.md.                  *
 **************************************************************************/
#pragma once
 
#include <tracking/trackFindingCDC/hough/trees/DynTree.h>
#include <tracking/trackFindingCDC/hough/baseelements/WithWeightedItems.h>
#include <tracking/trackFindingCDC/hough/baseelements/WithSharedMark.h>
 
#include <vector>
#include <memory>
#include <cassert>
#include <cfloat>
#include <cmath>
 
namespace Belle2 {
  namespace TrackFindingCDC {
 
    template<class T, class ADomain, class ADomainDivsion>
    class  WeightedParititioningDynTree :
      public DynTree<  WithWeightedItems<ADomain, T>, ADomainDivsion> {
    private:
      using Super = DynTree<  WithWeightedItems<ADomain, T>, ADomainDivsion>;
 
    public:
      WeightedParititioningDynTree(ADomain topDomain, ADomainDivsion domainDivsion) :
        Super(WithWeightedItems<ADomain, T>(std::move(topDomain)), std::move(domainDivsion))
      {}
    };
 
    template<class T, class ADomain, class ADomainDivsion>
    class WeightedFastHoughTree :
      public WeightedParititioningDynTree<WithSharedMark<T>, ADomain, ADomainDivsion> {
 
    private:
      using Super = WeightedParititioningDynTree<WithSharedMark<T>, ADomain, ADomainDivsion>;
 
    public:
      using WeightedParititioningDynTree<WithSharedMark<T>, ADomain, ADomainDivsion>::WeightedParititioningDynTree;
 
      using Node = typename Super::Node;
 
    public:
      template<class Ts>
      void seed(const Ts& items)
      {
        this->fell();
        Node& topNode = this->getTopNode();
        for (auto && item : items) {
          m_marks.push_back(false);
          bool& markOfItem = m_marks.back();
          Weight weight = DBL_MAX;
          topNode.insert(WithSharedMark<T>(T(item), &markOfItem), weight);
        }
      }
 
      template <class AItemInDomainMeasure>
      std::vector<std::pair<ADomain, std::vector<T>>>
      findHeavyLeavesDisjoint(AItemInDomainMeasure& weightItemInDomain,
                              int maxLevel,
                              double minWeight)
      {
        auto skipLowWeightNode = [minWeight](const Node * node) {
          return not(node->getWeight() >= minWeight);
        };
        return findLeavesDisjoint(weightItemInDomain, maxLevel, skipLowWeightNode);
      }
 
      template <class AItemInDomainMeasure, class ASkipNodePredicate>
      std::vector<std::pair<ADomain, std::vector<T>>>
      findLeavesDisjoint(AItemInDomainMeasure& weightItemInDomain,
                         int maxLevel,
                         ASkipNodePredicate& skipNode)
      {
        std::vector<std::pair<ADomain,  std::vector<T> > > found;
        auto isLeaf = [&found, &skipNode, maxLevel](Node * node) {
          // Skip the expansion and the filling of the children
          if (skipNode(node)) {
            return true;
          }
 
          // Node is a leaf at the maximum level
          // Save its content
          // Do not walk children
          if (node->getLevel() >= maxLevel) {
            const ADomain* domain = node;
            found.emplace_back(*domain, std::vector<T>(node->begin(), node->end()));
            for (WithSharedMark<T>& markableItem : *node) {
              markableItem.mark();
            }
            return true;
          }
 
          // Else to node has enough weight and is not at the lowest level
          // Signal that it is not a leaf
          // Continue to create and fill children.
          return false;
        };
        fillWalk(weightItemInDomain, isLeaf);
        return found;
      }
 
      template <class AItemInDomainMeasure>
      std::vector<std::pair<ADomain, std::vector<T>>>
      findHeaviestLeafRepeated(AItemInDomainMeasure& weightItemInDomain,
                               int maxLevel,
                               const Weight minWeight = NAN)
      {
        auto skipLowWeightNode = [minWeight](const Node * node) {
          return not(node->getWeight() >= minWeight);
        };
        return findHeaviestLeafRepeated(weightItemInDomain, maxLevel, skipLowWeightNode);
      }
 
      template <class AItemInDomainMeasure, class ASkipNodePredicate>
      std::vector<std::pair<ADomain, std::vector<T>>>
      findHeaviestLeafRepeated(AItemInDomainMeasure& weightItemInDomain,
                               int maxLevel,
                               ASkipNodePredicate& skipNode)
      {
        std::vector<std::pair<ADomain,  std::vector<T> > > found;
        Node* node = findHeaviestLeaf(weightItemInDomain, maxLevel, skipNode);
        while (node) {
          const ADomain* domain = node;
          found.emplace_back(*domain, std::vector<T>(node->begin(), node->end()));
          for (WithSharedMark<T>& markableItem : *node) {
            markableItem.mark();
          }
          node = findHeaviestLeaf(weightItemInDomain, maxLevel, skipNode);
        }
        return found;
      }
 
      template <class AItemInDomainMeasure, class ASkipNodePredicate>
      std::unique_ptr<std::pair<ADomain, std::vector<T>>>
      findHeaviestLeafSingle(AItemInDomainMeasure& weightItemInDomain,
                             int maxLevel,
                             ASkipNodePredicate& skipNode)
      {
        using Result = std::pair<ADomain,  std::vector<T> >;
        std::unique_ptr<Result> found = nullptr;
        Node* node = findHeaviestLeaf(weightItemInDomain, maxLevel, skipNode);
        if (node) {
          const ADomain* domain = node;
          found.reset(new Result(*domain, std::vector<T>(node->begin(), node->end())));
          for (WithSharedMark<T>& markableItem : *node) {
            markableItem.mark();
          }
        }
        return found;
      }
 
      template <class AItemInDomainMeasure, class ASkipNodePredicate>
      Node* findHeaviestLeaf(AItemInDomainMeasure& weightItemInDomain,
                             int maxLevel,
                             ASkipNodePredicate& skipNode)
      {
        Node* heaviestNode = nullptr;
        Weight heighestWeigth = NAN;
        auto isLeaf = [&heaviestNode, &heighestWeigth, maxLevel, &skipNode](Node * node) {
          // Skip the expansion and the filling of the children
          if (skipNode(node)) {
            return true;
          }
 
          Weight nodeWeight = node->getWeight();
          // Skip the expansion and filling of the children if the node has not enough weight
          if (not std::isnan(heighestWeigth) and not(nodeWeight > heighestWeigth)) {
            return true;
          }
 
          // Node is a leaf at the maximum level and is heavier than everything seen before.
          // Save its content
          // Do not walk children
          if (node->getLevel() >= maxLevel) {
            heaviestNode = node;
            heighestWeigth = nodeWeight;
            return true;
          }
          return false;
        };
        fillWalk(weightItemInDomain, isLeaf);
        return heaviestNode;
      }
 
    public:
      template<class AItemInDomainMeasure, class AIsLeafPredicate>
      void fillWalk(AItemInDomainMeasure& weightItemInDomain,
                    AIsLeafPredicate& isLeaf)
      {
        auto walker = [&weightItemInDomain, &isLeaf](Node * node) {
          // Check if node is a leaf
          // Do not create children in this case
          if (isLeaf(node)) {
            // Do not walk children.
            return false;
          }
 
          // Node is not a leaf.
          // Check if it has childen.
          // If children have not been created, create and fill them.
          typename Node::Children* children = node->getChildren();
          if (not children) {
            node->createChildren();
            children = node->getChildren();
            for (Node& childNode : *children) {
              assert(childNode.getChildren() == nullptr);
              assert(childNode.size() == 0);
              auto measure =
              [&childNode, &weightItemInDomain](WithSharedMark<T>& markableItem) -> Weight {
                // Weighting function should not see the mark, but only the item itself.
                T & item(markableItem);
                return weightItemInDomain(item, &childNode);
              };
              childNode.insert(*node, measure);
            }
          }
          // Continue to walk the children.
          return true;
        };
        walkHeighWeightFirst(walker);
      }
 
      template<class ATreeWalker>
      void walkHeighWeightFirst(ATreeWalker& walker)
      {
        auto priority = [](Node * node) -> float {
          auto isMarked = [](const WithSharedMark<T>& markableItem) -> bool {
            return markableItem.isMarked();
          };
          node->eraseIf(isMarked);
          return node->getWeight();
        };
 
        this->walk(walker, priority);
      }
 
      void fell()
      {
        this->getTopNode().clear();
        m_marks.clear();
        Super::fell();
      }
 
      void raze()
      {
        this->fell();
        Super::raze();
        m_marks.shrink_to_fit();
      }
 
    private:
      std::deque<bool> m_marks;
      // Note: Have to use a deque here because std::vector<bool> is special
      // std::vector<bool> m_marks;
    };
  }
}