EVEVisualization.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 <mdst/dataobjects/MCParticle.h>
#include <mdst/dataobjects/ECLCluster.h>
#include <mdst/dataobjects/KLMCluster.h>
#include <mdst/dataobjects/Track.h>
#include <cdc/dataobjects/CDCSimHit.h>
#include <cdc/dataobjects/CDCHit.h>
#include <trg/cdc/dataobjects/CDCTriggerSegmentHit.h>
#include <trg/cdc/dataobjects/CDCTriggerTrack.h>
#include <pxd/dataobjects/PXDSimHit.h>
#include <svd/dataobjects/SVDSimHit.h>
#include <svd/dataobjects/SVDCluster.h>
#include <klm/dataobjects/KLMHit2d.h>
#include <klm/dataobjects/KLMSimHit.h>
#include <arich/dataobjects/ARICHHit.h>
#include <top/dataobjects/TOPDigit.h>
#include <vxd/geometry/GeoCache.h>
#include <tracking/dataobjects/ROIid.h>
 
#include <framework/datastore/StoreArray.h>
 
#include <tracking/dataobjects/RecoTrack.h>
#include <genfit/GFRaveVertex.h>
 
#include <Math/Vector3D.h>
#include <TEveStraightLineSet.h>
#include <TEveTrack.h>
 
#include <string>
 
 
class TEveBox;
class TEveCaloDataVec;
class TEveCalo3D;
class TEveElementList;
class TEvePointSet;
class TEveTrackList;
class TEveTrackPropagator;
 
namespace Belle2 {
  class DisplayData;
  class EveVisBField;
 
  class EVEVisualization {
    enum eFitterType {
      SimpleKalman,
      RefKalman,
      DafSimple,
      DafRef,
      Gbl
    };
 
 
    struct MCTrack {
      TEveTrack* track; 
      TEvePointSet* simhits; 
      const MCParticle* parentParticle; 
    };
 
    struct ElementGroup {
      ElementGroup(): group(nullptr), visible(true) { }
      TEveElementList* group; 
      bool visible; 
    };
 
    const static int c_recoHitColor;
    const static int c_recoTrackColor;
    const static int c_trackColor;
    const static int c_trackMarkerColor;
    const static int c_unassignedHitColor;
    const static int c_klmClusterColor;
 
  public:
    explicit EVEVisualization();
 
    EVEVisualization(const EVEVisualization&) = delete;
    EVEVisualization& operator=(const EVEVisualization&) = delete;
    ~EVEVisualization();
 
    // === adding event data ===
 
    void addTrack(const Belle2::Track* belle2Track);
 
    void addTrackCandidate(const std::string& collectionName,
                           const RecoTrack& recoTrack);
 
    void addTrackCandidateImproved(const std::string& collectionName,
                                   const RecoTrack& recoTrack);
 
    void addCDCTriggerTrack(const std::string& collectionName,
                            const CDCTriggerTrack& track);
 
    template <class T> void addSimHits(const StoreArray<T>& hits)
    {
      const int numHits = hits.getEntries();
      for (int i = 0; i < numHits; i++) {
        const RelationsObject* rel = hits[i];
        const MCParticle* mcpart = rel->getRelatedFrom<MCParticle>();
 
        addSimHit(hits[i], mcpart);
      }
    }
 
    void addSimHit(const CDCSimHit* hit, const MCParticle* particle);
 
    void addSimHit(const PXDSimHit* hit, const MCParticle* particle);
 
    void addSimHit(const SVDSimHit* hit, const MCParticle* particle);
 
    void addSimHit(const KLMSimHit* hit, const MCParticle* particle);
 
    void addSimHit(const ROOT::Math::XYZVector& v, const MCParticle* particle);
 
    MCTrack* addMCParticle(const MCParticle* particle);
 
    void addVertex(const genfit::GFRaveVertex* vertex);
 
    void addECLCluster(const ECLCluster* cluster);
 
    void addKLMCluster(const KLMCluster* cluster);
 
    void addBKLMHit2d(const KLMHit2d* bklm2dhit);
 
    void addEKLMHit2d(const KLMHit2d* eklm2dhit);
 
    void addARICHHit(const ARICHHit* hit);
 
    void addROI(const ROIid* roi);
 
