packages/trackreco/PrelimDistortionCorrection.cc

0001 /*!
0002  *  \File PrelimDistortionCorrection.cc
0003  *  \brief  Makes preliminary TPC distortion corrections in pp running and replaces TPC seed parameters with new fits
0004  *  \author Tony Frawley
0005  */
0006
0007 #include "PrelimDistortionCorrection.h"
0008
0009 #include "ALICEKF.h"
0010
0011 #include <ffamodules/CDBInterface.h>
0012
0013 #include <fun4all/Fun4AllReturnCodes.h>
0014
0015 #include <phfield/PHFieldUtility.h>
0016 #include <phfield/PHFieldConfig.h>
0017
0018 #include <phool/PHTimer.h>
0019 #include <phool/getClass.h>
0020 #include <phool/phool.h>                       // for PHWHERE
0021
0022 // tpc distortion correction
0023 #include <tpc/TpcDistortionCorrectionContainer.h>
0024
0025 #include <trackbase/TrackFitUtils.h>
0026 #include <trackbase/TrkrCluster.h>
0027 #include <trackbase/TrkrClusterContainer.h>
0028 #include <trackbase/TrkrDefs.h>
0029 #include <trackbase/ActsGeometry.h>
0030
0031 #include <trackbase_historic/ActsTransformations.h>
0032 #include <trackbase_historic/TrackSeedContainer.h>
0033 #include <trackbase_historic/TrackSeed_v2.h>
0034 #include <trackbase_historic/TrackSeedHelper.h>
0035
0036 #include <Geant4/G4SystemOfUnits.hh>
0037
0038 #include <Eigen/Core>
0039 #include <Eigen/Dense>
0040
0041 #include <cmath>
0042 #include <iostream>                            // for operator<<, basic_ostream
0043 #include <vector>
0044
0045 //___________________________________________________________________________________________
0046 PrelimDistortionCorrection::PrelimDistortionCorrection(const std::string& name)
0047   : SubsysReco(name)
0048 {}
0049
0050 //___________________________________________________________________________________________
0051 int PrelimDistortionCorrection::End(PHCompositeNode* /*unused*/)
0052 {
0053   return Fun4AllReturnCodes::EVENT_OK;
0054 }
0055
0056 //___________________________________________________________________________________________
0057 int PrelimDistortionCorrection::InitRun(PHCompositeNode* topNode)
0058 {
0059
0060   int ret = get_nodes(topNode);
0061   if (ret != Fun4AllReturnCodes::EVENT_OK) { return ret; }
0062
0063   // load magnetic field from node tree
0064   /* note: if field is not found it is created with default configuration, as defined in PHFieldUtility */
0065   const auto field_map = PHFieldUtility::GetFieldMapNode(nullptr, topNode);
0066
0067   fitter = std::make_unique<ALICEKF>(_cluster_map,field_map, _min_clusters_per_track,_max_sin_phi,Verbosity());
0068   fitter->setNeonFraction(Ne_frac);
0069   fitter->setArgonFraction(Ar_frac);
0070   fitter->setCF4Fraction(CF4_frac);
0071   fitter->setNitrogenFraction(N2_frac);
0072   fitter->setIsobutaneFraction(isobutane_frac);
0073   fitter->useFixedClusterError(_use_fixed_clus_err);
0074   fitter->setFixedClusterError(0,_fixed_clus_err.at(0));
0075   fitter->setFixedClusterError(1,_fixed_clus_err.at(1));
0076   fitter->setFixedClusterError(2,_fixed_clus_err.at(2));
0077
0078   // properly set constField in ALICEKF, based on PHFieldConfig
0079   const auto field_config = PHFieldUtility::GetFieldConfigNode(nullptr, topNode);
0080   if( field_config->get_field_config() == PHFieldConfig::kFieldUniform )
0081   { fitter->setConstBField(field_config->get_field_mag_z()); }
0082
0083   return Fun4AllReturnCodes::EVENT_OK;
0084 }
0085
0086 //____________________________________________________________________________________________
0087 int PrelimDistortionCorrection::get_nodes(PHCompositeNode* topNode)
0088 {
0089   //---------------------------------
0090   // Get Objects off of the Node Tree
0091   //---------------------------------
0092
0093   _tgeometry = findNode::getClass<ActsGeometry>(topNode, "ActsGeometry");
0094   if(!_tgeometry)
0095     {
0096       std::cout << "No Acts tracking geometry, exiting." << std::endl;
0097       return Fun4AllReturnCodes::ABORTEVENT;
0098     }
0099
0100   // tpc distortion correction
0101   // input tpc distortion correction module edge
0102   m_dcc_module_edge = findNode::getClass<TpcDistortionCorrectionContainer>(topNode, "TpcDistortionCorrectionContainerModuleEdge");
0103   if (m_dcc_module_edge)
0104   {
0105     std::cout << "PrelimDistortionCorrection::get_nodes - found TPC distortion correction container module edge" << std::endl;
0106   }
0107
0108   // input tpc distortion correction static
0109   m_dcc_static = findNode::getClass<TpcDistortionCorrectionContainer>(topNode, "TpcDistortionCorrectionContainerStatic");
0110   if (m_dcc_static)
0111   {
0112     std::cout << "PrelimDistortionCorrection::get_nodes - found TPC distortion correction container static" << std::endl;
0113   }
0114
0115   // input tpc distortion correction average
0116   m_dcc_average = findNode::getClass<TpcDistortionCorrectionContainer>(topNode, "TpcDistortionCorrectionContainerAverage");
0117   if (m_dcc_average)
0118   {
0119     std::cout << "PrelimDistortionCorrection::get_nodes - found TPC distortion correction container average" << std::endl;
0120   }
0121
0122
0123   if(_use_truth_clusters) {
0124     _cluster_map = findNode::getClass<TrkrClusterContainer>(topNode, "TRKR_CLUSTER_TRUTH");
0125   } else {
0126     _cluster_map = findNode::getClass<TrkrClusterContainer>(topNode, "TRKR_CLUSTER");
0127 }
0128
0129   if (!_cluster_map)
0130   {
0131     std::cerr << PHWHERE << "PrelimDistortionCorrection::get_nodes - ERROR: Can't find node TRKR_CLUSTER" << std::endl;
0132     return Fun4AllReturnCodes::ABORTEVENT;
0133   }
0134
0135   _track_map = findNode::getClass<TrackSeedContainer>(topNode, "TpcTrackSeedContainer");
0136   if (!_track_map)
0137   {
0138     std::cerr << PHWHERE << "PrelimDistortionCorrection::get_nodes - ERROR: Can't find TrackSeedContainer " << std::endl;
0139     return Fun4AllReturnCodes::ABORTEVENT;
0140   }
0141
0142   return Fun4AllReturnCodes::EVENT_OK;
0143 }
0144
0145 //________________________________________________________________________________________________________________
0146 int PrelimDistortionCorrection::process_event(PHCompositeNode* /*topNode*/)
0147 {
0148   if(!_pp_mode)
0149   {
0150     std::cout << "PrelimDistortionCorrection called but not in pp_mode, do nothing!" << std::endl;
0151     return Fun4AllReturnCodes::EVENT_OK;
0152   }
0153
0154   PHTimer timer("PrelimDistortionCorrectionTimer");
0155   timer.stop();
0156   timer.restart();
0157
0158   if(Verbosity()>0)
0159   {
0160     std::cout << "starting PrelimDistortionCorrection process_event" << std::endl;
0161   }
0162
0163   // These accumulate trackseeds as we loop over tracks, keylist is a vector of vectors of seed cluskeys
0164   // The seeds are all given to the fitter at once
0165   std::vector<std::vector<TrkrDefs::cluskey>> keylist_A;
0166   PositionMap correctedOffsetTrackClusPositions;   //  this is an std::map<TrkrDefs::cluskey, Acts::Vector3>
0167
0168   for(unsigned int track_it = 0; track_it != _track_map->size(); ++track_it )
0169   {
0170     if(Verbosity()>0) { std::cout << "TPC seed " << track_it << std::endl; }
0171     // if not a TPC track, ignore
0172     TrackSeed* track = _track_map->get(track_it);
0173     if(!track)
0174       {
0175     continue;
0176       }
0177
0178     if(Verbosity() > 0)
0179       {
0180     std::cout << "Input seed pars for " << track_it
0181           << " q " << track->get_charge()
0182           << " qOverR " << std::abs(track->get_qOverR()) * track->get_charge()
0183           << " X0 " << TrackSeedHelper::get_x(track)
0184           << " Y0 " << TrackSeedHelper::get_y(track)
0185           << " Z0 " << TrackSeedHelper::get_z(track)
0186           << " eta " << track->get_eta()
0187           << " phi " << track->get_phi()
0188           << std::endl;
0189       }
0190
0191     const bool is_tpc = std::any_of(
0192       track->begin_cluster_keys(),
0193       track->end_cluster_keys(),
0194       []( const TrkrDefs::cluskey& key ) { return TrkrDefs::getTrkrId(key) == TrkrDefs::tpcId; } );
0195
0196     if(is_tpc)
0197     {
0198       // We want to move all clusters in this seed to point to Z0 = 0
0199       float Z0 = track->get_Z0();
0200       float offset_Z = 0.0 - Z0;
0201
0202       if(Verbosity() > 0)
0203     {
0204       std::cout << "  processing seed with offset_Z " << offset_Z
0205             << " eta " << track->get_eta()
0206         << " x " << TrackSeedHelper::get_x(track)
0207         << " y " << TrackSeedHelper::get_y(track)
0208         << " z " << TrackSeedHelper::get_z(track)
0209             << std::endl;
0210     }
0211
0212       // We want to make  distortion corrections to all clusters in this seed after offsetting the z values
0213       std::vector<TrkrDefs::cluskey> dumvec;
0214       for(TrackSeed::ConstClusterKeyIter iter = track->begin_cluster_keys();
0215       iter != track->end_cluster_keys();
0216       ++iter)
0217     {
0218       TrkrDefs::cluskey cluskey = *iter;
0219       TrkrCluster *cluster = _cluster_map->findCluster(cluskey);
0220
0221     // get global position
0222     const Acts::Vector3 pos = _tgeometry->getGlobalPosition(cluskey, cluster);
0223
0224     // apply z offset
0225     Acts::Vector3 offsetpos(pos.x(), pos.y(), pos.z() + offset_Z);
0226
0227     // apply distortion corrections
0228     if (m_dcc_module_edge)
0229     {
0230       offsetpos = m_distortionCorrection.get_corrected_position(offsetpos, m_dcc_module_edge);
0231     }
0232
0233     if (m_dcc_static)
0234     {
0235       offsetpos = m_distortionCorrection.get_corrected_position(offsetpos, m_dcc_static);
0236     }
0237
0238     if (m_dcc_average)
0239     {
0240       offsetpos = m_distortionCorrection.get_corrected_position(offsetpos, m_dcc_average);
0241     }
0242
0243       // now move the distortion corrected cluster back by offset_Z, to preserve the z measurement info
0244     const Acts::Vector3 corrpos(offsetpos(0), offsetpos(1), offsetpos(2) - offset_Z);
0245       correctedOffsetTrackClusPositions.emplace(cluskey,corrpos);
0246       dumvec.push_back(cluskey);
0247
0248       if(Verbosity() > 0)
0249         {
0250           std::cout << " cluskey " << cluskey << " input pos " << pos(0) << "  " << pos(1) << "  " << pos(2)
0251             << "   corr. pos " << corrpos(0) << "  " << corrpos(1) << "  " << corrpos(2) << std::endl
0252             << "distortion correction " <<  corrpos(0) - pos(0) << "  " << corrpos(1) - pos(1) << "  " << corrpos(2) - pos(2)
0253             << std::endl;
0254
0255         }
0256     }
0257
0258       /// Can't circle fit a seed with less than 3 clusters, skip it
0259       if(dumvec.size() < 3)
0260     { continue; }
0261       keylist_A.push_back(dumvec);
0262
0263       if(Verbosity() > 0)
0264     {
0265       std::cout << "Added  input seed " << track_it << "  becomes output seed " << keylist_A.size() - 1 << std::endl;
0266     }
0267
0268     } // end if TPC seed
0269
0270   }  // end loop over tracks
0271
0272   // reset the seed map for the TPC
0273   _track_map->Reset();
0274
0275   // refit the corrected clusters
0276   std::vector<float> trackChi2;
0277   auto seeds = fitter->ALICEKalmanFilter(keylist_A, true, correctedOffsetTrackClusPositions, trackChi2);
0278
0279   // update the seed parameters on the node tree
0280   // This calls circle fit and line fit (setting x, y, z, phi, eta) and sets qOverR explicitly using q from KF
0281   publishSeeds(seeds.first, correctedOffsetTrackClusPositions);
0282
0283   return Fun4AllReturnCodes::EVENT_OK;
0284 }
0285
0286 //____________________________________________________________________________________________________________
0287 void PrelimDistortionCorrection::publishSeeds(std::vector<TrackSeed_v2>& seeds, const PrelimDistortionCorrection::PositionMap& positions) const
0288 {
0289   int seed_index = 0;
0290   for(auto& seed: seeds )
0291   {
0292     if(seed.size_cluster_keys() < 3)
0293       {
0294     continue;   // ALICEKalmanFilter can drop clusters. Seeds require at least 3 clusters for circle fit
0295       }
0296
0297     /// The ALICEKF gives a better charge determination at high pT
0298     int q = seed.get_charge();
0299     TrackSeedHelper::circleFitByTaubin(&seed,positions, 7, 55);
0300     TrackSeedHelper::lineFit(&seed,positions, 7, 55);
0301
0302     seed.set_qOverR(std::abs(seed.get_qOverR()) * q);
0303     seed.set_phi(TrackSeedHelper::get_phi(&seed,positions));
0304     _track_map->insert(&seed);
0305
0306     if(Verbosity() > 0)
0307       {
0308     std::cout << "Publishing seed " << seed_index
0309           << " q " << q
0310           << " qOverR " << std::abs(seed.get_qOverR()) * q
0311           << " x " << TrackSeedHelper::get_x(&seed)
0312           << " y " << TrackSeedHelper::get_y(&seed)
0313           << " z " << TrackSeedHelper::get_z(&seed)
0314           << " pT " << seed.get_pt()
0315           << " eta " << seed.get_eta()
0316           << " phi " << seed.get_phi()
0317           << std::endl;
0318       }
0319     if(Verbosity() > 5)
0320       {
0321     TrackSeed* readseed = _track_map->get(seed_index);
0322     readseed->identify();
0323       }
0324
0325     seed_index++;
0326   }
0327 }