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 <trg/cdc/dataobjects/CDCTrigger3DHTrack.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);

    template <typename TriggerTrack>
    void addCDCTriggerTrack(const std::string& collectionName,
                            const TriggerTrack& trgTrack);

    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);

    static 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;

    static TEveBox* boxCreator(const ROOT::Math::XYZVector& o, ROOT::Math::XYZVector u, ROOT::Math::XYZVector v, float ud, float vd,
                               float depth);

    static 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;
  };

}