    template <class T> void addUnassignedRecoHits(const StoreArray<T>& hits)
    {
      if (hits.getEntries() == 0)
        return;
      if (!m_unassignedRecoHits) {
        m_unassignedRecoHits = new TEveStraightLineSet("Unassigned RecoHits");
        m_unassignedRecoHits->SetTitle("Unassigned RecoHits");
        m_unassignedRecoHits->SetMainColor(c_unassignedHitColor);
        m_unassignedRecoHits->SetMarkerColor(c_unassignedHitColor);
        m_unassignedRecoHits->SetMarkerStyle(6);
        //m_unassignedRecoHits->SetMainTransparency(60);
      }
      for (const T& hit : hits) {
        if (m_shownRecohits.count(&hit) == 0) {
          addRecoHit(&hit, m_unassignedRecoHits);
        }
      }
    }
 
    void addCDCHit(const CDCHit* hit, bool showTriggerHits = false);
 
    void addCDCTriggerSegmentHit(const std::string& collectionName, const CDCTriggerSegmentHit* hit);
 
    void addTOPDigits(const StoreArray<TOPDigit>& digits);
 
    void addObject(const TObject* dataStoreObject, TEveElement* visualRepresentation);
 
 
    void showUserData(const DisplayData& displayData);
 
 
    void makeTracks();
 
    void clearEvent();
 
    // === configuration options ===
 
    void setOptions(const std::string& opts);
 
    void setErrScale(double errScale = 1.);
 
    void setAssignToPrimaries(bool on) { m_assignToPrimaries = on; }
 
    void setHideSecondaries(bool on) { m_hideSecondaries = on; }
 
 
  private:
    TEveCalo3D* m_calo3d;
 
    TEveBox* boxCreator(const ROOT::Math::XYZVector& o, ROOT::Math::XYZVector u, ROOT::Math::XYZVector v, float ud, float vd,
                        float depth);
 
    void makeLines(TEveTrack* eveTrack, const genfit::StateOnPlane* prevState, const genfit::StateOnPlane* state,
                   const genfit::AbsTrackRep* rep,
                   TEvePathMark::EType_e markType, bool drawErrors, int markerPos = 1);
 
    template <class SomeVXDHit> void addRecoHit(const SomeVXDHit* hit, TEveStraightLineSet* lines)
    {
      static VXD::GeoCache& geo = VXD::GeoCache::getInstance();
 
      const ROOT::Math::XYZVector local_pos(hit->getU(), hit->getV(), 0.0); //z-component is height over the center of the detector plane
      const VXD::SensorInfoBase& sensor = geo.getSensorInfo(hit->getSensorID());
      const ROOT::Math::XYZVector global_pos = sensor.pointToGlobal(local_pos);
      lines->AddMarker(global_pos.X(), global_pos.Y(), global_pos.Z());
 
      m_shownRecohits.insert(hit);
    }
 
    void addRecoHit(const SVDCluster* hit, TEveStraightLineSet* lines);
 
    void addRecoHit(const CDCHit* hit, TEveStraightLineSet* lines);
 
    void addToGroup(const std::string& name, TEveElement* elem);
 
    double m_errorScale;
 
    std::string m_options;
 
    bool m_assignToPrimaries;
 
    bool m_hideSecondaries{false};
 
    std::map<const MCParticle*, MCTrack> m_mcparticleTracks;
 
    std::map<std::string, ElementGroup> m_groups;
 
    TEveTrackList* m_tracklist;
 
    TEveTrackPropagator* m_trackpropagator;
 
    TEveTrackPropagator* m_gftrackpropagator;
 
    TEveTrackPropagator* m_consttrackpropagator;
 
    TEveCaloDataVec* m_eclData;
 
    EveVisBField* m_bfield;
 
    std::set<const TObject*> m_shownRecohits;
 
    TEveStraightLineSet* m_unassignedRecoHits = nullptr;
    bool m_unassignedRecoHitsVisibility = true;
 
    static constexpr double c_minPCut = 0.00;
 
    bool m_drawCardinalRep = true;
 
    bool m_drawErrors = false;
 
    bool m_drawRefTrack = false;
 
    bool m_drawForward = false;
 
    bool m_drawBackward = false;
  };
}