Back to home page

sPhenix code displayed by LXR

 
 

    


File indexing completed on 2025-12-16 09:21:01

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 TrkrClusterContainer;
0035 class TrkrClusterIterationMapv1;
0036 class SvtxTrackMap;
0037 class TrackSeedContainer;
0038 class TrackSeed;
0039 
0040 using PositionMap = std::map<TrkrDefs::cluskey, Acts::Vector3>;
0041 
0042 class PHSimpleKFProp : public SubsysReco
0043 {
0044  public:
0045   PHSimpleKFProp(const std::string& name = "PHSimpleKFProp");
0046   ~PHSimpleKFProp() override = default;
0047 
0048   int InitRun(PHCompositeNode* topNode) override;
0049   int process_event(PHCompositeNode* topNode) override;
0050   int End(PHCompositeNode* topNode) override;
0051 
0052   // noop
0053   void set_field_dir(const double)
0054   { std::cout << "PHSimpleKFProp::set_field_dir - Warning - This function does nothing. Remove from macro" << std::endl; }
0055 
0056   // noop
0057   void magFieldFile(const std::string&)
0058   { std::cout << "PHSimpleKFProp::magFieldFile - Warning - This function does nothing. Remove from macro" << std::endl; }
0059 
0060   // noop
0061   void useConstBField(bool)
0062   { std::cout << "PHSimpleKFProp::useConstBField - Warning - This function does nothing. Remove from macro" << std::endl; }
0063 
0064   // noop
0065   void setConstBField(float)
0066   { std::cout << "PHSimpleKFProp::constBField - Warning - This function does nothing. Remove from macro" << std::endl; }
0067 
0068   void ghostRejection(bool set_value = true) { m_ghostrejection = set_value; }
0069   void set_max_window(double s) { _max_dist = s; }
0070   void useFixedClusterError(bool opt) { _use_fixed_clus_err = opt; }
0071   void setFixedClusterError(int i, double val) { _fixed_clus_err.at(i) = val; }
0072   void use_truth_clusters(bool truth)
0073   {
0074     _use_truth_clusters = truth;
0075   }
0076   void SetIteration(int iter) { _n_iteration = iter; }
0077   void set_pp_mode(bool mode) { _pp_mode = mode; }
0078   void set_max_seeds(unsigned int ui) { _max_seeds = ui; }
0079   enum class PropagationDirection
0080   {
0081     Outward,
0082     Inward
0083   };
0084 
0085   void setNeonFraction(double frac) { Ne_frac = frac; };
0086   void setArgonFraction(double frac) { Ar_frac = frac; };
0087   void setCF4Fraction(double frac) { CF4_frac = frac; };
0088   void setNitrogenFraction(double frac) { N2_frac = frac; };
0089   void setIsobutaneFraction(double frac) { isobutane_frac = frac; };
0090 
0091   void set_ghost_phi_cut(double d) { _ghost_phi_cut = d; }
0092   void set_ghost_eta_cut(double d) { _ghost_eta_cut = d; }
0093   void set_ghost_x_cut(double d) { _ghost_x_cut = d; }
0094   void set_ghost_y_cut(double d) { _ghost_y_cut = d; }
0095   void set_ghost_z_cut(double d) { _ghost_z_cut = d; }
0096 
0097   // number of threads
0098   void set_num_threads(int value) { m_num_threads = value; }
0099 
0100  private:
0101   bool _use_truth_clusters = false;
0102   bool m_ghostrejection = true;
0103   /// fetch node pointers
0104   int get_nodes(PHCompositeNode* topNode);
0105   std::vector<double> radii;
0106   std::vector<double> _vertex_x;
0107   std::vector<double> _vertex_y;
0108   std::vector<double> _vertex_z;
0109   std::vector<double> _vertex_xerr;
0110   std::vector<double> _vertex_yerr;
0111   std::vector<double> _vertex_zerr;
0112   std::vector<double> _vertex_ids;
0113   double _Bzconst = 10. * 0.000299792458f;
0114   // double _Bz = 1.4*_Bzconst;
0115   double _max_dist = .05;
0116   size_t _min_clusters_per_track = 3;
0117 
0118   //! internal field conversion to get the right particle sign
0119 
0120   double _max_sin_phi = 1.;
0121   bool _pp_mode = false;
0122   unsigned int _max_seeds = 0;
0123 
0124   TrkrClusterContainer* _cluster_map = nullptr;
0125 
0126   TrackSeedContainer* _track_map = nullptr;
0127 
0128   /// acts geometry
0129   ActsGeometry* m_tgeometry = nullptr;
0130 
0131   /// global position wrapper
0132   TpcGlobalPositionWrapper m_globalPositionWrapper;
0133 
0134   /// get global position for a given cluster
0135   /**
0136    * uses ActsTransformation to convert cluster local position into global coordinates
0137    * incorporates TPC distortion correction, if present
0138    */
0139   Acts::Vector3 getGlobalPosition(TrkrDefs::cluskey, TrkrCluster*) const;
0140 
0141   PositionMap PrepareKDTrees();
0142 
0143   bool TransportAndRotate(
0144     double old_layer,
0145     double new_layer,
0146     double& phi,
0147     GPUTPCTrackParam& kftrack,
0148     GPUTPCTrackParam::GPUTPCTrackFitParam& fp) const;
0149 
0150   bool PropagateStep(
0151     unsigned int& current_layer,
0152     double& current_phi,
0153     PropagationDirection& direction,
0154     std::vector<TrkrDefs::cluskey>& propagated_track,
0155     const std::vector<TrkrDefs::cluskey>& ckeys,
0156     GPUTPCTrackParam& kftrack,
0157     GPUTPCTrackParam::GPUTPCTrackFitParam& fp,
0158     const PositionMap& globalPositions) const;
0159 
0160   // TrackSeed objects store clusters in order of increasing cluster key (std::set<TrkrDefs::cluskey>),
0161   // which means we have to have a way to directly pass a list of clusters in order to extend looping tracks
0162   std::vector<TrkrDefs::cluskey> PropagateTrack(TrackSeed* track, PropagationDirection direction, GPUTPCTrackParam& aliceSeed, const PositionMap& globalPositions) const;
0163   std::vector<TrkrDefs::cluskey> PropagateTrack(TrackSeed* track, std::vector<TrkrDefs::cluskey>& ckeys, PropagationDirection direction, GPUTPCTrackParam& aliceSeed, const PositionMap& globalPositions) const;
0164   std::vector<std::vector<TrkrDefs::cluskey>> RemoveBadClusters(const std::vector<std::vector<TrkrDefs::cluskey>>& seeds, const PositionMap& globalPositions) const;
0165 
0166   template <typename T>
0167   struct KDPointCloud
0168   {
0169     KDPointCloud() = default;
0170 
0171     std::vector<std::vector<T>> pts;
0172 
0173     inline size_t kdtree_get_point_count() const
0174     {
0175       return pts.size();
0176     }
0177 
0178     inline T kdtree_distance(const T* p1, const size_t idx_p2, size_t /*size*/) const
0179     {
0180       const T d0 = p1[0] - pts[idx_p2][0];
0181       const T d1 = p1[1] - pts[idx_p2][1];
0182       const T d2 = p1[2] - pts[idx_p2][2];
0183       return d0 * d0 + d1 * d1 + d2 * d2;
0184     }
0185 
0186     inline T kdtree_get_pt(const size_t idx, int dim) const
0187     {
0188       if (dim == 0)
0189         return pts[idx][0];
0190       else if (dim == 1)
0191         return pts[idx][1];
0192       else
0193         return pts[idx][2];
0194     }
0195 
0196     template <class BBOX>
0197     bool kdtree_get_bbox(BBOX& /*bb*/) const
0198     { return false; }
0199 
0200   };
0201 
0202   std::vector<std::shared_ptr<KDPointCloud<double>>> _ptclouds;
0203 
0204   std::vector<std::shared_ptr<nanoflann::KDTreeSingleIndexAdaptor<nanoflann::L2_Simple_Adaptor<double, KDPointCloud<double>>, KDPointCloud<double>, 3>>> _kdtrees;
0205 
0206   std::unique_ptr<ALICEKF> fitter;
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 
0214   //! fixed cluster errors
0215   bool _use_fixed_clus_err = false;
0216 
0217   //! fixed cluster errors
0218   std::array<double, 3> _fixed_clus_err = {.2, .2, .5};
0219 
0220   //! keep track of used clusters in previous iteration for iterative tracking
0221   TrkrClusterIterationMapv1* _iteration_map = nullptr;
0222 
0223   //! iteration number
0224   int _n_iteration = 0;
0225 
0226   //!@name TPC gas. Needed for Kalman Filter
0227   //@{
0228   double Ne_frac = 0.00;
0229   double Ar_frac = 0.75;
0230   double CF4_frac = 0.20;
0231   double N2_frac = 0.00;
0232   double isobutane_frac = 0.05;
0233   //@}
0234 
0235   //!@name ghost rejection cuts
0236   //@{
0237   double _ghost_phi_cut = std::numeric_limits<double>::max();
0238   double _ghost_eta_cut = std::numeric_limits<double>::max();
0239   double _ghost_x_cut = std::numeric_limits<double>::max();
0240   double _ghost_y_cut = std::numeric_limits<double>::max();
0241   double _ghost_z_cut = std::numeric_limits<double>::max();
0242   //@}
0243 
0244   //! number of threads
0245   /**
0246    * default is 0. This corresponds to allocating as many threads as available on the host
0247    * on CONDOR, by default, this results in only one thread allocated
0248    * being allocated unless several cores are allocated to the job
0249    * via request_cpus directive in the jdf
0250    */
0251   int m_num_threads = 0;
0252 
0253 };
0254 
0255 #endif