TrackFindingCDCTestWithSimpleSimulation.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/display/EventDataPlotter.h>
#include <tracking/trackFindingCDC/sim/CDCSimpleSimulation.h>
 
#include <tracking/trackFindingCDC/eventdata/tracks/CDCTrack.h>
#include <tracking/trackFindingCDC/eventdata/segments/CDCSegment3D.h>
#include <tracking/trackFindingCDC/eventdata/segments/CDCSegment2D.h>
#include <tracking/trackFindingCDC/eventdata/trajectories/CDCTrajectory3D.h>
 
#include <tracking/trackFindingCDC/topology/CDCWireTopology.h>
 
#include <tracking/trackFindingCDC/geometry/Helix.h>
 
#include <tracking/trackFindingCDC/testFixtures/TrackFindingCDCTestWithTopology.h>
#include <tracking/trackFindingCDC/utilities/TimeIt.h>
 
#include <array>
 
namespace Belle2 {
  namespace TrackFindingCDC {
 
    class TrackFindingCDCTestWithSimpleSimulation :
      public TrackFindingCDCTestWithTopology {
 
    public:
      TrackFindingCDCTestWithSimpleSimulation() : m_simpleSimulation()
      {
      }
 
 
      void SetUp() override
      {
 
        m_savedLogLevel = Belle2::LogSystem::Instance().getLogConfig()->getLogLevel();
        m_savedDebugLogInfo = Belle2::LogSystem::Instance().getLogConfig()->getLogInfo(LogConfig::c_Debug);
        bool run_disabled = ::testing::GTEST_FLAG(also_run_disabled_tests);
        if (run_disabled) {
          Belle2::LogSystem::Instance().getLogConfig()->setLogLevel(LogConfig::ELogLevel::c_Debug);
          Belle2::LogSystem::Instance().getLogConfig()->setLogInfo(LogConfig::c_Debug, LogConfig::c_Level + LogConfig::c_Message);
        }
 
        m_plotter.clear();
        m_mcAxialSegment2Ds.clear();
        m_mcSegment2Ds.clear();
        m_mcTracks.clear();
        m_mcTrajectories.clear();
        m_axialWireHits.clear();
        m_wireHits.clear();
      }
 
      void simulate(const std::initializer_list<Helix>& helices)
      {
        simulate(std::vector<Helix>(helices));
      }
 
 
      void simulate(const std::vector<Helix>& helices)
      {
        std::vector<CDCTrajectory3D> trajectories;
        trajectories.reserve(helices.size());
        for (const Helix& helix : helices) {
          // had to make the implicit conversion to an explicit conversion. Maybe there is a more elegant way to do it
          trajectories.emplace_back(UncertainHelix(helix));
        }
        simulate(trajectories);
      }
 
      void simulate(const std::initializer_list<CDCTrajectory3D>& trajectories)
      {
        simulate(std::vector<CDCTrajectory3D>(trajectories));
      }
 
      void simulate(const std::vector<CDCTrajectory3D>& trajectories)
      {
        m_mcTrajectories = trajectories;
 
        // Construct tracks. Wire hits are stored in the simple simulation
        m_mcTracks = m_simpleSimulation.simulate(trajectories);
 
        fillCaches();
      }
 
      void loadPreparedEvent()
      {
        // No trajectory information present
        m_mcTrajectories.clear();
 
        // Load prepared tracks wire hits are stored in the simple simulation
        m_mcTracks = m_simpleSimulation.loadPreparedEvent();
 
        fillCaches();
      }
 
      void fillCaches()
      {
        for (size_t iTrack = 0; iTrack < m_mcTracks.size(); ++iTrack) {
          B2INFO("Size mc track " << iTrack << " : " << m_mcTracks[iTrack].size());
        }
 
        // Prepare the monte carlo segments
        for (const CDCTrack& mcTrack : m_mcTracks) {
          std::vector<CDCSegment3D> segment3DsInTrack = mcTrack.splitIntoSegments();
          for (const CDCSegment3D& segment3D :  segment3DsInTrack) {
            m_mcSegment2Ds.push_back(segment3D.stereoProjectToRef());
          }
        }
 
        // Filter the axial segments
        for (const CDCSegment2D& segment2D : m_mcSegment2Ds) {
          if (segment2D.getStereoKind() == EStereoKind::c_Axial) {
            m_mcAxialSegment2Ds.push_back(&segment2D);
          }
        }
 
        // Filter for axial hits
        for (const CDCWireHit& wireHit : m_simpleSimulation.getWireHits()) {
          if (wireHit.isAxial()) {
            m_axialWireHits.push_back(&wireHit);
          }
        }
 
        // Pick up points for all hits
        for (const CDCWireHit& wireHit : m_simpleSimulation.getWireHits()) {
          m_wireHits.push_back(&wireHit);
        }
 
        m_plotter.draw(CDCWireTopology::getInstance());
        for (const CDCWireHit& wireHit : m_simpleSimulation.getWireHits()) {
          m_plotter.draw(wireHit);
        }
      }
 
      void plotMCTracks()
      { for (const CDCTrack& mcTrack : m_mcTracks) m_plotter.draw(mcTrack); }
 
      void plotMCTrajectories()
      {
        for (const CDCTrajectory3D& mcTrajectory : m_mcTrajectories) {
          m_plotter.draw(mcTrajectory.getTrajectory2D());
        }
      }
 
      void saveDisplay(const std::string& svgFileName)
      {
        bool run_disabled = ::testing::GTEST_FLAG(also_run_disabled_tests);
        if (run_disabled) {
          m_plotter.save(svgFileName);
        } else {
          B2INFO("Not writing display file. To activate svg display output run with --gtest_also_run_disabled_tests");
        }
      }
 
      template<class AFunction >
      TimeItResult
      timeIt(size_t nExecutions,
             bool activateCallgrind,
             const AFunction& function,
             const std::function<void()>& setUp = doNothing,
             const std::function<void()>& tearDown = doNothing)
      {
        bool run_disabled = ::testing::GTEST_FLAG(also_run_disabled_tests);
        if (not run_disabled) {
          nExecutions = 1;
        }
        return Belle2::TrackFindingCDC::timeIt(nExecutions,
                                               activateCallgrind,
                                               function,
                                               setUp,
                                               tearDown);
 
      }
 
      template<class... Ts>
      void draw(Ts&& ... args)
      {
        m_plotter.draw(std::forward<Ts>(args) ...);
      }
 
      void TearDown() override
      {
        Belle2::LogSystem::Instance().getLogConfig()->setLogLevel(m_savedLogLevel);
        Belle2::LogSystem::Instance().getLogConfig()->setLogInfo(LogConfig::c_Debug, m_savedDebugLogInfo);
 
        m_plotter.clear();
        m_mcAxialSegment2Ds.clear();
        m_mcSegment2Ds.clear();
        m_mcTracks.clear();
        m_mcTrajectories.clear();
        m_axialWireHits.clear();
        m_wireHits.clear();
      }
 
    protected:
      LogConfig::ELogLevel m_savedLogLevel = LogConfig::c_Info;
 
      unsigned int m_savedDebugLogInfo = 100;
 
      const std::array<std::string, 6> m_colors{{ "red", "blue", "green", "yellow", "violet", "cyan" }};
 
      CDCSimpleSimulation m_simpleSimulation;
 
      std::vector<CDCTrajectory3D> m_mcTrajectories;
 
      std::vector<CDCTrack> m_mcTracks;
 
      std::vector<CDCSegment2D> m_mcSegment2Ds;
 
      std::vector<const CDCSegment2D*> m_mcAxialSegment2Ds;
 
      std::vector<const CDCWireHit*> m_axialWireHits;
 
      std::vector<const CDCWireHit*> m_wireHits;
 
    private:
      EventDataPlotter m_plotter;
 
    };
 
  }
}