|
30 | 30 | #include <array> |
31 | 31 | #include <cmath> |
32 | 32 | #include <cstdint> |
| 33 | +#include <utility> |
| 34 | +#include <vector> |
33 | 35 |
|
34 | 36 | namespace o2::aod |
35 | 37 | { |
@@ -165,6 +167,110 @@ DECLARE_SOA_DYNAMIC_COLUMN(IsTagged, isTagged, [](uint16_t bit, BJetTaggingMetho |
165 | 167 | JETTAGGING_TABLES_DEF(Charged, "C"); |
166 | 168 | JETTAGGING_TABLES_DEF(Full, "F"); |
167 | 169 |
|
| 170 | +// Tables for full inference of GNN-based b-jet tagger (including track origin, vertex predictions) |
| 171 | +namespace jetgnnindex |
| 172 | +{ |
| 173 | +DECLARE_SOA_INDEX_COLUMN(JTrack, track); //! index in JTracks of this GNN input track (not the parent AOD Tracks index) |
| 174 | +DECLARE_SOA_COLUMN(TrackIndex, trackIndex, uint8_t); //! node index (0..255) in the GNN input after track selection/truncation; need not match the constituent-array position |
| 175 | +// DECLARE_SOA_COLUMN(EdgeIndex, edgeIndex, uint8_t); //! unique in-jet edge index for the (i,j) constituent-track pair, i<j, triangular order i*(2N-i-1)/2+(j-i-1) |
| 176 | +// DECLARE_SOA_DYNAMIC_COLUMN(TrackIds, trackIds, //! inverse of the triangular edge indexing above: for n constituent tracks, recovers the (i,j) in-jet track index pair (i<j) this edge connects |
| 177 | +// [](uint8_t edgeIndex, int n) -> std::pair<uint8_t, uint8_t> { |
| 178 | +// int i = 0; |
| 179 | +// int offset = 0; |
| 180 | +// while (offset + (n - 1 - i) <= static_cast<int>(edgeIndex)) { |
| 181 | +// offset += (n - 1 - i); |
| 182 | +// i++; |
| 183 | +// } |
| 184 | +// int j = i + 1 + (static_cast<int>(edgeIndex) - offset); |
| 185 | +// return std::pair<uint8_t, uint8_t>{static_cast<uint8_t>(i), static_cast<uint8_t>(j)}; |
| 186 | +// }); |
| 187 | +} // namespace jetgnnindex |
| 188 | + |
| 189 | +namespace jetgnnlabel |
| 190 | +{ |
| 191 | +DECLARE_SOA_COLUMN(Origin, origin, uint8_t); //! trkOrigin from vertexClustering() (0: non-physical primary, 1: charm, 2: beauty, 3: primary vertex, 4: other secondary) |
| 192 | +DECLARE_SOA_COLUMN(VertexIndex, vertexIndex, int8_t); //! trkVtxIndex from vertexClustering() (-1: unassigned, 0: primary vertex, >0: secondary vertex, ranked by distance to PV) |
| 193 | +// DECLARE_SOA_COLUMN(SameVertex, sameVertex, bool); //! true if the (i,j) constituent-track pair of this edge share the same in-jet vertex cluster (same trkVtxIndex from vertexClustering()) |
| 194 | +} // namespace jetgnnlabel |
| 195 | + |
| 196 | +namespace jetgnnscore |
| 197 | +{ |
| 198 | +DECLARE_SOA_COLUMN(ScoreML, scoreML, float); //! GNN discriminant: one track-origin D_b per GnnTracks row |
| 199 | +} // namespace jetgnnscore |
| 200 | + |
| 201 | +// GnnTracks has one row per evaluated (jet, track) pair. Its jet and track indices |
| 202 | +// identify the pair independently of the ordering in JetConstituents. |
| 203 | +// GnnTrackIds is a positional jet extension: write one row for every jet in |
| 204 | +// Jets order, including an empty array for jets without evaluated tracks. Store |
| 205 | +// GnnTracks row indices (e.g. the backbone producer's lastIndex()), not JTrackIds. |
| 206 | +#define DECLARE_JET_GNNTRACKINDEX_TABLE(_jet_type_, _description_) \ |
| 207 | + namespace _jet_type_##gnntrackindextable \ |
| 208 | + { \ |
| 209 | + DECLARE_SOA_INDEX_COLUMN(_jet_type_, jet); \ |
| 210 | + } \ |
| 211 | + DECLARE_SOA_TABLE(_jet_type_##GnnTracks, "AOD", _description_, \ |
| 212 | + o2::soa::Index<>, \ |
| 213 | + _jet_type_##gnntrackindextable::_jet_type_##Id, \ |
| 214 | + jetgnnindex::JTrackId, \ |
| 215 | + jetgnnindex::TrackIndex); \ |
| 216 | + namespace _jet_type_##gnntrackindices \ |
| 217 | + { \ |
| 218 | + DECLARE_SOA_ARRAY_INDEX_COLUMN_FULL(_jet_type_##GnnTrack, tracksGnn, int32_t, _jet_type_##GnnTracks, ""); \ |
| 219 | + } \ |
| 220 | + DECLARE_SOA_TABLE(_jet_type_##GnnTrackIds, "AOD", _description_ "I", \ |
| 221 | + _jet_type_##gnntrackindices::_jet_type_##GnnTrackIds); |
| 222 | + |
| 223 | +/* |
| 224 | +#define DECLARE_JET_GNNEDGEINDEX_TABLE(_jet_type_, _description_) \ |
| 225 | + namespace _jet_type_##gnnedgeindextable \ |
| 226 | + { \ |
| 227 | + DECLARE_SOA_INDEX_COLUMN(_jet_type_, jet); \ |
| 228 | + } \ |
| 229 | + DECLARE_SOA_TABLE(_jet_type_##GnnEdges, "AOD", _description_, \ |
| 230 | + o2::soa::Index<>, \ |
| 231 | + _jet_type_##gnnedgeindextable::_jet_type_##Id, \ |
| 232 | + jetgnnindex::EdgeIndex, \ |
| 233 | + jetgnnindex::TrackIds<jetgnnindex::EdgeIndex>); \ |
| 234 | + namespace _jet_type_##gnnedgeindices \ |
| 235 | + { \ |
| 236 | + DECLARE_SOA_ARRAY_INDEX_COLUMN(_jet_type_##GnnEdge, edges); \ |
| 237 | + } \ |
| 238 | + DECLARE_SOA_TABLE(_jet_type_##GnnEdgeIds, "AOD", _description_ "I", \ |
| 239 | + _jet_type_##gnnedgeindices::_jet_type_##GnnEdgeIds); |
| 240 | +*/ |
| 241 | + |
| 242 | +// In-jet track/edge index backbone tables (own the in-jet ordering + the back-reference to the original track) |
| 243 | +DECLARE_JET_GNNTRACKINDEX_TABLE(ChargedJet, "CJGNT"); |
| 244 | +DECLARE_JET_GNNTRACKINDEX_TABLE(ChargedMCDetectorLevelJet, "CDJGNT"); |
| 245 | +// DECLARE_JET_GNNEDGEINDEX_TABLE(ChargedJet, "CJGNE"); |
| 246 | +// DECLARE_JET_GNNEDGEINDEX_TABLE(ChargedMCDetectorLevelJet, "CDJGNE"); |
| 247 | + |
| 248 | +#undef DECLARE_JET_GNNTRACKINDEX_TABLE |
| 249 | +// #undef DECLARE_JET_GNNEDGEINDEX_TABLE |
| 250 | + |
| 251 | +// Truth labels from vertexClustering() (MC detector level only). Thin tables: join positionally with |
| 252 | +// ChargedMCDetectorLevelJetGnnTracks / ChargedMCDetectorLevelJetGnnEdges (same producer loop, same row order). |
| 253 | +// TrackOriginDef must have exactly one row per GnnTracks backbone row, aligned |
| 254 | +// with GnnTrackOriginTags after GNN input selection/truncation. It is not a JTracks |
| 255 | +// extension: a JTrack used by multiple jets has a truth row for each backbone row. |
| 256 | +// tracksGnnIds() addresses this shared row position through a join containing |
| 257 | +// the backbone, the origin definition and the prediction scores. |
| 258 | +// VertexIndexDef uses the same GnnTracks row order and keeps the full-jet vertex |
| 259 | +// labels, without renumbering after GNN track selection/truncation. |
| 260 | +DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetTrackOriginDef, "AOD", "CDJTOD", jetgnnlabel::Origin); |
| 261 | +DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetVertexIndexDef, "AOD", "CDJTVD", jetgnnlabel::VertexIndex); |
| 262 | +// DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetEdgeDef, "AOD", "CDJEDF", jetgnnlabel::SameVertex); |
| 263 | + |
| 264 | +// GNN track-origin outputs, data and MC detector level. Thin tables: write one |
| 265 | +// scalar score per GnnTracks row, in the same row order as the backbone table. |
| 266 | +DECLARE_SOA_TABLE(ChargedJetGnnTrackOriginTags, "AOD", "CJTRKT", jetgnnscore::ScoreML); |
| 267 | +DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetGnnTrackOriginTags, "AOD", "CDJTRKT", jetgnnscore::ScoreML); |
| 268 | + |
| 269 | +// GNN ScoreML outputs, edge level, data and MC detector level. Thin tables: join positionally with |
| 270 | +// ChargedJetGnnEdges / ChargedMCDetectorLevelJetGnnEdges. |
| 271 | +// DECLARE_SOA_TABLE(ChargedJetGnnEdgeTags, "AOD", "CJEDGT", jetgnnscore::ScoreML); |
| 272 | +// DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetGnnEdgeTags, "AOD", "CDJEDGT", jetgnnscore::ScoreML); |
| 273 | + |
168 | 274 | } // namespace o2::aod |
169 | 275 |
|
170 | 276 | #endif // PWGJE_DATAMODEL_JETTAGGING_H_ |
0 commit comments