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