packages/trackreco/PHSimpleKFProp.h

0001 /*!
0002  *  \file PHSimpleKFProp.h
0003  *  \brief      kalman filter based propagator
0004  *  \author Michael Peters & Christof Roland
0005  */
0006
0007 #ifndef TRACKRECO_PHSIMPLEKFPROP_H
0008 #define TRACKRECO_PHSIMPLEKFPROP_H
0009
0010 #include "ALICEKF.h"
0011 #include "nanoflann.hpp"
0012
0013 // PHENIX includes
0014 #include <tpc/TpcGlobalPositionWrapper.h>
0015
0016 #include <trackbase/TrkrDefs.h>
0017
0018 #include <fun4all/SubsysReco.h>
0019
0020 #include <phool/PHTimer.h>
0021
0022 #include <Acts/MagneticField/MagneticFieldProvider.hpp>
0023
0024 #include <Eigen/Core>
0025
0026 // STL includes
0027 #include <limits>
0028 #include <memory>
0029 #include <string>
0030 #include <vector>
0031
0032 class ActsGeometry;
0033 class PHCompositeNode;
0034 class PHField;
0035 class TrkrClusterContainer;
0036 class TrkrClusterIterationMapv1;
0037 class SvtxTrackMap;
0038 class TrackSeedContainer;
0039 class TrackSeed;
0040
0041 using PositionMap = std::map<TrkrDefs::cluskey, Acts::Vector3>;
0042
0043 class PHSimpleKFProp : public SubsysReco
0044 {
0045  public:
0046   PHSimpleKFProp(const std::string& name = "PHSimpleKFProp");
0047   ~PHSimpleKFProp() override;
0048
0049   int InitRun(PHCompositeNode* topNode) override;
0050   int process_event(PHCompositeNode* topNode) override;
0051   int End(PHCompositeNode* topNode) override;
0052
0053   void set_field_dir(const double rescale)
0054   {
0055     _fieldDir = 1;
0056     if (rescale > 0)
0057     {
0058       _fieldDir = -1;
0059     }
0060   }
0061   void ghostRejection(bool set_value = true) { m_ghostrejection = set_value; }
0062   void magFieldFile(const std::string& fname) { m_magField = fname; }
0063   void set_max_window(double s) { _max_dist = s; }
0064   void useConstBField(bool opt) { _use_const_field = opt; }
0065   void setConstBField(float b) { _const_field = b; }
0066   void useFixedClusterError(bool opt) { _use_fixed_clus_err = opt; }
0067   void setFixedClusterError(int i, double val) { _fixed_clus_err.at(i) = val; }
0068   void use_truth_clusters(bool truth)
0069   {
0070     _use_truth_clusters = truth;
0071   }
0072   void SetIteration(int iter) { _n_iteration = iter; }
0073   void set_pp_mode(bool mode) { _pp_mode = mode; }
0074   void set_max_seeds(unsigned int ui) { _max_seeds = ui; }
0075   enum class PropagationDirection
0076   {
0077     Outward,
0078     Inward
0079   };
0080
0081   void setNeonFraction(double frac) { Ne_frac = frac; };
0082   void setArgonFraction(double frac) { Ar_frac = frac; };
0083   void setCF4Fraction(double frac) { CF4_frac = frac; };
0084   void setNitrogenFraction(double frac) { N2_frac = frac; };
0085   void setIsobutaneFraction(double frac) { isobutane_frac = frac; };
0086
0087   void set_ghost_phi_cut(double d) { _ghost_phi_cut = d; }
0088   void set_ghost_eta_cut(double d) { _ghost_eta_cut = d; }
0089   void set_ghost_x_cut(double d) { _ghost_x_cut = d; }
0090   void set_ghost_y_cut(double d) { _ghost_y_cut = d; }
0091   void set_ghost_z_cut(double d) { _ghost_z_cut = d; }
0092
0093   // number of threads
0094   void set_num_threads(int value) { m_num_threads = value; }
0095
0096  private:
0097   bool _use_truth_clusters = false;
0098   bool m_ghostrejection = true;
0099   /// fetch node pointers
0100   int get_nodes(PHCompositeNode* topNode);
0101   std::vector<double> radii;
0102   std::vector<double> _vertex_x;
0103   std::vector<double> _vertex_y;
0104   std::vector<double> _vertex_z;
0105   std::vector<double> _vertex_xerr;
0106   std::vector<double> _vertex_yerr;
0107   std::vector<double> _vertex_zerr;
0108   std::vector<double> _vertex_ids;
0109   double _Bzconst = 10. * 0.000299792458f;
0110   // double _Bz = 1.4*_Bzconst;
0111   double _max_dist = .05;
0112   size_t _min_clusters_per_track = 3;
0113   double _fieldDir = -1;
0114   double _max_sin_phi = 1.;
0115   bool _pp_mode = false;
0116   unsigned int _max_seeds = 0;
0117
0118   TrkrClusterContainer* _cluster_map = nullptr;
0119
0120   TrackSeedContainer* _track_map = nullptr;
0121
0122   //! magnetic field map
0123   PHField* _field_map = nullptr;
0124   bool m_own_fieldmap = false;
0125
0126   /// acts geometry
0127   ActsGeometry* m_tgeometry = nullptr;
0128
0129   /// global position wrapper
0130   TpcGlobalPositionWrapper m_globalPositionWrapper;
0131
0132   /// get global position for a given cluster
0133   /**
0134    * uses ActsTransformation to convert cluster local position into global coordinates
0135    * incorporates TPC distortion correction, if present
0136    */
0137   Acts::Vector3 getGlobalPosition(TrkrDefs::cluskey, TrkrCluster*) const;
0138
0139   PositionMap PrepareKDTrees();
0140
0141   bool TransportAndRotate(
0142     double old_layer,
0143     double new_layer,
0144     double& phi,
0145     GPUTPCTrackParam& kftrack,
0146     GPUTPCTrackParam::GPUTPCTrackFitParam& fp) const;
0147
0148   bool PropagateStep(
0149     unsigned int& current_layer,
0150     double& current_phi,
0151     PropagationDirection& direction,
0152     std::vector<TrkrDefs::cluskey>& propagated_track,
0153     const std::vector<TrkrDefs::cluskey>& ckeys,
0154     GPUTPCTrackParam& kftrack,
0155     GPUTPCTrackParam::GPUTPCTrackFitParam& fp,
0156     const PositionMap& globalPositions) const;
0157
0158   // TrackSeed objects store clusters in order of increasing cluster key (std::set<TrkrDefs::cluskey>),
0159   // which means we have to have a way to directly pass a list of clusters in order to extend looping tracks
0160   std::vector<TrkrDefs::cluskey> PropagateTrack(TrackSeed* track, PropagationDirection direction, GPUTPCTrackParam& aliceSeed, const PositionMap& globalPositions) const;
0161   std::vector<TrkrDefs::cluskey> PropagateTrack(TrackSeed* track, std::vector<TrkrDefs::cluskey>& ckeys, PropagationDirection direction, GPUTPCTrackParam& aliceSeed, const PositionMap& globalPositions) const;
0162   std::vector<std::vector<TrkrDefs::cluskey>> RemoveBadClusters(const std::vector<std::vector<TrkrDefs::cluskey>>& seeds, const PositionMap& globalPositions) const;
0163
0164   template <typename T>
0165   struct KDPointCloud
0166   {
0167     KDPointCloud() = default;
0168
0169     std::vector<std::vector<T>> pts;
0170
0171     inline size_t kdtree_get_point_count() const
0172     {
0173       return pts.size();
0174     }
0175
0176     inline T kdtree_distance(const T* p1, const size_t idx_p2, size_t /*size*/) const
0177     {
0178       const T d0 = p1[0] - pts[idx_p2][0];
0179       const T d1 = p1[1] - pts[idx_p2][1];
0180       const T d2 = p1[2] - pts[idx_p2][2];
0181       return d0 * d0 + d1 * d1 + d2 * d2;
0182     }
0183
0184     inline T kdtree_get_pt(const size_t idx, int dim) const
0185     {
0186       if (dim == 0)
0187         return pts[idx][0];
0188       else if (dim == 1)
0189         return pts[idx][1];
0190       else
0191         return pts[idx][2];
0192     }
0193
0194     template <class BBOX>
0195     bool kdtree_get_bbox(BBOX& /*bb*/) const
0196     { return false; }
0197
0198   };
0199
0200   std::vector<std::shared_ptr<KDPointCloud<double>>> _ptclouds;
0201
0202   std::vector<std::shared_ptr<nanoflann::KDTreeSingleIndexAdaptor<nanoflann::L2_Simple_Adaptor<double, KDPointCloud<double>>, KDPointCloud<double>, 3>>> _kdtrees;
0203
0204   std::unique_ptr<ALICEKF> fitter;
0205
0206   double get_Bz(double x, double y, double z) const;
0207
0208   void rejectAndPublishSeeds(std::vector<TrackSeed_v2>& seeds, const PositionMap& positions, std::vector<float>& trackChi2);
0209
0210   void publishSeeds(const std::vector<TrackSeed_v2>&);
0211
0212   int _max_propagation_steps = 200;
0213   std::string m_magField;
0214   bool _use_const_field = false;
0215   float _const_field = 1.4;
0216   bool _use_fixed_clus_err = false;
0217   std::array<double, 3> _fixed_clus_err = {.2, .2, .5};
0218   TrkrClusterIterationMapv1* _iteration_map = nullptr;
0219   int _n_iteration = 0;
0220
0221   double Ne_frac = 0.00;
0222   double Ar_frac = 0.75;
0223   double CF4_frac = 0.20;
0224   double N2_frac = 0.00;
0225   double isobutane_frac = 0.05;
0226
0227   double _ghost_phi_cut = std::numeric_limits<double>::max();
0228   double _ghost_eta_cut = std::numeric_limits<double>::max();
0229   double _ghost_x_cut = std::numeric_limits<double>::max();
0230   double _ghost_y_cut = std::numeric_limits<double>::max();
0231   double _ghost_z_cut = std::numeric_limits<double>::max();
0232
0233   // number of threads
0234   /*
0235    * default is 0. This corresponds to allocating as many threads as available on the host
0236    * on CONDOR, by default, this results in only one thread allocated
0237    * being allocated unless several cores are allocated to the job
0238    * via request_cpus directive in the jdf
0239    */
0240   int m_num_threads = 0;
0241
0242 };
0243
0244 #endif