dstmacro/data/Charge_Recalc.C

0001 #ifndef CHARGE_RECALC_C
0002 #define CHARGE_RECALC_C
0003
0004 #include <fun4all/SubsysReco.h>
0005 #include <fun4all/Fun4AllReturnCodes.h>
0006
0007 #include <phool/PHCompositeNode.h>
0008 #include <phool/getClass.h>
0009
0010 #include <trackbase/ActsGeometry.h>
0011 #include <trackbase/TrkrClusterContainer.h>
0012
0013 #include <trackbase_historic/SvtxTrackMap.h>
0014 #include <trackbase_historic/SvtxTrack.h>
0015 #include <trackbase_historic/TrackSeed.h> // explicit (avoid relying on indirect includes)
0016
0017 #include <globalvertex/SvtxVertexMap.h>
0018 #include <globalvertex/SvtxVertex.h>
0019
0020 #include <cmath>
0021 #include <limits>
0022 #include <iostream>
0023
0024 static inline float WrapPhi(float x)
0025 {
0026     while (x > M_PI)
0027         x -= 2.0f * M_PI;
0028     while (x < -M_PI)
0029         x += 2.0f * M_PI;
0030     return x;
0031 }
0032
0033 static inline float DeltaPhi(float a, float b)
0034 {
0035     return WrapPhi(a - b);
0036 }
0037
0038 class ChargeRecalc : public SubsysReco
0039 {
0040 public:
0041     ChargeRecalc(const std::string &name = "ChargeRecalc")
0042         : SubsysReco(name) {}
0043     void setVerbosity(int v) { m_verb = v; }
0044
0045     int process_event(PHCompositeNode *topNode) override
0046     {
0047         auto clustermap = findNode::getClass<TrkrClusterContainer>(topNode, m_clusterContainerName);
0048         auto geometry = findNode::getClass<ActsGeometry>(topNode, m_actsgeometryName);
0049         auto vertexmap = findNode::getClass<SvtxVertexMap>(topNode, m_vertexMapName);
0050         auto trackmap = findNode::getClass<SvtxTrackMap>(topNode, m_trackMapName);
0051         if (!trackmap)
0052         {
0053             std::cout << PHWHERE << "Missing trackmap, can't continue" << std::endl;
0054             return Fun4AllReturnCodes::EVENT_OK;
0055         }
0056         if (!clustermap)
0057         {
0058             std::cout << PHWHERE << "Missing clustermap, can't continue" << std::endl;
0059             return Fun4AllReturnCodes::EVENT_OK;
0060         }
0061         if (!geometry)
0062         {
0063             std::cout << PHWHERE << "Missing geometry, can't continue" << std::endl;
0064             return Fun4AllReturnCodes::EVENT_OK;
0065         }
0066         if (!vertexmap)
0067         {
0068             std::cout << PHWHERE << "NO vertex map, can't continue" << std::endl;
0069         }
0070
0071         // if (m_verb > 0)
0072         // {
0073         //     float vx = NAN, vy = NAN, vz = NAN;
0074         //     int vc = 999;
0075
0076         //     if (vertexmap && !vertexmap->empty())
0077         //     {
0078         //         SvtxVertex *vtx = vertexmap->begin()->second;
0079         //         if (vtx)
0080         //         {
0081         //             vx = vtx->get_x();
0082         //             vy = vtx->get_y();
0083         //             vz = vtx->get_z();
0084         //             vc = vtx->get_beam_crossing();
0085         //         }
0086         //     }
0087
0088         //     std::cout << "[ChargeRecalc] vtx(first): x=" << vx
0089         //               << " y=" << vy << " z=" << vz
0090         //               << " crossing=" << vc
0091         //               << std::endl;
0092         // }
0093
0094         std::cout << "vtx size : " << vertexmap->size() << " " << trackmap->size() << std::endl;
0095         if (vertexmap && !vertexmap->empty())
0096         {
0097             SvtxVertex *vtx0 = nullptr;
0098
0099             // crossing==0 has high priority
0100             int idx = 0;
0101             for (const auto &[k, v] : *vertexmap)
0102             {
0103                 // std::cout<<"vtx : "<<idx<<" "
0104                 //          <<v->get_id()<<" "
0105                 //          <<v->get_x()<<" "
0106                 //          <<v->get_y()<<" "
0107                 //          <<v->get_z()<<" "
0108                 //          <<v->get_beam_crossing()<<" "
0109                 //          <<std::endl;
0110                 idx++;
0111             }
0112         }
0113
0114         int n_changed = 0;
0115         int n_skipped = 0;
0116         int n_total = 0;
0117
0118         for (auto &iter : *trackmap)
0119         {
0120             SvtxTrack *track = iter.second;
0121             if (!track)
0122                 continue;
0123             n_total++;
0124
0125             TrackSeed *si_seed = track->get_silicon_seed();
0126             if (!si_seed)
0127             {
0128                 n_skipped++;
0129                 continue;
0130             }
0131
0132             float xvtx = 0.0f;
0133             float yvtx = 0.0f;
0134
0135             bool found_vtx = false;
0136             if (vertexmap && !vertexmap->empty())
0137             {
0138                 SvtxVertex *vtx0 = nullptr;
0139
0140                 // crossing==0 has high priority
0141                 for (const auto &[k, v] : *vertexmap)
0142                 {
0143                     if (v && v->get_beam_crossing() == 0)
0144                     {
0145                         vtx0 = v;
0146                         break;
0147                     }
0148                 }
0149
0150                 if (!vtx0)
0151                     vtx0 = vertexmap->begin()->second;
0152
0153                 if (vtx0)
0154                 {
0155                     xvtx = vtx0->get_x();
0156                     yvtx = vtx0->get_y();
0157                     found_vtx = true;
0158                 }
0159             }
0160
0161             if (m_verb > 0 && !found_vtx)
0162             {
0163                 std::cout << "[ChargeRecalc] WARNING: no usable vertex (vertexmap="
0164                           << (void *)vertexmap
0165                           << ", size=" << (vertexmap ? (int)vertexmap->size() : -1)
0166                           << ")" << std::endl;
0167                 n_skipped++;
0168                 continue;
0169             }
0170
0171             // find innermost/outermost cluster phi from vertex
0172             float rmin = std::numeric_limits<float>::max();
0173             float rmax = -1.0f;
0174             float phi_rmin = 0.0f;
0175             float phi_rmax = 0.0f;
0176
0177             for (auto key_iter = si_seed->begin_cluster_keys(); key_iter != si_seed->end_cluster_keys(); ++key_iter)
0178             {
0179                 const auto &cluster_key = *key_iter;
0180                 auto trkrCluster = clustermap->findCluster(cluster_key);
0181                 if (!trkrCluster)
0182                     continue;
0183
0184                 auto global = geometry->getGlobalPosition(cluster_key, trkrCluster);
0185
0186                 const float dx = global[0] - xvtx;
0187                 const float dy = global[1] - yvtx;
0188                 const float r = sqrt(dx * dx + dy * dy);
0189                 const float phi = std::atan2(dy, dx);
0190
0191                 if (r < rmin)
0192                 {
0193                     rmin = r;
0194                     phi_rmin = phi;
0195                 }
0196                 if (r > rmax)
0197                 {
0198                     rmax = r;
0199                     phi_rmax = phi;
0200                 }
0201             }
0202
0203             if (rmax < 0.0f)
0204             {
0205                 n_skipped++;
0206                 continue;
0207             }
0208
0209             const float dphi = DeltaPhi(phi_rmax, phi_rmin);
0210             const int q_new = (dphi > 0.0f) ? -1 : +1; // if official is flipped, swap -1/+1
0211             const int q_old = track->get_charge();
0212
0213             if (q_new != q_old)
0214             {
0215                 if (m_verb > 0)
0216                 {
0217                     std::cout << "[ChargeRecalc] trkid=" << track->get_id()
0218                               << " before=" << track->get_charge();
0219                 }
0220                 track->set_charge(q_new);
0221                 if (m_verb > 0)
0222                 {
0223                     std::cout << " after=" << track->get_charge() << std::endl;
0224                 }
0225                 n_changed++;
0226             }
0227         }
0228         if (m_verb > 0)
0229         {
0230             std::cout << "[ChargeRecalc] total=" << n_total
0231                       << " changed=" << n_changed
0232                       << " skipped=" << n_skipped
0233                       << std::endl;
0234         }
0235         return Fun4AllReturnCodes::EVENT_OK;
0236     }
0237
0238 private:
0239     // node names (match SiliconSeedsAna default)
0240     std::string m_clusterContainerName = "TRKR_CLUSTER";
0241     std::string m_actsgeometryName = "ActsGeometry";
0242     std::string m_trackMapName = "SvtxTrackMap";
0243     std::string m_vertexMapName = "SvtxVertexMap";
0244
0245     int m_verb = 0;
0246 };
0247
0248 #endif