Back to home page

sPhenix code displayed by LXR

 
 

    


File indexing completed on 2025-08-06 08:17:47

0001 /*****************/
0002 /* Cameron Dean  */
0003 /*   LANL 2020   */
0004 /* cdean@bnl.gov */
0005 /*****************/
0006 /*
0007  * Class to append reconstructed events to node tree
0008  */
0009 
0010 // Ideas taken from PHRaveVertexing
0011 
0012 #include "KFParticle_DST.h"
0013 
0014 #include "KFParticle_Container.h"
0015 #include "KFParticle_Tools.h"
0016 #include "KFParticle_truthAndDetTools.h"
0017 
0018 #include <trackbase_historic/SvtxTrack.h>     // for SvtxTrack
0019 #include <trackbase_historic/SvtxTrackMap.h>  // for SvtxTrackMap, SvtxTr...
0020 #include <trackbase_historic/SvtxTrackMap_v2.h>
0021 #include <trackbase_historic/SvtxTrack_v4.h>
0022 
0023 #include <fun4all/Fun4AllReturnCodes.h>
0024 
0025 #include <phool/PHCompositeNode.h>
0026 #include <phool/PHIODataNode.h>
0027 #include <phool/PHNode.h>  // for PHNode
0028 #include <phool/PHNodeIterator.h>
0029 #include <phool/PHObject.h>  // for PHObject
0030 #include <phool/getClass.h>
0031 
0032 #include <KFParticle.h>  // for KFParticle
0033 
0034 #include <cstdlib>   // for exit, size_t, abs
0035 #include <iostream>  // for operator<<, endl
0036 #include <map>       // for map, map<>::mapped_type
0037 #include <utility>   // for pair
0038 
0039 KFParticle_Tools kfpTupleTools_DST;
0040 KFParticle_truthAndDetTools kfpTruthTools_DST;
0041 
0042 int KFParticle_DST::createParticleNode(PHCompositeNode* topNode)
0043 {
0044   PHNodeIterator iter(topNode);
0045 
0046   PHCompositeNode* lowerNode = dynamic_cast<PHCompositeNode*>(iter.findFirst("PHCompositeNode", "DST"));
0047   if (!lowerNode)
0048   {
0049     lowerNode = new PHCompositeNode("DST");
0050     topNode->addNode(lowerNode);
0051     std::cout << "DST node added" << std::endl;
0052   }
0053 
0054   std::string baseName;
0055   std::string trackNodeName;
0056   std::string particleNodeName;
0057 
0058   if (m_container_name.empty())
0059   {
0060     baseName = "reconstructedParticles";
0061   }
0062   else
0063   {
0064     baseName = m_container_name;
0065   }
0066 
0067   // Cant have forward slashes in DST or else you make a subdirectory on save!!!
0068   std::string undrscr = "_";
0069   std::string nothing = "";
0070   std::map<std::string, std::string> forbiddenStrings;
0071   forbiddenStrings["/"] = undrscr;
0072   forbiddenStrings["("] = undrscr;
0073   forbiddenStrings[")"] = nothing;
0074   forbiddenStrings["+"] = "plus";
0075   forbiddenStrings["-"] = "minus";
0076   forbiddenStrings["*"] = "star";
0077   for (auto const& [badString, goodString] : forbiddenStrings)
0078   {
0079     size_t pos;
0080     while ((pos = baseName.find(badString)) != std::string::npos)
0081     {
0082       baseName.replace(pos, 1, goodString);
0083     }
0084   }
0085 
0086   trackNodeName = baseName + "_SvtxTrackMap";
0087   particleNodeName = baseName + "_KFParticle_Container";
0088 
0089   if (m_write_track_container)
0090   {
0091     m_recoTrackMap = new SvtxTrackMap_v2();
0092     PHIODataNode<PHObject>* trackNode = new PHIODataNode<PHObject>(m_recoTrackMap, trackNodeName.c_str(), "PHObject");
0093     lowerNode->addNode(trackNode);
0094     std::cout << trackNodeName << " node added" << std::endl;
0095   }
0096 
0097   if (m_write_particle_container)
0098   {
0099     m_recoParticleMap = new KFParticle_Container();
0100     PHIODataNode<PHObject>* particleNode = new PHIODataNode<PHObject>(m_recoParticleMap, particleNodeName.c_str(), "PHObject");
0101     lowerNode->addNode(particleNode);
0102     std::cout << particleNodeName << " node added" << std::endl;
0103   }
0104 
0105   if (!m_write_track_container && !m_write_particle_container)
0106   {
0107     std::cout << "You have asked to put your selection on the node tree but disabled both the SvtxTrackMap and KFParticle_Container\n";
0108     std::cout << "Check your options" << std::endl;
0109     exit(0);
0110   }
0111 
0112   return Fun4AllReturnCodes::EVENT_OK;
0113 }
0114 
0115 void KFParticle_DST::fillParticleNode(PHCompositeNode* topNode, KFParticle& motherParticle,
0116                                       KFParticle& PV,
0117                                       const std::vector<KFParticle>& daughters,
0118                                       const std::vector<KFParticle>& intermediates)
0119 {
0120   if (m_write_track_container)
0121   {
0122     fillParticleNode_Track(topNode, motherParticle, daughters, intermediates);
0123   }
0124   if (m_write_particle_container)
0125   {
0126     fillParticleNode_Particle(topNode, motherParticle, PV, daughters, intermediates);
0127   }
0128 }
0129 
0130 void KFParticle_DST::fillParticleNode_Track(PHCompositeNode* topNode, KFParticle& motherParticle,
0131                                             std::vector<KFParticle> daughters,
0132                                             std::vector<KFParticle> intermediates)
0133 {
0134   //Make keys for daughters, mothers and intermediates
0135   unsigned int daughterCounter = 0;
0136   unsigned int resonanceCounter = UINT_MAX;
0137 
0138   std::string baseName;
0139   std::string trackNodeName;
0140 
0141   if (m_container_name.empty())
0142   {
0143     baseName = "reconstructedParticles";
0144   }
0145   else
0146   {
0147     baseName = m_container_name;
0148   }
0149 
0150   // Cant have forward slashes in DST or else you make a subdirectory on save!!!
0151   std::string undrscr = "_";
0152   std::string nothing = "";
0153   std::map<std::string, std::string> forbiddenStrings;
0154   forbiddenStrings["/"] = undrscr;
0155   forbiddenStrings["("] = undrscr;
0156   forbiddenStrings[")"] = nothing;
0157   forbiddenStrings["+"] = "plus";
0158   forbiddenStrings["-"] = "minus";
0159   forbiddenStrings["*"] = "star";
0160   for (auto const& [badString, goodString] : forbiddenStrings)
0161   {
0162     size_t pos;
0163     while ((pos = baseName.find(badString)) != std::string::npos)
0164     {
0165       baseName.replace(pos, 1, goodString);
0166     }
0167   }
0168 
0169   trackNodeName = baseName + "_SvtxTrackMap";
0170 
0171   m_recoTrackMap = findNode::getClass<SvtxTrackMap>(topNode, trackNodeName.c_str());
0172 
0173   SvtxTrack* m_recoTrack = new SvtxTrack_v4();
0174 
0175   m_recoTrack = buildSvtxTrack(motherParticle);
0176 
0177   SvtxTrack *dummyMother = nullptr;  
0178   while (!dummyMother)
0179   {
0180     if (!m_recoTrackMap->get(resonanceCounter))
0181     {
0182       dummyMother = m_recoTrackMap->insertWithKey(m_recoTrack, resonanceCounter);
0183     }
0184     --resonanceCounter;
0185   }
0186   m_recoTrack->Reset();
0187 
0188   if (m_has_intermediates_DST)
0189   {
0190     KFParticle* intermediateArray = &intermediates[0];
0191 
0192     for (unsigned int k = 0; k < intermediates.size(); ++k)
0193     {
0194       m_recoTrack = buildSvtxTrack(intermediateArray[k]);
0195       SvtxTrack *dummyIntermediate = nullptr;  
0196       while (!dummyIntermediate)
0197       {
0198         if (!m_recoTrackMap->get(resonanceCounter))
0199     {
0200       dummyIntermediate = m_recoTrackMap->insertWithKey(m_recoTrack, resonanceCounter);
0201     }
0202         --resonanceCounter;
0203       }
0204       m_recoTrack->Reset();
0205     }
0206   }
0207 
0208   SvtxTrackMap* originalTrackMap = findNode::getClass<SvtxTrackMap>(topNode, m_origin_track_map_node_name.c_str());
0209   KFParticle* daughterArray = &daughters[0];
0210 
0211   for (unsigned int k = 0; k < daughters.size(); ++k)
0212   {
0213     if (originalTrackMap->size() == 0)
0214     {
0215       std::cout << "There was no original track map found, the tracks will have no cluster information!" << std::endl;
0216       m_recoTrack = buildSvtxTrack(daughterArray[k]);
0217 
0218       SvtxTrack *dummyDaughter = nullptr;  
0219       while (!dummyDaughter)
0220       {
0221         if (!m_recoTrackMap->get(daughterCounter))
0222     {
0223           dummyDaughter = m_recoTrackMap->insertWithKey(m_recoTrack, daughterCounter);
0224     }
0225         ++daughterCounter;
0226       }
0227     }
0228     else
0229     {
0230       m_recoTrack = kfpTruthTools_DST.getTrack(daughterArray[k].Id(), originalTrackMap);
0231       if (!m_recoTrackMap->get(daughterArray[k].Id()))
0232       {
0233         m_recoTrackMap->insertWithKey(m_recoTrack, daughterArray[k].Id());
0234       }
0235     }
0236   }
0237 }
0238 
0239 void KFParticle_DST::fillParticleNode_Particle(PHCompositeNode* topNode, KFParticle& motherParticle,
0240                                                KFParticle& PV,
0241                                                std::vector<KFParticle> daughters,
0242                                                std::vector<KFParticle> intermediates)
0243 {
0244   std::string baseName;
0245   std::string particleNodeName;
0246 
0247   if (m_container_name.empty())
0248   {
0249     baseName = "reconstructedParticles";
0250   }
0251   else
0252   {
0253     baseName = m_container_name;
0254   }
0255 
0256   // Cant have forward slashes in DST or else you make a subdirectory on save!!!
0257   std::string undrscr = "_";
0258   std::string nothing = "";
0259   std::map<std::string, std::string> forbiddenStrings;
0260   forbiddenStrings["/"] = undrscr;
0261   forbiddenStrings["("] = undrscr;
0262   forbiddenStrings[")"] = nothing;
0263   forbiddenStrings["+"] = "plus";
0264   forbiddenStrings["-"] = "minus";
0265   forbiddenStrings["*"] = "star";
0266   for (auto const& [badString, goodString] : forbiddenStrings)
0267   {
0268     size_t pos;
0269     while ((pos = baseName.find(badString)) != std::string::npos)
0270     {
0271       baseName.replace(pos, 1, goodString);
0272     }
0273   }
0274 
0275   particleNodeName = baseName + "_KFParticle_Container";
0276 
0277   m_recoParticleMap = findNode::getClass<KFParticle_Container>(topNode, particleNodeName.c_str());
0278 
0279   motherParticle.SetProductionVertex(PV);
0280   motherParticle.TransportToDecayVertex();
0281 
0282   KFParticle* intermediateArray = &intermediates[0];
0283   if (m_has_intermediates_DST)
0284   {
0285     for (unsigned int k = 0; k < intermediates.size(); ++k)
0286     {
0287       intermediateArray[k].SetProductionVertex(motherParticle);
0288       intermediateArray[k].TransportToDecayVertex();
0289     }
0290   }
0291 
0292   KFParticle* daughterArray = &daughters[0];
0293   for (unsigned int k = 0; k < daughters.size(); ++k)
0294   {
0295       bool didntSetTrackToIntermediate = true;
0296       for (auto& intermediate : intermediates)
0297       {
0298         const std::vector<int> daughterIDs = intermediate.DaughterIds();
0299         for (auto& id : daughterIDs) 
0300         {
0301           if (daughterArray[k].Id() == id)
0302           {
0303             didntSetTrackToIntermediate = false;
0304             daughterArray[k].SetProductionVertex(intermediate);
0305           }
0306         }
0307       }
0308 
0309       if (didntSetTrackToIntermediate)
0310       {
0311         daughterArray[k].SetProductionVertex(motherParticle);
0312       }
0313 
0314     m_recoParticleMap->insert(&daughterArray[k]);
0315   }
0316 
0317   if (m_has_intermediates_DST)
0318   {
0319     for (unsigned int k = 0; k < intermediates.size(); ++k)
0320     {
0321       intermediateArray[k].TransportToProductionVertex();
0322       m_recoParticleMap->insert(&intermediateArray[k]);
0323     }
0324   }
0325 
0326   motherParticle.TransportToProductionVertex();
0327   m_recoParticleMap->insert(&motherParticle);
0328 }
0329 
0330 SvtxTrack* KFParticle_DST::buildSvtxTrack(const KFParticle& particle)
0331 {
0332   SvtxTrack* track = new SvtxTrack_v4();
0333 
0334   track->set_id(std::abs(particle.GetPDG()));
0335   track->set_charge((int) particle.GetQ());
0336   track->set_chisq(particle.GetChi2());
0337   track->set_ndf(particle.GetNDF());
0338 
0339   track->set_x(particle.GetX());
0340   track->set_y(particle.GetY());
0341   track->set_z(particle.GetZ());
0342 
0343   track->set_px(particle.GetPx());
0344   track->set_py(particle.GetPy());
0345   track->set_pz(particle.GetPz());
0346 
0347   for (int i = 0; i < 6; ++i)
0348   {
0349     for (int j = 0; j < 6; ++j)
0350     {
0351       track->set_error(i, j, particle.GetCovariance(i, j));
0352     }
0353   }
0354 
0355   return track;
0356 }
0357 
0358 void KFParticle_DST::printNode(PHCompositeNode* topNode)
0359 {
0360   std::string baseName;
0361   std::string trackNodeName;
0362   std::string particleNodeName;
0363 
0364   if (m_container_name.empty())
0365   {
0366     baseName = "reconstructedParticles";
0367   }
0368   else
0369   {
0370     baseName = m_container_name;
0371   }
0372 
0373   // Cant have forward slashes in DST or else you make a subdirectory on save!!!
0374   std::string undrscr = "_";
0375   std::string nothing = "";
0376   std::map<std::string, std::string> forbiddenStrings;
0377   forbiddenStrings["/"] = undrscr;
0378   forbiddenStrings["("] = undrscr;
0379   forbiddenStrings[")"] = nothing;
0380   forbiddenStrings["+"] = "plus";
0381   forbiddenStrings["-"] = "minus";
0382   forbiddenStrings["*"] = "star";
0383   for (auto const& [badString, goodString] : forbiddenStrings)
0384   {
0385     size_t pos;
0386     while ((pos = baseName.find(badString)) != std::string::npos)
0387     {
0388       baseName.replace(pos, 1, goodString);
0389     }
0390   }
0391 
0392   if (m_write_track_container)
0393   {
0394     trackNodeName = baseName + "_SvtxTrackMap";
0395     std::cout << "----------------";
0396     std::cout << " KFParticle_DST: " << trackNodeName << " information ";
0397     std::cout << "----------------" << std::endl;
0398     SvtxTrackMap* trackmap = findNode::getClass<SvtxTrackMap>(topNode, trackNodeName.c_str());
0399     for (auto& iter : *trackmap)
0400     {
0401       SvtxTrack* track = iter.second;
0402       track->identify();
0403     }
0404     std::cout << "--------------------------------------------------------------------------------------------------" << std::endl;
0405   }
0406 
0407   if (m_write_particle_container)
0408   {
0409     particleNodeName = baseName + "_KFParticle_Container";
0410     std::cout << "----------------";
0411     std::cout << " KFParticle_DST: " << particleNodeName << " information ";
0412     std::cout << "----------------" << std::endl;
0413     KFParticle_Container* particlemap = findNode::getClass<KFParticle_Container>(topNode, particleNodeName.c_str());
0414     for (auto& iter : *particlemap)
0415     {
0416       KFParticle* particle = iter.second;
0417       kfpTupleTools_DST.identify(*particle);
0418     }
0419     std::cout << "--------------------------------------------------------------------------------------------------" << std::endl;
0420   }
0421 }