diff --git a/Common/DataModel/CollisionAssociation.h b/Common/DataModel/CollisionAssociation.h index 958f97cb6ae..73ff1d19c77 100644 --- a/Common/DataModel/CollisionAssociation.h +++ b/Common/DataModel/CollisionAssociation.h @@ -22,13 +22,17 @@ namespace o2::aod namespace track_association { -DECLARE_SOA_INDEX_COLUMN(Collision, collision); //! Collision index -DECLARE_SOA_INDEX_COLUMN(Track, track); //! Track index +DECLARE_SOA_INDEX_COLUMN(Collision, collision); //! Collision index +DECLARE_SOA_INDEX_COLUMN(Track, track); //! Track index +DECLARE_SOA_SELF_ARRAY_INDEX_COLUMN(CompatibleColl, compatibleColl); //! Array of collision indices } // namespace track_association DECLARE_SOA_TABLE(TrackAssoc, "AOD", "TRACKASSOC", //! Table for track-to-collision association for e.g. HF vertex finding - tracks can appear for several collisions track_association::CollisionId, track_association::TrackId); + +DECLARE_SOA_TABLE(TrackCompColls, "AOD", "TRACKCOMPCOLLS", //! Table with vectors of collision indices stored per track + track_association::CompatibleCollIds); } // namespace o2::aod #endif // COMMON_DATAMODEL_COLLISIONASSOCIATION_H_ diff --git a/Common/TableProducer/trackToCollisionAssociator.cxx b/Common/TableProducer/trackToCollisionAssociator.cxx index 727027dad5a..ee9b4bd0ea0 100644 --- a/Common/TableProducer/trackToCollisionAssociator.cxx +++ b/Common/TableProducer/trackToCollisionAssociator.cxx @@ -30,6 +30,7 @@ using namespace o2::aod; struct TrackToCollisionAssociation { Produces association; + Produces reverseIndices; Configurable nSigmaForTimeCompat{"nSigmaForTimeCompat", 4.f, "number of sigmas for time compatibility"}; Configurable timeMargin{"timeMargin", 0.f, "time margin in ns added to uncertainty because of uncalibrated TPC"}; @@ -37,6 +38,7 @@ struct TrackToCollisionAssociation { Configurable applyIsGlobalTrackWoDCA{"applyIsGlobalTrackWoDCA", true, "flag to apply global track w/o DCA selection"}; Configurable usePVAssociation{"usePVAssociation", true, "if the track is a PV contributor, use the collision time for it"}; Configurable includeUnassigned{"includeUnassigned", false, "consider also tracks which are not assigned to any collision"}; + Configurable fillTableOfCollIdsPerTrack{"fillTableOfCollIdsPerTrack", false, "fill additional table with vector of collision ids per track"}; Filter trackFilter = (applyIsGlobalTrackWoDCA == false) || requireGlobalTrackWoDCAInFilter(); using TracksWithSel = soa::Join; @@ -54,6 +56,7 @@ struct TrackToCollisionAssociation { } void processAssocWithTime(Collisions const& collisions, + TracksWithSel const& tracksUnfiltered, TracksWithSelFilter const& tracks, AmbiguousTracks const& ambiguousTracks, BCs const& bcs) @@ -73,6 +76,9 @@ struct TrackToCollisionAssociation { } } + // define vector of vectors to store indices of compatible collisions per track + std::vector>> collsPerTrack(tracksUnfiltered.size()); + // loop over collisions to find time-compatible tracks auto trackBegin = tracks.begin(); constexpr auto bOffsetMax = 241; // 6 mus (ITS) @@ -128,6 +134,29 @@ struct TrackToCollisionAssociation { const auto trackIdx = track.globalIndex(); LOGP(debug, "Filling track id {} for coll id {}", trackIdx, collIdx); association(collIdx, trackIdx); + if (fillTableOfCollIdsPerTrack) { + if (collsPerTrack[trackIdx] == nullptr) { + collsPerTrack[trackIdx] = std::make_unique>(); + } + collsPerTrack[trackIdx].get()->push_back(collIdx); + } + } + } + } + + // create reverse index track to collisions if enabled + if (fillTableOfCollIdsPerTrack) { + std::vector empty{}; + for (const auto& track : tracksUnfiltered) { + + const auto trackId = track.globalIndex(); + if (collsPerTrack[trackId] == nullptr) { + reverseIndices(empty); + } else { + // for (const auto& collId : *collsPerTrack[trackId]) { + // LOGP(info, " -> Coll id {}", collId); + // } + reverseIndices(*collsPerTrack[trackId].get()); } } } @@ -135,22 +164,39 @@ struct TrackToCollisionAssociation { PROCESS_SWITCH(TrackToCollisionAssociation, processAssocWithTime, "Use track-to-collision association based on time", false); - void processStandardAssoc(Collision const& collision, - TracksWithSel const& tracksPerCollision) + Preslice tracksPerCollisions = aod::track::collisionId; + + void processStandardAssoc(Collisions const& collisions, + TracksWithSel const& tracks) { // we do it for all tracks, to be compatible with Run 2 analyses - for (const auto& track : tracksPerCollision) { - bool hasGoodQuality = true; - if (applyIsGlobalTrackWoDCA && !track.isGlobalTrackWoDCA()) { - hasGoodQuality = false; - } else if (applyMinimalTrackSelForRun2) { - unsigned char itsClusterMap = track.itsClusterMap(); - if (!(track.tpcNClsFound() >= 50 && track.flags() & o2::aod::track::ITSrefit && track.flags() & o2::aod::track::TPCrefit && (TESTBIT(itsClusterMap, 0) || TESTBIT(itsClusterMap, 1)))) { + for (const auto& collision : collisions) { + auto tracksThisCollision = tracks.sliceBy(tracksPerCollisions, collision.globalIndex()); + for (const auto& track : tracksThisCollision) { + bool hasGoodQuality = true; + if (applyIsGlobalTrackWoDCA && !track.isGlobalTrackWoDCA()) { hasGoodQuality = false; + } else if (applyMinimalTrackSelForRun2) { + unsigned char itsClusterMap = track.itsClusterMap(); + if (!(track.tpcNClsFound() >= 50 && track.flags() & o2::aod::track::ITSrefit && track.flags() & o2::aod::track::TPCrefit && (TESTBIT(itsClusterMap, 0) || TESTBIT(itsClusterMap, 1)))) { + hasGoodQuality = false; + } + } + if (hasGoodQuality) { + association(collision.globalIndex(), track.globalIndex()); } } - if (hasGoodQuality) { - association(collision.globalIndex(), track.globalIndex()); + } + + // create reverse index track to collisions if enabled + std::vector empty{}; + if (fillTableOfCollIdsPerTrack) { + for (const auto& track : tracks) { + if (track.has_collision()) { + reverseIndices(std::vector{track.collisionId()}); + } else { + reverseIndices(empty); + } } } }