233 {
234 graph = SegmentEdgeGraph{};
235 graph.segments.reserve(segments.
size());
237
238
239
240
241
242
243
244 std::map<SegmentBucketKey, std::vector<const xAOD::MuonSegment*>>
245 segmentsByBucket;
247 segmentsByBucket[segmentBucketKey(*segment)].push_back(segment);
248 }
249
250 const InferenceUtils::SegmentQualityOrder betterSegment{};
251
252 std::unordered_set<const xAOD::MuonSegment*> retainedSegments;
253 retainedSegments.reserve(segments.size());
254 for (auto& [_, bucketSegments] : segmentsByBucket) {
255 std::ranges::sort(bucketSegments, betterSegment);
257 ? bucketSegments.size()
258 : std::min<std::size_t>(
259 bucketSegments.size(),
261 retainedSegments.insert(bucketSegments.begin(),
262 bucketSegments.begin() + nKeep);
263 }
264
265 std::vector<Amg::Vector3D>
pos;
266 std::vector<Amg::Vector3D>
dir;
267 std::vector<BucketSegmentFeatures> bucket;
268 pos.reserve(retainedSegments.size());
269 dir.reserve(retainedSegments.size());
270 bucket.reserve(retainedSegments.size());
271
273 if (!retainedSegments.contains(segment)) continue;
274
277 const SegmentBucketKey
key = segmentBucketKey(*segment);
278 const auto bucketIt = segmentsByBucket.find(key);
279 const int multiplicity =
280 bucketIt == segmentsByBucket.end()
281 ? 1
282 : static_cast<int>(bucketIt->second.size());
283
284 const int chamberIndex = static_cast<int>(segment->chamberIndex());
286 const int sector = segment->sector();
287
288 graph.segments.push_back(segment);
289 pos.emplace_back(position / Gaudi::Units::m);
290 dir.emplace_back(direction);
291 bucket.emplace_back(BucketSegmentFeatures{
292 chamberIndex,
layers, sector, multiplicity});
294 graph.nodeFeatures.push_back(
295 nodeFeatureValue(featureId,
pos.back(),
dir.back(), bucket.back()));
296 }
297 }
298 graph.nNodes = graph.segments.size();
299
300 if (
pos.size() != graph.nNodes ||
dir.size() != graph.nNodes ||
301 bucket.size() != graph.nNodes) {
302 ATH_MSG_ERROR(
"Inconsistent vector sizes during graph building: nodes="
303 << graph.nNodes <<
", pos=" <<
pos.size()
304 <<
", dir=" <<
dir.size() <<
", bucket=" << bucket.size());
305 return StatusCode::FAILURE;
306 }
307
308 if (graph.nNodes < 2) {
309 graph.nEdges = 0;
310 return StatusCode::SUCCESS;
311 }
312
313 const auto wrapRegularSector = [&](int sector) {
314
315
319 }
320 return sector;
321 };
322
323
324
325
326
327
328 std::unordered_map<int, std::vector<std::size_t>> nodesBySector;
329 nodesBySector.reserve(graph.nNodes);
330 for (std::size_t node = 0; node < graph.nNodes; ++node) {
331 nodesBySector[wrapRegularSector(bucket[node].sector)].push_back(node);
332 }
333
334 std::unordered_map<int, std::vector<int>> targetSectorsBySourceSector;
335 targetSectorsBySourceSector.reserve(nodesBySector.size());
336 std::size_t sectorLocalEdgeUpperBound = 0;
337 for (const auto& [sourceSector, sourceNodes] : nodesBySector) {
338 std::vector<int> targetSectors;
342 targetSectors.push_back(wrapRegularSector(sourceSector + delta));
343 }
344 for (const int targetSector : targetSectors) {
345 const auto found = nodesBySector.find(targetSector);
346 if (found == nodesBySector.end()) continue;
347 sectorLocalEdgeUpperBound += sourceNodes.size() *
found->second.size();
348 if (targetSector == sourceSector) {
349 sectorLocalEdgeUpperBound -= sourceNodes.size();
350 }
351 }
352 targetSectorsBySourceSector.emplace(sourceSector,
353 std::move(targetSectors));
354 }
355
356
357
358
359
360
361
362
363 struct UndirectedEdge {
364 std::size_t
first{0};
368 float dz{0.f};
370 float cosAngle{0.f};
371 };
372 const auto betterEdge = [](
const UndirectedEdge&
first,
373 const UndirectedEdge&
second) {
376 if (cosOrder != 0) {
377 return cosOrder < 0;
378 }
379
380 const int distanceOrder =
382 if (distanceOrder != 0) {
383 return distanceOrder < 0;
384 }
387 };
388 const auto edgeKey = [](const UndirectedEdge& edge) {
389 return (static_cast<std::uint64_t>(edge.first) << 32) |
390 static_cast<std::uint64_t>(edge.second);
391 };
392
393 const unsigned int maxEdgesPerNode =
395 const unsigned int maxEdgesPerTargetChamber =
397 const bool usePreInferenceSelection =
398 maxEdgesPerNode != 0 || maxEdgesPerTargetChamber != 0;
399 std::vector<std::vector<UndirectedEdge>> bestEdgesByNode;
400 if (usePreInferenceSelection) {
401 bestEdgesByNode.resize(graph.nNodes);
402 const unsigned int reservePerNode =
403 maxEdgesPerNode != 0 ? maxEdgesPerNode : maxEdgesPerTargetChamber;
404 for (std::vector<UndirectedEdge>& edges : bestEdgesByNode) {
405 edges.reserve(reservePerNode);
406 }
407 } else {
408 graph.edgeIndex.reserve(2 * sectorLocalEdgeUpperBound);
410 }
411 const auto appendDirectedPair = [&](const UndirectedEdge& edge) {
412 graph.edgeIndex.push_back(static_cast<int64_t>(edge.first));
413 graph.edgeIndex.push_back(static_cast<int64_t>(edge.second));
414 graph.edgeFeatures.insert(
415 graph.edgeFeatures.end(),
416 {edge.dx, edge.dy, edge.dz, edge.distance, edge.cosAngle,
417 float(bucket[edge.first].chamberIndex ==
418 bucket[edge.second].chamberIndex),
419 float(bucket[edge.first].sector == bucket[edge.second].sector)});
420
421 graph.edgeIndex.push_back(static_cast<int64_t>(edge.second));
422 graph.edgeIndex.push_back(static_cast<int64_t>(edge.first));
423 graph.edgeFeatures.insert(
424 graph.edgeFeatures.end(),
425 {-edge.dx, -edge.dy, -edge.dz, edge.distance, edge.cosAngle,
426 float(bucket[edge.first].chamberIndex ==
427 bucket[edge.second].chamberIndex),
428 float(bucket[edge.first].sector == bucket[edge.second].sector)});
429 };
430
431 const auto retainForNode = [&](std::size_t node,
432 const UndirectedEdge& candidate) {
433 std::vector<UndirectedEdge>& retained = bestEdgesByNode[node];
434 const std::size_t
other = candidate.first == node ? candidate.second
435 : candidate.first;
436 const int targetChamber = bucket[
other].chamberIndex;
437
438 if (maxEdgesPerTargetChamber != 0) {
439 unsigned int sameChamberCount = 0;
440 auto worstSameChamber = retained.end();
441 for (auto it = retained.begin(); it != retained.end(); ++it) {
442 const std::size_t retainedOther =
443 it->first == node ?
it->second :
it->first;
444 if (bucket[retainedOther].chamberIndex != targetChamber) continue;
445 ++sameChamberCount;
446 if (worstSameChamber == retained.end() ||
447 betterEdge(*worstSameChamber, *it)) {
448 worstSameChamber =
it;
449 }
450 }
451 if (sameChamberCount >= maxEdgesPerTargetChamber) {
452 if (!betterEdge(candidate, *worstSameChamber)) return;
453 *worstSameChamber = candidate;
454 } else {
455 retained.push_back(candidate);
456 }
457 } else {
458 retained.push_back(candidate);
459 }
460
461 if (maxEdgesPerNode != 0 && retained.size() > maxEdgesPerNode) {
462 auto worst = retained.begin();
463 for (auto it = std::next(retained.begin()); it != retained.end(); ++it) {
464 if (betterEdge(*worst, *it))
worst =
it;
465 }
466 retained.erase(worst);
467 }
468 };
469
470 std::size_t candidatePairs = 0;
471 for (std::size_t first = 0;
first < graph.nNodes; ++
first) {
472 const auto sectorsIt =
473 targetSectorsBySourceSector.find(wrapRegularSector(bucket[first].sector));
474 if (sectorsIt == targetSectorsBySourceSector.end()) continue;
475 for (const int sector : sectorsIt->second) {
476 const auto targetIt = nodesBySector.find(sector);
477 if (targetIt == nodesBySector.end()) continue;
478
479 for (const std::size_t second : targetIt->second) {
480
481 if (second <= first) continue;
482 if (sectorDistance(bucket[first].sector, bucket[second].sector,
485 continue;
486 }
488 bucket[first].chamberIndex == bucket[second].chamberIndex) {
489 continue;
490 }
491 const float cosAngle =
static_cast<float>(
dir[
first].dot(dir[second]));
493
495 const UndirectedEdge candidate{
498 static_cast<float>(delta.x()),
499 static_cast<float>(delta.y()),
500 static_cast<float>(delta.z()),
501 static_cast<float>(delta.mag()),
502 cosAngle};
503 ++candidatePairs;
504
505 if (!usePreInferenceSelection) {
506 appendDirectedPair(candidate);
507 } else {
508 retainForNode(first, candidate);
509 retainForNode(second, candidate);
510 }
511 }
512 }
513 }
514 std::size_t retainedPairs = candidatePairs;
515 if (usePreInferenceSelection) {
516 const unsigned int selectedReservePerNode =
517 maxEdgesPerNode != 0 ? maxEdgesPerNode : maxEdgesPerTargetChamber;
518 std::unordered_set<std::uint64_t> selectedKeys;
519 selectedKeys.reserve(graph.nNodes * selectedReservePerNode);
520 std::vector<UndirectedEdge> selectedEdges;
521 selectedEdges.reserve(graph.nNodes * selectedReservePerNode);
522
523 for (const std::vector<UndirectedEdge>& nodeEdges : bestEdgesByNode) {
524 for (const UndirectedEdge& edge : nodeEdges) {
525 if (selectedKeys.insert(edgeKey(edge)).second) {
526 selectedEdges.push_back(edge);
527 }
528 }
529 }
530 std::sort(selectedEdges.begin(), selectedEdges.end(),
531 [](const UndirectedEdge& first,
532 const UndirectedEdge& second) {
533 if (first.first != second.first) {
534 return first.first < second.first;
535 }
537 });
538
539 retainedPairs = selectedEdges.size();
540 graph.edgeIndex.reserve(4 * retainedPairs);
542 for (const UndirectedEdge& edge : selectedEdges) {
543 appendDirectedPair(edge);
544 }
545 }
546 graph.nEdges = graph.edgeIndex.size() / 2;
547 const std::size_t nodesBeforeIsolatedNodeDrop = graph.nNodes;
549 std::vector<unsigned char>
active(graph.nNodes, 0);
550 for (const int64_t index : graph.edgeIndex) {
552 }
553 const std::size_t activeNodes =
554 std::count(
active.begin(),
active.end(),
static_cast<unsigned char>(1));
555 if (activeNodes != graph.nNodes) {
556 std::vector<std::size_t> oldToNew(graph.nNodes, graph.nNodes);
557 std::vector<const xAOD::MuonSegment*> compactedSegments;
558 std::vector<float> compactedNodeFeatures;
559 compactedSegments.reserve(activeNodes);
561 for (std::size_t oldNode = 0; oldNode < graph.nNodes; ++oldNode) {
562 if (!active[oldNode]) continue;
563 oldToNew[oldNode] = compactedSegments.size();
564 compactedSegments.push_back(graph.segments[oldNode]);
565 const auto featureBegin = graph.nodeFeatures.begin() +
567 compactedNodeFeatures.insert(compactedNodeFeatures.end(),
568 featureBegin,
570 }
571 for (int64_t& index : graph.edgeIndex) {
572 index =
static_cast<int64_t
>(oldToNew[
static_cast<std::size_t
>(
index)]);
573 }
574 graph.segments = std::move(compactedSegments);
575 graph.nodeFeatures = std::move(compactedNodeFeatures);
576 graph.nNodes = activeNodes;
577 }
578 }
579 ATH_MSG_DEBUG(
"buildGraph: input segments=" << segments.size()
580 << ", kept nodes=" << graph.nNodes
581 << ", nodes before isolated-node drop=" << nodesBeforeIsolatedNodeDrop
583 << ", candidate pairs=" << candidatePairs
584 << ", retained pairs=" << retainedPairs
585 << ", built directed edges=" << graph.nEdges
587 << ", per-target-chamber cap=" << maxEdgesPerTargetChamber
590 << ", sector-local reserve=" << sectorLocalEdgeUpperBound);
591 return StatusCode::SUCCESS;
592}
const SpacePointBucket * parentBucket() const
Returns the bucket out of which the seed was formed.
const SegmentSeed * parent() const
Returns the seed out of which the segment was built.
float distance(const Amg::Vector3D &p1, const Amg::Vector3D &p2)
calculates the distance between two point in 3D space
Eigen::Matrix< double, 3, 1 > Vector3D
int compareFloat(float first, float second)
Three-way float comparison which orders NaN after all numeric values.
int compareFloatDescending(float first, float second)
Three-way descending comparison which also orders NaN last.
SegmentNodeFeatureId
Identifier for each node feature in segment-based GNNs.
const Segment * detailedSegment(const xAOD::MuonSegment &seg)
Helper function to navigate from the xAOD::MuonSegment to the MuonR4::Segment.
const Amg::Vector3D & direction() const
Method to retrieve the direction at the Intersection.
const Amg::Vector3D & position() const
Method to retrieve the position of the Intersection.
void sort(typename DataModel_detail::iterator< DVL > beg, typename DataModel_detail::iterator< DVL > end)
Specialization of sort for DataVector/List.
MuonSegment_v1 MuonSegment
Reference the current persistent version: