packages/trackreco/PrelimDistortionCorrectionAuAu.cc

0001 /*!
0002  *  \File PrelimDistortionCorrectionAuAu.cc
0003  *  \brief  Refits AuAu TPC seeds with distortion corrections and current calibrations
0004  *  \author Tony Frawley
0005  */
0006
0007 #include "PrelimDistortionCorrectionAuAu.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 PrelimDistortionCorrectionAuAu::PrelimDistortionCorrectionAuAu(const std::string& name)
0047   : SubsysReco(name)
0048 {}
0049
0050 //___________________________________________________________________________________________
0051 int PrelimDistortionCorrectionAuAu::End(PHCompositeNode* /*unused*/)
0052 {
0053   return Fun4AllReturnCodes::EVENT_OK;
0054 }
0055
0056 //___________________________________________________________________________________________
0057 int PrelimDistortionCorrectionAuAu::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 PrelimDistortionCorrectionAuAu::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 << "PrelimDistortionCorrectionAuAu::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 << "PrelimDistortionCorrectionAuAu::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 << "PrelimDistortionCorrectionAuAu::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 << "PrelimDistortionCorrectionAuAu::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 << "PrelimDistortionCorrectionAuAu::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 PrelimDistortionCorrectionAuAu::process_event(PHCompositeNode* /*topNode*/)
0147 {
0148   if(_pp_mode)
0149   {
0150     std::cout << "PrelimDistortionCorrectionAuAu called but in pp_mode, do nothing!" << std::endl;
0151     return Fun4AllReturnCodes::EVENT_OK;
0152   }
0153
0154   PHTimer timer("PrelimDistortionCorrectionAuAuTimer");
0155   timer.stop();
0156   timer.restart();
0157
0158   if(Verbosity()>0)
0159   {
0160     std::cout << "starting PrelimDistortionCorrectionAuAu 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();   // for pp
0200       float Z0 = 0;                                // for AuAu
0201       float offset_Z = 0.0 - Z0;
0202
0203       if(Verbosity() > 0)
0204     {
0205       std::cout << "  processing seed with offset_Z " << offset_Z
0206             << " eta " << track->get_eta()
0207         << " x " << TrackSeedHelper::get_x(track)
0208         << " y " << TrackSeedHelper::get_y(track)
0209         << " z " << TrackSeedHelper::get_z(track)
0210             << std::endl;
0211     }
0212
0213       // We want to make  distortion corrections to all clusters in this seed after offsetting the z values
0214       std::vector<TrkrDefs::cluskey> dumvec;
0215       for(TrackSeed::ConstClusterKeyIter iter = track->begin_cluster_keys();
0216       iter != track->end_cluster_keys();
0217       ++iter)
0218     {
0219       TrkrDefs::cluskey cluskey = *iter;
0220       TrkrCluster *cluster = _cluster_map->findCluster(cluskey);
0221
0222     // get global position
0223     const Acts::Vector3 pos = _tgeometry->getGlobalPosition(cluskey, cluster);
0224
0225     // apply z offset
0226     Acts::Vector3 offsetpos(pos.x(), pos.y(), pos.z() + offset_Z);
0227
0228     // apply distortion corrections
0229     if (m_dcc_module_edge)
0230     {
0231       offsetpos = m_distortionCorrection.get_corrected_position(offsetpos, m_dcc_module_edge);
0232     }
0233
0234     if (m_dcc_static)
0235     {
0236       offsetpos = m_distortionCorrection.get_corrected_position(offsetpos, m_dcc_static);
0237     }
0238
0239     if (m_dcc_average)
0240     {
0241       offsetpos = m_distortionCorrection.get_corrected_position(offsetpos, m_dcc_average);
0242     }
0243
0244       // now move the distortion corrected cluster back by offset_Z, to preserve the z measurement info
0245     const Acts::Vector3 corrpos(offsetpos(0), offsetpos(1), offsetpos(2) - offset_Z);
0246       correctedOffsetTrackClusPositions.emplace(cluskey,corrpos);
0247       dumvec.push_back(cluskey);
0248
0249       if(Verbosity() > 0)
0250         {
0251           std::cout << " cluskey " << cluskey << " input pos " << pos(0) << "  " << pos(1) << "  " << pos(2)
0252             << "   corr. pos " << corrpos(0) << "  " << corrpos(1) << "  " << corrpos(2) << std::endl
0253             << "distortion correction " <<  corrpos(0) - pos(0) << "  " << corrpos(1) - pos(1) << "  " << corrpos(2) - pos(2)
0254             << std::endl;
0255
0256         }
0257     }
0258
0259       /// Can't circle fit a seed with less than 3 clusters, skip it
0260       if(dumvec.size() < 3)
0261     { continue; }
0262       keylist_A.push_back(dumvec);
0263
0264       if(Verbosity() > 0)
0265     {
0266       std::cout << "Added  input seed " << track_it << "  becomes output seed " << keylist_A.size() - 1 << std::endl;
0267     }
0268
0269     } // end if TPC seed
0270
0271   }  // end loop over tracks
0272
0273   // reset the seed map for the TPC
0274   _track_map->Reset();
0275
0276   // refit the corrected clusters
0277   std::vector<float> trackChi2;
0278   auto seeds = fitter->ALICEKalmanFilter(keylist_A, true, correctedOffsetTrackClusPositions, trackChi2);
0279
0280   // update the seed parameters on the node tree
0281   // This calls circle fit and line fit (setting x, y, z, phi, eta) and sets qOverR explicitly using q from KF
0282   publishSeeds(seeds.first, correctedOffsetTrackClusPositions);
0283
0284   return Fun4AllReturnCodes::EVENT_OK;
0285 }
0286
0287 //____________________________________________________________________________________________________________
0288 void PrelimDistortionCorrectionAuAu::publishSeeds(std::vector<TrackSeed_v2>& seeds, const PrelimDistortionCorrectionAuAu::PositionMap& positions) const
0289 {
0290   int seed_index = 0;
0291   for(auto& seed: seeds )
0292   {
0293     if(seed.size_cluster_keys() < 3)
0294       {
0295     continue;   // ALICEKalmanFilter can drop clusters. Seeds require at least 3 clusters for circle fit
0296       }
0297
0298     /// The ALICEKF gives a better charge determination at high pT
0299     int q = seed.get_charge();
0300     TrackSeedHelper::circleFitByTaubin(&seed,positions, 7, 55);
0301     TrackSeedHelper::lineFit(&seed,positions, 7, 55);
0302
0303     seed.set_qOverR(std::abs(seed.get_qOverR()) * q);
0304     seed.set_phi(TrackSeedHelper::get_phi(&seed,positions));
0305     _track_map->insert(&seed);
0306
0307     if(Verbosity() > 0)
0308       {
0309     std::cout << "Publishing seed " << seed_index
0310           << " q " << q
0311           << " qOverR " << std::abs(seed.get_qOverR()) * q
0312           << " x " << TrackSeedHelper::get_x(&seed)
0313           << " y " << TrackSeedHelper::get_y(&seed)
0314           << " z " << TrackSeedHelper::get_z(&seed)
0315           << " pT " << seed.get_pt()
0316           << " eta " << seed.get_eta()
0317           << " phi " << seed.get_phi()
0318           << std::endl;
0319       }
0320     if(Verbosity() > 5)
0321       {
0322     TrackSeed* readseed = _track_map->get(seed_index);
0323     readseed->identify();
0324       }
0325
0326     seed_index++;
0327   }
0328 }