TrackFindingExaTrkX/src/createFeatures.cpp

0001 // This file is part of the ACTS project.
0002 //
0003 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008
0009 #include "createFeatures.hpp"
0010
0011 #include "Acts/Utilities/AngleHelpers.hpp"
0012 #include "Acts/Utilities/VectorHelpers.hpp"
0013
0014 namespace ActsExamples {
0015
0016 std::vector<float> createFeatures(
0017     const SimSpacePointContainer& spacepoints,
0018     const std::optional<ClusterContainer>& clusters,
0019     const std::vector<TrackFindingAlgorithmExaTrkX::NodeFeature>& nodeFeatures,
0020     const std::vector<float>& featureScales) {
0021   using namespace ActsExamples;
0022
0023   assert(nodeFeatures.size() == featureScales.size());
0024   std::vector<float> features(spacepoints.size() * nodeFeatures.size());
0025
0026   for (auto isp = 0ul; isp < spacepoints.size(); ++isp) {
0027     const auto& sp = spacepoints[isp];
0028
0029     // For now just take the first index since does require one single index
0030     // per spacepoint
0031     // TODO does it work for the module map construction to use only the first
0032     // sp?
0033     const auto& sl1 = sp.sourceLinks()[0].template get<IndexSourceLink>();
0034
0035     // This should be fine, because check in constructor
0036     const Cluster* cl1 = clusters ? &clusters->at(sl1.index()) : nullptr;
0037     const Cluster* cl2 = cl1;
0038
0039     if (sp.sourceLinks().size() == 2) {
0040       const auto& sl2 = sp.sourceLinks()[1].template get<IndexSourceLink>();
0041       cl2 = clusters ? &clusters->at(sl2.index()) : nullptr;
0042     }
0043
0044     // I would prefer to use a std::span or boost::span here once available
0045     float* f = features.data() + isp * nodeFeatures.size();
0046
0047     using NF = TrackFindingAlgorithmExaTrkX::NodeFeature;
0048
0049     using namespace Acts::VectorHelpers;
0050     using namespace Acts::AngleHelpers;
0051
0052     // clang-format off
0053 #define MAKE_CLUSTER_FEATURES(n) \
0054     break; case NF::eCluster##n##X:   f[ift] = cl##n->globalPosition[Acts::ePos0]; \
0055     break; case NF::eCluster##n##Y:   f[ift] = cl##n->globalPosition[Acts::ePos1]; \
0056     break; case NF::eCluster##n##R:   f[ift] = perp(cl##n->globalPosition); \
0057     break; case NF::eCluster##n##Phi: f[ift] = phi(cl##n->globalPosition); \
0058     break; case NF::eCluster##n##Z:   f[ift] = cl##n->globalPosition[Acts::ePos2]; \
0059     break; case NF::eCluster##n##Eta: f[ift] = eta(cl##n->globalPosition); \
0060     break; case NF::eCellCount##n:    f[ift] = cl##n->channels.size(); \
0061     break; case NF::eChargeSum##n:    f[ift] = cl##n->sumActivations(); \
0062     break; case NF::eLocDir0##n:      f[ift] = cl##n->localDirection[0]; \
0063     break; case NF::eLocDir1##n:      f[ift] = cl##n->localDirection[1]; \
0064     break; case NF::eLocDir2##n:      f[ift] = cl##n->localDirection[2]; \
0065     break; case NF::eLengthDir0##n:   f[ift] = cl##n->lengthDirection[0]; \
0066     break; case NF::eLengthDir1##n:   f[ift] = cl##n->lengthDirection[1]; \
0067     break; case NF::eLengthDir2##n:   f[ift] = cl##n->lengthDirection[2]; \
0068     break; case NF::eLocEta##n:       f[ift] = cl##n->localEta; \
0069     break; case NF::eLocPhi##n:       f[ift] = cl##n->localPhi; \
0070     break; case NF::eGlobEta##n:      f[ift] = cl##n->globalEta; \
0071     break; case NF::eGlobPhi##n:      f[ift] = cl##n->globalPhi; \
0072     break; case NF::eEtaAngle##n:     f[ift] = cl##n->etaAngle; \
0073     break; case NF::ePhiAngle##n:     f[ift] = cl##n->phiAngle;
0074     // clang-format on
0075
0076     Acts::Vector3 spPos{sp.x(), sp.y(), sp.z()};
0077
0078     for (auto ift = 0ul; ift < nodeFeatures.size(); ++ift) {
0079       // clang-format off
0080       switch(nodeFeatures[ift]) {
0081         // Spacepoint features
0082         break; case NF::eR:           f[ift] = perp(spPos);
0083         break; case NF::ePhi:         f[ift] = phi(spPos);
0084         break; case NF::eZ:           f[ift] = sp.z();
0085         break; case NF::eX:           f[ift] = sp.x();
0086         break; case NF::eY:           f[ift] = sp.y();
0087         break; case NF::eEta:         f[ift] = eta(spPos);
0088         // Single cluster features
0089         break; case NF::eClusterLoc0: f[ift] = cl1->sizeLoc0;
0090         break; case NF::eClusterLoc1: f[ift] = cl1->sizeLoc1;
0091         break; case NF::eCellCount:   f[ift] = cl1->channels.size();
0092         break; case NF::eChargeSum:   f[ift] = cl1->sumActivations();
0093         // Features for split clusters
0094         MAKE_CLUSTER_FEATURES(1)
0095         MAKE_CLUSTER_FEATURES(2)
0096       }
0097       // clang-format on
0098
0099       assert(std::isfinite(f[ift]));
0100       f[ift] /= featureScales[ift];
0101     }
0102 #undef MAKE_CLUSTER_FEATURES
0103   }
0104
0105   return features;
0106 }
0107
0108 }  // namespace ActsExamples