Clusterizer.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/ca/AutomatonCell.h>
 
#include <tracking/trackFindingCDC/utilities/VectorRange.h>
#include <tracking/trackFindingCDC/utilities/WeightedRelation.h>
 
#include <tracking/trackFindingCDC/findlets/base/Findlet.h>
 
#include <framework/logging/Logger.h>
 
namespace Belle2 {
  namespace TrackFindingCDC {
 
    template <class ACellHolder, class ACluster = std::vector<ACellHolder*> >
    class Clusterizer
      : public Findlet<ACellHolder* const, WeightedRelation<ACellHolder> const, ACluster> {
 
    public:
      void apply(std::vector<ACellHolder*> const& cellHolders,
                 std::vector<WeightedRelation<ACellHolder> > const& cellHolderRelations,
                 std::vector<ACluster>& clusters) override
      {
        // Expect the relations to be sorted for lookup
        assert(std::is_sorted(cellHolderRelations.begin(),
                              cellHolderRelations.end()));
 
        // Prepare some output clusters
        clusters.reserve(30);
 
        // Prepare states
        for (ACellHolder* cellHolder : cellHolders) {
          setCellState(cellHolder, -1);
        }
 
        // Work horse cluster
        std::vector<ACellHolder*> cluster;
 
        // Go through each cell holder and start a cluster on each that is unassigned yet
        int iCluster = -1;
        for (ACellHolder* cellHolder : cellHolders) {
          if (getCellState(cellHolder) != -1) continue;
 
          cluster.clear();
 
          ++iCluster;
          setCellState(cellHolder, iCluster);
          cluster.push_back(cellHolder);
 
          expandCluster(cellHolderRelations, cluster);
 
          clusters.emplace_back(std::move(cluster));
          cluster.clear();
        }
      }
 
    private:
      void expandCluster(std::vector<WeightedRelation<ACellHolder>> const& cellHolderRelations,
                         std::vector<ACellHolder*>& cluster) const
      {
        ACellHolder* seedCellHolder = cluster.front();
        int iCluster = getCellState(seedCellHolder);
 
        // Grow the cluster iterativelly
        std::vector<ACellHolder*> checkNow;
        std::vector<ACellHolder*> checkNext;
 
        checkNow.reserve(10);
        checkNext.reserve(10);
 
        checkNext.push_back(seedCellHolder);
 
        while (not checkNext.empty()) {
 
          checkNow.swap(checkNext);
          checkNext.clear();
 
          for (ACellHolder* cellHolder : checkNow) {
 
            ConstVectorRange<WeightedRelation<ACellHolder> > neighborRelations(
              std::equal_range(cellHolderRelations.begin(),
                               cellHolderRelations.end(),
                               cellHolder));
 
            // Setting the cell weight to the number of neighbors
            size_t nNeighbors = neighborRelations.size();
            setCellWeight(cellHolder, nNeighbors);
 
            for (const WeightedRelation<ACellHolder>& neighborRelation : neighborRelations) {
              ACellHolder* neighborCellHolder = neighborRelation.getTo();
 
              Weight neighborICluster = getCellState(neighborCellHolder);
              if (neighborICluster == -1) {
                // Neighbor not yet in cluster
                setCellState(neighborCellHolder, iCluster);
                cluster.push_back(neighborCellHolder);
 
                // Register neighbor for further expansion
                checkNext.push_back(neighborCellHolder);
                continue;
              }
 
              if (neighborICluster != iCluster) {
                B2WARNING("Clusterizer: Neighboring item was already assigned to different "
                          "cluster. Check if the neighborhood is symmetric.");
                continue;
              }
            }
          }
        }
      }
 
      void setCellState(ACellHolder* cellHolder, Weight cellState) const
      {
        AutomatonCell& automatonCell = cellHolder->getAutomatonCell();
        automatonCell.setCellState(cellState);
      }
 
      Weight getCellState(ACellHolder* cellHolder) const
      {
        const AutomatonCell& automatonCell = cellHolder->getAutomatonCell();
        return automatonCell.getCellState();
      }
 
      void setCellWeight(ACellHolder* cellHolder, Weight cellWeight) const
      {
        AutomatonCell& automatonCell = cellHolder->getAutomatonCell();
        automatonCell.setCellWeight(cellWeight);
      }
 
    };
  }
}